###### tags: `frontend`, `shaping`
[TOC]
GT4Py Frontend design - Evaluation test cases
=============================================
## Examples overview
A collection of code snippets to cover all patterns we would like to express in the frontend. The snippets should be expressed using the *semantic IR* as well as the proposed frontend syntax flavors.
__[WIP]__ This document is work-in-progress. We are not trying to design "the" frontend, but a user-friendly (see [goals](https://hackmd.io/ZGMjwyALQW2iW0784Fp5ew#Goals) in the main document) proposal that maps easily to the *semantic IR*.
## Driver code example
```python=
T = TypeVar("T")
DT = TypeVar("DT", bound=gt.ScalarType)
# TODO: Field typing conventions
field: Field[[*Dimensions], DT] #
# Axes names: V or Vertex, V2V or VV or V_V
# ...
# TODO(tehrengruber): tentative
# mapping from axis to connectivity / neighbor table
# user can provide something different (e.g. diamond connectivity)
connectivities = {
# Connectivity[From, To, max_neighbors, has_skip_values]
E2V: Connectivity[E, V, 2, False](range(0, 10), np.array([[0, 1], ...]))
# ...
}
# -- Beautified iterator view --
# Decorator to define local operations (i.e. mapping iterators to values)
@gt.beautified_iterator_view
def my_local_op(field: Iterator[Vertex, Vertex]) -> Value:
if field() < 0:
return -2. * field()
else:
return 0.5 * field()
# Feature for debugging & testing: direct execution of stencils
new_field = my_local_op[my_vertex_domain](input_field)
# -- Lifted iterator view --
# Decorator to define operations from fields to a field
@gt.lazy_field_view
def my_field_op(
field_a: Field[Vertex], field_b: Field[Vertex],
) -> Field[Vertex]:
return field_a + field_b
# Program/fencil on the fly feature:
# outputs = my_local_op[vertex_domain](*inputs)
# note that you can't pass the out kwarg here
# -- Entry point to the toolchain from Python ('fencil') --
# Input fields can only be updated inside fencils
@gt.program
def my_program(
field_a: Field[Vertex],
field_b: Field[Vertex],
out_field: Field[Vertex],
):
my_field_op(field_a, field_b, out=[out_field])
vertex_domain = gt.IndexRange(Vertex, 0, 3)
my_local_op[vertex_domain](*inputs, out=[outputs])
@gt.program
def my_dynamic_program(
field_a: Field[Vertex],
field_b: Field[Vertex],
out_field: Field[Vertex],
point_a: int,
point_b: int
):
vertex_domain = gt.IndexRange(Vertex, point_a-3, point_a + 2)
my_local_op[other_domain](*inputs, out=[outputs])
# If we later want to allow the creation of temporaries
# inside the program, we can add the feature in a more or less
# backward compatible way like this:
# tmp_field = new_temp()
# my_local_op[vertex_domain](*inputs, out=[tmp_field])
```
## Frontends overview
Possible frontends
- Plain iterator: _deref()_, _shift()_ and _lift()_ as explicit functions
- Beautified iterator:
+ `name()` -> _deref(name)_
+ `name(offset)` -> _deref(shift(offset)(name))_
+ `func_name[...]` -> _lift(func_name)_
- Lifted iterator (aka _Lazy Field_): everything within an expression is dereferenced, put in lambda and lifted (i.e. everything is an iterator).
### Beautified Iterator
Features:
- Trivial translation to Iterator IR as there are only syntactic changes
How to lower:
- combined offset `O[0] == (O, 0) == O + 0`
- `it(offsets, ...)` expands to `deref(shift(offsets, ...)(it))`
- trivial: `it()` -> `deref(shift()(it)) == deref(it)`
- `fun[...]` -> `lift(fun)`
### Lifted Iterator (aka (Lazy)Field)
Features:
- can be constructed from the Iterator Model
- no need for `lift()` in the syntax as stencils take iterators and return iterators
How to lower:
- in every expression:
- expand `it(offsets, ...)` to `deref(shift(offsets.reverse(), ...)(it))`
<!-- discuss(till): it is not sufficient to just deref as it was written before, as you need to be able to shift/remap inside the expression. The reverse it is needed for (C2E, E2V) to read as from vertex to edge to cell. Alternatively we could disallow such shifts.-->
- put the expression in a lambda that has all iterators as parameters
- _lift_ the lambda
- call the lifted lambda with the original iterators as arguments
### Optional syntactic sugar
- reductions with list comprehension:
- `sum(fun(it1(n), ..., itN(n)) for n in Offset)` -> `reduce(lambda acc, val1, ..., valN: acc + fun(va1, ..., valN), 0)(shift(Offset)(it1), ..., shift(Offset)(itN))`
## Examples
### Regular horizontal stencil (id: S-V4)
+ _Plain Iterator:_
```python=
@fundef
def sum_of_all_vertex_neighbors_of_a_vertex(inp):
# length(V2V) == 4
return (
deref(shift(V2V, 0)(inp))
+ deref(shift(V2V, 1)(inp))
+ deref(shift(V2V, 2)(inp))
+ deref(shift(V2V, 3)(inp))
)
```
+ _Beautified iterator:_
```python=
@gt.beautified_iterator_view
def sum_of_all_vertex_neighbors_of_a_vertex(inp: Iterator[Vertex, Vertex]) -> DT:
return inp(V2V[0]) + inp(V2V[1]) + inp(V2V[2]) + inp(V2V[3])
```
+ _Lifted iterator view:_
```python=
def sum_of_all_vertex_neighbors_of_a_vertex(inp: Field[Vertex]) -> Field[Vertex]:
return (
inp(V2V[0]) + inp(V2V[1]) + inp(V2V[2]) + inp(V2V[3])
) # It would be an out-of-bounds error if `V2V[i]` is not pointing to a valid position
def sum_of_all_vertex_neighbors_of_a_vertex(inp: Field[Vertex]):
# not all length(V2V) have to be the same
return sum(inp(v) for v in V2V)
```
### Regular horizontal stencil with shifting (id: S-C2E2V)
Semantic: Compute the value of the first vertex neighbor of the first edge neighbor of the current cell.
```
Mesh Iterator Field
v3
* * *
/ \ / o\ / X \
e3 / c1 \ e2 / ↓ \ / ↑ \
*------* *------* *-------*
v1 e1 v2 X ← o →
```
+ _Plain iterator:_
```python=
@fundef
def c2e2v_stencil(inp): # iterator pointing to Cell for Vertex field
return deref(shift(E2V, 0)(shift(C2E, 0)(inp)))
```
+ _Beautified iterator:_
```python=
@gt.beautified_iterator_view
def c2e2v_stencil(inp: Iterator[Cell, Vertex]) -> DT:
# note: it is not possible to write this as two seperate shift
# without a lift operation, as the call operator automatically
# dereferences
# return inp(C2E[0])(E2V[0]) # not allowed (trying to shift a value)
# return inp(C2E[0], E2V[0]) # don't do it now, because need extra explanation
```
+ _Lifted iterator view:_
```python=
@fundef
def c2e2v_stencil(inp: Field[Vertex]) -> Field[Cell]:
# inp(E2V[0])(C2E[0])
# inp(E2V[0], C2E[0])
tmp1: Field[Edge] = inp(E2V[0]) # (lambda inp: inp(E2V[0]))[...]
tmp2: Field[Cell] = tmp1(C2E[0])
# todo: discuss
# inp(C2E[0], E2V[0])
return tmp2
```
```python=
def c2e2v_stencil(inp: Iterator[Cell, Vertex]):
edge_neigh_of_cell: Iterator[Edge, Vertex] = shift(C2E, 0)(inp)
vertex_neigh_of_edge: Iterator[Vertex, Vertex] = shift(E2V, 0)(edge_neigh_of_cell)
return deref(vertex_neigh_of_edge)
def c2e2v_stencil(inp):
vertex_neigh_of_edge = lift(lambda x: deref(shift(E2V, 0)(x)))(inp)
edge_neigh_of_cell = lift(lambda x: deref(shift(C2E, 0)(x)))(vertex_neigh_of_edge)
return deref(edge_neigh_of_cell)
def c2e2v_stencil(inp: Field[Vertex]) -> Field[Cell]:
vertex_neigh_of_edge: Field[Edge] = inp(E2V[0])
edge_neigh_of_cell: Field[Cell] = vertex_neigh_or_edge(C2E[0])
return edge_neigh_of_cell
```
### Regular horizontal stencil with shifting and lifting (id: S-C2E2V-lifted)
Note: this is the same computation as above, but with an additional lifting operation.
+ _Plain iterator:_
```python=
# inp is an iterator with E-position
@fundef
def first_vertex_nb(inp):
# shift(E2V, 0)(inp) returns an iterator on a Vertex position
# what type of buffer inp refers to is undefined at this point
return deref(shift(E2V, 0)(inp))
@fundef
def c2e2v_stencil(inp):
# lift(first_vertex_nb) takes an iterator on an edge position
# an iterator on an edge position
# note: The following can not be expressed in the beautified and
# lifted iterator view as a shift always involves a deref.
return first_vertex_nb(shift(C2E,0)(inp))
return deref(lift(first_vertex_nb)(shift(C2E, 0)(inp)))
return deref(shift(C2E, 0)(lift(first_vertex_nb)(inp)))
```
+ _Beautified iterator view_:
```python=
@gt.beautified_iterator_view
def first_vertex_nb(inp: Iterator[Edge, Vertex]) -> DT:
return inp(E2V[0])
@gt.beautified_iterator_view
def c2e2v_stencil(inp: Iterator[Cell, Vertex]) -> DT:
# first_vertex_nb[...](inp) returns an iterator of type
# Iterator[Cell, Edge], where Cell originates from the
# inp argument and Edge comes from first_vertex_nb signature.
# Note that the position type of all arguments of a shiftable
# stencil need to agree.
# first_cell_nb[...](inp) can also be seen as an edge field
# (originating from inp) that we can access by shifting our
# current position to an edge in the vicinity.
return first_vertex_nb[...](inp)(C2E[0])
```
+ _Lifted iterator view:_
```python=
def first_vertex_nb(inp: Field[Vertex]) -> Field[Edge]:
return inp(E2V[0])
def c2e2v_stencil(inp: Field[Vertex]) -> Field[Cell]:
return first_vertex_nb(inp)(C2E[0])
```
<!--
```python=
def stencil(vertex_field: Iterator[Edge, Vertex]) -> DT:
return field(E2V[0])
def bar(vertex_field: Iterator[Vertex, Vertex]):
field_edge = stencil[...](vertex_field)
field_vertex = field_edge(V2E[0])
return field_vertex
def bar_without_lift():
# note how the ordering changes
return field_vertex.shift(V2E[0]).shift(E2V[0])
```
```python=
def stencil(vertex_field: Field[Vertex]) -> Field[Edge]:
# the value at the 0-th edge for all vertices that E2V maps to
return field(E2V[0])
def bar(vertex_field: Field[Vertex]):
field_edge = stencil(vertex_field)
field_vertex = field_edge(V2E[0])
return field_vertex
```
-->
### Regular horizontal stencil with reduction (id: S-VRED)
+ _Plain iterator:_
```python=
@fundef
def sum_of_all_vertex_neighbors_of_a_vertex(inp):
return reduce(lambda a, b: a+b, shift(V2V)(inp))
```
+ _Beautified iterator view_:
```python=
@gt.beautified_iterator_view
def sum_of_all_vertex_neighbors_of_a_vertex(inp: Iterator[Vertex, Vertex]) -> DT:
return sum(inp(v) for v in V2V)
```
+ _Lifted iterator view:_
```python=
def sum_of_all_vertex_neighbors_of_a_vertex(inp: Field[Vertex]):
return sum(inp(v) for v in V2V)
```
### Sparse-field (Hannes) (id: S-SF)
+ _Plain iterator:_
```python=
def sum_of_all_vertex_neighbors_of_a_vertex(inp): # inp is an iterator in the sparse field state
return reduce(lambda a, b: a+b, 0)(inp)
def driver(inp) # non sparse field
return sum_of_all_vertex_neighbors_of_a_vertex(shift(V2V)(inp))
```
+ _Lifted iterator view:_
```python=
def sum_of_all_vertex_neighbors_of_a_vertex(inp: Iterator[Vertex, ?]) -> Iterator[Vertex]:
return reduce(lambda a, b: a+b, 0)(inp)
def driver(inp: Iterator[Vertex]) -> Iterator[Vertex]
return sum_of_all_vertex_neighbors_of_a_vertex(inp(V2V))
```
### Nabla (id: S-FVM-NBL)
+ _Plain iterator:_
```python=
@fundef
def compute_zavgS(pp, S_M):
zavg = 0.5 * (deref(shift(E2V, 0)(pp)) + deref(shift(E2V, 1)(pp)))
# zavg = 0.5 * reduce(lambda a, b: a + b, 0)(shift(E2V)(pp))
# zavg = 0.5 * library.sum()(shift(E2V)(pp))
return deref(S_M) * zavg
@fundef
def compute_pnabla(pp, S_M, sign, vol):
zavgS = lift(compute_zavgS)(pp, S_M)
# pnabla_M = reduce(lambda a, b, c: a + b * c, 0)(shift(V2E)(zavgS), sign)
# pnabla_M = library.sum(lambda a, b: a * b)(shift(V2E)(zavgS), sign)
pnabla_M = library.dot(shift(V2E)(zavgS), sign)
return pnabla_M / deref(vol)
```
+ _Beautified iterator view:_
```python=
@gt.beautified_iterator_view
def compute_zavgS(pp: Iterator[Edge, Vertex], S_M: Iterator[Edge, Edge]) -> DT:
zavg = 0.5 * (pp(E2V[0]) + pp(E2V[1]))
return S_M() * zavg
@gt.beautified_iterator_view
def compute_pnabla(
pp: Iterator[Vertex, Vertex],
S_M: Iterator[Vertex, Edge],
sign: Iterator[Vertex, ?],
vol: Iterator[Vertex, Vertex]
) -> DT:
zavgS: Iterator[Vertex, Edge] = compute_zavgS[...](pp, S_M)
# todo: sign() could return Iterator[Vertex, Edge]
pnabla_M = sum(zavgS(e)*sign()(e) for e in V2E)
return pnabla_M / vol()
```
+ _Lifted iterator view:_
```python=
def compute_zavgS(pp: Field[Vertex], S_M: Field[Edge]) -> Field[Edge]:
zavg: Field[Edge] = 0.5 * (pp(E2V[0]) + pp(E2V[1]))
# zavg = lift(lambda a, b: 0.5 * (deref(shift(E2V, 0)(a)) + deref(shift(E2V, 1)(a))))(pp)
return S_M * zavg
# return lift(lambda a,b: deref(a)*deref(b))(S_M, zavg)
def compute_pnabla(pp: Field[Vertex], S_M: Field[Edge], sign, vol: Field[Vertex]) -> Field[Vertex]:
zavgS: Field[Edge] = compute_zavgS(pp, S_M)
pnabla_M: Field[Vertex] = sum(
zavgS(v) * sign[i] for i, v in enumerate(V2E)
#zavgS(e) * sign[i] for i, e in enumerate(V2E)
) # TODO I don't like adressing the sign field with `i`, dusk solves this by inprecise syntax
# pnabla_M = lift(reduce( lambda a, b, c: a+b*c, 0))(shift(V2E)(zavgS), sign)
return pnabla_M / vol
# return lift(lambda a,b: deref(a)/deref(b))(pnabla_M, vol)
```
### ICON reductions & dusk/dawn weights (id: S-ICO-DRED)
+ _Beautified iterator:_
```python=
@gt.beautified_iterator_view
def reductions_and_weights(
cells_a: Iterator[Edge, Cell],
vertices_a: Iterator[Edge, Vertex]
) -> DT:
# simple difference between neighboring cells of an edge:
edges_a: DT = cells_a(E2C[1]) - cells_a(E2C[0])
# difference between the upper and lower vertex of the diamond
edges_a: DT = vertices_a(E2C[1], C2V[0]) - vertices_a(E2C[0], C2V[0])
# sum of the upper and lower vertex of the diamond (selective/filtered sum)
edges_a: DT = vertices_a(E2C[0], C2V[0]) + vertices_a(E2C[1], C2V[0])
return edges_a
```
+ _Lifted iterator_
```python=
def reductions_and_weights(
cells_a: Field[Cell],
vertices_a: Field[Vertex]
) -> Field[Edge]:
# simple difference between neighboring cells of an edge:
edges_a: Field[Edge] = cells_a(E2C[1]) - cells_a(E2C[0])
# difference between the upper and lower vertex of the diamond
edges_a: Field[Edge] = vertices_a(E2C[1], C2V[0]) - vertices_a(E2C[0], C2V[0])
# sum of the upper and lower vertex of the diamond (selective/filtered sum)
edges_a: Field[Edge] = vertices_a(E2C[0], C2V[0]) + vertices_a(E2C[1], C2V[0])
return edges_a
```
+ dusk syntax:
```python=
@stencil
def reductions_and_weights(
edges_a: Field[Edge],
cells_a: Field[Cell],
vertices_a: Field[Vertex],
):
with domain.upward:
# These examples all involve the following two triangle shape:
# (we call this the diamond shape, orientation is irrelevant)
#
# *
# / \
# / \
# *-----*
# \ /
# \ /
# *
# simple difference between neighboring cells of an edge:
edges_a = sum_over(Edge > Cell, cells_a, weights=[-1, 1])
# difference between the upper and lower vertex of the diamond
edges_a = sum_over(Edge > Cell > Vertex, vertices_a, weights=[-1, 1, 0, 0])
# sum of the upper and lower vertex of the diamond (selective/filtered sum)
edges_a = sum_over(Edge > Cell > Vertex, vertices_a, weights=[1, 1, 0, 0])
```
### [WIP] Scan pass (id: S-KSUM)
+ _Plain iterator:_
```python=
@fundef
def sum_scanpass(state, inp):
return if_(is_none(state), deref(inp), state + deref(inp))
@fundef
def ksum(inp):
return scan(sum_scanpass, True, None)(inp)
@fendef(column_axis=KDim)
def ksum_fencil(i_size, k_size, inp, out):
closure(
domain(named_range(IDim, 0, i_size), named_range(KDim, 0, k_size)),
ksum,
[out],
[inp],
)
```
+ _Beautified iterator:_
```python=
@gt.beautified_iterator_view
def sum_scanpass(state: Optional[DT], inp: Iterator[Vertex]):
if state is None:
return inp()
return state+inp()
@gt.column_operator(column_axis=K)
def ksum(inp: Iterator[[V, K], [V, K]]):
c: Iterator[V] = scan(sum_scanpass, forward=True, init=None)(inp)
return c
```
+ _Indexing proposal:_
```python=
# option 1
@gt.local_operator
def sum_scanpass(state: Optional[DT], inp: Field[], *, position):
(v, k) = position
if state is None:
return inp[v, k]
return state+inp[v, k]
# option 2
@gt.scan_pass(forward=True, init=None, column_axis=K)
def sum_solver(state, inp, *, position):
# ... same as before
@gt.column_operator(column_axis=K)
def ksum(inp: Field[[V, K]], inp_2d: Field[[V]]):
(v,) = position
# option 1
return scan(sum_scanpass, forward=True, init=None)(inp, position=position)
# option 2
return sum_solver(inp, position=position)
@gt.program(...)
def ksum_fencil(i_size, k_size, inp, out):
domain = IndexRange[IDim](0, i_size)*IndexRange[IDim](0, k_size)
return ksum[domain](inp, out=out)
```
### Tridiag solver (id: S-TRI)
+ _Plain iterator:_
```python=
@fundef
def tridiag_forward(state, a, b, c, d):
# not tracable
# if is_none(state):
# cp_k = deref(c) / deref(b)
# dp_k = deref(d) / deref(b)
# else:
# cp_km1, dp_km1 = state
# cp_k = deref(c) / (deref(b) - deref(a) * cp_km1)
# dp_k = (deref(d) - deref(a) * dp_km1) / (deref(b) - deref(a) * cp_km1)
# return make_tuple(cp_k, dp_k)
# variant a
# return if_(
# is_none(state),
# make_tuple(deref(c) / deref(b), deref(d) / deref(b)),
# make_tuple(
# deref(c) / (deref(b) - deref(a) * nth(0, state)),
# (deref(d) - deref(a) * nth(1, state))
# / (deref(b) - deref(a) * nth(0, state)),
# ),
# )
# variant b
def initial():
return make_tuple(deref(c) / deref(b), deref(d) / deref(b))
def step():
return make_tuple(
deref(c) / (deref(b) - deref(a) * nth(0, state)),
(deref(d) - deref(a) * nth(1, state)) / (deref(b) - deref(a) * nth(0, state)),
)
return if_(is_none(state), initial, step)()
@fundef
def tridiag_backward(x_kp1, cp, dp):
# if is_none(x_kp1):
# x_k = deref(dp)
# else:
# x_k = deref(dp) - deref(cp) * x_kp1
# return x_k
return if_(is_none(x_kp1), deref(dp), deref(dp) - deref(cp) * x_kp1)
@fundef
def solve_tridiag(a, b, c, d):
tup = lift(scan(tridiag_forward, True, None))(a, b, c, d)
cp = nth(0, tup)
dp = nth(1, tup)
return scan(tridiag_backward, False, None)(cp, dp)
@fendef
def fen_solve_tridiag(i_size, j_size, k_size, a, b, c, d, x):
closure(
domain(
named_range(IDim, 0, i_size),
named_range(JDim, 0, j_size),
named_range(KDim, 0, k_size),
),
solve_tridiag,
[x],
[a, b, c, d],
)
```
+ _Beautified iterator:_
```python=
@gt.beautified_iterator_view
def tridiag_forward(state: Optional[DT], a, b, c, d):
if state == None:
cp_k = c() / b()
dp_k = d() / b()
else:
cp_km1, dp_km1 = state
cp_k = c() / (b() - a() * cp_km1)
dp_k = (d() - a() * dp_km1) / (b() - a() * cp_km1)
return (cp_k, dp_k)
@gt.beautified_iterator_view
def tridiag_backward(x_kp1: Optional[DT], cp, dp):
if x_kp1 == None:
x_k = dp()
else:
x_k = dp() - cp() * x_kp1
return x_k
@gt.column_operator(column_axis=K)
def ksum(inp: Iterator[[V, K], [V, K]]):
c: Iterator[V] = scan(sum_scanpass, forward=True, init=None)(inp)
return c
```
### Laplacian Cartesian (id: S-L4)
+ _Plain iterator:_
```python=
@fundef
def laplacian(inp):
return -4.0 * deref(inp) + (
deref(shift(I, 1)(inp))
+ deref(shift(I, -1)(inp))
+ deref(shift(J, 1)(inp))
+ deref(shift(J, -1)(inp))
)
```
+ _Beautified Iterator:_
```python=
@gt.beautified_iterator_view
@fundef
def laplacian(inp: Iterator[(I, J), (I, J)]) -> DT:
return -4.0 * inp() + inp(I+1) + inp(I-1) + inp(J+1) + inp(J-1)
```
+ _Lifted iterator:_
```python=
@fundef
def laplacian(inp: Field[I, J]) -> Field[I, J]:
return -4.0 * inp() + inp(I+1) + inp(I-1) + inp(J+1) + inp(J-1)
```
### Horizontal diff Cartesian (id: S-HD)
+ _Plain iterator:_
```python=
@fundef
def flux_x(inp):
lap = lift(laplacian)(inp)
flux = deref(shift(I, -1/2)(lap)) - deref(shift(I, 1/2)(lap))
return if_(flux * deref(shift((I, 1/2))(inp)) - deref(shift((I, -1/2))(inp)) > 0, 0, flux)
# code without staggered dimensions looks like this:
def flux_x(inp):
lap = lift(laplacian)(inp)
flux = deref(lap) - deref(shift(I, 1)(lap))
return if_(flux * (deref(shift(I, 1)(inp)) - deref(inp)) > 0.0, 0.0, flux)
```
+ _Beautified Iterator:_
```python=
@gt.beautified_iterator_view
def laplacian(inp: Iterator[(I, J), (I, J)]):
return ...
@gt.beautified_iterator_view
def flux_x(inp: Iterator[(I+1/2, J), (I, J)]):
lap = laplacian[...](inp)
flux = lap(I-1/2) - lap(I+1/2)
if flux * inp(I+1/2) - I(I-1/2) > 0.0:
return 0
return flux
flux_x[i_staggered_domain](inp)
```
+ _Lifted Iterator:_
```python=
def flux_x(inp: Field[(I, J)]):
lap = laplacian[...](inp)
flux = lap(I-1/2) - lap(I+1/2)
if flux * inp(I+1/2) - I(I-1/2) > 0.0:
return 0
return flux
```
Note: it is not possible in this proposal to specify that
this operator was formulated to be posed on the staggered domain.
### [WIP] BCs Cartesian (id: S-L4-BC)
+ _Accessor proposal:_
```python=
@gt.local_operator(domain=(I, J))
def laplacian(boundary: IndexRange[I, J], inp: Accessor[[I, J], DT]):
if i, j in boundary:
return 0.
return -4.0 * inp() + (inp(I+1) + inp(I-1) + inp(J+1) + inp(J-1))
@gt.local_operator(domain=(I, J))
def laplacian(inp: Accessor[[I, J], DT]):
if i, j in gt.domain[1:-1, 1:-1]:
return 0.
return -4.0 * inp() + (inp(I+1) + inp(I-1) + inp(J+1) + inp(J-1))
```
### [WIP] BCs Unstructured
```python=
@gt.local_operator
def laplacian(inp: Field[[V], DT], on_boundary: Field[[V], bool], *, position: Position[V]):
(v,) = position
if on_boundary:
return 0.
return -4.0 * inp[v] + sum(inp[v2] for v2 in neighbors(v, V2V))
```
```python=
@gt.local_operator
def laplacian(inp: Field[[V], DT], *, position: Position[V]):
(v,) = position
if v in boundary_vertices:
return 0.
return -4.0 * inp[v] + sum(inp[v2] for v2 in neighbors(v, V2V))
```
### [WIP] PBCs Cartesian
### [WIP] Elliptic solver (id: S-ES)
*Positional with field composition*
```python=
@gt.local_operator
def laplacian(domain, phi: Field[(I, J), DT], *, position: Position[I, J]):
if position not in domain[1:-1, 1:-1]:
return -(phi[i-1, j]-2*phi[i, j]+phi[i+1, j])/(dx*dx) - (phi[i, j-1]-2*phi[i, j]+phi[i, j+1])/(dy*dy)
return 0.
@gt.field_operator
def step(domain, β: DT, L: Field[(I, J), DT], ϕ: Field[(I, J), DT], r: Field[(I, J), DT]):
ϕ = ϕ + β * r # update solution
r = r + β * L_ # update residual
return ϕ, r
# driver code
def steepest_descent(ϕ_0, L, R, ϵ=0.01, max_it=10000):
ϕ = ϕ_0 # solution
r = L(ϕ)-R # residual
# ensure BCs are respected
r[0, :]=r[:, 0]=r[-1, :]=r[:, -1]=0
for it in range(0, max_it):
L = laplacian[domain](domain, r)
β: DT = - sum(r*r)/sum(r * L_) # step size
step(β, L, ϕ, r, out=(ϕ, r))
# debug output
print(f"it: {it}, ||r||_L2 = {np.linalg.norm(r)}")
#if plotting_enabled:
# plot(ϕ)
# break when converged
if np.linalg.norm(r) <= ϵ:
break
return ϕ
```
### [WIP] Average (id: S-AVG)
*Positional with field composition*
```python=
def backward_diff(inp):
return forward_diff(inp)[i-1]
@gt.local_operator
def laplacian(inp):
return -4*inp[i, j]+inp[i-1, j]+inp[i+1, j]+inp[i, j-1]+inp[i, j+1]
return (inp[i-1] - inp) - (+inp - inp[i+1]) + (...)
return forward_diff(inp)[i-1] - forward_diff(inp)
return backward_diff(inp) - forward_diff(inp)
@gt.local_operator
def forward_diff_i(inp):
return inp[i+1, j]-inp[i, j]
@gt.local_operator
def central_diff_i(inp):
return inp[i+1, j]-inp[i-1, j]
@gt.field_operator
def laplacian(inp):
@gt.local_operator
def average(inp: Field[(Edge, K), DT], *, position: Position[Vertex, K]) -> DT:
v, k = position
return sum(inp[e, k] for e in neighbors(v, V2E))
@gt.local_operator
def average_same(inp: Field[(Vertex, K), DT], *, position: Position[Vertex, K]):
v, k = position
return sum(inp[v, k] for v in neighbors(v, V2V))
@gt.field_operator
def average_composition(vertex_domain, inp: Field[(Edge, K), DT]):
tmp: Field[Vertex, DT] = average[vertex_domain](inp)
return average_same[vertex_domain](tmp)
def driver():
vertex_domain = Domain((Vertex, 0, 10), (K, 0, 10))
edge_domain = Domain((Edge, 0, 10), (K, 0, 10))
inp_edge = Field.from_function(edge_domain, lambda e, k: 42)
out_vertex = Field.from_function(vertex_domain, lambda v, k: 0)
average[vertex_domain](inp_edge, out=out_vertex)
average_composition(vertex_domain, inp_edge, out=out_vertex)
```
### Different iterator positions in the same stencil (id: S-EY)
*Beautified Iterator*
```python=
@gt.beautified_iterator_view
def foo(it0, it1):
return it0() + it1()
@gt.beautified_iterator_view
def bar(inp):
foo(inp, shift(I+1)(inp))
```
*Positional*
```python=
@local_operator
def foo(inp, inp2, *, position):
return inp[position] + inp2[position]
@local_operator
def shift(inp, dim, offset, *, position):
pos = position.shift(dim, offset)
return inp[pos]
@field_operator
def bar(inp):
foo[...](inp, shift[...](inp, I, 1 ))
```
### "Chiasm" stencil (id: S-CHI)
This example illustrates a valid stencil (call) that can not be lifted and shifted. Namely stencils for which the position type of some argument iterators is different.
*Beautified Iterator*
```python=
@gt.beautified_iterator_view
def foo(it0: Iterator[Vertex, Edge], it1: Iterator[Edge, Edge]):
return it0(V2E(0)) + it1()
@gt.beautified_iterator_view
def bar(it0: Iterator[Vertex, Edge]):
# we have not introduced a shift syntax-wise for this model
# but we can always read it as (lambda inp: inp(E2V, 0))[...]
it1: Iterator[Edge, Edge] = it0.shift(E2V[0])
tmp = foo[...](it0, it1)
return tmp() # perfectly valid
# not possible with whatever shift:
# return tmp(some_shift)
# Proof by contradiction: Let it0 and it1 be iterators
# located on different positions. By definition of the
# semantics of a call to a lifted stencil this must
# be equal to:
# foo(it0.shift(some_shift), it1.shift(some_shift))
# However by the definition of shift it0 and it1 must
# be located at the "from" dimension of some_shift.
# This contradics the assumption.
# This can also be seen by inling the above example:
# it0(some_shift, V2E(0)) + it0(E2V(0), some_shift)
# Consequently some_shift is from Vertex (2. call) to
# Vertex (1. call). As such the first call shifts it0
# from a vertex to an edge, but the second call shifts
# it from an edge to a vertex. An iterator can however
# either point to a vertex or an edge, but not both.
# "fencil"
bar[vertex_domain](edge_field)
```
### [WIP] Nested reductions (id: S-NR)
### [WIP] Sign field FVM (id: S-FVM-SGN)
### Conditional
- _Plain iterator_:
```python=
def conditional(cond, a, b):
return if_(deref(cond), deref(a), deref(b))
```
- _Beautified iterator (with if beautification):_
```python=
def conditional(cond, a,b):
if cond():
tmp2 = a()
tmp = tmp2
else:
tmp = b()
return tmp
```
- _Lifted iterator:_
```python=
def conditional(cond: Field, a: Field, b: Field):
return if_(cond, a, b)
```
Sign in with Wallet
Connect another wallet