Merge remote-tracking branch 'origin/generic-part-bdry' into generic-part-bdry

Author: majosm

doc/conf.py

@@ -41,3 +41,16 @@
     "https://mpi4py.readthedocs.io/en/stable": None,
     "h5py": ("https://docs.h5py.org/en/stable", None),
 }
+
+
+# Some modules need to import things just so that sphinx can resolve symbols in
+# type annotations. Often, we do not want these imports (e.g. of PyOpenCL) when
+# in normal use (because they would introduce unintended side effects or hard
+# dependencies). This flag exists so that these imports only occur during doc
+# build. Since sphinx appears to resolve type hints lexically (as it should),
+# this needs to be cross-module (since, e.g. an inherited arraycontext
+# docstring can be read by sphinx when building meshmode, a dependent package),
+# this needs a setting of the same name across all packages involved, that's
+# why this name is as global-sounding as it is.
+import sys
+sys._BUILDING_SPHINX_DOCS = True

examples/moving-geometry.py

@@ -25,7 +25,6 @@
 
 from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.transform_metadata import FirstAxisIsElementsTag
-from arraycontext import thaw
 
 from pytools import keyed_memoize_in
 from pytools.obj_array import make_obj_array
@@ -186,15 +185,15 @@ def velocity_field(nodes, alpha=1.0):
 
     def source(t, x):
         discr = reconstruct_discr_from_nodes(actx, discr0, x)
-        u = velocity_field(thaw(discr.nodes(), actx))
+        u = velocity_field(actx.thaw(discr.nodes()))
 
         # {{{
 
         # NOTE: these are just here because this was at some point used to
         # profile some more operators (turned out well!)
 
         from meshmode.discretization import num_reference_derivative
-        x = thaw(discr.nodes()[0], actx)
+        x = actx.thaw(discr.nodes()[0])
         gradx = sum(
                 num_reference_derivative(discr, (i,), x)
                 for i in range(discr.dim))
@@ -214,7 +213,7 @@ def source(t, x):
     maxiter = int(tmax // timestep) + 1
     dt = tmax / maxiter + 1.0e-15
 
-    x = thaw(discr0.nodes(), actx)
+    x = actx.thaw(discr0.nodes())
     t = 0.0
 
     if visualize:

examples/parallel-vtkhdf.py

@@ -63,9 +63,8 @@ def main(*, ambient_dim: int) -> None:
 
     logger.info("[%4d] discretization: finished", mpirank)
 
-    from arraycontext import thaw
-    vector_field = thaw(discr.nodes(), actx)
-    scalar_field = actx.np.sin(thaw(discr.nodes()[0], actx))
+    vector_field = actx.thaw(discr.nodes())
+    scalar_field = actx.np.sin(vector_field[0])
     part_id = 1 + mpirank + discr.zeros(actx)
     logger.info("[%4d] fields: finished", mpirank)
 

examples/plot-connectivity.py

@@ -2,7 +2,6 @@
 import pyopencl as cl
 
 from meshmode.array_context import PyOpenCLArrayContext
-from arraycontext import thaw
 
 order = 4
 
@@ -30,7 +29,7 @@ def main():
     vis = make_visualizer(actx, discr, order)
 
     vis.write_vtk_file("geometry.vtu", [
-        ("f", thaw(discr.nodes()[0], actx)),
+        ("f", actx.thaw(discr.nodes()[0])),
         ])
 
     from meshmode.discretization.visualization import \

examples/simple-dg.py

@@ -35,7 +35,6 @@
 from meshmode.array_context import (PyOpenCLArrayContext,
                                     PytatoPyOpenCLArrayContext)
 from arraycontext import (
-        freeze, thaw,
         ArrayContainer,
         map_array_container,
         with_container_arithmetic,
@@ -57,7 +56,7 @@
 # {{{ discretization
 
 def parametrization_derivative(actx, discr):
-    thawed_nodes = thaw(discr.nodes(), actx)
+    thawed_nodes = actx.thaw(discr.nodes())
 
     from meshmode.discretization import num_reference_derivative
     result = np.zeros((discr.ambient_dim, discr.dim), dtype=object)
@@ -175,17 +174,17 @@ def get_discr(self, where):
 
     @memoize_method
     def parametrization_derivative(self):
-        return freeze(
+        return self._setup_actx.freeze(
                 parametrization_derivative(self._setup_actx, self.volume_discr))
 
     @memoize_method
     def vol_jacobian(self):
-        [a, b], [c, d] = thaw(self.parametrization_derivative(), self._setup_actx)
-        return freeze(a*d-b*c)
+        [a, b], [c, d] = self._setup_actx.thaw(self.parametrization_derivative())
+        return self._setup_actx.freeze(a*d - b*c)
 
     @memoize_method
     def inverse_parametrization_derivative(self):
-        [a, b], [c, d] = thaw(self.parametrization_derivative(), self._setup_actx)
+        [a, b], [c, d] = self._setup_actx.thaw(self.parametrization_derivative())
 
         result = np.zeros((2, 2), dtype=object)
         det = a*d-b*c
@@ -194,13 +193,13 @@ def inverse_parametrization_derivative(self):
         result[1, 0] = -c/det
         result[1, 1] = a/det
 
-        return freeze(result)
+        return self._setup_actx.freeze(result)
 
     def zeros(self, actx):
         return self.volume_discr.zeros(actx)
 
     def grad(self, vec):
-        ipder = thaw(self.inverse_parametrization_derivative(), vec.array_context)
+        ipder = vec.array_context.thaw(self.inverse_parametrization_derivative())
 
         from meshmode.discretization import num_reference_derivative
         dref = [
@@ -222,15 +221,15 @@ def normal(self, where):
         ((a,), (b,)) = parametrization_derivative(self._setup_actx, bdry_discr)
 
         nrm = 1/(a**2+b**2)**0.5
-        return freeze(flat_obj_array(b*nrm, -a*nrm))
+        return self._setup_actx.freeze(flat_obj_array(b*nrm, -a*nrm))
 
     @memoize_method
     def face_jacobian(self, where):
         bdry_discr = self.get_discr(where)
 
         ((a,), (b,)) = parametrization_derivative(self._setup_actx, bdry_discr)
 
-        return freeze((a**2 + b**2)**0.5)
+        return self._setup_actx.freeze((a**2 + b**2)**0.5)
 
     @memoize_method
     def get_inverse_mass_matrix(self, grp, dtype):
@@ -261,7 +260,7 @@ def inverse_mass(self, vec):
                     tagged=(FirstAxisIsElementsTag(),)
                 ) for grp, vec_i in zip(discr.groups, vec)
             )
-        ) / thaw(self.vol_jacobian(), actx)
+        ) / actx.thaw(self.vol_jacobian())
 
     @memoize_method
     def get_local_face_mass_matrix(self, afgrp, volgrp, dtype):
