def nested(mesh: Mesh, f_1: Field[Edge, dtype], f_2: Field[Vertex, dtype], f_3: Field[Edge, dtype]):
with computation(FORWARD), interval(0, None):
with location(Edge) as e:
f_1 = 1
with location(Vertex) as v:
f_2 = 2
with computation(FORWARD), interval(0, None), location(Edge) as e:
f_3 = 3
def sten(e2v: E2V, in_field: Field[Vertex, dtype], out_field: Field[Edge,
dtype]):
with computation(FORWARD), interval(0, None):
with location(Edge) as e:
# TODO: ints don't work right now
weights = LocalField[E2V, dtype]([-1.0, 1.0])
out_field = sum(in_field[v] * weights[e, v] for v in e2v[e])
def nabla(
mesh: Mesh,
S_MXX: Field[Edge, dtype],
S_MYY: Field[Edge, dtype],
pp: Field[Vertex, dtype],
pnabla_MXX: Field[Vertex, dtype],
pnabla_MYY: Field[Vertex, dtype],
vol: Field[Vertex, dtype],
sign: Field[Vertex, Local[Edge], dtype],
):
with computation(FORWARD):
with location(Edge) as e:
zavg = 0.5 * sum(pp[v] for v in vertices(e))
zavgS_MXX = S_MXX * zavg
zavgS_MYY = S_MYY * zavg
with location(Vertex) as v:
pnabla_MXX = sum(zavgS_MXX[e] * sign[v, e] for e in edges(v))
pnabla_MYY = sum(zavgS_MYY[e] * sign[v, e] for e in edges(v))
pnabla_MXX = pnabla_MXX / vol
pnabla_MYY = pnabla_MYY / vol
def sten(
c2c: C2C,
field_in: Field[[Cell, K], dtype],
field_out: Field[[Cell, K], dtype],
field_sparse: SparseField[[C2C, K], dtype],
field_k: Field[K, dtype],
):
with computation(FORWARD), location(Cell) as c1:
field_out[c1] = 2.0 * field_k + sum(
field_in[c1] + field_in[c2] + field_sparse[c1, c2] for c2 in c2c[c1]
)
def fvm_nabla(
v2e: V2E,
e2v: E2V,
S_MXX: Field[Edge, dtype],
S_MYY: Field[Edge, dtype],
pp: Field[Vertex, dtype],
pnabla_MXX: Field[Vertex, dtype],
pnabla_MYY: Field[Vertex, dtype],
vol: Field[Vertex, dtype],
sign: SparseField[V2E, dtype],
):
with computation(FORWARD), interval(0, None):
with location(Edge) as e:
zavg = 0.5 * sum(pp[v] for v in e2v[e])
zavgS_MXX = S_MXX * zavg
zavgS_MYY = S_MYY * zavg
with location(Vertex) as v:
pnabla_MXX = sum(zavgS_MXX[e] * sign[v, e] for e in v2e[v])
pnabla_MYY = sum(zavgS_MYY[e] * sign[v, e] for e in v2e[v])
pnabla_MXX = pnabla_MXX / vol
pnabla_MYY = pnabla_MYY / vol
def copy2(field_in: Field[Vertex, dtype], field_out: Field[Vertex, dtype]):
with computation(FORWARD), interval(0, None), location(Vertex) as v:
field_in[v] = field_out[v]
def native_functions(c2c: C2C, field_in: Field[Cell, dtype],
field_out: Field[Cell, dtype]):
with computation(FORWARD), location(Cell) as c1:
field_out[c1] = (
sqrt(field_in) + max(1, 2) +
sum(max(field_in[c1], field_in[c2]) for c2 in c2c[c1]))
def sparse_field_assign(e2v: E2V, in_sparse_field: SparseField[E2V, dtype],
out_sparse_field: SparseField[E2V, dtype]):
with computation(FORWARD), interval(0, None):
with location(Edge) as e:
# TODO: maybe support slicing for lhs: out_sparse_field[e,:]
out_sparse_field = (in_sparse_field[e, v] for v in e2v[e])
def weights(e2v: E2V, in_field: Field[Vertex, dtype], out_field: Field[Edge,
dtype]):
with computation(FORWARD), interval(0, None):
with location(Edge) as e:
weights = LocalField[E2V, dtype]([2, 1])
out_field = sum(in_field[v] * weights[e, v] for v in e2v[e])
def sten(field_in: Field[Edge, dtype], field_out: Field[Edge, dtype]):
with computation(FORWARD), location(Edge) as e:
field_out[e] = field_in[e]
def temporary_field(out: Field[Vertex, dtype]):
with computation(FORWARD), interval(0, None), location(Vertex) as e:
tmp = 1
with computation(FORWARD), interval(0, None), location(Vertex) as v:
out = tmp
def sten(mesh: Mesh, field_in: Field[Cell, dtype], field_out: Field[Cell,
dtype]):
with computation(FORWARD), location(Cell) as c1:
field_out[c1] = sum(field_in[c1] + field_in[c2] for c2 in cells(c1))
def sten(c2c: C2C, field_in: Field[Cell, dtype], field_out: Field[Cell, dtype]):
with computation(FORWARD), location(Cell) as c1:
field_out[c1] = sum(field_in[c1] + field_in[c2] for c2 in c2c[c1])
def sparse_ex(
mesh: Mesh, edge_field: Field[Edge, dtype], sparse_field: Field[Edge, Local[Vertex], dtype]
):
with computation(FORWARD), interval(0, None), location(Edge) as e:
edge_field = sum(sparse_field[e, v] for v in vertices(e))
def edge_reduction(mesh: Mesh, edge_field: Field[Edge, dtype], vertex_field: Field[Vertex, dtype]):
with computation(FORWARD), interval(0, None), location(Edge) as e:
edge_field = 0.5 * sum(vertex_field[v] for v in vertices(e))
def edge_reduction(e2v: E2V, edge_field: Field[Edge, dtype],
vertex_field: Field[Vertex, dtype]):
with computation(FORWARD), interval(0, None), location(Edge) as e:
edge_field = 0.5 * sum(vertex_field[v] for v in e2v[e])
def sparse_ex(e2v: E2V, edge_field: Field[Edge, dtype],
sparse_field: SparseField[E2V, dtype]):
with computation(FORWARD), interval(0, None), location(Edge) as e:
edge_field = sum(sparse_field[e, v] for v in e2v[e])
def sten(e2v: E2V, field_in: Field[Vertex, dtype], field_out: Field[Edge,
dtype]):
with computation(FORWARD), location(Edge) as e:
field_out[e] = sum(field_in[v] for v in e2v[e])
def sten(mesh: Mesh, field_in: Field[Cell, dtype], field_out: Field[Cell,
dtype]):
with computation(FORWARD), location(Cell):
tmp = field_in
with computation(FORWARD), location(Cell):
field_out = tmp # noqa: F841
def sten(mesh: Mesh, field_in: Field[Vertex, dtype], field_out: Field[Edge,
dtype]):
with computation(FORWARD), location(Edge) as e:
field_out[e] = sum(field_in[v] for v in vertices(e))
