update some docs; fix dim from num faces computation

grudge/array_context.py

@@ -53,8 +53,8 @@
 from pytools.tag import Tag
 
 from grudge.transform.mappers import (
-    InverseMassPropagator,
-    InverseMassRemover,
+    InverseMassDistributor,
+    RedundantMassTimesMassInverseRemover,
     MassInverseTimesStiffnessSimplifier,
 )
 from grudge.transform.metadata import (
@@ -273,10 +273,10 @@ def transform_dag(self, dag):
         self.dot_codes_before.append(pt.get_dot_graph(dag))
         if 1:
             # step 1: distribute mass inverse through DAG, across index lambdas
-            dag = InverseMassPropagator()(dag)
+            dag = InverseMassDistributor()(dag)
 
             # step 2: remove mass-times-mass-inverse einsums
-            dag = InverseMassRemover()(dag)
+            dag = RedundantMassTimesMassInverseRemover()(dag)
 
             # step 3: create new operator out of inverse mass times stiffness
             dag = MassInverseTimesStiffnessSimplifier()(dag)
@@ -289,7 +289,6 @@ def transform_dag(self, dag):
 
         # }}}
 
-        # dag = pt.transform.materialize_with_mpms(dag)
         dag = deduplicate_data_wrappers(dag)
         num_nodes_after = get_num_nodes(dag)
         self.dot_codes_after.append(pt.get_dot_graph(dag))
@@ -312,34 +311,10 @@ def eliminate_reshapes_of_data_wrappers(ary):
         dag = pt.transform.map_and_copy(dag,
                                         eliminate_reshapes_of_data_wrappers)
 
-        def materialize_all_einsums_or_reduces(expr):
-            from pytato.raising import ReduceOp, index_lambda_to_high_level_op
-
-            if isinstance(expr, pt.Einsum):
-                return expr.tagged(pt.tags.ImplStored())
-            elif (isinstance(expr, pt.IndexLambda)
-                    and isinstance(index_lambda_to_high_level_op(expr), ReduceOp)):
-                return expr.tagged(pt.tags.ImplStored())
-            else:
-                return expr
-
-        # logger.info("transform_dag.materialize_all_einsums_or_reduces")
-        # dag = pt.transform.map_and_copy(dag, materialize_all_einsums_or_reduces)
-
         logger.info("transform_dag.infer_axes_tags")
         from grudge.transform.metadata import unify_discretization_entity_tags
         dag = unify_discretization_entity_tags(dag)
 
-        def untag_loopy_call_results(expr):
-            from pytato.loopy import LoopyCallResult
-            if isinstance(expr, LoopyCallResult):
-                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)
-
         def _untag_impl_stored(expr):
             if isinstance(expr, pt.InputArgumentBase):
                 return expr

grudge/transform/mappers.py

@@ -74,6 +74,15 @@ def map_einsum(self, expr):
 # {{{ algebraic dag rewrites
 
 class MassInverseTimesStiffnessSimplifier(CopyMapperWithExtraArgs):
+    """
+    This is part three of a three part transformation pipeline.
+
+    Creates a new operator that is the result of a mass inverse times weak
+    derivative operator and replaces all instances of these two operators in
+    each einsum with the new operator.
+
+    See `InverseMassDistributor` for background.
+    """
     def map_einsum(self, expr, *args, **kwargs):
         iarg_stiffness = None
         iarg_mass_inverse = None
@@ -121,7 +130,13 @@ def map_einsum(self, expr, *args, **kwargs):
         return expr.copy(args=tuple(new_args))
 
 
-class InverseMassRemover(CopyMapperWithExtraArgs):
+class RedundantMassTimesMassInverseRemover(CopyMapperWithExtraArgs):
+    """
+    This is part two of a three-part transformation pipeline. Removes einsum
+    nodes that contain both a mass inverse and mass operator.
+
+    See `InverseMassDistributor` for more information.
+    """
     def map_einsum(self, expr, *args, **kwargs):
         found_mass = False
         found_mass_inverse = False
@@ -142,37 +157,54 @@ def map_einsum(self, expr, *args, **kwargs):
         return expr.copy(args=tuple(new_args))
 
 
