Tag axes of discretization-related arrays

examples/advection/surface.py

@@ -31,7 +31,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw
 from grudge.array_context import PyOpenCLArrayContext
 
 from meshmode.dof_array import flatten
@@ -62,7 +61,7 @@ def __init__(self, actx, dcoll, order, visualize=True):
             import matplotlib.pyplot as pt
             self.fig = pt.figure(figsize=(8, 8), dpi=300)
 
-            x = thaw(dcoll.discr_from_dd(dof_desc.DD_VOLUME).nodes(), actx)
+            x = actx.thaw(dcoll.discr_from_dd(dof_desc.DD_VOLUME).nodes())
             self.x = actx.to_numpy(flatten(actx.np.arctan2(x[1], x[0])))
         elif self.ambient_dim == 3:
             from grudge.shortcuts import make_visualizer
@@ -174,7 +173,7 @@ def main(ctx_factory, dim=2, order=4, use_quad=False, visualize=False):
     # {{{ Surface advection operator
 
     # velocity field
-    x = thaw(dcoll.nodes(), actx)
+    x = actx.thaw(dcoll.nodes())
     c = make_obj_array([-x[1], x[0], 0.0])[:dim]
 
     def f_initial_condition(x):
@@ -238,7 +237,7 @@ def rhs(t, u):
 
         df = dof_desc.DOFDesc(FACE_RESTR_INTERIOR)
         face_discr = dcoll.discr_from_dd(df)
-        face_normal = thaw(dcoll.normal(dd=df), actx)
+        face_normal = actx.thaw(dcoll.normal(dd=df))
 
         from meshmode.discretization.visualization import make_visualizer
         vis = make_visualizer(actx, face_discr)

examples/advection/var-velocity.py

@@ -29,7 +29,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw
 from grudge.array_context import PyOpenCLArrayContext
 
 from meshmode.dof_array import flatten
@@ -61,7 +60,7 @@ def __init__(self, actx, dcoll, order, visualize=True, ylim=None):
             self.ylim = ylim
 
             volume_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)
-            self.x = actx.to_numpy(flatten(thaw(volume_discr.nodes()[0], actx)))
+            self.x = actx.to_numpy(flatten(actx.thaw(volume_discr.nodes()[0])))
         else:
             from grudge.shortcuts import make_visualizer
             self.vis = make_visualizer(dcoll)
@@ -168,7 +167,7 @@ def zero_inflow_bc(dtag, t=0):
 
     from grudge.models.advection import VariableCoefficientAdvectionOperator
 
-    x = thaw(dcoll.nodes(), actx)
+    x = actx.thaw(dcoll.nodes())
 
     # velocity field
     if dim == 1:

examples/advection/weak.py

@@ -30,7 +30,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw
 from grudge.array_context import PyOpenCLArrayContext
 
 from meshmode.dof_array import flatten
@@ -60,7 +59,7 @@ def __init__(self, actx, dcoll, order, visualize=True, ylim=None):
             self.ylim = ylim
 
             volume_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)
-            self.x = actx.to_numpy(flatten(thaw(volume_discr.nodes()[0], actx)))
+            self.x = actx.to_numpy(flatten(actx.thaw(volume_discr.nodes()[0])))
         else:
             from grudge.shortcuts import make_visualizer
             self.vis = make_visualizer(dcoll)
@@ -152,13 +151,13 @@ def u_analytic(x, t=0):
         dcoll,
         c,
         inflow_u=lambda t: u_analytic(
-            thaw(dcoll.nodes(dd=BTAG_ALL), actx),
+            actx.thaw(dcoll.nodes(dd=BTAG_ALL)),
             t=t
         ),
         flux_type=flux_type
     )
 
-    nodes = thaw(dcoll.nodes(), actx)
+    nodes = actx.thaw(dcoll.nodes())
     u = u_analytic(nodes, t=0)
 
     def rhs(t, u):