@@ -300,7 +299,7 @@ def face_mass(self, vec):
         all_faces_discr = all_faces_conn.to_discr
         vol_discr = all_faces_conn.from_discr
 
-        fj = thaw(self.face_jacobian("all_faces"), vec.array_context)
+        fj = vec.array_context.thaw(self.face_jacobian("all_faces"))
         vec = vec*fj
 
         assert len(all_faces_discr.groups) == len(vol_discr.groups)
@@ -367,7 +366,7 @@ def wave_flux(actx, discr, c, q_tpair):
     u = q_tpair.u
     v = q_tpair.v
 
-    normal = thaw(discr.normal(q_tpair.where), actx)
+    normal = actx.thaw(discr.normal(q_tpair.where))
 
     flux_weak = WaveState(
             u=np.dot(v.avg, normal),
@@ -422,7 +421,7 @@ def bump(actx, discr, t=0):
     source_width = 0.05
     source_omega = 3
 
-    nodes = thaw(discr.volume_discr.nodes(), actx)
+    nodes = actx.thaw(discr.volume_discr.nodes())
     center_dist = flat_obj_array([
         nodes[0] - source_center[0],
         nodes[1] - source_center[1],
@@ -492,8 +491,8 @@ def rhs(t, q):
     compiled_rhs = actx_rhs.compile(rhs)
 
     def rhs_wrapper(t, q):
-        r = compiled_rhs(t, thaw(freeze(q, actx_outer), actx_rhs))
-        return thaw(freeze(r, actx_rhs), actx_outer)
+        r = compiled_rhs(t, actx_rhs.thaw(actx_outer.freeze(q)))
+        return actx_outer.thaw(actx_rhs.freeze(r))
 
     t = np.float64(0)
     t_final = 3

examples/to_firedrake.py

@@ -25,7 +25,6 @@
 import pyopencl as cl
 
 from meshmode.array_context import PyOpenCLArrayContext
-from arraycontext import thaw
 
 
 # Nb: Some of the initial setup was adapted from meshmode/examplse/simple-dg.py
@@ -75,7 +74,7 @@ def main():
     #           = e^x cos(y)
     nodes = discr.nodes()
     for i in range(len(nodes)):
-        nodes[i] = thaw(nodes[i], actx)
+        nodes[i] = actx.thaw(nodes[i])
     # First index is dimension
     candidate_sol = actx.np.exp(nodes[0]) * actx.np.cos(nodes[1])
 

meshmode/array_context.py

@@ -41,9 +41,7 @@ def thaw(actx, ary):
             "meshmode.array_context.thaw will continue to work until 2022.",
             DeprecationWarning, stacklevel=2)
 
-    from arraycontext import thaw as _thaw
-    # /!\ arg order flipped
-    return _thaw(ary, actx)
+    return actx.thaw(ary)
 
 
 # {{{ kernel transform function
@@ -235,6 +233,27 @@ def transform_loopy_program(self, t_unit):
 # {{{ pytato pyopencl array context subclass
 
 class PytatoPyOpenCLArrayContext(PytatoPyOpenCLArrayContextBase):
+    def transform_dag(self, dag):
+        dag = super().transform_dag(dag)
+
+        # {{{ /!\ Remove tags from NamedArrays
+        # See <https://www.github.com/inducer/pytato/issues/195>
+
+        import pytato as pt
+
+        def untag_loopy_call_results(expr):
+            if isinstance(expr, pt.NamedArray):
+                return expr.copy(tags=frozenset(),
+                                 axes=(pt.Axis(frozenset()),)*expr.ndim)
+            else:
+                return expr
+
+        dag = pt.transform.map_and_copy(dag, untag_loopy_call_results)
+
+        # }}}
+
+        return dag
+
     def transform_loopy_program(self, t_unit):
         # FIXME: Do not parallelize for now.
         return t_unit

meshmode/discretization/__init__.py

@@ -36,11 +36,13 @@
 import numpy as np
 
 import loopy as lp
-from arraycontext import ArrayContext, make_loopy_program
+from arraycontext import ArrayContext, make_loopy_program, tag_axes
 from pytools import memoize_in, memoize_method, keyed_memoize_in
 from pytools.obj_array import make_obj_array
 from meshmode.transform_metadata import (
-        ConcurrentElementInameTag, ConcurrentDOFInameTag, FirstAxisIsElementsTag)
+        ConcurrentElementInameTag, ConcurrentDOFInameTag,
+        FirstAxisIsElementsTag, DiscretizationElementAxisTag,
+        DiscretizationDOFAxisTag)
 
 # underscored because it shouldn't be imported from here.
 from meshmode.dof_array import DOFArray as _DOFArray
@@ -542,9 +544,14 @@ def _new_array(self, actx, creation_func, dtype=None):
         else:
             dtype = np.dtype(dtype)
 
-        return _DOFArray(actx, tuple(
-            creation_func(shape=(grp.nelements, grp.nunit_dofs), dtype=dtype)
-            for grp in self.groups))
+        return tag_axes(actx, {
+                    0: DiscretizationElementAxisTag(),
+                    1: DiscretizationDOFAxisTag()},
+                _DOFArray(actx,
+                           tuple(creation_func(shape=(grp.nelements,
+                                                      grp.nunit_dofs),
+                                               dtype=dtype)
+                                 for grp in self.groups)))
 
     def empty(self, actx: ArrayContext,
               dtype: Optional[np.dtype] = None) -> _DOFArray:
@@ -642,21 +649,32 @@ def nodes(self, cached: bool = True) -> np.ndarray:
             raise ElementGroupTypeError("Element groups must be nodal.")
 
         def resample_mesh_nodes(grp, iaxis):
+            name_hint = f"nodes{iaxis}_{self.ambient_dim}d"
             # TODO: would be nice to have the mesh use an array context already
-            nodes = actx.from_numpy(grp.mesh_el_group.nodes[iaxis])
+            nodes = tag_axes(actx,
+                    {0: DiscretizationElementAxisTag(),
+                        1: DiscretizationDOFAxisTag()},
+                    actx.from_numpy(grp.mesh_el_group.nodes[iaxis]))
 
             grp_unit_nodes = grp.unit_nodes.reshape(-1)
             meg_unit_nodes = grp.mesh_el_group.unit_nodes.reshape(-1)
 
+            from arraycontext.metadata import NameHint
+
             tol = 10 * np.finfo(grp_unit_nodes.dtype).eps
             if (grp_unit_nodes.shape == meg_unit_nodes.shape
                     and np.linalg.norm(grp_unit_nodes - meg_unit_nodes) < tol):
-                return nodes
+                return actx.tag(NameHint(name_hint), nodes)
 
             return actx.einsum("ij,ej->ei",
-                               actx.from_numpy(grp.from_mesh_interp_matrix()),
+                               actx.tag_axis(
+                                   0,
+                                   DiscretizationDOFAxisTag(),
+                                   actx.from_numpy(grp.from_mesh_interp_matrix())),
                                nodes,
-                               tagged=(FirstAxisIsElementsTag(),))
+                               tagged=(
+                                   FirstAxisIsElementsTag(),
+                                   NameHint(name_hint)))
 
         result = make_obj_array([
             _DOFArray(None, tuple([
@@ -714,7 +732,9 @@ def get_mat(grp, gref_axes):
 
     return _DOFArray(actx, tuple(
             actx.einsum("ij,ej->ei",
-                        get_mat(grp, ref_axes),
+                        actx.tag_axis(0,
+                                      DiscretizationDOFAxisTag(),
+                                      get_mat(grp, ref_axes)),
                         vec[igrp],
                         tagged=(FirstAxisIsElementsTag(),))
             for igrp, grp in enumerate(discr.groups)))