Add hooks for constructing connections to and from modal representations

grudge/discretization.py

@@ -21,7 +21,8 @@
 """
 
 from pytools import memoize_method
-from grudge.dof_desc import QTAG_NONE, DTAG_BOUNDARY, DOFDesc, as_dofdesc
+from grudge.dof_desc import \
+    QTAG_NONE, QTAG_MODAL, DTAG_BOUNDARY, DOFDesc, as_dofdesc
 import numpy as np  # noqa: F401
 from meshmode.array_context import ArrayContext
 from meshmode.discretization.connection import \
@@ -82,6 +83,12 @@ def __init__(self, array_context, mesh, order=None,
                 quad_tag_to_group_factory[QTAG_NONE] = \
                         PolynomialWarpAndBlendGroupFactory(order=order)
 
+        # Modal discr should always comes from the base discretization
+        quad_tag_to_group_factory[QTAG_MODAL] = \
+            _generate_modal_group_factory(
+                quad_tag_to_group_factory[QTAG_NONE]
+            )
+
         self.quad_tag_to_group_factory = quad_tag_to_group_factory
 
         from meshmode.discretization import Discretization
@@ -164,6 +171,9 @@ def discr_from_dd(self, dd):
 
         qtag = dd.quadrature_tag
 
+        if qtag is QTAG_MODAL:
+            return self._modal_discr(dd.domain_tag)
+
         if dd.is_volume():
             if qtag is not QTAG_NONE:
                 return self._quad_volume_discr(qtag)
@@ -195,10 +205,23 @@ def connection_from_dds(self, from_dd, to_dd):
         to_dd = as_dofdesc(to_dd)
 
         to_qtag = to_dd.quadrature_tag
+        from_qtag = from_dd.quadrature_tag
+
+        # {{{ mapping between modal and nodal representations
+
+        if (from_qtag is QTAG_MODAL and to_qtag is not QTAG_MODAL):
+            return self._modal_to_nodal_connection(to_dd)
+
+        if (to_qtag is QTAG_MODAL and from_qtag is not QTAG_MODAL):
+            return self._nodal_to_modal_connection(from_dd)
+
+        # }}}
+
+        assert (to_qtag is not QTAG_MODAL and from_qtag is not QTAG_MODAL)
 
         if (
                 not from_dd.is_volume()
-                and from_dd.quadrature_tag == to_dd.quadrature_tag
+                and from_qtag == to_qtag
                 and to_dd.domain_tag is FACE_RESTR_ALL):
             faces_conn = self.connection_from_dds(
                     DOFDesc("vol"),
@@ -214,10 +237,9 @@ def connection_from_dds(self, from_dd, to_dd):
 
         # {{{ simplify domain + qtag change into chained
 
-        if (
-                from_dd.domain_tag != to_dd.domain_tag
-                and from_dd.quadrature_tag is QTAG_NONE
-                and to_dd.quadrature_tag is not QTAG_NONE):
+        if (from_dd.domain_tag != to_dd.domain_tag
+                and from_qtag is QTAG_NONE
+                and to_qtag is not QTAG_NONE):
 
             from meshmode.discretization.connection import \
                     ChainedDiscretizationConnection
@@ -239,10 +261,11 @@ def connection_from_dds(self, from_dd, to_dd):
 
         # {{{ generic to-quad
 
-        if (
-                from_dd.domain_tag == to_dd.domain_tag
-                and from_dd.quadrature_tag is QTAG_NONE
-                and to_dd.quadrature_tag is not QTAG_NONE):
+        # QTAG_MODAL is handled above
+        if (from_dd.domain_tag == to_dd.domain_tag
+                and from_qtag is QTAG_NONE
+                and to_qtag is not QTAG_NONE):
+
             from meshmode.discretization.connection.same_mesh import \
                     make_same_mesh_connection
 
@@ -253,7 +276,7 @@ def connection_from_dds(self, from_dd, to_dd):
 
         # }}}
 
-        if from_dd.quadrature_tag is not QTAG_NONE:
+        if from_qtag is not QTAG_NONE:
             raise ValueError("cannot interpolate *from* a "
                     "(non-interpolatory) quadrature grid")
 
@@ -265,12 +288,12 @@ def connection_from_dds(self, from_dd, to_dd):
             if to_dd.domain_tag is FACE_RESTR_INTERIOR:
                 return self._interior_faces_connection()
             elif to_dd.is_boundary_or_partition_interface():
-                assert from_dd.quadrature_tag is QTAG_NONE
+                assert from_qtag is QTAG_NONE
                 return self._boundary_connection(to_dd.domain_tag.tag)
             elif to_dd.is_volume():
                 from meshmode.discretization.connection import \
                         make_same_mesh_connection
-                to_discr = self._quad_volume_discr(to_dd.quadrature_tag)
+                to_discr = self._quad_volume_discr(to_qtag)
                 from_discr = self._volume_discr
                 return make_same_mesh_connection(self._setup_actx, to_discr,
                             from_discr)
@@ -298,6 +321,50 @@ def _quad_volume_discr(self, quadrature_tag):
         return Discretization(self._setup_actx, self._volume_discr.mesh,
                 self.group_factory_for_quadrature_tag(quadrature_tag))
 
+    # {{{ modal to nodal connections
+
+    @memoize_method
+    def _modal_discr(self, domain_tag):
+        from meshmode.discretization import Discretization
+
+        discr_base = self.discr_from_dd(DOFDesc(domain_tag, QTAG_NONE))
+        return Discretization(
+            self._setup_actx, discr_base.mesh,
+            self.group_factory_for_quadrature_tag(QTAG_MODAL)
+        )
+
+    @memoize_method
+    def _modal_to_nodal_connection(self, to_dd):
+        """
+        :arg to_dd: a :class:`grudge.dof_desc.DOFDesc`
+            describing the dofs corresponding to the
+            *to_discr*
+        """
+        from meshmode.discretization.connection import \
+            ModalToNodalDiscretizationConnection
+
+        return ModalToNodalDiscretizationConnection(
+            from_discr=self._modal_discr(to_dd.domain_tag),
+            to_discr=self.discr_from_dd(to_dd)
+        )
+
+    @memoize_method
+    def _nodal_to_modal_connection(self, from_dd):
+        """
+        :arg from_dd: a :class:`grudge.dof_desc.DOFDesc`
+            describing the dofs corresponding to the
+            *from_discr*
+        """
+        from meshmode.discretization.connection import \
+            NodalToModalDiscretizationConnection
+
+        return NodalToModalDiscretizationConnection(
+            from_discr=self.discr_from_dd(from_dd),
+            to_discr=self._modal_discr(from_dd.domain_tag)
+        )
+
+    # }}}
+
     # {{{ boundary
 
     @memoize_method
@@ -403,4 +470,24 @@ def __init__(self, *args, **kwargs):
 
         super().__init__(*args, **kwargs)
 
+
+def _generate_modal_group_factory(nodal_group_factory):
+    from meshmode.discretization.poly_element import (
+        ModalSimplexGroupFactory,
+        ModalTensorProductGroupFactory
+    )
+    from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup
+
+    order = nodal_group_factory.order
+    mesh_group_cls = nodal_group_factory.mesh_group_class
+
+    if mesh_group_cls is SimplexElementGroup:
+        return ModalSimplexGroupFactory(order=order)
+    elif mesh_group_cls is TensorProductElementGroup:
+        return ModalTensorProductGroupFactory(order=order)
+    else:
+        raise ValueError(
+            f"Unknown mesh element group: {mesh_group_cls}"
+        )
+
 # vim: foldmethod=marker

grudge/dof_desc.py

@@ -36,12 +36,14 @@
 .. autoclass:: DTAG_VOLUME_ALL
 .. autoclass:: DTAG_BOUNDARY
 .. autoclass:: QTAG_NONE
+.. autoclass:: QTAG_MODAL
 
 .. autoclass:: DOFDesc
 .. autofunction:: as_dofdesc
 
 .. data:: DD_SCALAR
 .. data:: DD_VOLUME
+.. data:: DD_VOLUME_MODAL
 """
 
 
