Tensor product differentiation operators

grudge/op.py

@@ -166,6 +166,38 @@
     )
 
 
+# {{{ Temporary tools for tensor product operators
+# NOTE: Will possibly be removed in a future version of tensor product operator
+# development since (I think) it is not entirely necessary
+from pytools.tag import Tag
+class OutputIsTensorProductDOFArrayOrdered(Tag):
+    pass
+
+
+from grudge.array_context import PyOpenCLArrayContext
+class TensorProductArrayContext(PyOpenCLArrayContext):
+    def transform_loopy_program(self, t_unit):
+        if len(t_unit.callables_table) == 1:
+            knl = t_unit.default_entrypoint
+            if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered):
+                new_args = []
+                for arg in knl.args:
+                    if arg.is_output:
+                        arg = arg.copy(dim_tags=(
+                            f"N{len(arg.shape)-1},"
+                            + ",".join(f"N{i}"
+                                       for i in range(len(arg.shape)-1))
+                            ))
+
+                    new_args.append(arg)
+
+                knl = knl.copy(args=new_args)
+                t_unit = t_unit.with_kernel(knl)
+
+        return super().transform_loopy_program(t_unit)
+# }}}
+
+
 # {{{ common derivative "kernels"
 
 def _single_axis_derivative_kernel(
@@ -202,19 +234,97 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec,
         *, metric_in_matvec):
     # See _single_axis_derivative_kernel for comments on the usage scenarios
     # (both strong and weak derivative) and their differences.
+
+    def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm):
+        """Exploits tensor product structure to differentiate each coordinate
+        axis using a single differentiation matrix of shape (nnodes1d, nnodes1d)
+        """
+
+        actx_tp = TensorProductArrayContext(
+                actx.queue,
+                allocator=actx.allocator,
+                force_device_scalars=actx._force_device_scalars)
+
+        from modepy.tools import (
+                reshape_array_for_tensor_product_space,
+                unreshape_array_for_tensor_product_space)
+
+        # reshape u to expose tensor product structure
+        vec = reshape_array_for_tensor_product_space(grp.space, vec)
+
+        # apply differentiation matrix to vec
+        # check len(vec.shape) since shape is expected to be
+        # (nelements, nnodes1d, nnodes1d)
+        if len(vec.shape) == 3:
+            specs = ["il,elj->eij",
+                     "jl,eil->eij"]
+        elif len(vec.shape) == 4:
+            specs = ["il,eljk->eijk",
+                     "jl,eilk->eijk",
+                     "kl,eijl->eijk"]
+        else:
+            specs = None
+        assert specs is not None
+
+        diff_mat = get_diff_mat(actx, grp, grp)
+        grad = make_obj_array([
+                actx_tp.einsum(
+                    spec,
+                    diff_mat,
+                    vec,
+                    arg_names=("diff_mat", "vec"),
+                    tagged=(FirstAxisIsElementsTag(),
+                            OutputIsTensorProductDOFArrayOrdered()))
+            for spec in specs
+            ])
+
+        # unreshape grad to apply geometric factors
+        # NOTE: In a future version, do not reshape before application of
+        # geometric factors. Can possibly "chain" the einsum. For example, the
+        # simplicial case below has einsum with spec
+        #                       ("xrei,rij,ei->ei")
+        # for the strong local gradient case
+        grad = make_obj_array([
+            unreshape_array_for_tensor_product_space(grp.space, grad[i])
+            for i in range(grad.shape[0])
+            ])
+
+        # apply geometric factors to current grad
+        # FIXME: using einsum spec ("xrei,xei->xei") throws error:
+        # "Loopy does not directly support object arrays"
+        grad = make_obj_array([
+            actx_tp.einsum(
+                "rei,ei->ei",
+                ijm[i],
+                grad[i],
+                tagged=(FirstAxisIsElementsTag(),)),
+                arg_names=("inv_jac_t", "vec")
+            for i in range(grad.shape[0])
+            ])
+
+        return grad
+
+    from meshmode.discretization.poly_element import \
+        TensorProductElementGroupBase
     per_group_grads = [
+
+        compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i)
+        if isinstance(in_grp, TensorProductElementGroupBase)
+
         # r for rst axis
         # x for xyz axis
-        actx.einsum("xrej,rij,ej->xei" if metric_in_matvec else "xrei,rij,ej->xei",
-                    ijm_i,
-                    get_diff_mat(
-                        actx,
-                        out_element_group=out_grp,
-                        in_element_group=in_grp
-                    ),
-                    vec_i,
-                    arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"),
-                    tagged=(FirstAxisIsElementsTag(),))
+        else actx.einsum(
+            "xrej,rij,ej->xei" if metric_in_matvec else "xrei,rij,ej->xei",
+            ijm_i,
+            get_diff_mat(
+                actx,
+                out_element_group=out_grp,
+                in_element_group=in_grp
+            ),
+            vec_i,
+            arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"),
+            tagged=(FirstAxisIsElementsTag(),))
+
         for out_grp, in_grp, vec_i, ijm_i in zip(
             out_discr.groups, in_discr.groups, vec,
             inv_jac_mat)]