examples/euler/acoustic_pulse.py

@@ -28,7 +28,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw, freeze
 from grudge.array_context import (
     PyOpenCLArrayContext,
     PytatoPyOpenCLArrayContext
@@ -175,7 +174,7 @@ def rhs(t, q):
     cn = 0.5*(order + 1)**2
     dt = cfl * actx.to_numpy(h_min_from_volume(dcoll)) / cn
 
-    fields = acoustic_pulse_condition(thaw(dcoll.nodes(), actx))
+    fields = acoustic_pulse_condition(actx.thaw(dcoll.nodes()))
 
     logger.info("Timestep size: %g", dt)
 
@@ -204,7 +203,7 @@ def rhs(t, q):
                 )
             assert norm_q < 5
 
-        fields = thaw(freeze(fields, actx), actx)
+        fields = actx.thaw(actx.freeze(fields))
         fields = rk4_step(fields, t, dt, compiled_rhs)
         t += dt
         step += 1

examples/euler/vortex.py

@@ -26,8 +26,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw, freeze
-
 from grudge.array_context import PytatoPyOpenCLArrayContext, PyOpenCLArrayContext
 from grudge.models.euler import (
     vortex_initial_condition,
@@ -111,7 +109,7 @@ def rhs(t, q):
 
     compiled_rhs = actx.compile(rhs)
 
-    fields = vortex_initial_condition(thaw(dcoll.nodes(), actx))
+    fields = vortex_initial_condition(actx.thaw(dcoll.nodes()))
 
     from grudge.dt_utils import h_min_from_volume
 
@@ -147,7 +145,7 @@ def rhs(t, q):
                 )
             assert norm_q < 200
 
-        fields = thaw(freeze(fields, actx), actx)
+        fields = actx.thaw(actx.freeze(fields))
         fields = rk4_step(fields, t, dt, compiled_rhs)
         t += dt
         step += 1

examples/geometry.py

@@ -30,7 +30,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw
 from grudge.array_context import PyOpenCLArrayContext
 
 from grudge import DiscretizationCollection, shortcuts
@@ -51,9 +50,9 @@ def main(write_output=True):
 
     dcoll = DiscretizationCollection(actx, mesh, order=4)
 
-    nodes = thaw(dcoll.nodes(), actx)
-    bdry_nodes = thaw(dcoll.nodes(dd=BTAG_ALL), actx)
-    bdry_normals = thaw(dcoll.normal(dd=BTAG_ALL), actx)
+    nodes = actx.thaw(dcoll.nodes())
+    bdry_nodes = actx.thaw(dcoll.nodes(dd=BTAG_ALL))
+    bdry_normals = actx.thaw(dcoll.normal(dd=BTAG_ALL))
 
     if write_output:
         vis = shortcuts.make_visualizer(dcoll)

examples/hello-grudge.py

@@ -14,7 +14,6 @@
 import grudge.op as op
 from meshmode.mesh.generation import generate_box_mesh
 from meshmode.array_context import PyOpenCLArrayContext
-from arraycontext import thaw
 from grudge.dof_desc import DTAG_BOUNDARY, FACE_RESTR_INTERIOR
 
 
@@ -43,7 +42,7 @@ def left_boundary_condition(x, t):
 def flux(dcoll, u_tpair):
     dd = u_tpair.dd
     velocity = np.array([2 * np.pi])