@@ -102,6 +104,16 @@ class QTAG_NONE:  # noqa: N801
     """
 
 
+class QTAG_MODAL:  # noqa: N801
+    """A quadrature tag indicating the use of a
+    basic discretization grid with modal degrees of
+    freedom. This tag is used to distinguish the
+    modal discretization (`QTAG_MODAL`) from
+    the base (nodal) discretization (e.g. `QTAG_NONE`)
+    or discretizations on quadrature grids.
+    """
+
+
 class DOFDesc:
     """Describes the meaning of degrees of freedom.
 
@@ -236,6 +248,8 @@ def fmt(s):
 
 DD_VOLUME = DOFDesc(DTAG_VOLUME_ALL, None)
 
+DD_VOLUME_MODAL = DOFDesc(DTAG_VOLUME_ALL, QTAG_MODAL)
+
 
 def as_dofdesc(dd):
     if isinstance(dd, DOFDesc):

test/test_modal_connections.py

@@ -0,0 +1,133 @@
+__copyright__ = "Copyright (C) 2021 University of Illinois Board of Trustees"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+
+from meshmode.array_context import (  # noqa
+    pytest_generate_tests_for_pyopencl_array_context
+    as pytest_generate_tests
+)
+from meshmode.discretization.poly_element import (
+    # Simplex group factories
+    InterpolatoryQuadratureSimplexGroupFactory,
+    PolynomialWarpAndBlendGroupFactory,
+    PolynomialEquidistantSimplexGroupFactory,
+    # Tensor product group factories
+    LegendreGaussLobattoTensorProductGroupFactory,
+    # Quadrature-based (non-interpolatory) group factories
+    QuadratureSimplexGroupFactory
+    )
+from meshmode.dof_array import thaw
+import meshmode.mesh.generation as mgen
+
+from grudge import DiscretizationCollection
+import grudge.dof_desc as dof_desc
+
+import pytest
+
+
+@pytest.mark.parametrize("nodal_group_factory", [
+    InterpolatoryQuadratureSimplexGroupFactory,
+    PolynomialWarpAndBlendGroupFactory,
+    PolynomialEquidistantSimplexGroupFactory,
+    LegendreGaussLobattoTensorProductGroupFactory,
+    ]
+)
+def test_inverse_modal_connections(actx_factory, nodal_group_factory):
+    actx = actx_factory()
+    order = 4
+
+    def f(x):
+        return 2*actx.np.sin(20*x) + 0.5*actx.np.cos(10*x)
+
+    # Make a regular rectangle mesh
+    mesh = mgen.generate_regular_rect_mesh(
+        a=(0, 0), b=(5, 3), npoints_per_axis=(10, 6), order=order,
+        group_cls=nodal_group_factory.mesh_group_class
+    )
+
+    dcoll = DiscretizationCollection(
+        actx, mesh,
+        quad_tag_to_group_factory={
+            dof_desc.QTAG_NONE: nodal_group_factory(order)
+        }
+    )
+
+    dd_modal = dof_desc.DD_VOLUME_MODAL
+    dd_volume = dof_desc.DD_VOLUME
+
+    x_nodal = thaw(actx, dcoll.discr_from_dd(dd_volume).nodes()[0])
+    nodal_f = f(x_nodal)
+
+    # Map nodal coefficients of f to modal coefficients
+    forward_conn = dcoll.connection_from_dds(dd_volume, dd_modal)
+    modal_f = forward_conn(nodal_f)
+    # Now map the modal coefficients back to nodal
+    backward_conn = dcoll.connection_from_dds(dd_modal, dd_volume)
+    nodal_f_2 = backward_conn(modal_f)
+
+    # This error should be small since we composed a map with
+    # its inverse
+    err = actx.np.linalg.norm(nodal_f - nodal_f_2)
+
+    assert err <= 1e-13
+
+
+def test_inverse_modal_connections_quadgrid(actx_factory):
+    actx = actx_factory()
+    order = 5
+
+    def f(x):
+        return 1 + 2*x + 3*x**2
+
+    # Make a regular rectangle mesh
+    mesh = mgen.generate_regular_rect_mesh(
+        a=(0, 0), b=(5, 3), npoints_per_axis=(10, 6), order=order,
+        group_cls=QuadratureSimplexGroupFactory.mesh_group_class
+    )
+
+    dcoll = DiscretizationCollection(
+        actx, mesh,
+        quad_tag_to_group_factory={
+            dof_desc.QTAG_NONE: PolynomialWarpAndBlendGroupFactory(order),
+            "quad": QuadratureSimplexGroupFactory(2*order)
+        }
+    )
+
+    # Use dof descriptors on the quadrature grid
+    dd_modal = dof_desc.DD_VOLUME_MODAL
+    dd_quad = dof_desc.DOFDesc(dof_desc.DTAG_VOLUME_ALL, "quad")
+
+    x_quad = thaw(actx, dcoll.discr_from_dd(dd_quad).nodes()[0])
+    quad_f = f(x_quad)
+
+    # Map nodal coefficients of f to modal coefficients
+    forward_conn = dcoll.connection_from_dds(dd_quad, dd_modal)
+    modal_f = forward_conn(quad_f)
+    # Now map the modal coefficients back to nodal
+    backward_conn = dcoll.connection_from_dds(dd_modal, dd_quad)
+    quad_f_2 = backward_conn(modal_f)
+
+    # This error should be small since we composed a map with
+    # its inverse
+    err = actx.np.linalg.norm(quad_f - quad_f_2)
+
+    assert err <= 1e-11