algebraic transforms v0.1

grudge/array_context.py

@@ -221,14 +221,6 @@ class PytatoPyOpenCLArrayContext(_PytatoPyOpenCLArrayContextBase):
     which there isn't any, for now.)
     """
 
-    dot_codes_before: list[str] = []
-    dot_codes_after: list[str] = []
-
-    num_nodes_before: int = 0
-    num_nodes_after: int = 0
-
-    device_code: str = ""
-
     def __init__(self, queue, allocator=None,
             *,
             compile_trace_callback: Callable[[Any, str, Any], None] | None
@@ -251,21 +243,23 @@ def __init__(self, queue, allocator=None,
 
     def transform_dag(self, dag):
 
-        self.num_nodes_before += get_num_nodes(dag)
-        self.dot_codes_before.append(pt.visualization.get_dot_graph(dag))
-
         # {{{ tensor-product algebraic DAG rewrites
 
-        # FIXME: remove after enough cases have been tested
-        if 1:
-            # step 1: distribute mass inverse through DAG, across index lambdas
-            dag = InverseMassPropagator()(dag)
+        num_nodes_before = get_num_nodes(dag)
+        # step 1: distribute mass inverse through DAG, across index lambdas
+        dag = InverseMassPropagator()(dag)
+
+        # step 2: remove mass-times-mass-inverse einsums
+        dag = InverseMassRemover()(dag)
 
-            # step 2: remove mass-times-mass-inverse einsums
-            dag = InverseMassRemover()(dag)
+        # step 3: create new operator out of inverse mass times stiffness
+        dag = MassInverseTimesStiffnessSimplifier()(dag)
+        num_nodes_after = get_num_nodes(dag)
 
-            # step 3: create new operator out of inverse mass times stiffness
-            dag = MassInverseTimesStiffnessSimplifier()(dag)
+        if num_nodes_before != num_nodes_after:
+            logger.info("tensor-product xforms: %d nodes via tensor-product "
+                        "algebraic transformations ",
+                        num_nodes_before - num_nodes_after)
 
         # }}}
 
@@ -283,8 +277,6 @@ def eliminate_reshapes_of_data_wrappers(ary):
 
         dag = pt.transform.map_and_copy(dag,
                                         eliminate_reshapes_of_data_wrappers)
-        self.dot_codes_after.append(pt.visualization.get_dot_graph(dag))
-        self.num_nodes_after += get_num_nodes(dag)
 
         return dag
 

grudge/bilinear_forms.py

@@ -474,7 +474,7 @@ def _compute_reference_gradient(self, vec: DOFArray) -> DOFArray:
                     self.output_group.space,
                     reference_gradient[rst_axis])
 
-        return DOFArray(self.actx, data=tuple(reference_gradient))
+        return reference_gradient
 
     def mass_operator(self, vec: DOFArray,
                       exclude_axis: int | None = None,
@@ -541,20 +541,31 @@ def divergence_operator(self, vec: DOFArray) -> DOFArray:
         derivative in the field.
         """
 
-        vec = vec * self.area_element
+        vec = vec * self.area_element * self.metric_terms
 
         if self.input_group.dim != 1:
             vec = fold(self.input_group.space, vec)
 
-        reference_gradients = DOFArray(self.actx, data=tuple([
-            self._compute_reference_gradient(vec[func_axis])
-            for func_axis in range(self.input_group.dim)]))
+        div = 0.0
+        for xyz_axis in range(self.input_group.dim):
+            partial = 0.0
+            for rst_axis in range(self.input_group.dim):
+                reference_partial = self.mass_operator(
+                    vec[xyz_axis, rst_axis],
+                    exclude_axis=rst_axis,
+                    exclude_metric=True)
 
-        return self.actx.einsum(
-            "xrej,xrej->ej",
-            self.metric_terms,
-            reference_gradients,
-            arg_names=("metrics", "vec"))
+                partial += single_axis_contraction(
+                        self.actx, self.input_group.dim, rst_axis,
+                        self._stiffness_operator, reference_partial,
+                        arg_names=("stiffness_1d",
+                                   f"ref_partial_{rst_axis}"))
+
+            if self.input_group.dim != 1:
+                partial = unfold(self.output_group.space, partial)
+            div += partial
+
+        return div
 
     def derivative_operator(self, vec: DOFArray, derivative: int) -> DOFArray:
         """

test/test_op.py

@@ -233,8 +233,7 @@ def get_flux(u_tpair, dcoll=dcoll):
             vis.write_vtk_file(filename, [
                 ("u", u),
                 ("grad_u", grad_u),
-                ("expected_grad_u", expected_grad_u),
-                ], overwrite=True)
+                ("expected_grad_u", expected_grad_u), ], overwrite=True)
 
         rel_linf_err = actx.to_numpy(
             op.norm(dcoll, grad_u - expected_grad_u, np.inf)
@@ -254,7 +253,6 @@ def get_flux(u_tpair, dcoll=dcoll):
 @pytest.mark.parametrize("form", ["strong", "weak", "weak-overint"])
 @pytest.mark.parametrize("dim", [1, 2, 3])
 @pytest.mark.parametrize("order", [2, 3])
-@pytest.mark.parametrize("warp_mesh", [True])
 @pytest.mark.parametrize(("vectorize", "nested"), [
     (False, False),
     (True, False),
@@ -265,7 +263,7 @@ def get_flux(u_tpair, dcoll=dcoll):
     TensorProductElementGroup
 ])
 def test_divergence(actx_factory, form, dim, order, vectorize, nested,
-                    group_cls, warp_mesh, visualize=False):
+                    group_cls, visualize=False):
 
     actx = actx_factory()