-    normal = thaw(dcoll.normal(dd), actx)
+    normal = actx.thaw(dcoll.normal(dd))
 
     v_dot_n = np.dot(velocity, normal)
     u_upwind = actx.np.where(v_dot_n > 0,
@@ -55,8 +54,8 @@ def flux(dcoll, u_tpair):
 left_bndry = DTAG_BOUNDARY("left")
 right_bndry = DTAG_BOUNDARY("right")
 
-x_vol = thaw(dcoll.nodes(), actx)
-x_bndry = thaw(dcoll.discr_from_dd(left_bndry).nodes(), actx)
+x_vol = actx.thaw(dcoll.nodes())
+x_bndry = actx.thaw(dcoll.discr_from_dd(left_bndry).nodes())
 
 uh = initial_condition(x_vol)
 

examples/maxwell/cavities.py

@@ -28,7 +28,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw
 from grudge.array_context import PyOpenCLArrayContext
 
 from grudge.shortcuts import set_up_rk4
@@ -84,7 +83,7 @@ def cavity_mode(x, t=0):
         else:
             return get_rectangular_cavity_mode(actx, x, t, 1, (2, 3))
 
-    fields = cavity_mode(thaw(dcoll.nodes(), actx), t=0)
+    fields = cavity_mode(actx.thaw(dcoll.nodes()), t=0)
 
     maxwell_operator.check_bc_coverage(mesh)
 

examples/wave/var-propagation-speed.py

@@ -28,7 +28,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw
 from grudge.array_context import PyOpenCLArrayContext
 
 from grudge.shortcuts import set_up_rk4
@@ -63,7 +62,7 @@ def source_f(actx, dcoll, t=0):
         source_center = np.array([0.1, 0.22, 0.33])[:dcoll.dim]
         source_width = 0.05
         source_omega = 3
-        nodes = thaw(dcoll.nodes(), actx)
+        nodes = actx.thaw(dcoll.nodes())
         source_center_dist = flat_obj_array(
             [nodes[i] - source_center[i] for i in range(dcoll.dim)]
         )
@@ -75,7 +74,7 @@ def source_f(actx, dcoll, t=0):
             )
         )
 
-    x = thaw(dcoll.nodes(), actx)
+    x = actx.thaw(dcoll.nodes())
     ones = dcoll.zeros(actx) + 1
     c = actx.np.where(np.dot(x, x) < 0.15, 0.1 * ones, 0.2 * ones)
 

examples/wave/wave-min-mpi.py

@@ -28,7 +28,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw
 from grudge.array_context import MPIPyOpenCLArrayContext
 
 from grudge.shortcuts import set_up_rk4
@@ -88,7 +87,7 @@ def source_f(actx, dcoll, t=0):
         source_center = np.array([0.1, 0.22, 0.33])[:dcoll.dim]
         source_width = 0.05
         source_omega = 3
-        nodes = thaw(dcoll.nodes(), actx)
+        nodes = actx.thaw(dcoll.nodes())
         source_center_dist = flat_obj_array(
             [nodes[i] - source_center[i] for i in range(dcoll.dim)]
         )

examples/wave/wave-op-mpi.py

@@ -30,7 +30,6 @@
 import pyopencl.tools as cl_tools
 
 from arraycontext import (
-    thaw,
     with_container_arithmetic,
     dataclass_array_container
 )
@@ -73,12 +72,12 @@ def array_context(self):
         return self.u.array_context
 
 
-def wave_flux(dcoll, c, w_tpair):
+def wave_flux(actx, dcoll, c, w_tpair):
     u = w_tpair.u
     v = w_tpair.v
     dd = w_tpair.dd
 