@@ -259,16 +369,37 @@ def _reference_derivative_matrices(actx: ArrayContext,
     # _reference_stiffness_transpose_matrices.
     assert out_element_group is in_element_group
 
+    from meshmode.mesh import TensorProductElementGroup
+
     @keyed_memoize_in(
         actx, _reference_derivative_matrices,
         lambda grp: grp.discretization_key())
     def get_ref_derivative_mats(grp):
-        from meshmode.discretization.poly_element import diff_matrices
-        return actx.freeze(
-                actx.tag_axis(
-                    1, DiscretizationDOFAxisTag(),
-                    actx.from_numpy(
-                        np.asarray(diff_matrices(grp)))))
+
+        from meshmode.discretization.poly_element import \
+                TensorProductElementGroupBase
+        if isinstance(grp, TensorProductElementGroupBase):
+            import modepy as mp
+            import numpy.linalg as la
+
+            space1d = grp.space.bases[0]
+            shape1d = grp.shape.bases[0]
+
+            nodes1d = mp.edge_clustered_nodes_for_space(space1d, shape1d)
+            basis1d = mp.basis_for_space(space1d, shape1d)
+
+            vdm1d = mp.vandermonde(basis1d.functions, nodes1d)
+            vdm_p1d = mp.vandermonde(basis1d.gradients, nodes1d)[0]
+
+            return actx.freeze(actx.from_numpy(vdm_p1d @ la.inv(vdm1d)))
+
+        else:
+            from meshmode.discretization.poly_element import diff_matrices
+            return actx.freeze(
+                    actx.tag_axis(
+                        1, DiscretizationDOFAxisTag(),
+                        actx.from_numpy(
+                            np.asarray(diff_matrices(grp)))))
     return get_ref_derivative_mats(out_element_group)
 
 

test/test_op.py

@@ -32,6 +32,7 @@
 
 import pytest
 
+from grudge.discretization import make_discretization_collection
 from grudge.array_context import PytestPyOpenCLArrayContextFactory
 from arraycontext import pytest_generate_tests_for_array_contexts
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
@@ -159,6 +160,81 @@ def get_flux(u_tpair):
     assert (eoc_rec.order_estimate() >= order - 0.5
                 or eoc_rec.max_error() < 1e-11)
 
+
+@pytest.mark.parametrize("form", ["strong"])
+@pytest.mark.parametrize("dim", [2, 3])
+@pytest.mark.parametrize("order", [2, 3])
+@pytest.mark.parametrize(("vectorize", "nested"), [
+    (False, False)
+    ])
+def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize,
+                                 nested, visualize=False):
+    """A "one-dimensional tensor product element" does not make sense, so the
+    one-dimensional case is excluded from this test.
+    """
+    actx = actx_factory()
+    from pytools.convergence import EOCRecorder
+    eoc_rec = EOCRecorder()
+
+    from meshmode.mesh import TensorProductElementGroup
+    from meshmode.discretization.poly_element import \
+        LegendreGaussLobattoTensorProductGroupFactory as LGL
+    for n in [4, 6, 8]:
+        mesh = mgen.generate_regular_rect_mesh(
+                a=(-1,)*dim,
+                b=(1,)*dim,
+                nelements_per_axis=(n,)*dim,
+                group_cls=TensorProductElementGroup)
+
+        import grudge.dof_desc as dd
+        dcoll = make_discretization_collection(
+                actx,
+                mesh,
+                discr_tag_to_group_factory={
+                    dd.DISCR_TAG_BASE: LGL(order)})
+
+
+        def f(x):
+            if dim == 2:
+                ret = actx.np.cos(np.pi*x[0]) + actx.np.sin(np.pi*x[1])
+            elif dim == 3:
+                ret = actx.np.cos(np.pi*x[0]) + actx.np.sin(np.pi*x[1]) \
+                        + actx.np.sin(np.pi*x[2])
+            else:
+                ret = None
+            assert ret is not None
+
+            return ret
+
+
+        def grad_f(x):
+            ret = make_obj_array([dcoll.zeros(actx) for _ in range(dim)])
+
+            if dim == 2:
+                ret[0] = -np.pi*actx.np.sin(np.pi*x[0])
+                ret[1] = np.pi*actx.np.cos(np.pi*x[1])
+            elif dim == 3:
+                ret[0] = -np.pi*actx.np.sin(np.pi*x[0])
+                ret[1] = np.pi*actx.np.cos(np.pi*x[1])
+                ret[2] = np.pi*actx.np.cos(np.pi*x[2])
+
+            return ret
+
+
+        x = actx.thaw(dcoll.nodes())
+        u = f(x)
+        ref_grad = grad_f(x)
+        grad = op.local_grad(dcoll, u)
+
+        rel_linf_error = actx.to_numpy(op.norm(dcoll, ref_grad - grad, np.inf) /
+                                       op.norm(dcoll, ref_grad, np.inf))
+        eoc_rec.add_data_point(1./n, rel_linf_error)
+
+    print("L^inf error:")
+    print(eoc_rec)
+    assert (eoc_rec.order_estimate() >= order - 0.5 or
+            eoc_rec.max_error() < 1e-11)
+
 # }}}