-class InverseMassPropagator(CopyMapperWithExtraArgs):
+class InverseMassDistributor(CopyMapperWithExtraArgs):
     r"""
-    Applying a full mass inverse operator when using a tensor-product
-    discretization results in redundant work. For example, suppose we have a 3D
-    tensor-product discretization. Then the weak differentiation operator
-    associated with the r-coordinate direction can be expressed as
-
-    .. math::
+    Implements one part of a three-part transformation pipeline to realize an
+    algebraic simplification arising in tensor-product discretizations.
 
-        K_r = \hat{K} \otimes \hat{M} \otimes \hat{M},
+    Specifically, one can represent a weak derivative operator associated with
+    a tensor-product discretization as a Kronecker product of a 1D weak
+    derivative operator with a variable number of 1D mass matrices, which
+    depends on the dimension of the problem.
 
-    where :math:`\hat{M}` is a 1D mass operator and :math:`\hat{K}` is a 1D weak
-    differentiation operator. Similarly, the inverse mass operator can be
+    Let $S$ be the full weak operator, $\hat{S}$ as the 1D weak derivative
+    operator and $\hat{M}$ as the 1D mass operator. For example, consider a 2D
+    tensor-product discretization. The weak $x$-derivative operator can be
     expressed as
 
-    .. math::
+    ..math::
+
+        S_x = \hat{S} \otimes \hat{M}
+
+    Since we are using a tensor-product discretization, the mass operator can be
+    decomposed as a tensor-product of a variable number of mass operators.
+    Hence, the mass inverse operator can also be expressed via a tensor-product.
 
-        M^{-1} = \hat{M}^{-1} \otimes \hat{M}^{-1} \otimes \hat{M}^{-1},
+    ..math::
 
-    By the properties of the tensor-product, multiplying the weak derivative
-    operator on the left by the inverse mass operator results in
+    M^{-1} S_x^k = (\hat{M}^{-1} \otimes \hat{M}^{-1})(\hat{S} \otimes \hat{M})
+
+    By associativity of the tensor-product,
 
     .. math::
 
-        M^{-1} K_r = \hat{M}^{-1} \hat{K}_r \otimes \hat{I} \otimes \hat{I},
+        M^{-1} S_x^k = \hat{M}^{-1}\hat{S} \otimes \hat{M}^{-1}\hat{M}
+
+    Thus, we can instead apply the operator as
 
-    where :math:`\hat{I}` is a 1D identity operator.
+    ..math::
 
-    The goal of this mapper is to remove redundant mass-times-mass inverse
-    operations from an expression graph of operations involved with a
-    tensor-product discretization.
+        M^{-1} S_x^k = \hat{M}^{-1}\hat{S} \otimes \hat{I}
+
+    where $\hat{I}$ is a 1D identity operator. This results in both a reduction
+    in the total number of operations and the required storage for the
+    operators.
+
+    This transformation takes care of the distribution of the mass inverse
+    operator through the DAG to other tensor-product application routines
+    Moreover, the mass inverse is distribtued to the face mass terms (if
+    included in the original einsum), and properly reshapes to and from
+    tensor-product form to apply the 1D mass operator.
     """
     def map_einsum(self, expr, *args, **kwargs):
         new_args = []
@@ -195,18 +227,9 @@ def map_einsum(self, expr, *args, **kwargs):
                 nfaces, nelts, _ = expr.args[1].shape
                 ndofs = expr.shape[1]
 
-                if nfaces == 2:
-                    dim = 1
-                elif nfaces == 4:
-                    dim = 2
-                elif nfaces == 6:
-                    dim = 3
-                else:
-                    raise ValueError("Unable to determine the dimension of the",
-                                     " hypercube to apply transformations.")
-
                 from math import ceil
-                ndofs_1d = int(ceil(ndofs**(1/dim)))
+                dim = ceil(nfaces/2)
+                ndofs_1d = ceil(ndofs**(1/dim))
                 expr = expr.reshape(nelts, *(ndofs_1d,)*dim)
 
                 for axis in range(dim):