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 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 fvm_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), interval(0, None):
with location(Edge) as e:
zavg = 0.5 * sum(pp[v] for v in vertices(e))
zavg = 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 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 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 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 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 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 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))