Skip to content

Commit

Permalink
use lazy eval array context for setup actx
Browse files Browse the repository at this point in the history
  • Loading branch information
kaushikcfd committed Mar 6, 2024
1 parent 7359d59 commit ce9b9fa
Showing 1 changed file with 12 additions and 9 deletions.
21 changes: 12 additions & 9 deletions test/test_pytato_transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ def _map_index_base(self, expr: IndexBase) -> bool:
from grudge.pytato_transforms.pytato_indirection_transforms import (
_is_materialized)
return self.combine(
_is_materialized(expr.array) or isinstance(expr.array, BasicIndex),
_is_materialized(expr.array) or isinstance(expr, BasicIndex),
self.rec(expr.array)
)

Expand Down Expand Up @@ -227,6 +227,7 @@ def _compute_flux_2(dcoll, actx, u):
normal_on_bdry_faces = actx.thaw(dcoll.normal(BTAG_ALL))
flux_on_interior_faces = u_interior_tpair.avg * normal_on_interior_faces
flux_on_bdry = op.project(dcoll, "vol", BTAG_ALL, u) * normal_on_bdry_faces

flux_on_all_faces = (
op.project(dcoll,
FACE_RESTR_INTERIOR,
Expand All @@ -245,21 +246,23 @@ def test_resampling_indirections_are_fused_2(ctx_factory):
cl_ctx = ctx_factory()
cq = cl.CommandQueue(cl_ctx)

ref_actx = PyOpenCLArrayContext(cq)
from grudge.array_context import get_reasonable_array_context_class

ref_actx = get_reasonable_array_context_class(lazy=True, distributed=False)(cq)
actx = FluxOptimizerActx(cq)

dim = 2
nel_1d = 4
dim = 3
nel_1d = 16
order = 4
mesh = generate_regular_rect_mesh(
a=(-0.5,)*dim,
b=(0.5,)*dim,
nelements_per_axis=(nel_1d,)*dim,
boundary_tag_to_face={"bdry": ["-x", "+x",
"-y", "+y"]}
)
dcoll = grudge.make_discretization_collection(ref_actx, mesh, order=2)

x, _ = dcoll.nodes()
dcoll = grudge.make_discretization_collection(
ref_actx, mesh,
order=order)
x, _, _ = dcoll.nodes()
compiled_flux_2 = actx.compile(lambda ary: _compute_flux_2(dcoll, actx, ary))

ref_output = ref_actx.to_numpy(
Expand Down

0 comments on commit ce9b9fa

Please sign in to comment.