-    normal = thaw(dcoll.normal(dd), u.int.array_context)
+    normal = actx.thaw(dcoll.normal(dd))
 
     flux_weak = WaveState(
         u=v.avg @ normal,
@@ -99,7 +98,7 @@ class _WaveStateTag:
     pass
 
 
-def wave_operator(dcoll, c, w, quad_tag=None):
+def wave_operator(actx, dcoll, c, w, quad_tag=None):
     dd_base = as_dofdesc("vol")
     dd_vol = DOFDesc("vol", quad_tag)
     dd_faces = DOFDesc("all_faces", quad_tag)
@@ -135,13 +134,14 @@ def interp_to_surf_quad(utpair):
                 dcoll,
                 dd_faces,
                 wave_flux(
+                    actx,
                     dcoll, c=c,
                     w_tpair=op.bdry_trace_pair(dcoll,
                                                dd_btag,
                                                interior=dir_bval,
                                                exterior=dir_bc)
                 ) + sum(
-                    wave_flux(dcoll, c=c, w_tpair=interp_to_surf_quad(tpair))
+                    wave_flux(actx, dcoll, c=c, w_tpair=interp_to_surf_quad(tpair))
                     for tpair in op.interior_trace_pairs(dcoll, w,
                         comm_tag=_WaveStateTag)
                 )
@@ -165,7 +165,7 @@ def bump(actx, dcoll, t=0):
     source_width = 0.05
     source_omega = 3
 
-    nodes = thaw(dcoll.nodes(), actx)
+    nodes = actx.thaw(dcoll.nodes())
     center_dist = flat_obj_array([
         nodes[i] - source_center[i]
         for i in range(dcoll.dim)
@@ -258,7 +258,7 @@ def main(ctx_factory, dim=2, order=3,
     vis = make_visualizer(dcoll)
 
     def rhs(t, w):
-        return wave_operator(dcoll, c=c, w=w, quad_tag=quad_tag)
+        return wave_operator(actx, dcoll, c=c, w=w, quad_tag=quad_tag)
 
     compiled_rhs = actx.compile(rhs)
 

examples/wave/wave-op-var-velocity.py

@@ -29,7 +29,6 @@
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 
-from arraycontext import thaw
 from grudge.array_context import PyOpenCLArrayContext
 
 from pytools.obj_array import flat_obj_array
@@ -48,14 +47,14 @@
 
 # {{{ wave equation bits
 
-def wave_flux(dcoll, c, w_tpair):
+def wave_flux(actx, dcoll, c, w_tpair):
     dd = w_tpair.dd
     dd_quad = dd.with_discr_tag(DISCR_TAG_QUAD)
 
     u = w_tpair[0]
     v = w_tpair[1:]
 
-    normal = thaw(dcoll.normal(dd), u.int.array_context)
+    normal = actx.thaw(dcoll.normal(dd))
 
     flux_weak = flat_obj_array(
             np.dot(v.avg, normal),
@@ -76,7 +75,7 @@ def wave_flux(dcoll, c, w_tpair):
     return op.project(dcoll, dd_quad, dd_allfaces_quad, c_quad*flux_quad)
 
 
-def wave_operator(dcoll, c, w):
+def wave_operator(actx, dcoll, c, w):
     u = w[0]
     v = w[1:]
 
@@ -104,13 +103,14 @@ def wave_operator(dcoll, c, w):
                 dcoll,
                 dd_allfaces_quad,
                 wave_flux(
+                    actx,
                     dcoll, c=c,
                     w_tpair=op.bdry_trace_pair(dcoll,
                                                BTAG_ALL,
                                                interior=dir_bval,
                                                exterior=dir_bc)
                 ) + sum(
-                    wave_flux(dcoll, c=c, w_tpair=tpair)
+                    wave_flux(actx, dcoll, c=c, w_tpair=tpair)
                     for tpair in op.interior_trace_pairs(dcoll, w)
                 )
             )
@@ -135,7 +135,7 @@ def bump(actx, dcoll, t=0, width=0.05, center=None):
     center = center[:dcoll.dim]
     source_omega = 3
 
-    nodes = thaw(dcoll.nodes(), actx)
+    nodes = actx.thaw(dcoll.nodes())
     center_dist = flat_obj_array([
         nodes[i] - center[i]
         for i in range(dcoll.dim)
@@ -189,7 +189,7 @@ def main(ctx_factory, dim=2, order=3, visualize=False):
     vis = make_visualizer(dcoll)
 
     def rhs(t, w):
-        return wave_operator(dcoll, c=c, w=w)
+        return wave_operator(actx, dcoll, c=c, w=w)
 
     logger.info("dt = %g", dt)