From 836d0af9db46f054990577943ed9ffdfb6100ea4 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 15 Mar 2022 09:23:36 -0500 Subject: [PATCH 01/12] Actually remove quad_tag_to_group_factory --- grudge/discretization.py | 30 +----------------------------- 1 file changed, 1 insertion(+), 29 deletions(-) diff --git a/grudge/discretization.py b/grudge/discretization.py index 43bd24226..2af050d58 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -86,9 +86,7 @@ class DiscretizationCollection: def __init__(self, array_context: ArrayContext, mesh: Mesh, order=None, - discr_tag_to_group_factory=None, mpi_communicator=None, - # FIXME: `quad_tag_to_group_factory` is deprecated - quad_tag_to_group_factory=None): + discr_tag_to_group_factory=None, mpi_communicator=None): """ :arg discr_tag_to_group_factory: A mapping from discretization tags (typically one of: :class:`grudge.dof_desc.DISCR_TAG_BASE`, @@ -101,22 +99,6 @@ def __init__(self, array_context: ArrayContext, mesh: Mesh, discretization. """ - if (quad_tag_to_group_factory is not None - and discr_tag_to_group_factory is not None): - raise ValueError( - "Both `quad_tag_to_group_factory` and `discr_tag_to_group_factory` " - "are specified. Use `discr_tag_to_group_factory` instead." - ) - - # FIXME: `quad_tag_to_group_factory` is deprecated - if (quad_tag_to_group_factory is not None - and discr_tag_to_group_factory is None): - warn("`quad_tag_to_group_factory` is a deprecated kwarg and will " - "be dropped in version 2022.x. Use `discr_tag_to_group_factory` " - "instead.", - DeprecationWarning, stacklevel=2) - discr_tag_to_group_factory = quad_tag_to_group_factory - self._setup_actx = array_context.clone() from meshmode.discretization.poly_element import \ @@ -196,16 +178,6 @@ def mpi_communicator(self): return self._mpi_communicator - @property - def quad_tag_to_group_factory(self): - warn("`DiscretizationCollection.quad_tag_to_group_factory` " - "is deprecated and will go away in 2022. Use " - "`DiscretizationCollection.discr_tag_to_group_factory` " - "instead.", - DeprecationWarning, stacklevel=2) - - return self.discr_tag_to_group_factory - def get_management_rank_index(self): return 0 From b46d3078425a561e0e2fbfa52f0f84d663266fc1 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 15 Mar 2022 09:25:40 -0500 Subject: [PATCH 02/12] Remove get_distributed_boundary_swap_connection, DGDiscretizationWithBoundaries --- grudge/discretization.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/grudge/discretization.py b/grudge/discretization.py index 2af050d58..41bdfd1d6 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -224,14 +224,6 @@ def _set_up_distributed_communication(self, mpi_communicator, array_context): return boundary_connections - def get_distributed_boundary_swap_connection(self, dd): - warn("`DiscretizationCollection.get_distributed_boundary_swap_connection` " - "is deprecated and will go away in 2022. Use " - "`DiscretizationCollection.distributed_boundary_swap_connection` " - "instead.", - DeprecationWarning, stacklevel=2) - return self.distributed_boundary_swap_connection(dd) - def distributed_boundary_swap_connection(self, dd): """Provides a mapping from the base volume discretization to the exterior boundary restriction on a parallel boundary @@ -713,15 +705,6 @@ def normal(self, dd): # }}} -class DGDiscretizationWithBoundaries(DiscretizationCollection): - def __init__(self, *args, **kwargs): - warn("DGDiscretizationWithBoundaries is deprecated and will go away " - "in 2022. Use DiscretizationCollection instead.", - DeprecationWarning, stacklevel=2) - - super().__init__(*args, **kwargs) - - def _generate_modal_group_factory(nodal_group_factory): from meshmode.discretization.poly_element import ( ModalSimplexGroupFactory, From 1badecdc1a2620d65385812587037b37e95cb174 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 15 Mar 2022 09:31:25 -0500 Subject: [PATCH 03/12] Fix sectioning, remove DColl.order --- grudge/discretization.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/grudge/discretization.py b/grudge/discretization.py index 41bdfd1d6..49e14efe3 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -497,6 +497,8 @@ def group_factory_for_discretization_tag(self, discretization_tag): # }}} + # {{{ (internal) discretization getters + @memoize_method def _discr_tag_volume_discr(self, discretization_tag): assert discretization_tag is not None @@ -521,6 +523,8 @@ def _modal_discr(self, domain_tag): self.group_factory_for_discretization_tag(DISCR_TAG_MODAL) ) + # }}} + # {{{ connection factories: modal<->nodal @memoize_method @@ -617,6 +621,8 @@ def _all_faces_volume_connection(self): # }}} + # {{{ properties + @property def dim(self): """Return the topological dimension.""" @@ -644,6 +650,10 @@ def mesh(self): """ return self._volume_discr.mesh + # }}} + + # {{{ array creators + def empty(self, array_context: ArrayContext, dtype=None): """Return an empty :class:`~meshmode.dof_array.DOFArray` defined at the volume nodes: :class:`grudge.dof_desc.DD_VOLUME`. @@ -669,17 +679,9 @@ def zeros(self, array_context: ArrayContext, dtype=None): def is_volume_where(self, where): return where is None or as_dofdesc(where).is_volume() - @property - def order(self): - warn("DiscretizationCollection.order is deprecated, " - "consider using the orders of element groups instead. " - "'order' will go away in 2021.", - DeprecationWarning, stacklevel=2) - - from pytools import single_valued - return single_valued(egrp.order for egrp in self._volume_discr.groups) + # }}} - # {{{ Discretization-specific geometric properties + # {{{ discretization-specific geometric fields def nodes(self, dd=None): r"""Return the nodes of a discretization specified by *dd*. From 4794f1bb62f3a2fe548d86b4608afed5ca3c352c Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 17 Mar 2022 00:42:37 -0500 Subject: [PATCH 04/12] Drop removed stuff from pylintrc --- .pylintrc-local.yml | 8 -------- 1 file changed, 8 deletions(-) diff --git a/.pylintrc-local.yml b/.pylintrc-local.yml index 08f36d4c9..13303e6cb 100644 --- a/.pylintrc-local.yml +++ b/.pylintrc-local.yml @@ -1,14 +1,6 @@ - arg: ignore val: - mappers - - gas_dynamics - - burgers.py - - diffusion.py - - dt_finding.py - - nd_calculus.py - - pml.py - - poisson.py - - second_order.py - arg: ignored-modules val: - sympy From efb95141206f11e6eced66a74ab3b89aad23f3d3 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 29 Aug 2022 15:24:15 -0500 Subject: [PATCH 05/12] add more sectioning --- grudge/discretization.py | 5 +++++ test/test_grudge.py | 8 ++++++++ 2 files changed, 13 insertions(+) diff --git a/grudge/discretization.py b/grudge/discretization.py index 49e14efe3..de06aca62 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -707,6 +707,8 @@ def normal(self, dd): # }}} +# {{{ modal group factory + def _generate_modal_group_factory(nodal_group_factory): from meshmode.discretization.poly_element import ( ModalSimplexGroupFactory, @@ -726,4 +728,7 @@ def _generate_modal_group_factory(nodal_group_factory): f"Unknown mesh element group: {mesh_group_cls}" ) +# }}} + + # vim: foldmethod=marker diff --git a/test/test_grudge.py b/test/test_grudge.py index ce0b199b8..916bd6673 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -979,6 +979,8 @@ def bessel_j(actx, n, r): # }}} +# {{{ test norms + @pytest.mark.parametrize("p", [2, np.inf]) def test_norm_real(actx_factory, p): actx = actx_factory() @@ -1042,6 +1044,10 @@ def test_norm_obj_array(actx_factory, p): logger.info("norm: %.5e %.5e", norm, ref_norm) assert abs(norm-ref_norm) / abs(ref_norm) < 1e-14 +# }}} + + +# {{{ empty boundaries def test_empty_boundary(actx_factory): # https://github.com/inducer/grudge/issues/54 @@ -1061,6 +1067,8 @@ def test_empty_boundary(actx_factory): assert isinstance(component, DOFArray) assert len(component) == len(dcoll.discr_from_dd(BTAG_NONE).groups) +# }}} + # You can test individual routines by typing # $ python test_grudge.py 'test_routine()' From 07b0b568d79847c5e02f2511d7b8c41978a57466 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 29 Aug 2022 15:18:35 -0500 Subject: [PATCH 06/12] remove deprecated group_factory_for_quadrature_tag --- grudge/discretization.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/grudge/discretization.py b/grudge/discretization.py index de06aca62..862e3af4e 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -477,15 +477,6 @@ def connection_from_dds(self, from_dd, to_dd): # {{{ group_factory_for_discretization_tag - def group_factory_for_quadrature_tag(self, discretization_tag): - warn("`DiscretizationCollection.group_factory_for_quadrature_tag` " - "is deprecated and will go away in 2022. Use " - "`DiscretizationCollection.group_factory_for_discretization_tag` " - "instead.", - DeprecationWarning, stacklevel=2) - - return self.group_factory_for_discretization_tag(discretization_tag) - def group_factory_for_discretization_tag(self, discretization_tag): """ OK to override in user code to control mode/node choice. From 5c9556e8e7db4078af57b19108c666908ac409e6 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 29 Aug 2022 15:29:09 -0500 Subject: [PATCH 07/12] correct characteristic_lengthscales docstring --- grudge/dt_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 6c846cc49..b949a2399 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -79,8 +79,8 @@ def characteristic_lengthscales( node distance on the reference cell (see :func:`dt_non_geometric_factors`), and :math:`r_D` is the inradius of the cell (see :func:`dt_geometric_factors`). - :returns: a frozen :class:`~meshmode.dof_array.DOFArray` containing a - characteristic lengthscale for each element, at each nodal location. + :returns: a :class:`~meshmode.dof_array.DOFArray` containing a characteristic + lengthscale for each element, at each nodal location. .. note:: From 098d57244c9e1e50c60cc02245d84dd54600fc7f Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Fri, 26 Aug 2022 13:35:56 -0500 Subject: [PATCH 08/12] add support for multiple independent volumes --- doc/conf.py | 1 + examples/hello-grudge.py | 6 +- grudge/__init__.py | 5 +- grudge/discretization.py | 494 +++++++++++++++++++++----------- grudge/dof_desc.py | 500 +++++++++++++++++++-------------- grudge/dt_utils.py | 35 ++- grudge/eager.py | 16 +- grudge/geometry/metrics.py | 94 ++++--- grudge/models/advection.py | 25 +- grudge/op.py | 129 ++++++--- grudge/projection.py | 32 ++- grudge/reductions.py | 14 +- grudge/shortcuts.py | 9 +- grudge/trace_pair.py | 247 +++++++++++----- test/test_grudge.py | 30 +- test/test_mpi_communication.py | 6 +- test/test_reductions.py | 4 +- 17 files changed, 1051 insertions(+), 596 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index 0cd5ba6d6..1ff3ca070 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -35,6 +35,7 @@ def get_version(): "https://documen.tician.de/arraycontext/": None, "https://documen.tician.de/meshmode/": None, "https://documen.tician.de/loopy/": None, + "https://mpi4py.readthedocs.io/en/stable": None, } # index-page demo uses pyopencl via plot_directive diff --git a/examples/hello-grudge.py b/examples/hello-grudge.py index cfb724115..017430d86 100644 --- a/examples/hello-grudge.py +++ b/examples/hello-grudge.py @@ -14,7 +14,7 @@ import grudge.op as op from meshmode.mesh.generation import generate_box_mesh from meshmode.array_context import PyOpenCLArrayContext -from grudge.dof_desc import DTAG_BOUNDARY, FACE_RESTR_INTERIOR +from grudge.dof_desc import BoundaryDomainTag, FACE_RESTR_INTERIOR ctx = cl.create_some_context() @@ -51,8 +51,8 @@ def flux(dcoll, u_tpair): vol_discr = dcoll.discr_from_dd("vol") -left_bndry = DTAG_BOUNDARY("left") -right_bndry = DTAG_BOUNDARY("right") +left_bndry = BoundaryDomainTag("left") +right_bndry = BoundaryDomainTag("right") x_vol = actx.thaw(dcoll.nodes()) x_bndry = actx.thaw(dcoll.discr_from_dd(left_bndry).nodes()) diff --git a/grudge/__init__.py b/grudge/__init__.py index aad8dbd1c..fa4a2b3b7 100644 --- a/grudge/__init__.py +++ b/grudge/__init__.py @@ -20,8 +20,9 @@ THE SOFTWARE. """ -from grudge.discretization import DiscretizationCollection +from grudge.discretization import ( + DiscretizationCollection, make_discretization_collection) __all__ = [ - "DiscretizationCollection" + "DiscretizationCollection", "make_discretization_collection" ] diff --git a/grudge/discretization.py b/grudge/discretization.py index 862e3af4e..8e57ca503 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -1,7 +1,12 @@ """ -.. currentmodule:: grudge +.. autoclass:: DiscretizationTag + +.. currentmodule:: grudge .. autoclass:: DiscretizationCollection +.. autofunction:: make_discretization_collection + +.. currentmodule:: grudge.discretization """ __copyright__ = """ @@ -29,30 +34,89 @@ THE SOFTWARE. """ -from pytools import memoize_method +from typing import Mapping, Optional, Union, TYPE_CHECKING, Any + +from pytools import memoize_method, single_valued from grudge.dof_desc import ( - DD_VOLUME, + VTAG_ALL, + DD_VOLUME_ALL, DISCR_TAG_BASE, DISCR_TAG_MODAL, - DTAG_BOUNDARY, + VolumeDomainTag, BoundaryDomainTag, DOFDesc, - as_dofdesc + VolumeTag, DomainTag, + DiscretizationTag, + as_dofdesc, + ConvertibleToDOFDesc ) import numpy as np # noqa: F401 from arraycontext import ArrayContext +from meshmode.discretization import ElementGroupFactory, Discretization from meshmode.discretization.connection import ( FACE_RESTR_INTERIOR, FACE_RESTR_ALL, - make_face_restriction + make_face_restriction, + DiscretizationConnection ) from meshmode.mesh import Mesh, BTAG_PARTITION +from meshmode.dof_array import DOFArray from warnings import warn +if TYPE_CHECKING: + import mpi4py.MPI + + +# {{{ discr_tag_to_group_factory normalization + +def _normalize_discr_tag_to_group_factory( + dim: int, + discr_tag_to_group_factory: Optional[ + Mapping[DiscretizationTag, ElementGroupFactory]], + order: Optional[int] + ) -> Mapping[DiscretizationTag, ElementGroupFactory]: + from meshmode.discretization.poly_element import \ + default_simplex_group_factory + + if discr_tag_to_group_factory is None: + if order is None: + raise TypeError( + "one of 'order' and 'discr_tag_to_group_factory' must be given" + ) + + discr_tag_to_group_factory = { + DISCR_TAG_BASE: default_simplex_group_factory( + base_dim=dim, order=order)} + else: + discr_tag_to_group_factory = dict(discr_tag_to_group_factory) + + if order is not None: + if DISCR_TAG_BASE in discr_tag_to_group_factory: + raise ValueError( + "if 'order' is given, 'discr_tag_to_group_factory' must " + "not have a key of DISCR_TAG_BASE" + ) + + discr_tag_to_group_factory[DISCR_TAG_BASE] = \ + default_simplex_group_factory(base_dim=dim, order=order) + + assert discr_tag_to_group_factory is not None + + # Modal discr should always come from the base discretization + if DISCR_TAG_MODAL not in discr_tag_to_group_factory: + discr_tag_to_group_factory[DISCR_TAG_MODAL] = \ + _generate_modal_group_factory( + discr_tag_to_group_factory[DISCR_TAG_BASE] + ) + + return discr_tag_to_group_factory + +# }}} + class DiscretizationCollection: """A collection of discretizations, defined on the same underlying @@ -60,11 +124,13 @@ class DiscretizationCollection: (volume, interior facets, boundaries) and associated element groups. - .. automethod:: __init__ + .. note:: + + Do not call the constructor directly. Use + :func:`make_discretization_collection` instead. .. autoattribute:: dim .. autoattribute:: ambient_dim - .. autoattribute:: mesh .. autoattribute:: real_dtype .. autoattribute:: complex_dtype @@ -84,9 +150,13 @@ class DiscretizationCollection: # {{{ constructor - def __init__(self, array_context: ArrayContext, mesh: Mesh, - order=None, - discr_tag_to_group_factory=None, mpi_communicator=None): + def __init__(self, array_context: ArrayContext, + volume_discrs: Union[Mesh, Mapping[VolumeTag, Discretization]], + order: Optional[int] = None, + discr_tag_to_group_factory: Optional[ + Mapping[DiscretizationTag, ElementGroupFactory]] = None, + mpi_communicator: Optional["mpi4py.MPI.Intracomm"] = None, + ) -> None: """ :arg discr_tag_to_group_factory: A mapping from discretization tags (typically one of: :class:`grudge.dof_desc.DISCR_TAG_BASE`, @@ -101,45 +171,6 @@ def __init__(self, array_context: ArrayContext, mesh: Mesh, self._setup_actx = array_context.clone() - from meshmode.discretization.poly_element import \ - default_simplex_group_factory - - if discr_tag_to_group_factory is None: - if order is None: - raise TypeError( - "one of 'order' and 'discr_tag_to_group_factory' must be given" - ) - - discr_tag_to_group_factory = { - DISCR_TAG_BASE: default_simplex_group_factory( - base_dim=mesh.dim, order=order)} - else: - if order is not None: - discr_tag_to_group_factory = discr_tag_to_group_factory.copy() - if DISCR_TAG_BASE in discr_tag_to_group_factory: - raise ValueError( - "if 'order' is given, 'discr_tag_to_group_factory' must " - "not have a key of DISCR_TAG_BASE" - ) - - discr_tag_to_group_factory[DISCR_TAG_BASE] = \ - default_simplex_group_factory(base_dim=mesh.dim, order=order) - - # Modal discr should always come from the base discretization - discr_tag_to_group_factory[DISCR_TAG_MODAL] = \ - _generate_modal_group_factory( - discr_tag_to_group_factory[DISCR_TAG_BASE] - ) - - self.discr_tag_to_group_factory = discr_tag_to_group_factory - - from meshmode.discretization import Discretization - - self._volume_discr = Discretization( - array_context, mesh, - self.group_factory_for_discretization_tag(DISCR_TAG_BASE) - ) - # {{{ process mpi_communicator argument if mpi_communicator is not None: @@ -163,9 +194,42 @@ def __init__(self, array_context: ArrayContext, mesh: Mesh, # }}} - self._dist_boundary_connections = \ - self._set_up_distributed_communication( - mpi_communicator, array_context) + from meshmode.discretization import Discretization + + if isinstance(volume_discrs, Mesh): + # {{{ deprecated backward compatibility yuck + + warn("Calling the DiscretizationCollection constructor directly " + "is deprecated, call make_discretization_collection " + "instead. This will stop working in 2023.", + DeprecationWarning, stacklevel=2) + + mesh = volume_discrs + + discr_tag_to_group_factory = _normalize_discr_tag_to_group_factory( + dim=mesh.dim, + discr_tag_to_group_factory=discr_tag_to_group_factory, + order=order) + self._discr_tag_to_group_factory = discr_tag_to_group_factory + + volume_discr = Discretization( + array_context, mesh, + self.group_factory_for_discretization_tag(DISCR_TAG_BASE)) + volume_discrs = {VTAG_ALL: volume_discr} + + del mesh + + # }}} + else: + assert discr_tag_to_group_factory is not None + self._discr_tag_to_group_factory = discr_tag_to_group_factory + + self._volume_discrs = volume_discrs + + self._dist_boundary_connections = { + vtag: self._set_up_distributed_communication( + vtag, mpi_communicator, array_context) + for vtag in self._volume_discrs.keys()} # }}} @@ -190,13 +254,14 @@ def is_management_rank(self): # {{{ distributed - def _set_up_distributed_communication(self, mpi_communicator, array_context): - from_dd = DOFDesc("vol", DISCR_TAG_BASE) + def _set_up_distributed_communication( + self, vtag, mpi_communicator, array_context): + from_dd = DOFDesc(VolumeDomainTag(vtag), DISCR_TAG_BASE) boundary_connections = {} from meshmode.distributed import get_connected_partitions - connected_parts = get_connected_partitions(self._volume_discr.mesh) + connected_parts = get_connected_partitions(self._volume_discrs[vtag].mesh) if connected_parts: if mpi_communicator is None: @@ -209,8 +274,7 @@ def _set_up_distributed_communication(self, mpi_communicator, array_context): local_boundary_connections = {} for i_remote_part in connected_parts: local_boundary_connections[i_remote_part] = self.connection_from_dds( - from_dd, DOFDesc(BTAG_PARTITION(i_remote_part), - DISCR_TAG_BASE)) + from_dd, from_dd.trace(BTAG_PARTITION(i_remote_part))) from meshmode.distributed import MPIBoundaryCommSetupHelper with MPIBoundaryCommSetupHelper(mpi_communicator, array_context, @@ -233,7 +297,7 @@ def distributed_boundary_swap_connection(self, dd): :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. The domain tag must be a subclass - of :class:`grudge.dof_desc.DTAG_BOUNDARY` with an + of :class:`grudge.dof_desc.BoundaryDomainTag` with an associated :class:`meshmode.mesh.BTAG_PARTITION` corresponding to a particular communication rank. """ @@ -244,22 +308,21 @@ def distributed_boundary_swap_connection(self, dd): f"{dd.discretization_tag} is not implemented." ) - assert isinstance(dd.domain_tag, DTAG_BOUNDARY) + assert isinstance(dd.domain_tag, BoundaryDomainTag) assert isinstance(dd.domain_tag.tag, BTAG_PARTITION) - return self._dist_boundary_connections[dd.domain_tag.tag.part_nr] + vtag = dd.domain_tag.volume_tag + + return self._dist_boundary_connections[vtag][dd.domain_tag.tag.part_nr] # }}} # {{{ discr_from_dd @memoize_method - def discr_from_dd(self, dd): + def discr_from_dd(self, dd: "ConvertibleToDOFDesc") -> Discretization: """Provides a :class:`meshmode.discretization.Discretization` object from *dd*. - - :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value - convertible to one. """ dd = as_dofdesc(dd) @@ -269,45 +332,43 @@ def discr_from_dd(self, dd): return self._modal_discr(dd.domain_tag) if dd.is_volume(): - if discr_tag is not DISCR_TAG_BASE: - return self._discr_tag_volume_discr(discr_tag) - return self._volume_discr + return self._volume_discr_from_dd(dd) if discr_tag is not DISCR_TAG_BASE: - no_quad_discr = self.discr_from_dd(DOFDesc(dd.domain_tag)) + base_discr = self.discr_from_dd(dd.with_discr_tag(DISCR_TAG_BASE)) from meshmode.discretization import Discretization return Discretization( self._setup_actx, - no_quad_discr.mesh, + base_discr.mesh, self.group_factory_for_discretization_tag(discr_tag) ) assert discr_tag is DISCR_TAG_BASE - if dd.domain_tag is FACE_RESTR_ALL: - return self._all_faces_volume_connection().to_discr - elif dd.domain_tag is FACE_RESTR_INTERIOR: - return self._interior_faces_connection().to_discr - elif dd.is_boundary_or_partition_interface(): - return self._boundary_connection(dd.domain_tag.tag).to_discr + if isinstance(dd.domain_tag, BoundaryDomainTag): + if dd.domain_tag.tag in [FACE_RESTR_ALL, FACE_RESTR_INTERIOR]: + return self._faces_connection(dd.domain_tag).to_discr + else: + return self._boundary_connection(dd.domain_tag).to_discr else: - raise ValueError("DOF desc tag not understood: " + str(dd)) + raise ValueError(f"DOF desc not understood: {dd}") # }}} # {{{ _base_to_geoderiv_connection @memoize_method - def _has_affine_groups(self): + def _has_affine_groups(self, domain_tag: DomainTag) -> bool: from modepy.shapes import Simplex + discr = self.discr_from_dd(DOFDesc(domain_tag, DISCR_TAG_BASE)) return any( megrp.is_affine and issubclass(megrp._modepy_shape_cls, Simplex) - for megrp in self._volume_discr.mesh.groups) + for megrp in discr.mesh.groups) @memoize_method - def _base_to_geoderiv_connection(self, dd: DOFDesc): + def _base_to_geoderiv_connection(self, dd: DOFDesc) -> DiscretizationConnection: r"""The "geometry derivatives" discretization for a given *dd* is typically identical to the one returned by :meth:`discr_from_dd`, however for affinely-mapped simplicial elements, it will use a @@ -320,7 +381,7 @@ def _base_to_geoderiv_connection(self, dd: DOFDesc): :mod:`grudge`. """ base_discr = self.discr_from_dd(dd) - if not self._has_affine_groups(): + if not self._has_affine_groups(dd.domain_tag): # no benefit to having another discretization that takes # advantage of affine-ness from meshmode.discretization.connection import \ @@ -357,7 +418,9 @@ def geo_group_factory(megrp, index): # {{{ connection_from_dds @memoize_method - def connection_from_dds(self, from_dd, to_dd): + def connection_from_dds( + self, from_dd: "ConvertibleToDOFDesc", to_dd: "ConvertibleToDOFDesc" + ) -> DiscretizationConnection: """Provides a mapping (connection) from one discretization to another, e.g. from the volume to the boundary, or from the base to the an overintegrated quadrature discretization, or from @@ -389,12 +452,15 @@ def connection_from_dds(self, from_dd, to_dd): assert (to_discr_tag is not DISCR_TAG_MODAL and from_discr_tag is not DISCR_TAG_MODAL) - if (not from_dd.is_volume() + if (isinstance(from_dd.domain_tag, BoundaryDomainTag) and from_discr_tag == to_discr_tag - and to_dd.domain_tag is FACE_RESTR_ALL): + and isinstance(to_dd.domain_tag, BoundaryDomainTag) + and to_dd.domain_tag.tag is FACE_RESTR_ALL): faces_conn = self.connection_from_dds( - DOFDesc("vol"), - DOFDesc(from_dd.domain_tag)) + DOFDesc( + VolumeDomainTag(from_dd.domain_tag.volume_tag), + DISCR_TAG_BASE), + from_dd.with_discr_tag(DISCR_TAG_BASE)) from meshmode.discretization.connection import \ make_face_to_all_faces_embedding @@ -412,7 +478,7 @@ def connection_from_dds(self, from_dd, to_dd): from meshmode.discretization.connection import \ ChainedDiscretizationConnection - intermediate_dd = DOFDesc(to_dd.domain_tag) + intermediate_dd = to_dd.with_discr_tag(DISCR_TAG_BASE) return ChainedDiscretizationConnection( [ # first change domain @@ -446,66 +512,79 @@ def connection_from_dds(self, from_dd, to_dd): # }}} if from_discr_tag is not DISCR_TAG_BASE: - raise ValueError("cannot interpolate *from* a " - "(non-interpolatory) quadrature grid") + raise ValueError("cannot get a connection *from* a " + f"(non-interpolatory) quadrature grid: '{from_dd}'") assert to_discr_tag is DISCR_TAG_BASE - if from_dd.is_volume(): - if to_dd.domain_tag is FACE_RESTR_ALL: - return self._all_faces_volume_connection() - if to_dd.domain_tag is FACE_RESTR_INTERIOR: - return self._interior_faces_connection() - elif to_dd.is_boundary_or_partition_interface(): - assert from_discr_tag is DISCR_TAG_BASE - return self._boundary_connection(to_dd.domain_tag.tag) + if isinstance(from_dd.domain_tag, VolumeDomainTag): + if isinstance(to_dd.domain_tag, BoundaryDomainTag): + if to_dd.domain_tag.volume_tag != from_dd.domain_tag.tag: + raise ValueError("cannot get a connection from one volume " + f"('{from_dd.domain_tag.tag}') " + "to the boundary of another volume " + f"('{to_dd.domain_tag.volume_tag}') ") + if to_dd.domain_tag.tag in [FACE_RESTR_ALL, FACE_RESTR_INTERIOR]: + return self._faces_connection(to_dd.domain_tag) + elif isinstance(to_dd.domain_tag, BoundaryDomainTag): + assert from_discr_tag is DISCR_TAG_BASE + return self._boundary_connection(to_dd.domain_tag) elif to_dd.is_volume(): + if to_dd.domain_tag != from_dd.domain_tag: + raise ValueError("cannot get a connection between " + "volumes of different tags: requested " + f"'{from_dd.domain_tag}' -> '{to_dd.domain_tag}'") + from meshmode.discretization.connection import \ make_same_mesh_connection - to_discr = self._discr_tag_volume_discr(to_discr_tag) - from_discr = self._volume_discr - return make_same_mesh_connection(self._setup_actx, to_discr, - from_discr) + return make_same_mesh_connection( + self._setup_actx, + self._volume_discr_from_dd(to_dd), + self._volume_discr_from_dd(from_dd)) else: - raise ValueError("cannot interpolate from volume to: " + str(to_dd)) + raise ValueError( + f"cannot get a connection from volume to: '{to_dd}'") else: - raise ValueError("cannot interpolate from: " + str(from_dd)) + raise ValueError(f"cannot get a connection from: '{from_dd}'") # }}} # {{{ group_factory_for_discretization_tag def group_factory_for_discretization_tag(self, discretization_tag): - """ - OK to override in user code to control mode/node choice. - """ if discretization_tag is None: discretization_tag = DISCR_TAG_BASE - return self.discr_tag_to_group_factory[discretization_tag] + return self._discr_tag_to_group_factory[discretization_tag] # }}} # {{{ (internal) discretization getters @memoize_method - def _discr_tag_volume_discr(self, discretization_tag): - assert discretization_tag is not None + def _volume_discr_from_dd(self, dd: DOFDesc) -> Discretization: + assert isinstance(dd.domain_tag, VolumeDomainTag) + + try: + base_volume_discr = self._volume_discrs[dd.domain_tag.tag] + except KeyError: + raise ValueError("a volume discretization with volume tag " + f"'{dd.domain_tag.tag}' is not known") # Refuse to re-make the volume discretization - if discretization_tag is DISCR_TAG_BASE: - return self._volume_discr + if dd.discretization_tag is DISCR_TAG_BASE: + return base_volume_discr from meshmode.discretization import Discretization return Discretization( - self._setup_actx, self._volume_discr.mesh, - self.group_factory_for_discretization_tag(discretization_tag) + self._setup_actx, base_volume_discr.mesh, + self.group_factory_for_discretization_tag(dd.discretization_tag) ) @memoize_method - def _modal_discr(self, domain_tag): + def _modal_discr(self, domain_tag) -> Discretization: from meshmode.discretization import Discretization discr_base = self.discr_from_dd(DOFDesc(domain_tag, DISCR_TAG_BASE)) @@ -519,7 +598,7 @@ def _modal_discr(self, domain_tag): # {{{ connection factories: modal<->nodal @memoize_method - def _modal_to_nodal_connection(self, to_dd): + def _modal_to_nodal_connection(self, to_dd: DOFDesc) -> DiscretizationConnection: """ :arg to_dd: a :class:`grudge.dof_desc.DOFDesc` describing the dofs corresponding to the @@ -534,7 +613,8 @@ def _modal_to_nodal_connection(self, to_dd): ) @memoize_method - def _nodal_to_modal_connection(self, from_dd): + def _nodal_to_modal_connection( + self, from_dd: DOFDesc) -> DiscretizationConnection: """ :arg from_dd: a :class:`grudge.dof_desc.DOFDesc` describing the dofs corresponding to the @@ -553,25 +633,31 @@ def _nodal_to_modal_connection(self, from_dd): # {{{ connection factories: boundary @memoize_method - def _boundary_connection(self, boundary_tag): + def _boundary_connection( + self, domain_tag: BoundaryDomainTag) -> DiscretizationConnection: return make_face_restriction( - self._setup_actx, - self._volume_discr, - self.group_factory_for_discretization_tag(DISCR_TAG_BASE), - boundary_tag=boundary_tag - ) + self._setup_actx, + self._volume_discr_from_dd( + DOFDesc(VolumeDomainTag(domain_tag.volume_tag), DISCR_TAG_BASE)), + self.group_factory_for_discretization_tag(DISCR_TAG_BASE), + boundary_tag=domain_tag.tag) # }}} - # {{{ connection factories: interior faces + # {{{ connection factories: faces @memoize_method - def _interior_faces_connection(self): + def _faces_connection( + self, domain_tag: BoundaryDomainTag) -> DiscretizationConnection: + assert domain_tag.tag in [FACE_RESTR_INTERIOR, FACE_RESTR_ALL] + return make_face_restriction( self._setup_actx, - self._volume_discr, + self._volume_discr_from_dd( + DOFDesc( + VolumeDomainTag(domain_tag.volume_tag), DISCR_TAG_BASE)), self.group_factory_for_discretization_tag(DISCR_TAG_BASE), - FACE_RESTR_INTERIOR, + domain_tag.tag, # FIXME: This will need to change as soon as we support # pyramids or other elements with non-identical face @@ -580,7 +666,8 @@ def _interior_faces_connection(self): ) @memoize_method - def opposite_face_connection(self): + def opposite_face_connection( + self, domain_tag: BoundaryDomainTag) -> DiscretizationConnection: """Provides a mapping from the base volume discretization to the exterior boundary restriction on a neighboring element. This does not take into account parallel partitions. @@ -588,84 +675,71 @@ def opposite_face_connection(self): from meshmode.discretization.connection import \ make_opposite_face_connection + assert domain_tag.tag is FACE_RESTR_INTERIOR + return make_opposite_face_connection( self._setup_actx, - self._interior_faces_connection()) - - # }}} - - # {{{ connection factories: all-faces - - @memoize_method - def _all_faces_volume_connection(self): - return make_face_restriction( - self._setup_actx, - self._volume_discr, - self.group_factory_for_discretization_tag(DISCR_TAG_BASE), - FACE_RESTR_ALL, - - # FIXME: This will need to change as soon as we support - # pyramids or other elements with non-identical face - # types. - per_face_groups=False - ) + self._faces_connection(domain_tag)) # }}} # {{{ properties @property - def dim(self): + def dim(self) -> int: """Return the topological dimension.""" - return self._volume_discr.dim + return single_valued(discr.dim for discr in self._volume_discrs.values()) @property - def ambient_dim(self): + def ambient_dim(self) -> int: """Return the dimension of the ambient space.""" - return self._volume_discr.ambient_dim + return single_valued( + discr.ambient_dim for discr in self._volume_discrs.values()) @property - def real_dtype(self): + def real_dtype(self) -> "np.dtype[Any]": """Return the data type used for real-valued arithmetic.""" - return self._volume_discr.real_dtype + return single_valued( + discr.real_dtype for discr in self._volume_discrs.values()) @property - def complex_dtype(self): + def complex_dtype(self) -> "np.dtype[Any]": """Return the data type used for complex-valued arithmetic.""" - return self._volume_discr.complex_dtype - - @property - def mesh(self): - """Return the :class:`meshmode.mesh.Mesh` over which the discretization - collection is built. - """ - return self._volume_discr.mesh + return single_valued( + discr.complex_dtype for discr in self._volume_discrs.values()) # }}} # {{{ array creators - def empty(self, array_context: ArrayContext, dtype=None): + def empty(self, array_context: ArrayContext, dtype=None, + *, dd: Optional[DOFDesc] = None) -> DOFArray: """Return an empty :class:`~meshmode.dof_array.DOFArray` defined at - the volume nodes: :class:`grudge.dof_desc.DD_VOLUME`. + the volume nodes: :class:`grudge.dof_desc.DD_VOLUME_ALL`. :arg array_context: an :class:`~arraycontext.context.ArrayContext`. :arg dtype: type special value 'c' will result in a vector of dtype :attr:`complex_dtype`. If *None* (the default), a real vector will be returned. """ - return self._volume_discr.empty(array_context, dtype) + if dd is None: + dd = DD_VOLUME_ALL + return self.discr_from_dd(dd).empty(array_context, dtype) - def zeros(self, array_context: ArrayContext, dtype=None): + def zeros(self, array_context: ArrayContext, dtype=None, + *, dd: Optional[DOFDesc] = None) -> DOFArray: """Return a zero-initialized :class:`~meshmode.dof_array.DOFArray` - defined at the volume nodes, :class:`grudge.dof_desc.DD_VOLUME`. + defined at the volume nodes, :class:`grudge.dof_desc.DD_VOLUME_ALL`. :arg array_context: an :class:`~arraycontext.context.ArrayContext`. :arg dtype: type special value 'c' will result in a vector of dtype :attr:`complex_dtype`. If *None* (the default), a real vector will be returned. """ - return self._volume_discr.zeros(array_context, dtype) + if dd is None: + dd = DD_VOLUME_ALL + + return self.discr_from_dd(dd).zeros(array_context, dtype) def is_volume_where(self, where): return where is None or as_dofdesc(where).is_volume() @@ -682,7 +756,7 @@ def nodes(self, dd=None): :returns: an object array of frozen :class:`~meshmode.dof_array.DOFArray`\ s """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL return self.discr_from_dd(dd).nodes() def normal(self, dd): @@ -722,4 +796,88 @@ def _generate_modal_group_factory(nodal_group_factory): # }}} +# {{{ make_discretization_collection + +MeshOrDiscr = Union[Mesh, Discretization] + + +def make_discretization_collection( + array_context: ArrayContext, + volumes: Union[ + MeshOrDiscr, + Mapping[VolumeTag, MeshOrDiscr]], + order: Optional[int] = None, + discr_tag_to_group_factory: Optional[ + Mapping[DiscretizationTag, ElementGroupFactory]] = None, + ) -> DiscretizationCollection: + """ + :arg discr_tag_to_group_factory: A mapping from discretization tags + (typically one of: :class:`~grudge.dof_desc.DISCR_TAG_BASE`, + :class:`~grudge.dof_desc.DISCR_TAG_MODAL`, or + :class:`~grudge.dof_desc.DISCR_TAG_QUAD`) to a + :class:`~meshmode.discretization.ElementGroupFactory` + indicating with which type of discretization the operations are + to be carried out, or *None* to indicate that operations with this + discretization tag should be carried out with the standard volume + discretization. + + .. note:: + + If passing a :class:`~meshmode.discretization.Discretization` in + *volumes*, it must be nodal and unisolvent, consistent with + :class:`~grudge.dof_desc.DISCR_TAG_BASE`. + + .. note:: + + To use the resulting :class:`DiscretizationCollection` in a + distributed-memory manner, the *array_context* passed in + must be one of the distributed-memory array contexts + from :mod:`grudge.array_context`. Unlike the (now-deprecated, + for direct use) constructor of :class:`DiscretizationCollection`, + this function no longer accepts a separate MPI communicator. + + .. note:: + + If the resulting :class:`DiscretizationCollection` is distributed + across multiple ranks, then this is an MPI-collective operation, + i.e. all ranks in the communicator must enter this function at the same + time. + """ + if isinstance(volumes, (Mesh, Discretization)): + volumes = {VTAG_ALL: volumes} + + from pytools import is_single_valued + + assert len(volumes) > 0 + assert is_single_valued(mesh_or_discr.ambient_dim + for mesh_or_discr in volumes.values()) + + discr_tag_to_group_factory = _normalize_discr_tag_to_group_factory( + dim=single_valued( + mesh_or_discr.dim for mesh_or_discr in volumes.values()), + discr_tag_to_group_factory=discr_tag_to_group_factory, + order=order) + + del order + + if any( + isinstance(mesh_or_discr, Discretization) + for mesh_or_discr in volumes.values()): + raise NotImplementedError("Doesn't work at the moment") + + volume_discrs = { + vtag: Discretization( + array_context, + mesh, + discr_tag_to_group_factory[DISCR_TAG_BASE]) + for vtag, mesh in volumes.items()} + + return DiscretizationCollection( + array_context=array_context, + volume_discrs=volume_discrs, + discr_tag_to_group_factory=discr_tag_to_group_factory) + +# }}} + + # vim: foldmethod=marker diff --git a/grudge/dof_desc.py b/grudge/dof_desc.py index 267e4f56e..cf285a30e 100644 --- a/grudge/dof_desc.py +++ b/grudge/dof_desc.py @@ -1,4 +1,55 @@ -"""Degree of freedom (DOF) descriptions""" +""" +Volume tags +----------- + +.. autoclass:: VolumeTag +.. autoclass:: VTAG_ALL + +:mod:`grudge`-specific boundary tags +------------------------------------ + +Domain tags +----------- + +A domain tag identifies a geometric part (or whole) of the domain described +by a :class:`grudge.DiscretizationCollection`. This can be a volume or a boundary. + +.. autoclass:: DTAG_SCALAR +.. autoclass:: DTAG_VOLUME_ALL +.. autoclass:: VolumeDomainTag +.. autoclass:: BoundaryDomainTag + +Discretization tags +------------------- + +A discretization tag serves as a symbolic identifier of the manner in which +meaning is assigned to degrees of freedom. + +.. autoclass:: DISCR_TAG_BASE +.. autoclass:: DISCR_TAG_QUAD +.. autoclass:: DISCR_TAG_MODAL + +DOF Descriptor +-------------- + +.. autoclass:: DOFDesc +.. autofunction:: as_dofdesc + +Shortcuts +--------- + +.. data:: DD_SCALAR +.. data:: DD_VOLUME_ALL +.. data:: DD_VOLUME_ALL_MODAL + +Internal things that are visble due to type annotations +------------------------------------------------------- + +.. class:: _DiscretizationTag +.. class:: ConvertibleToDOFDesc + + Anything that is convertible to a :class:`DOFDesc` via :func:`as_dofdesc`. +""" __copyright__ = """ Copyright (C) 2008 Andreas Kloeckner @@ -25,31 +76,18 @@ THE SOFTWARE. """ -from meshmode.discretization.connection import \ - FACE_RESTR_INTERIOR, FACE_RESTR_ALL -from meshmode.mesh import \ - BTAG_PARTITION, BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE -from warnings import warn import sys +from warnings import warn +from typing import Hashable, Union, Type, Optional, Any, Tuple +from dataclasses import dataclass, replace +from meshmode.discretization.connection import ( + FACE_RESTR_INTERIOR, FACE_RESTR_ALL) +from meshmode.mesh import ( + BTAG_PARTITION, BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE, BoundaryTag) -__doc__ = """ -.. autoclass:: DTAG_SCALAR -.. autoclass:: DTAG_VOLUME_ALL -.. autoclass:: DTAG_BOUNDARY - -.. autoclass:: DISCR_TAG_BASE -.. autoclass:: DISCR_TAG_QUAD -.. autoclass:: DISCR_TAG_MODAL - -.. autoclass:: DOFDesc -.. autofunction:: as_dofdesc - -.. data:: DD_SCALAR -.. data:: DD_VOLUME -.. data:: DD_VOLUME_MODAL -""" +# {{{ _to_identifier def _to_identifier(name: str) -> str: if not name.isidentifier(): @@ -57,71 +95,91 @@ def _to_identifier(name: str) -> str: else: return name +# }}} + + +# {{{ volume tags -# {{{ DOF description +class VTAG_ALL: # noqa: N801 + pass -class DTAG_SCALAR: # noqa: N801 + +VolumeTag = Hashable + +# }}} + + +# {{{ domain tag + +@dataclass(frozen=True, eq=True) +class ScalarDomainTag: # noqa: N801 """A domain tag denoting scalar values.""" -class DTAG_VOLUME_ALL: # noqa: N801 - """ - A domain tag denoting values defined - in all cell volumes. +DTAG_SCALAR = ScalarDomainTag() + + +@dataclass(frozen=True, eq=True, init=True) +class VolumeDomainTag: + """A domain tag referring to a volume identified by the + volume tag :attr:`tag`. These volume identifiers are only used + when the :class:`~grudge.discretization.DiscretizationCollection` contains + more than one volume. + + .. attribute:: tag + + .. automethod:: __init__ """ + tag: VolumeTag -class DTAG_BOUNDARY: # noqa: N801 - """A domain tag describing the values on element - boundaries which are adjacent to elements - of another :class:`~meshmode.mesh.Mesh`. +DTAG_VOLUME_ALL = VolumeDomainTag(VTAG_ALL) + + +@dataclass(frozen=True, eq=True, init=True) +class BoundaryDomainTag: + """A domain tag referring to a boundary identified by the + boundary tag :attr:`tag`. .. attribute:: tag + .. attribute:: volume_tag .. automethod:: __init__ - .. automethod:: __eq__ - .. automethod:: __ne__ - .. automethod:: __hash__ """ + tag: BoundaryTag + volume_tag: VolumeTag = VTAG_ALL - def __init__(self, tag): - """ - :arg tag: One of the following: - :class:`~meshmode.mesh.BTAG_ALL`, - :class:`~meshmode.mesh.BTAG_NONE`, - :class:`~meshmode.mesh.BTAG_REALLY_ALL`, - :class:`~meshmode.mesh.BTAG_PARTITION`. - """ - self.tag = tag - def __eq__(self, other): - return isinstance(other, DTAG_BOUNDARY) and self.tag == other.tag +DomainTag = Union[ScalarDomainTag, VolumeDomainTag, BoundaryDomainTag] - def __ne__(self, other): - return not self.__eq__(other) +# }}} + + +# {{{ discretization tag + +class _DiscretizationTag: # noqa: N801 + pass - def __hash__(self): - return hash(type(self)) ^ hash(self.tag) - def __repr__(self): - return "<{}({})>".format(type(self).__name__, repr(self.tag)) +DiscretizationTag = Type[_DiscretizationTag] -class DISCR_TAG_BASE: # noqa: N801 +class DISCR_TAG_BASE(_DiscretizationTag): # noqa: N801 """A discretization tag indicating the use of a - basic discretization grid. This tag is used + nodal and unisolvent discretization. This tag is used to distinguish the base discretization from quadrature (e.g. overintegration) or modal (:class:`DISCR_TAG_MODAL`) discretizations. """ -class DISCR_TAG_QUAD: # noqa: N801 - """A discretization tag indicating the use of a - quadrature discretization grid. This tag is used - to distinguish the quadrature discretization - (e.g. overintegration) from modal (:class:`DISCR_TAG_MODAL`) - or base (:class:`DISCR_TAG_BASE`) discretizations. +class DISCR_TAG_QUAD(_DiscretizationTag): # noqa: N801 + """A discretization tag indicating the use of a quadrature discretization + grid, which typically affords higher quadrature accuracy (e.g. for + nonlinear terms) at the expense of unisolvency. This tag is used to + distinguish the quadrature discretization (e.g. overintegration) from modal + (:class:`DISCR_TAG_MODAL`) or base (:class:`DISCR_TAG_BASE`) + discretizations. For working with multiple quadrature grids, it is recommended to create appropriate subclasses of @@ -135,20 +193,22 @@ class CustomQuadTag(DISCR_TAG_QUAD): "A custom quadrature discretization tag." dd = DOFDesc(DTAG_VOLUME_ALL, CustomQuadTag) - """ -class DISCR_TAG_MODAL: # noqa: N801 - """A discretization tag indicating the use of a - basic discretization grid with modal degrees of - freedom. This tag is used to distinguish the - modal discretization from the base (nodal) - discretization (e.g. :class:`DISCR_TAG_BASE`) or +class DISCR_TAG_MODAL(_DiscretizationTag): # noqa: N801 + """A discretization tag indicating the use of unisolvent modal degrees of + freedom. This tag is used to distinguish the modal discretization from the + base (nodal) discretization (e.g. :class:`DISCR_TAG_BASE`) or discretizations on quadrature grids (:class:`DISCR_TAG_QUAD`). """ +# }}} + + +# {{{ DOF descriptor +@dataclass(frozen=True, eq=True) class DOFDesc: """Describes the meaning of degrees of freedom. @@ -165,8 +225,9 @@ class DOFDesc: .. automethod:: uses_quadrature + .. automethod:: with_domain_tag .. automethod:: with_discr_tag - .. automethod:: with_dtag + .. automethod:: trace .. automethod:: __eq__ .. automethod:: __ne__ @@ -174,159 +235,87 @@ class DOFDesc: .. automethod:: as_identifier """ - def __init__(self, domain_tag, discretization_tag=None, - # FIXME: `quadrature_tag` is deprecated - quadrature_tag=None): - """ - :arg domain_tag: One of the following: - :class:`DTAG_SCALAR` (or the string ``"scalar"``), - :class:`DTAG_VOLUME_ALL` (or the string ``"vol"``) - for the default volume discretization, - :data:`~meshmode.discretization.connection.FACE_RESTR_ALL` - (or the string ``"all_faces"``), or - :data:`~meshmode.discretization.connection.FACE_RESTR_INTERIOR` - (or the string ``"int_faces"``), or one of - :class:`~meshmode.mesh.BTAG_ALL`, - :class:`~meshmode.mesh.BTAG_NONE`, - :class:`~meshmode.mesh.BTAG_REALLY_ALL`, - :class:`~meshmode.mesh.BTAG_PARTITION`, - or *None* to indicate that the geometry is not yet known. - - :arg discretization_tag: - *None* or :class:`DISCR_TAG_BASE` to indicate the use of the basic - discretization grid, :class:`DISCR_TAG_MODAL` to indicate a - modal discretization, or :class:`DISCR_TAG_QUAD` to indicate - the use of a quadrature grid. - """ + domain_tag: DomainTag + discretization_tag: DiscretizationTag - if domain_tag is None: - pass - elif domain_tag in [DTAG_SCALAR, "scalar"]: - domain_tag = DTAG_SCALAR - elif domain_tag in [DTAG_VOLUME_ALL, "vol"]: - domain_tag = DTAG_VOLUME_ALL - elif domain_tag in [FACE_RESTR_ALL, "all_faces"]: - domain_tag = FACE_RESTR_ALL - elif domain_tag in [FACE_RESTR_INTERIOR, "int_faces"]: - domain_tag = FACE_RESTR_INTERIOR - elif isinstance(domain_tag, BTAG_PARTITION): - domain_tag = DTAG_BOUNDARY(domain_tag) - elif domain_tag in [BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE]: - domain_tag = DTAG_BOUNDARY(domain_tag) - elif isinstance(domain_tag, DTAG_BOUNDARY): - pass - else: - raise ValueError("domain tag not understood: %s" % domain_tag) + def __init__(self, domain_tag: Any, + discretization_tag: Optional[type[DiscretizationTag]] = None): - if (quadrature_tag is not None and discretization_tag is not None): - raise ValueError( - "Both `quadrature_tag` and `discretization_tag` are specified. " - "Use `discretization_tag` instead." - ) - - # FIXME: `quadrature_tag` is deprecated - if (quadrature_tag is not None and discretization_tag is None): - warn("`quadrature_tag` is a deprecated kwarg and will be dropped " - "in version 2022.x. Use `discretization_tag` instead.", - DeprecationWarning, stacklevel=2) - discretization_tag = quadrature_tag + if ( + not (isinstance(domain_tag, + (ScalarDomainTag, BoundaryDomainTag, VolumeDomainTag))) + or discretization_tag is None + or ( + not isinstance(discretization_tag, type) + or not issubclass(discretization_tag, _DiscretizationTag))): + warn("Sloppy construction of DOFDesc is deprecated. " + "This will stop working in 2023. " + "Call as_dofdesc instead, with the same arguments. ", + DeprecationWarning, stacklevel=2) - if domain_tag is DTAG_SCALAR and discretization_tag is not None: - raise ValueError("cannot have nontrivial discretization tag on scalar") + domain_tag, discretization_tag = _normalize_domain_and_discr_tag( + domain_tag, discretization_tag) - if discretization_tag is None: - discretization_tag = DISCR_TAG_BASE - - # FIXME: String tags are deprecated - if isinstance(discretization_tag, str): - warn("Support for string values of `discretization_tag` will " - "be dropped in version 2022.x. Use one of the `DISCR_TAG_` " - "tags instead.", - DeprecationWarning, stacklevel=2) - - self.domain_tag = domain_tag - self.discretization_tag = discretization_tag - - @property - def quadrature_tag(self): - warn("`DOFDesc.quadrature_tag` is deprecated and will be dropped " - "in version 2022.x. Use `DOFDesc.discretization_tag` instead.", - DeprecationWarning, stacklevel=2) - return self.discretization_tag + object.__setattr__(self, "domain_tag", domain_tag) + object.__setattr__(self, "discretization_tag", discretization_tag) - def is_scalar(self): - return self.domain_tag is DTAG_SCALAR + def is_scalar(self) -> bool: + return isinstance(self.domain_tag, ScalarDomainTag) - def is_discretized(self): + def is_discretized(self) -> bool: return not self.is_scalar() - def is_volume(self): - return self.domain_tag is DTAG_VOLUME_ALL + def is_volume(self) -> bool: + return isinstance(self.domain_tag, VolumeDomainTag) - def is_boundary_or_partition_interface(self): - return isinstance(self.domain_tag, DTAG_BOUNDARY) - - def is_trace(self): - return (self.is_boundary_or_partition_interface() - or self.domain_tag in [ + def is_boundary_or_partition_interface(self) -> bool: + return (isinstance(self.domain_tag, BoundaryDomainTag) + and self.domain_tag.tag not in [ FACE_RESTR_ALL, FACE_RESTR_INTERIOR]) - def uses_quadrature(self): + def is_trace(self) -> bool: + return isinstance(self.domain_tag, BoundaryDomainTag) + + def uses_quadrature(self) -> bool: # FIXME: String tags are deprecated - # Check for string first, otherwise - # `issubclass` will raise an exception whenever - # its first argument is not a class. - # This can go away once support for strings is dropped - # completely. if isinstance(self.discretization_tag, str): # All strings are interpreted as quadrature-related tags return True - elif issubclass(self.discretization_tag, DISCR_TAG_QUAD): - return True - elif issubclass(self.discretization_tag, - (DISCR_TAG_BASE, DISCR_TAG_MODAL)): - return False - else: - raise ValueError( - f"Unsure how to interpret tag: {self.discretization_tag}" - ) - - def with_qtag(self, discr_tag): - warn("`DOFDesc.with_qtag` is deprecated and will be dropped " - "in version 2022.x. Use `DOFDesc.with_discr_tag` instead.", - DeprecationWarning, stacklevel=2) - return self.with_discr_tag(discr_tag) - - def with_discr_tag(self, discr_tag): - return type(self)(domain_tag=self.domain_tag, - discretization_tag=discr_tag) - - def with_dtag(self, dtag): - return type(self)(domain_tag=dtag, - discretization_tag=self.discretization_tag) - - def __eq__(self, other): - return (type(self) == type(other) - and self.domain_tag == other.domain_tag - and self.discretization_tag == other.discretization_tag) - - def __ne__(self, other): - return not self.__eq__(other) - - def __hash__(self): - return hash((type(self), self.domain_tag, self.discretization_tag)) - - def __repr__(self): - def fmt(s): - if isinstance(s, type): - return s.__name__ - else: - return repr(s) + elif isinstance(self.discretization_tag, type): + if issubclass(self.discretization_tag, DISCR_TAG_QUAD): + return True + elif issubclass(self.discretization_tag, + (DISCR_TAG_BASE, DISCR_TAG_MODAL)): + return False + + raise ValueError( + f"Invalid discretization tag: {self.discretization_tag}") + + def with_dtag(self, dtag) -> "DOFDesc": + from warnings import warn + warn("'with_dtag' is deprecated. Use 'with_domain_tag' instead. " + "This will stop working in 2023", + DeprecationWarning, stacklevel=2) + return replace(self, domain_tag=dtag) + + def with_domain_tag(self, dtag) -> "DOFDesc": + return replace(self, domain_tag=dtag) + + def trace(self, btag: BoundaryTag) -> "DOFDesc": + """Return a :class:`DOFDesc` for the restriction of the volume + descriptor *self* to the boundary named by *btag*. + + An error is raised if this method is called on a non-volume instance of + :class:`DOFDesc`. + """ + if not isinstance(self.domain_tag, VolumeDomainTag): + raise ValueError(f"must originate on volume, got '{self.domain_tag}'") + return replace(self, + domain_tag=BoundaryDomainTag(btag, volume_tag=self.domain_tag.tag)) - return "DOFDesc({}, {})".format( - fmt(self.domain_tag), - fmt(self.discretization_tag)) + def with_discr_tag(self, discr_tag) -> "DOFDesc": + return replace(self, discretization_tag=discr_tag) def as_identifier(self) -> str: """Returns a descriptive string for this :class:`DOFDesc` that is usable @@ -341,7 +330,16 @@ def as_identifier(self) -> str: dom_id = "f_all" elif self.domain_tag is FACE_RESTR_INTERIOR: dom_id = "f_int" - elif isinstance(self.domain_tag, DTAG_BOUNDARY): + elif isinstance(self.domain_tag, VolumeDomainTag): + vtag = self.domain_tag.tag + if isinstance(vtag, type): + vtag = vtag.__name__.replace("VTAG_", "").lower() + elif isinstance(vtag, str): + vtag = _to_identifier(vtag) + else: + vtag = _to_identifier(str(vtag)) + dom_id = f"v_{vtag}" + elif isinstance(self.domain_tag, BoundaryDomainTag): btag = self.domain_tag.tag if isinstance(btag, type): btag = btag.__name__.replace("BTAG_", "").lower() @@ -369,31 +367,101 @@ def as_identifier(self) -> str: return f"{dom_id}{discr_id}" -DD_SCALAR = DOFDesc(DTAG_SCALAR, None) +DD_SCALAR = DOFDesc(DTAG_SCALAR, DISCR_TAG_BASE) +DD_VOLUME_ALL = DOFDesc(DTAG_VOLUME_ALL, DISCR_TAG_BASE) +DD_VOLUME_ALL_MODAL = DOFDesc(DTAG_VOLUME_ALL, DISCR_TAG_MODAL) + + +def _normalize_domain_and_discr_tag( + domain: Any, + discretization_tag: Optional[DiscretizationTag] = None, + *, _contextual_volume_tag: Optional[VolumeTag] = None + ) -> Tuple[DomainTag, DiscretizationTag]: + + if _contextual_volume_tag is None: + _contextual_volume_tag = VTAG_ALL + + if domain == "scalar": + domain = DTAG_SCALAR + elif isinstance(domain, (ScalarDomainTag, BoundaryDomainTag, VolumeDomainTag)): + pass + elif domain in [VTAG_ALL, "vol"]: + domain = DTAG_VOLUME_ALL + elif domain in [FACE_RESTR_ALL, "all_faces"]: + domain = BoundaryDomainTag(FACE_RESTR_ALL, _contextual_volume_tag) + elif domain in [FACE_RESTR_INTERIOR, "int_faces"]: + domain = BoundaryDomainTag(FACE_RESTR_INTERIOR, _contextual_volume_tag) + elif isinstance(domain, BTAG_PARTITION): + domain = BoundaryDomainTag(domain, _contextual_volume_tag) + elif domain in [BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE]: + domain = BoundaryDomainTag(domain, _contextual_volume_tag) + else: + raise ValueError("domain tag not understood: %s" % domain) + + if domain is DTAG_SCALAR and discretization_tag is not None: + raise ValueError("cannot have nontrivial discretization tag on scalar") + + if discretization_tag is None: + discretization_tag = DISCR_TAG_BASE -DD_VOLUME = DOFDesc(DTAG_VOLUME_ALL, None) + return domain, discretization_tag + + +ConvertibleToDOFDesc = Any + + +def as_dofdesc( + domain: "ConvertibleToDOFDesc", + discretization_tag: Optional[DiscretizationTag] = None, + *, _contextual_volume_tag: Optional[VolumeTag] = None) -> DOFDesc: + """ + :arg domain_tag: One of the following: + :class:`DTAG_SCALAR` (or the string ``"scalar"``), + :class:`DTAG_VOLUME_ALL` (or the string ``"vol"``) + for the default volume discretization, + :data:`~meshmode.discretization.connection.FACE_RESTR_ALL` + (or the string ``"all_faces"``), or + :data:`~meshmode.discretization.connection.FACE_RESTR_INTERIOR` + (or the string ``"int_faces"``), or one of + :class:`~meshmode.mesh.BTAG_ALL`, + :class:`~meshmode.mesh.BTAG_NONE`, + :class:`~meshmode.mesh.BTAG_REALLY_ALL`, + :class:`~meshmode.mesh.BTAG_PARTITION`, + or *None* to indicate that the geometry is not yet known. + + :arg discretization_tag: + *None* or :class:`DISCR_TAG_BASE` to indicate the use of the basic + discretization grid, :class:`DISCR_TAG_MODAL` to indicate a + modal discretization, or :class:`DISCR_TAG_QUAD` to indicate + the use of a quadrature grid. + """ -DD_VOLUME_MODAL = DOFDesc(DTAG_VOLUME_ALL, DISCR_TAG_MODAL) + if isinstance(domain, DOFDesc): + return domain + domain, discretization_tag = _normalize_domain_and_discr_tag( + domain, discretization_tag, + _contextual_volume_tag=_contextual_volume_tag) -def as_dofdesc(dd): - if isinstance(dd, DOFDesc): - return dd - return DOFDesc(dd, discretization_tag=None) + return DOFDesc(domain, discretization_tag) # }}} -# {{{ Deprecated tags +# {{{ deprecations -_deprecated_name_to_new_name = {"QTAG_NONE": "DISCR_TAG_BASE", - "QTAG_MODAL": "DISCR_TAG_MODAL"} +_deprecated_name_to_new_name = { + "DTAG_VOLUME": "VolumeDomainTag", + "DTAG_BOUNDARY": "BoundaryDomainTag", + "DD_VOLUME": "DD_VOLUME_ALL", + "DD_VOLUME_MODAL": "DD_VOLUME_ALL_MODAL" + } def __getattr__(name): if name in _deprecated_name_to_new_name: warn(f"'{name}' is deprecated and will be dropped " - f"in version 2022.x. Use '{_deprecated_name_to_new_name[name]}' " + f"in version 2023.x. Use '{_deprecated_name_to_new_name[name]}' " "instead.", DeprecationWarning, stacklevel=2) return globals()[_deprecated_name_to_new_name[name]] diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index b949a2399..64f22cd34 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -43,6 +43,7 @@ """ +from typing import Optional, Sequence import numpy as np from arraycontext import ArrayContext, Scalar, tag_axes @@ -52,7 +53,8 @@ DiscretizationFaceAxisTag, DiscretizationElementAxisTag) -from grudge.dof_desc import DD_VOLUME, DOFDesc, as_dofdesc +from grudge.dof_desc import ( + DD_VOLUME_ALL, DOFDesc, as_dofdesc, BoundaryDomainTag, FACE_RESTR_ALL) from grudge.discretization import DiscretizationCollection import grudge.op as op @@ -63,7 +65,8 @@ def characteristic_lengthscales( - actx: ArrayContext, dcoll: DiscretizationCollection) -> DOFArray: + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None) -> DOFArray: r"""Computes the characteristic length scale :math:`h_{\text{loc}}` at each node. The characteristic length scale is mainly useful for estimating the stable time step size. E.g. for a hyperbolic system, an estimate of the @@ -91,7 +94,7 @@ def characteristic_lengthscales( methods has been used as a guide. Any concrete time integrator will likely require scaling of the values returned by this routine. """ - @memoize_in(dcoll, (characteristic_lengthscales, + @memoize_in(dcoll, (characteristic_lengthscales, dd, "compute_characteristic_lengthscales")) def _compute_characteristic_lengthscales(): return actx.freeze( @@ -103,15 +106,16 @@ def _compute_characteristic_lengthscales(): # corresponding group non-geometric factor cng * geo_facts for cng, geo_facts in zip( - dt_non_geometric_factors(dcoll), - actx.thaw(dt_geometric_factors(dcoll))))))) + dt_non_geometric_factors(dcoll, dd), + actx.thaw(dt_geometric_factors(dcoll, dd))))))) return actx.thaw(_compute_characteristic_lengthscales()) @memoize_on_first_arg def dt_non_geometric_factors( - dcoll: DiscretizationCollection, dd=None) -> list: + dcoll: DiscretizationCollection, dd: Optional[DOFDesc] = None + ) -> Sequence[float]: r"""Computes the non-geometric scale factors following [Hesthaven_2008]_, section 6.4, for each element group in the *dd* discretization: @@ -128,7 +132,7 @@ def dt_non_geometric_factors( node distance on the reference element for each group. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL discr = dcoll.discr_from_dd(dd) min_delta_rs = [] @@ -160,7 +164,8 @@ def dt_non_geometric_factors( @memoize_on_first_arg def h_max_from_volume( - dcoll: DiscretizationCollection, dim=None, dd=None) -> Scalar: + dcoll: DiscretizationCollection, dim=None, + dd: Optional[DOFDesc] = None) -> Scalar: """Returns a (maximum) characteristic length based on the volume of the elements. This length may not be representative if the elements have very high aspect ratios. @@ -175,7 +180,7 @@ def h_max_from_volume( from grudge.reductions import nodal_max, elementwise_sum if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dd = as_dofdesc(dd) if dim is None: @@ -191,7 +196,8 @@ def h_max_from_volume( @memoize_on_first_arg def h_min_from_volume( - dcoll: DiscretizationCollection, dim=None, dd=None) -> Scalar: + dcoll: DiscretizationCollection, dim=None, + dd: Optional[DOFDesc] = None) -> Scalar: """Returns a (minimum) characteristic length based on the volume of the elements. This length may not be representative if the elements have very high aspect ratios. @@ -206,7 +212,7 @@ def h_min_from_volume( from grudge.reductions import nodal_min, elementwise_sum if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dd = as_dofdesc(dd) if dim is None: @@ -221,7 +227,7 @@ def h_min_from_volume( def dt_geometric_factors( - dcoll: DiscretizationCollection, dd=None) -> DOFArray: + dcoll: DiscretizationCollection, dd: Optional[DOFDesc] = None) -> DOFArray: r"""Computes a geometric scaling factor for each cell following [Hesthaven_2008]_, section 6.4, defined as the inradius (radius of an inscribed circle/sphere). @@ -245,7 +251,7 @@ def dt_geometric_factors( from meshmode.discretization.poly_element import SimplexElementGroupBase if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL actx = dcoll._setup_actx volm_discr = dcoll.discr_from_dd(dd) @@ -272,7 +278,8 @@ def dt_geometric_factors( # Inscribed "circle" radius is half the cell size return actx.freeze(cell_vols/2) - dd_face = DOFDesc("all_faces", dd.discretization_tag) + dd_face = dd.with_domain_tag( + BoundaryDomainTag(FACE_RESTR_ALL, dd.domain_tag.tag)) face_discr = dcoll.discr_from_dd(dd_face) # Compute areas of each face diff --git a/grudge/eager.py b/grudge/eager.py index 1886cfd04..626e15592 100644 --- a/grudge/eager.py +++ b/grudge/eager.py @@ -47,14 +47,14 @@ def __init__(self, *args, **kwargs): def project(self, src, tgt, vec): return op.project(self, src, tgt, vec) - def grad(self, vec): - return op.local_grad(self, vec) + def grad(self, *args): + return op.local_grad(self, *args) - def d_dx(self, xyz_axis, vec): - return op.local_d_dx(self, xyz_axis, vec) + def d_dx(self, xyz_axis, *args): + return op.local_d_dx(self, xyz_axis, *args) - def div(self, vecs): - return op.local_div(self, vecs) + def div(self, *args): + return op.local_div(self, *args) def weak_grad(self, *args): return op.weak_local_grad(self, *args) @@ -68,8 +68,8 @@ def weak_div(self, *args): def mass(self, *args): return op.mass(self, *args) - def inverse_mass(self, vec): - return op.inverse_mass(self, vec) + def inverse_mass(self, *args): + return op.inverse_mass(self, *args) def face_mass(self, *args): return op.face_mass(self, *args) diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py index 89e1f1f2c..f5c622cd5 100644 --- a/grudge/geometry/metrics.py +++ b/grudge/geometry/metrics.py @@ -58,6 +58,7 @@ """ +from typing import Optional, Tuple, Union import numpy as np from arraycontext import ArrayContext, tag_axes @@ -68,7 +69,7 @@ import grudge.dof_desc as dof_desc from grudge.dof_desc import ( - DD_VOLUME, DOFDesc, DISCR_TAG_BASE + DD_VOLUME_ALL, DOFDesc, DISCR_TAG_BASE ) from meshmode.transform_metadata import (DiscretizationAmbientDimAxisTag, @@ -115,7 +116,8 @@ def to_quad(vec): def forward_metric_nth_derivative( actx: ArrayContext, dcoll: DiscretizationCollection, - xyz_axis, ref_axes, dd=None, + xyz_axis: int, ref_axes: Union[int, Tuple[Tuple[int, int], ...]], + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> DOFArray: r"""Pointwise metric derivatives representing repeated derivatives of the physical coordinate enumerated by *xyz_axis*: :math:`x_{\mathrm{xyz\_axis}}` @@ -150,7 +152,7 @@ def forward_metric_nth_derivative( metric derivative at each nodal coordinate. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL inner_dd = dd.with_discr_tag(DISCR_TAG_BASE) @@ -182,8 +184,10 @@ def forward_metric_nth_derivative( def forward_metric_derivative_vector( - actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, dd=None, - *, _use_geoderiv_connection=False) -> np.ndarray: + actx: ArrayContext, dcoll: DiscretizationCollection, + rst_axis: Union[int, Tuple[Tuple[int, int], ...]], + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False + ) -> np.ndarray: r"""Computes an object array containing the forward metric derivatives of each physical coordinate. @@ -207,7 +211,9 @@ def forward_metric_derivative_vector( def forward_metric_derivative_mv( - actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + rst_axis: Union[int, Tuple[Tuple[int, int], ...]], + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> MultiVector: r"""Computes a :class:`pymbolic.geometric_algebra.MultiVector` containing the forward metric derivatives of each physical coordinate. @@ -230,7 +236,8 @@ def forward_metric_derivative_mv( def forward_metric_derivative_mat( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> np.ndarray: r"""Computes the forward metric derivative matrix, also commonly called the Jacobian matrix, with entries defined as the @@ -257,7 +264,7 @@ def forward_metric_derivative_mat( ambient_dim = dcoll.ambient_dim if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim @@ -271,7 +278,8 @@ def forward_metric_derivative_mat( def first_fundamental_form(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None, *, _use_geoderiv_connection=False) -> np.ndarray: + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False + ) -> np.ndarray: r"""Computes the first fundamental form using the Jacobian matrix: .. math:: @@ -297,7 +305,7 @@ def first_fundamental_form(actx: ArrayContext, dcoll: DiscretizationCollection, form. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL mder = forward_metric_derivative_mat( actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection) @@ -306,7 +314,8 @@ def first_fundamental_form(actx: ArrayContext, dcoll: DiscretizationCollection, def inverse_metric_derivative_mat( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> np.ndarray: r"""Computes the inverse metric derivative matrix, which is the inverse of the Jacobian (forward metric derivative) matrix. @@ -320,7 +329,7 @@ def inverse_metric_derivative_mat( ambient_dim = dcoll.ambient_dim if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim @@ -336,7 +345,8 @@ def inverse_metric_derivative_mat( def inverse_first_fundamental_form( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> np.ndarray: r"""Computes the inverse of the first fundamental form: @@ -360,7 +370,7 @@ def inverse_first_fundamental_form( first fundamental form. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim @@ -387,7 +397,8 @@ def inverse_first_fundamental_form( def inverse_metric_derivative( - actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, xyz_axis, dd, + actx: ArrayContext, dcoll: DiscretizationCollection, + rst_axis: int, xyz_axis: int, dd: DOFDesc, *, _use_geoderiv_connection=False ) -> DOFArray: r"""Computes the inverse metric derivative of the physical @@ -446,7 +457,7 @@ def outprod_with_unit(i, at): def inverse_surface_metric_derivative( actx: ArrayContext, dcoll: DiscretizationCollection, - rst_axis, xyz_axis, dd=None, + rst_axis, xyz_axis, dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False): r"""Computes the inverse surface metric derivative of the physical coordinate enumerated by *xyz_axis* with respect to the @@ -468,7 +479,7 @@ def inverse_surface_metric_derivative( ambient_dim = dcoll.ambient_dim if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dd = dof_desc.as_dofdesc(dd) if ambient_dim == dim: @@ -488,7 +499,8 @@ def inverse_surface_metric_derivative( def inverse_surface_metric_derivative_mat( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, times_area_element=False, _use_geoderiv_connection=False): r"""Computes the matrix of inverse surface metric derivatives, indexed by ``(xyz_axis, rst_axis)``. It returns all values of @@ -509,7 +521,7 @@ def inverse_surface_metric_derivative_mat( """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dd = dof_desc.as_dofdesc(dd) @memoize_in(dcoll, (inverse_surface_metric_derivative_mat, dd, @@ -542,7 +554,7 @@ def _inv_surf_metric_deriv(): def _signed_face_ones( - actx: ArrayContext, dcoll: DiscretizationCollection, dd + actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc ) -> DOFArray: assert dd.is_trace() @@ -550,7 +562,7 @@ def _signed_face_ones( # NOTE: ignore quadrature_tags on dd, since we only care about # the face_id here all_faces_conn = dcoll.connection_from_dds( - DD_VOLUME, DOFDesc(dd.domain_tag) + DD_VOLUME_ALL, DOFDesc(dd.domain_tag, DISCR_TAG_BASE) ) signed_ones = dcoll.discr_from_dd(dd.with_discr_tag(DISCR_TAG_BASE)).zeros( actx, dtype=dcoll.real_dtype @@ -571,7 +583,7 @@ def _signed_face_ones( def parametrization_derivative( - actx: ArrayContext, dcoll: DiscretizationCollection, dd, + actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc, *, _use_geoderiv_connection=False) -> MultiVector: r"""Computes the product of forward metric derivatives spanning the tangent space with topological dimension *dim*. @@ -584,7 +596,7 @@ def parametrization_derivative( the product of metric derivatives. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim if dim == 0: @@ -605,8 +617,10 @@ def parametrization_derivative( ) -def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None, *, _use_geoderiv_connection=False) -> MultiVector: +def pseudoscalar( + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False + ) -> MultiVector: r"""Computes the field of pseudoscalars for the domain/discretization identified by *dd*. @@ -618,7 +632,7 @@ def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection, :class:`~meshmode.dof_array.DOFArray`\ s. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL return parametrization_derivative( actx, dcoll, dd, @@ -626,7 +640,8 @@ def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection, def area_element( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False ) -> DOFArray: r"""Computes the scale factor used to transform integrals from reference @@ -642,7 +657,7 @@ def area_element( volumes for each element. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL @memoize_in(dcoll, (area_element, dd, _use_geoderiv_connection)) def _area_elements(): @@ -662,7 +677,8 @@ def _area_elements(): # {{{ surface normal vectors def rel_mv_normal( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> MultiVector: r"""Computes surface normals at each nodal location as a :class:`~pymbolic.geometric_algebra.MultiVector` relative to the @@ -688,7 +704,7 @@ def rel_mv_normal( def mv_normal( - actx: ArrayContext, dcoll: DiscretizationCollection, dd, + actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc, *, _use_geoderiv_connection=False ) -> MultiVector: r"""Exterior unit normal as a :class:`~pymbolic.geometric_algebra.MultiVector`. @@ -744,10 +760,10 @@ def _normal(): from grudge.op import project volm_normal = MultiVector( - project(dcoll, dof_desc.DD_VOLUME, dd, + project(dcoll, DD_VOLUME_ALL, dd, rel_mv_normal( actx, dcoll, - dd=dof_desc.DD_VOLUME, + dd=DD_VOLUME_ALL, _use_geoderiv_connection=_use_geoderiv_connection ).as_vector(dtype=object)) ) @@ -768,7 +784,7 @@ def _normal(): return actx.thaw(_normal()) -def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd, +def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc, *, _use_geoderiv_connection=None): """Get the unit normal to the specified surface discretization, *dd*. This supports both volume discretizations @@ -798,8 +814,8 @@ def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd, # {{{ Curvature computations def second_fundamental_form( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None - ) -> np.ndarray: + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None) -> np.ndarray: r"""Computes the second fundamental form: .. math:: @@ -817,7 +833,7 @@ def second_fundamental_form( """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim normal = rel_mv_normal(actx, dcoll, dd=dd).as_vector(dtype=object) @@ -846,7 +862,7 @@ def second_fundamental_form( def shape_operator(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None) -> np.ndarray: + dd: Optional[DOFDesc] = None) -> np.ndarray: r"""Computes the shape operator (also called the curvature tensor) containing second order derivatives: @@ -871,7 +887,7 @@ def shape_operator(actx: ArrayContext, dcoll: DiscretizationCollection, def summed_curvature(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None) -> DOFArray: + dd: Optional[DOFDesc] = None) -> DOFArray: r"""Computes the sum of the principal curvatures: .. math:: @@ -888,7 +904,7 @@ def summed_curvature(actx: ArrayContext, dcoll: DiscretizationCollection, """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.ambient_dim - 1 diff --git a/grudge/models/advection.py b/grudge/models/advection.py index cfe1a4920..c5d35e2d1 100644 --- a/grudge/models/advection.py +++ b/grudge/models/advection.py @@ -214,13 +214,13 @@ def __init__(self, dcoll, v, inflow_u, flux_type="central", quad_tag=None): self.quad_tag = quad_tag def flux(self, u_tpair): - from grudge.dof_desc import DD_VOLUME + from grudge.dof_desc import DD_VOLUME_ALL - surf_v = op.project(self.dcoll, DD_VOLUME, u_tpair.dd, self.v) + surf_v = op.project(self.dcoll, DD_VOLUME_ALL, u_tpair.dd, self.v) return advection_weak_flux(self.dcoll, self.flux_type, u_tpair, surf_v) def operator(self, t, u): - from grudge.dof_desc import DOFDesc, DD_VOLUME, DTAG_VOLUME_ALL + from grudge.dof_desc import DOFDesc, DD_VOLUME_ALL, DTAG_VOLUME_ALL from meshmode.mesh import BTAG_ALL from meshmode.discretization.connection import FACE_RESTR_ALL @@ -234,7 +234,7 @@ def flux(tpair): return op.project(dcoll, tpair.dd, face_dd, self.flux(tpair)) def to_quad(arg): - return op.project(dcoll, DD_VOLUME, quad_dd, arg) + return op.project(dcoll, DD_VOLUME_ALL, quad_dd, arg) if self.inflow_u is not None: inflow_flux = flux(op.bv_trace_pair(dcoll, @@ -279,7 +279,7 @@ def to_quad(arg): # {{{ closed surface advection def v_dot_n_tpair(actx, dcoll, velocity, trace_dd): - from grudge.dof_desc import DTAG_BOUNDARY + from grudge.dof_desc import BoundaryDomainTag from grudge.trace_pair import TracePair from meshmode.discretization.connection import FACE_RESTR_INTERIOR @@ -287,10 +287,9 @@ def v_dot_n_tpair(actx, dcoll, velocity, trace_dd): v_dot_n = velocity.dot(normal) i = op.project(dcoll, trace_dd.with_discr_tag(None), trace_dd, v_dot_n) - if trace_dd.domain_tag is FACE_RESTR_INTERIOR: - e = dcoll.opposite_face_connection()(i) - elif isinstance(trace_dd.domain_tag, DTAG_BOUNDARY): - e = dcoll.distributed_boundary_swap_connection(trace_dd)(i) + assert isinstance(trace_dd.domain_tag, BoundaryDomainTag) + if trace_dd.domain_tag.tag is FACE_RESTR_INTERIOR: + e = dcoll.opposite_face_connection(trace_dd.domain_tag)(i) else: raise ValueError("Unrecognized domain tag: %s" % trace_dd.domain_tag) @@ -325,9 +324,9 @@ def __init__(self, dcoll, v, flux_type="central", quad_tag=None): self.quad_tag = quad_tag def flux(self, u_tpair): - from grudge.dof_desc import DD_VOLUME + from grudge.dof_desc import DD_VOLUME_ALL - surf_v = op.project(self.dcoll, DD_VOLUME, + surf_v = op.project(self.dcoll, DD_VOLUME_ALL, u_tpair.dd.with_discr_tag(None), self.v) return surface_advection_weak_flux(self.dcoll, self.flux_type, @@ -335,7 +334,7 @@ def flux(self, u_tpair): surf_v) def operator(self, t, u): - from grudge.dof_desc import DOFDesc, DD_VOLUME, DTAG_VOLUME_ALL + from grudge.dof_desc import DOFDesc, DD_VOLUME_ALL, DTAG_VOLUME_ALL from meshmode.discretization.connection import FACE_RESTR_ALL face_dd = DOFDesc(FACE_RESTR_ALL, self.quad_tag) @@ -347,7 +346,7 @@ def flux(tpair): return op.project(dcoll, tpair.dd, face_dd, self.flux(tpair)) def to_quad(arg): - return op.project(dcoll, DD_VOLUME, quad_dd, arg) + return op.project(dcoll, DD_VOLUME_ALL, quad_dd, arg) quad_v = to_quad(self.v) quad_u = to_quad(u) diff --git a/grudge/op.py b/grudge/op.py index 015e0718b..ff40ae3fd 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -81,6 +81,7 @@ DiscretizationFaceAxisTag) from grudge.discretization import DiscretizationCollection +from grudge.dof_desc import as_dofdesc from pytools import keyed_memoize_in from pytools.obj_array import make_obj_array @@ -88,6 +89,10 @@ import numpy as np import grudge.dof_desc as dof_desc +from grudge.dof_desc import ( + DD_VOLUME_ALL, FACE_RESTR_ALL, DISCR_TAG_BASE, + DOFDesc, VolumeDomainTag +) from grudge.interpolation import interp from grudge.projection import project @@ -244,14 +249,14 @@ def get_ref_derivative_mats(grp): def _strong_scalar_grad(dcoll, dd_in, vec): - assert dd_in == dof_desc.as_dofdesc(dof_desc.DD_VOLUME) + assert isinstance(dd_in.domain_tag, VolumeDomainTag) from grudge.geometry import inverse_surface_metric_derivative_mat - discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + discr = dcoll.discr_from_dd(dd_in) actx = vec.array_context - inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, + inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, dd=dd_in, _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) return _gradient_kernel(actx, discr, discr, _reference_derivative_matrices, inverse_jac_mat, vec, @@ -259,7 +264,7 @@ def _strong_scalar_grad(dcoll, dd_in, vec): def local_grad( - dcoll: DiscretizationCollection, vec, *, nested=False) -> ArrayOrContainer: + dcoll: DiscretizationCollection, *args, nested=False) -> ArrayOrContainer: r"""Return the element-local gradient of a function :math:`f` represented by *vec*: @@ -268,15 +273,26 @@ def local_grad( \nabla|_E f = \left( \partial_x|_E f, \partial_y|_E f, \partial_z|_E f \right) + May be called with ``(vec)`` or ``(dd_in, vec)``. + :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` of them. + :arg dd_in: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. + Defaults to the base volume discretization if not provided. :arg nested: return nested object arrays instead of a single multidimensional array if *vec* is non-scalar. :returns: an object array (possibly nested) of :class:`~meshmode.dof_array.DOFArray`\ s or :class:`~arraycontext.ArrayContainer` of object arrays. """ - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + if len(args) == 1: + vec, = args + dd_in = DD_VOLUME_ALL + elif len(args) == 2: + dd_in, vec = args + else: + raise TypeError("invalid number of arguments") + from grudge.tools import rec_map_subarrays return rec_map_subarrays( partial(_strong_scalar_grad, dcoll, dd_in), @@ -285,7 +301,7 @@ def local_grad( def local_d_dx( - dcoll: DiscretizationCollection, xyz_axis, vec) -> ArrayOrContainer: + dcoll: DiscretizationCollection, xyz_axis, *args) -> ArrayOrContainer: r"""Return the element-local derivative along axis *xyz_axis* of a function :math:`f` represented by *vec*: @@ -293,22 +309,34 @@ def local_d_dx( \frac{\partial f}{\partial \lbrace x,y,z\rbrace}\Big|_E + May be called with ``(vec)`` or ``(dd, vec)``. + :arg xyz_axis: an integer indicating the axis along which the derivative is taken. + :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. + Defaults to the base volume discretization if not provided. :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` of them. :returns: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` of them. """ + if len(args) == 1: + vec, = args + dd = DD_VOLUME_ALL + elif len(args) == 2: + dd, vec = args + else: + raise TypeError("invalid number of arguments") + if not isinstance(vec, DOFArray): - return map_array_container(partial(local_d_dx, dcoll, xyz_axis), vec) + return map_array_container(partial(local_d_dx, dcoll, xyz_axis, dd), vec) - discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + discr = dcoll.discr_from_dd(dd) actx = vec.array_context from grudge.geometry import inverse_surface_metric_derivative_mat - inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, - _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, dd=dd, + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) return _single_axis_derivative_kernel( actx, discr, discr, @@ -316,7 +344,7 @@ def local_d_dx( metric_in_matvec=False) -def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainer: +def local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: r"""Return the element-local divergence of the vector function :math:`\mathbf{f}` represented by *vecs*: @@ -324,6 +352,10 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainer: \nabla|_E \cdot \mathbf{f} = \sum_{i=1}^d \partial_{x_i}|_E \mathbf{f}_i + May be called with ``(vec)`` or ``(dd, vec)``. + + :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. + Defaults to the base volume discretization if not provided. :arg vecs: an object array of :class:`~meshmode.dof_array.DOFArray`\s or an :class:`~arraycontext.ArrayContainer` object @@ -332,13 +364,21 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainer: :returns: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` of them. """ + if len(args) == 1: + vec, = args + dd = DD_VOLUME_ALL + elif len(args) == 2: + dd, vec = args + else: + raise TypeError("invalid number of arguments") + from grudge.tools import rec_map_subarrays return rec_map_subarrays( lambda vec: sum( - local_d_dx(dcoll, i, vec_i) + local_d_dx(dcoll, i, dd, vec_i) for i, vec_i in enumerate(vec)), (dcoll.ambient_dim,), (), - vecs, scalar_cls=DOFArray) + vec, scalar_cls=DOFArray) # }}} @@ -391,8 +431,9 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): def _weak_scalar_grad(dcoll, dd_in, vec): from grudge.geometry import inverse_surface_metric_derivative_mat + dd_in = as_dofdesc(dd_in) in_discr = dcoll.discr_from_dd(dd_in) - out_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + out_discr = dcoll.discr_from_dd(dd_in.with_discr_tag(DISCR_TAG_BASE)) actx = vec.array_context inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, dd=dd_in, @@ -429,7 +470,7 @@ def weak_local_grad( """ if len(args) == 1: vecs, = args - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = DD_VOLUME_ALL elif len(args) == 2: dd_in, vecs = args else: @@ -474,7 +515,7 @@ def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: """ if len(args) == 2: xyz_axis, vec = args - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = dof_desc.DD_VOLUME_ALL elif len(args) == 3: dd_in, xyz_axis, vec = args else: @@ -488,8 +529,9 @@ def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: from grudge.geometry import inverse_surface_metric_derivative_mat + dd_in = as_dofdesc(dd_in) in_discr = dcoll.discr_from_dd(dd_in) - out_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + out_discr = dcoll.discr_from_dd(dd_in.with_discr_tag(DISCR_TAG_BASE)) actx = vec.array_context inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, dd=dd_in, @@ -533,7 +575,7 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: """ if len(args) == 1: vecs, = args - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = DD_VOLUME_ALL elif len(args) == 2: dd_in, vecs = args else: @@ -628,7 +670,7 @@ def mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: *vec* being an :class:`~arraycontext.ArrayContainer`, the mass operator is applied component-wise. - May be called with ``(vec)`` or ``(dd, vec)``. + May be called with ``(vec)`` or ``(dd_in, vec)``. Specifically, this function applies the mass matrix elementwise on a vector of coefficients :math:`\mathbf{f}` via: @@ -640,7 +682,7 @@ def mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: where :math:`\phi_i` are local polynomial basis functions on :math:`E`. - :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. + :arg dd_in: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. Defaults to the base volume discretization if not provided. :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` of them. @@ -650,13 +692,15 @@ def mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: if len(args) == 1: vec, = args - dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = dof_desc.DD_VOLUME_ALL elif len(args) == 2: - dd, vec = args + dd_in, vec = args else: raise TypeError("invalid number of arguments") - return _apply_mass_operator(dcoll, dof_desc.DD_VOLUME, dd, vec) + dd_out = dd_in.with_discr_tag(DISCR_TAG_BASE) + + return _apply_mass_operator(dcoll, dd_out, dd_in, vec) # }}} @@ -714,7 +758,7 @@ def _apply_inverse_mass_operator( return DOFArray(actx, data=tuple(group_data)) -def inverse_mass(dcoll: DiscretizationCollection, vec) -> ArrayOrContainer: +def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: r"""Return the action of the DG mass matrix inverse on a vector (or vectors) of :class:`~meshmode.dof_array.DOFArray`\ s, *vec*. In the case of *vec* being an :class:`~arraycontext.ArrayContainer`, @@ -744,15 +788,24 @@ def inverse_mass(dcoll: DiscretizationCollection, vec) -> ArrayOrContainer: where :math:`\widehat{\mathbf{M}}` is the reference mass matrix on :math:`\widehat{E}`. + May be called with ``(vec)`` or ``(dd, vec)``. + :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` of them. + :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. + Defaults to the base volume discretization if not provided. :returns: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` like *vec*. """ + if len(args) == 1: + vec, = args + dd = DD_VOLUME_ALL + elif len(args) == 2: + dd, vec = args + else: + raise TypeError("invalid number of arguments") - return _apply_inverse_mass_operator( - dcoll, dof_desc.DD_VOLUME, dof_desc.DD_VOLUME, vec - ) + return _apply_inverse_mass_operator(dcoll, dd, dd, vec) # }}} @@ -850,21 +903,25 @@ def get_ref_face_mass_mat(face_grp, vol_grp): return get_ref_face_mass_mat(face_element_group, vol_element_group) -def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec): +def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd_in, vec): if not isinstance(vec, DOFArray): return map_array_container( - partial(_apply_face_mass_operator, dcoll, dd), vec + partial(_apply_face_mass_operator, dcoll, dd_in), vec ) from grudge.geometry import area_element - volm_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) - face_discr = dcoll.discr_from_dd(dd) + dd_out = DOFDesc( + VolumeDomainTag(dd_in.domain_tag.volume_tag), + DISCR_TAG_BASE) + + volm_discr = dcoll.discr_from_dd(dd_out) + face_discr = dcoll.discr_from_dd(dd_in) dtype = vec.entry_dtype actx = vec.array_context assert len(face_discr.groups) == len(volm_discr.groups) - surf_area_elements = area_element(actx, dcoll, dd=dd, + surf_area_elements = area_element(actx, dcoll, dd=dd_in, _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) return DOFArray( @@ -901,7 +958,7 @@ def face_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: *vec* being an arbitrary :class:`~arraycontext.ArrayContainer`, the face mass operator is applied component-wise. - May be called with ``(vec)`` or ``(dd, vec)``. + May be called with ``(vec)`` or ``(dd_in, vec)``. Specifically, this function applies the face mass matrix elementwise on a vector of coefficients :math:`\mathbf{f}` as the sum of contributions for @@ -932,13 +989,13 @@ def face_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: if len(args) == 1: vec, = args - dd = dof_desc.DOFDesc("all_faces", dof_desc.DISCR_TAG_BASE) + dd_in = DD_VOLUME_ALL.trace(FACE_RESTR_ALL) elif len(args) == 2: - dd, vec = args + dd_in, vec = args else: raise TypeError("invalid number of arguments") - return _apply_face_mass_operator(dcoll, dd, vec) + return _apply_face_mass_operator(dcoll, dd_in, vec) # }}} diff --git a/grudge/projection.py b/grudge/projection.py index 425239591..e21e02295 100644 --- a/grudge/projection.py +++ b/grudge/projection.py @@ -37,13 +37,19 @@ from arraycontext import ArrayOrContainer from grudge.discretization import DiscretizationCollection -from grudge.dof_desc import as_dofdesc +from grudge.dof_desc import ( + as_dofdesc, + VolumeDomainTag, + BoundaryDomainTag, + ConvertibleToDOFDesc) from numbers import Number def project( - dcoll: DiscretizationCollection, src, tgt, vec) -> ArrayOrContainer: + dcoll: DiscretizationCollection, + src: "ConvertibleToDOFDesc", + tgt: "ConvertibleToDOFDesc", vec) -> ArrayOrContainer: """Project from one discretization to another, e.g. from the volume to the boundary, or from the base to the an overintegrated quadrature discretization. @@ -55,10 +61,24 @@ def project( :returns: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` like *vec*. """ - src = as_dofdesc(src) - tgt = as_dofdesc(tgt) + # {{{ process dofdesc arguments - if isinstance(vec, Number) or src == tgt: + src_dofdesc = as_dofdesc(src) + + contextual_volume_tag = None + if isinstance(src_dofdesc.domain_tag, VolumeDomainTag): + contextual_volume_tag = src_dofdesc.domain_tag.tag + elif isinstance(src_dofdesc.domain_tag, BoundaryDomainTag): + contextual_volume_tag = src_dofdesc.domain_tag.volume_tag + + tgt_dofdesc = as_dofdesc(tgt, _contextual_volume_tag=contextual_volume_tag) + + del src + del tgt + + # }}} + + if isinstance(vec, Number) or src_dofdesc == tgt_dofdesc: return vec - return dcoll.connection_from_dds(src, tgt)(vec) + return dcoll.connection_from_dds(src_dofdesc, tgt_dofdesc)(vec) diff --git a/grudge/reductions.py b/grudge/reductions.py index 95ed44726..6087b5725 100644 --- a/grudge/reductions.py +++ b/grudge/reductions.py @@ -94,7 +94,7 @@ def norm(dcoll: DiscretizationCollection, vec, p, dd=None) -> Scalar: :returns: a nonegative scalar denoting the norm. """ if dd is None: - dd = dof_desc.DD_VOLUME + dd = dof_desc.DD_VOLUME_ALL from arraycontext import get_container_context_recursively actx = get_container_context_recursively(vec) @@ -128,7 +128,7 @@ def nodal_sum(dcoll: DiscretizationCollection, dd, vec) -> Scalar: if comm is None: return nodal_sum_loc(dcoll, dd, vec) - # NOTE: Don't move this + # NOTE: Do not move, we do not want to import mpi4py in single-rank computations from mpi4py import MPI from arraycontext import get_container_context_recursively @@ -174,7 +174,7 @@ def nodal_min(dcoll: DiscretizationCollection, dd, vec, *, initial=None) -> Scal if comm is None: return nodal_min_loc(dcoll, dd, vec, initial=initial) - # NOTE: Don't move this + # NOTE: Do not move, we do not want to import mpi4py in single-rank computations from mpi4py import MPI actx = vec.array_context @@ -231,7 +231,7 @@ def nodal_max(dcoll: DiscretizationCollection, dd, vec, *, initial=None) -> Scal if comm is None: return nodal_max_loc(dcoll, dd, vec, initial=initial) - # NOTE: Don't move this + # NOTE: Do not move, we do not want to import mpi4py in single-rank computations from mpi4py import MPI actx = vec.array_context @@ -320,7 +320,7 @@ def _apply_elementwise_reduction( """ if len(args) == 1: vec, = args - dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd = dof_desc.DD_VOLUME_ALL elif len(args) == 2: dd, vec = args else: @@ -344,7 +344,7 @@ def _apply_elementwise_reduction( ) ) else: - @memoize_in(actx, (_apply_elementwise_reduction, + @memoize_in(actx, (_apply_elementwise_reduction, dd, "elementwise_%s_prg" % op_name)) def elementwise_prg(): # FIXME: This computes the reduction value redundantly for each @@ -485,7 +485,7 @@ def elementwise_integral( """ if len(args) == 1: vec, = args - dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd = dof_desc.DD_VOLUME_ALL elif len(args) == 2: dd, vec = args else: diff --git a/grudge/shortcuts.py b/grudge/shortcuts.py index 0aca64a58..e6e62cc55 100644 --- a/grudge/shortcuts.py +++ b/grudge/shortcuts.py @@ -20,6 +20,8 @@ THE SOFTWARE. """ +from grudge.dof_desc import DD_VOLUME_ALL + from pytools import memoize_in @@ -76,11 +78,14 @@ def set_up_rk4(field_var_name, dt, fields, rhs, t_start=0.0): return dt_stepper -def make_visualizer(dcoll, vis_order=None, **kwargs): +def make_visualizer(dcoll, vis_order=None, volume_dd=None, **kwargs): from meshmode.discretization.visualization import make_visualizer + if volume_dd is None: + volume_dd = DD_VOLUME_ALL + return make_visualizer( dcoll._setup_actx, - dcoll.discr_from_dd("vol"), vis_order, **kwargs) + dcoll.discr_from_dd(volume_dd), vis_order, **kwargs) def make_boundary_visualizer(dcoll, vis_order=None, **kwargs): diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index cb2de38f6..3b1f10821 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -51,6 +51,7 @@ """ +from warnings import warn from typing import List, Hashable, Optional, Type, Any from pytools.persistent_dict import KeyBuilder @@ -70,7 +71,6 @@ from numbers import Number from pytools import memoize_on_first_arg -from pytools.obj_array import obj_array_vectorize from grudge.discretization import DiscretizationCollection from grudge.projection import project @@ -78,7 +78,13 @@ from meshmode.mesh import BTAG_PARTITION import numpy as np + import grudge.dof_desc as dof_desc +from grudge.dof_desc import ( + DOFDesc, DD_VOLUME_ALL, FACE_RESTR_INTERIOR, DISCR_TAG_BASE, + VolumeDomainTag, + ConvertibleToDOFDesc, + ) # {{{ trace pair container class @@ -107,12 +113,22 @@ class TracePair: .. automethod:: __len__ """ - dd: dof_desc.DOFDesc + dd: DOFDesc interior: ArrayContainer exterior: ArrayContainer - def __init__(self, dd, *, interior, exterior): - object.__setattr__(self, "dd", dof_desc.as_dofdesc(dd)) + def __init__(self, dd: DOFDesc, *, + interior: ArrayOrContainer, + exterior: ArrayOrContainer): + if not isinstance(dd, DOFDesc): + warn("Constructing a TracePair with a first argument that is not " + "exactly a DOFDesc (but convertible to one) is deprecated. " + "This will stop working in July 2022. " + "Pass an actual DOFDesc instead.", + DeprecationWarning, stacklevel=2) + dd = dof_desc.as_dofdesc(dd) + + object.__setattr__(self, "dd", dd) object.__setattr__(self, "interior", interior) object.__setattr__(self, "exterior", exterior) @@ -178,7 +194,8 @@ def diff(self): # {{{ boundary trace pairs def bdry_trace_pair( - dcoll: DiscretizationCollection, dd, interior, exterior) -> TracePair: + dcoll: DiscretizationCollection, dd: "ConvertibleToDOFDesc", + interior, exterior) -> TracePair: """Returns a trace pair defined on the exterior boundary. Input arguments are assumed to already be defined on the boundary denoted by *dd*. If the input arguments *interior* and *exterior* are @@ -197,11 +214,19 @@ def bdry_trace_pair( be used for the flux. :returns: a :class:`TracePair` on the boundary. """ + if not isinstance(dd, DOFDesc): + warn("Calling bdry_trace_pair with a first argument that is not " + "exactly a DOFDesc (but convertible to one) is deprecated. " + "This will stop working in July 2022. " + "Pass an actual DOFDesc instead.", + DeprecationWarning, stacklevel=2) + dd = dof_desc.as_dofdesc(dd) return TracePair(dd, interior=interior, exterior=exterior) def bv_trace_pair( - dcoll: DiscretizationCollection, dd, interior, exterior) -> TracePair: + dcoll: DiscretizationCollection, dd: "ConvertibleToDOFDesc", + interior, exterior) -> TracePair: """Returns a trace pair defined on the exterior boundary. The interior argument is assumed to be defined on the volume discretization, and will therefore be restricted to the boundary *dd* prior to creating a @@ -223,21 +248,29 @@ def bv_trace_pair( be used for the flux. :returns: a :class:`TracePair` on the boundary. """ + if not isinstance(dd, DOFDesc): + warn("Calling bv_trace_pair with a first argument that is not " + "exactly a DOFDesc (but convertible to one) is deprecated. " + "This will stop working in July 2022. " + "Pass an actual DOFDesc instead.", + DeprecationWarning, stacklevel=2) + dd = dof_desc.as_dofdesc(dd) return bdry_trace_pair( - dcoll, dd, project(dcoll, "vol", dd, interior), exterior - ) + dcoll, dd, project(dcoll, dd.domain_tag.volume_tag, dd, interior), exterior) # }}} # {{{ interior trace pairs -def local_interior_trace_pair(dcoll: DiscretizationCollection, vec) -> TracePair: +def local_interior_trace_pair( + dcoll: DiscretizationCollection, vec, *, + volume_dd: Optional[DOFDesc] = None, + ) -> TracePair: r"""Return a :class:`TracePair` for the interior faces of *dcoll* with a discretization tag specified by *discr_tag*. This does not include interior faces on different MPI ranks. - :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` of them. @@ -250,21 +283,33 @@ def local_interior_trace_pair(dcoll: DiscretizationCollection, vec) -> TracePair computation. :returns: a :class:`TracePair` object. """ - i = project(dcoll, "vol", "int_faces", vec) + if volume_dd is None: + volume_dd = DD_VOLUME_ALL + + assert isinstance(volume_dd.domain_tag, VolumeDomainTag) + trace_dd = volume_dd.trace(FACE_RESTR_INTERIOR) + + interior = project(dcoll, volume_dd, trace_dd, vec) - def get_opposite_face(el): - if isinstance(el, Number): - return el + opposite_face_conn = dcoll.opposite_face_connection(trace_dd.domain_tag) + + def get_opposite_trace(ary): + if isinstance(ary, Number): + return ary else: - return dcoll.opposite_face_connection()(el) + return opposite_face_conn(ary) - e = obj_array_vectorize(get_opposite_face, i) + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + exterior = rec_map_array_container( + get_opposite_trace, + interior, + leaf_class=DOFArray) - return TracePair("int_faces", interior=i, exterior=e) + return TracePair(trace_dd, interior=interior, exterior=exterior) def interior_trace_pair(dcoll: DiscretizationCollection, vec) -> TracePair: - from warnings import warn warn("`grudge.op.interior_trace_pair` is deprecated and will be dropped " "in version 2022.x. Use `local_interior_trace_pair` " "instead, or `interior_trace_pairs` which also includes contributions " @@ -274,7 +319,8 @@ def interior_trace_pair(dcoll: DiscretizationCollection, vec) -> TracePair: def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *, - comm_tag: Hashable = None, tag: Hashable = None) -> List[TracePair]: + comm_tag: Hashable = None, tag: Hashable = None, + volume_dd: Optional[DOFDesc] = None) -> List[TracePair]: r"""Return a :class:`list` of :class:`TracePair` objects defined on the interior faces of *dcoll* and any faces connected to a parallel boundary. @@ -293,7 +339,6 @@ def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *, """ if tag is not None: - from warnings import warn warn("Specifying 'tag' is deprecated and will stop working in July of 2022. " "Specify 'comm_tag' instead.", DeprecationWarning, stacklevel=2) if comm_tag is not None: @@ -302,46 +347,97 @@ def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *, comm_tag = tag del tag + if volume_dd is None: + volume_dd = DD_VOLUME_ALL + return ( - [local_interior_trace_pair(dcoll, vec)] - + cross_rank_trace_pairs(dcoll, vec, comm_tag=comm_tag) + [local_interior_trace_pair( + dcoll, vec, volume_dd=volume_dd)] + + cross_rank_trace_pairs( + dcoll, vec, comm_tag=comm_tag, volume_dd=volume_dd) ) # }}} -# {{{ distributed-memory functionality +# {{{ distributed: helper functions + +class _TagKeyBuilder(KeyBuilder): + def update_for_type(self, key_hash, key: Type[Any]): + self.rec(key_hash, (key.__module__, key.__name__, key.__name__,)) + @memoize_on_first_arg -def connected_ranks(dcoll: DiscretizationCollection): +def connected_ranks( + dcoll: DiscretizationCollection, + volume_dd: Optional[DOFDesc] = None): + if volume_dd is None: + volume_dd = DD_VOLUME_ALL + from meshmode.distributed import get_connected_partitions - return get_connected_partitions(dcoll._volume_discr.mesh) + return get_connected_partitions( + dcoll._volume_discrs[volume_dd.domain_tag.tag].mesh) + + +def _sym_tag_to_num_tag(comm_tag: Optional[Hashable]) -> Optional[int]: + if comm_tag is None: + return comm_tag + + if isinstance(comm_tag, int): + return comm_tag + + # FIXME: This isn't guaranteed to be correct. + # See here for discussion: + # - https://github.com/illinois-ceesd/mirgecom/issues/617#issuecomment-1057082716 # noqa + # - https://github.com/inducer/grudge/pull/222 + + from mpi4py import MPI + tag_ub = MPI.COMM_WORLD.Get_attr(MPI.TAG_UB) + key_builder = _TagKeyBuilder() + digest = key_builder(comm_tag) + + num_tag = sum(ord(ch) << i for i, ch in enumerate(digest)) % tag_ub + warn("Encountered unknown symbolic tag " + f"'{comm_tag}', assigning a value of '{num_tag}'. " + "This is a temporary workaround, please ensure that " + "tags are sufficiently distinct for your use case.") -class _RankBoundaryCommunication: + return num_tag + +# }}} + + +# {{{ eager rank-boundary communication + +class _RankBoundaryCommunicationEager: base_comm_tag = 1273 def __init__(self, dcoll: DiscretizationCollection, array_container: ArrayOrContainer, - remote_rank, comm_tag: Optional[int] = None): + remote_rank, comm_tag: Optional[int] = None, + volume_dd=DD_VOLUME_ALL): actx = get_container_context_recursively(array_container) - btag = BTAG_PARTITION(remote_rank) + bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_rank)) - local_bdry_data = project(dcoll, "vol", btag, array_container) + local_bdry_data = project(dcoll, volume_dd, bdry_dd, array_container) comm = dcoll.mpi_communicator + assert comm is not None self.dcoll = dcoll self.array_context = actx - self.remote_btag = btag - self.bdry_discr = dcoll.discr_from_dd(btag) + self.remote_bdry_dd = bdry_dd + self.bdry_discr = dcoll.discr_from_dd(bdry_dd) self.local_bdry_data = local_bdry_data self.local_bdry_data_np = \ to_numpy(flatten(self.local_bdry_data, actx), actx) self.comm_tag = self.base_comm_tag + comm_tag = _sym_tag_to_num_tag(comm_tag) if comm_tag is not None: self.comm_tag += comm_tag + del comm_tag # Here, we initialize both send and recieve operations through # mpi4py `Request` (MPI_Request) instances for comm.Isend (MPI_Isend) @@ -382,36 +478,42 @@ def finish(self): remote_bdry_data = unflatten(self.local_bdry_data, remote_bdry_data_flat, actx) bdry_conn = self.dcoll.distributed_boundary_swap_connection( - dof_desc.as_dofdesc(dof_desc.DTAG_BOUNDARY(self.remote_btag))) + self.remote_bdry_dd) swapped_remote_bdry_data = bdry_conn(remote_bdry_data) # Complete the nonblocking send request associated with communicating # `self.local_bdry_data_np` self.send_req.Wait() - return TracePair(self.remote_btag, + return TracePair(self.remote_bdry_dd, interior=self.local_bdry_data, exterior=swapped_remote_bdry_data) +# }}} -from pytato import make_distributed_recv, staple_distributed_send +# {{{ lazy rank-boundary communication class _RankBoundaryCommunicationLazy: def __init__(self, dcoll: DiscretizationCollection, array_container: ArrayOrContainer, - remote_rank: int, comm_tag: Hashable): + remote_rank: int, comm_tag: Hashable, + volume_dd=DD_VOLUME_ALL): if comm_tag is None: raise ValueError("lazy communication requires 'tag' to be supplied") + bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_rank)) + self.dcoll = dcoll self.array_context = get_container_context_recursively(array_container) - self.remote_btag = BTAG_PARTITION(remote_rank) - self.bdry_discr = dcoll.discr_from_dd(self.remote_btag) + self.remote_bdry_dd = bdry_dd + self.bdry_discr = dcoll.discr_from_dd(self.remote_bdry_dd) self.local_bdry_data = project( - dcoll, "vol", self.remote_btag, array_container) + dcoll, volume_dd, bdry_dd, array_container) + + from pytato import make_distributed_recv, staple_distributed_send def communicate_single_array(key, local_bdry_ary): ary_tag = (comm_tag, key) @@ -428,22 +530,22 @@ def communicate_single_array(key, local_bdry_ary): def finish(self): bdry_conn = self.dcoll.distributed_boundary_swap_connection( - dof_desc.as_dofdesc(dof_desc.DTAG_BOUNDARY(self.remote_btag))) + self.remote_bdry_dd) - return TracePair(self.remote_btag, + return TracePair(self.remote_bdry_dd, interior=self.local_bdry_data, exterior=bdry_conn(self.remote_data)) +# }}} -class _TagKeyBuilder(KeyBuilder): - def update_for_type(self, key_hash, key: Type[Any]): - self.rec(key_hash, (key.__module__, key.__name__, key.__name__,)) +# {{{ cross_rank_trace_pairs def cross_rank_trace_pairs( - dcoll: DiscretizationCollection, ary, - comm_tag: Hashable = None, - tag: Hashable = None) -> List[TracePair]: + dcoll: DiscretizationCollection, ary: ArrayOrContainer, + tag: Hashable = None, + *, comm_tag: Hashable = None, + volume_dd: Optional[DOFDesc] = None) -> List[TracePair]: r"""Get a :class:`list` of *ary* trace pairs for each partition boundary. For each partition boundary, the field data values in *ary* are @@ -472,57 +574,48 @@ def cross_rank_trace_pairs( :returns: a :class:`list` of :class:`TracePair` objects. """ + # {{{ process arguments + + if volume_dd is None: + volume_dd = DD_VOLUME_ALL + + if not isinstance(volume_dd.domain_tag, VolumeDomainTag): + raise TypeError(f"expected a volume DOFDesc, got '{volume_dd}'") + if volume_dd.discretization_tag != DISCR_TAG_BASE: + raise TypeError(f"expected a base-discretized DOFDesc, got '{volume_dd}'") + if tag is not None: - from warnings import warn warn("Specifying 'tag' is deprecated and will stop working in July of 2022. " - "Specify 'comm_tag' instead.", DeprecationWarning, stacklevel=2) + "Specify 'comm_tag' (keyword-only) instead.", + DeprecationWarning, stacklevel=2) if comm_tag is not None: raise TypeError("may only specify one of 'tag' and 'comm_tag'") else: comm_tag = tag del tag + # }}} + if isinstance(ary, Number): # NOTE: Assumed that the same number is passed on every rank - return [TracePair(BTAG_PARTITION(remote_rank), interior=ary, exterior=ary) - for remote_rank in connected_ranks(dcoll)] + return [TracePair( + volume_dd.trace(BTAG_PARTITION(remote_rank)), + interior=ary, exterior=ary) + for remote_rank in connected_ranks(dcoll, volume_dd=volume_dd)] actx = get_container_context_recursively(ary) from grudge.array_context import MPIPytatoArrayContextBase if isinstance(actx, MPIPytatoArrayContextBase): - rbc = _RankBoundaryCommunicationLazy + rbc_class = _RankBoundaryCommunicationLazy else: - rbc = _RankBoundaryCommunication - if comm_tag is not None: - num_tag: Optional[int] = None - if isinstance(comm_tag, int): - num_tag = comm_tag - - if num_tag is None: - # FIXME: This isn't guaranteed to be correct. - # See here for discussion: - # - https://github.com/illinois-ceesd/mirgecom/issues/617#issuecomment-1057082716 # noqa - # - https://github.com/inducer/grudge/pull/222 - from mpi4py import MPI - tag_ub = MPI.COMM_WORLD.Get_attr(MPI.TAG_UB) - key_builder = _TagKeyBuilder() - digest = key_builder(comm_tag) - num_tag = sum(ord(ch) << i for i, ch in enumerate(digest)) % tag_ub - - from warnings import warn - warn("Encountered unknown symbolic tag " - f"'{comm_tag}', assigning a value of '{num_tag}'. " - "This is a temporary workaround, please ensure that " - "tags are sufficiently distinct for your use case.") - - comm_tag = num_tag + rbc_class = _RankBoundaryCommunicationEager # Initialize and post all sends/receives rank_bdry_communcators = [ - rbc(dcoll, ary, remote_rank, comm_tag=comm_tag) - for remote_rank in connected_ranks(dcoll) + rbc_class(dcoll, ary, remote_rank, comm_tag=comm_tag, volume_dd=volume_dd) + for remote_rank in connected_ranks(dcoll, volume_dd=volume_dd) ] # Complete send/receives and return communicated data diff --git a/test/test_grudge.py b/test/test_grudge.py index 916bd6673..c32c2ca2c 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -38,7 +38,7 @@ from pytools.obj_array import flat_obj_array -from grudge import DiscretizationCollection +from grudge import DiscretizationCollection, make_discretization_collection import grudge.dof_desc as dof_desc import grudge.op as op @@ -341,7 +341,10 @@ def test_face_normal_surface(actx_factory, mesh_name): surf_normal = surf_normal / actx.np.sqrt(sum(surf_normal**2)) face_normal_i = actx.thaw(dcoll.normal(df)) - face_normal_e = dcoll.opposite_face_connection()(face_normal_i) + face_normal_e = dcoll.opposite_face_connection( + dof_desc.BoundaryDomainTag( + dof_desc.FACE_RESTR_INTERIOR, dof_desc.VTAG_ALL) + )(face_normal_i) if mesh.ambient_dim == 3: from grudge.geometry import pseudoscalar, area_element @@ -1070,6 +1073,29 @@ def test_empty_boundary(actx_factory): # }}} +# {{{ multi-volume + +def test_multiple_independent_volumes(actx_factory): + dim = 2 + actx = actx_factory() + + mesh1 = mgen.generate_regular_rect_mesh( + a=(-2,)*dim, b=(-1,)*dim, + nelements_per_axis=(4,)*dim, order=4) + + mesh2 = mgen.generate_regular_rect_mesh( + a=(1,)*dim, b=(2,)*dim, + nelements_per_axis=(8,)*dim, order=4) + + volume_to_mesh = { + "vol1": mesh1, + "vol2": mesh2} + + make_discretization_collection(actx, volume_to_mesh, order=4) + +# }}} + + # You can test individual routines by typing # $ python test_grudge.py 'test_routine()' diff --git a/test/test_mpi_communication.py b/test/test_mpi_communication.py index 47100b415..424b78199 100644 --- a/test/test_mpi_communication.py +++ b/test/test_mpi_communication.py @@ -45,6 +45,7 @@ from pytools.obj_array import flat_obj_array import grudge.op as op +import grudge.dof_desc as dof_desc class SimpleTag: @@ -153,7 +154,10 @@ def hopefully_zero(): return ( op.project( dcoll, "int_faces", "all_faces", - dcoll.opposite_face_connection()(int_faces_func) + dcoll.opposite_face_connection( + dof_desc.BoundaryDomainTag( + dof_desc.FACE_RESTR_INTERIOR, dof_desc.VTAG_ALL) + )(int_faces_func) ) + sum(op.project(dcoll, tpair.dd, "all_faces", tpair.ext) for tpair in op.cross_rank_trace_pairs(dcoll, myfunc, diff --git a/test/test_reductions.py b/test/test_reductions.py index 9e7387bad..e0e424fe4 100644 --- a/test/test_reductions.py +++ b/test/test_reductions.py @@ -149,7 +149,7 @@ def f(x): min_res = np.empty(grp_f.shape) max_res = np.empty(grp_f.shape) sum_res = np.empty(grp_f.shape) - for eidx in range(dcoll.mesh.nelements): + for eidx in range(mesh.nelements): element_data = actx.to_numpy(grp_f[eidx]) min_res[eidx, :] = np.min(element_data) max_res[eidx, :] = np.max(element_data) @@ -272,7 +272,7 @@ def _get_ref_data(field): min_res = np.empty(grp_f.shape) max_res = np.empty(grp_f.shape) sum_res = np.empty(grp_f.shape) - for eidx in range(dcoll.mesh.nelements): + for eidx in range(mesh.nelements): element_data = actx.to_numpy(grp_f[eidx]) min_res[eidx, :] = np.min(element_data) max_res[eidx, :] = np.max(element_data) From baab4bd98291fcb487eac3697b4a32652e91b318 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Tue, 30 Aug 2022 10:12:42 -0500 Subject: [PATCH 09/12] move with_discr_tag to match docstring order --- grudge/dof_desc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/grudge/dof_desc.py b/grudge/dof_desc.py index cf285a30e..246a5228b 100644 --- a/grudge/dof_desc.py +++ b/grudge/dof_desc.py @@ -302,6 +302,9 @@ def with_dtag(self, dtag) -> "DOFDesc": def with_domain_tag(self, dtag) -> "DOFDesc": return replace(self, domain_tag=dtag) + def with_discr_tag(self, discr_tag) -> "DOFDesc": + return replace(self, discretization_tag=discr_tag) + def trace(self, btag: BoundaryTag) -> "DOFDesc": """Return a :class:`DOFDesc` for the restriction of the volume descriptor *self* to the boundary named by *btag*. @@ -314,9 +317,6 @@ def trace(self, btag: BoundaryTag) -> "DOFDesc": return replace(self, domain_tag=BoundaryDomainTag(btag, volume_tag=self.domain_tag.tag)) - def with_discr_tag(self, discr_tag) -> "DOFDesc": - return replace(self, discretization_tag=discr_tag) - def as_identifier(self) -> str: """Returns a descriptive string for this :class:`DOFDesc` that is usable in Python identifiers. From 42fa3ad5d38d6c6eac09a5671f4d0446ae58a5a4 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Tue, 30 Aug 2022 10:13:06 -0500 Subject: [PATCH 10/12] add untrace and with_boundary_tag methods to DOFDesc --- grudge/dof_desc.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/grudge/dof_desc.py b/grudge/dof_desc.py index 246a5228b..cb8cd810b 100644 --- a/grudge/dof_desc.py +++ b/grudge/dof_desc.py @@ -228,6 +228,8 @@ class DOFDesc: .. automethod:: with_domain_tag .. automethod:: with_discr_tag .. automethod:: trace + .. automethod:: untrace + .. automethod:: with_boundary_tag .. automethod:: __eq__ .. automethod:: __ne__ @@ -317,6 +319,30 @@ def trace(self, btag: BoundaryTag) -> "DOFDesc": return replace(self, domain_tag=BoundaryDomainTag(btag, volume_tag=self.domain_tag.tag)) + def untrace(self) -> "DOFDesc": + """Return a :class:`DOFDesc` for the volume associated with the boundary + descriptor *self*. + + An error is raised if this method is called on a non-boundary instance of + :class:`DOFDesc`. + """ + if not isinstance(self.domain_tag, BoundaryDomainTag): + raise ValueError(f"must originate on boundary, got '{self.domain_tag}'") + return replace(self, + domain_tag=VolumeDomainTag(self.domain_tag.volume_tag)) + + def with_boundary_tag(self, btag: BoundaryTag) -> "DOFDesc": + """Return a :class:`DOFDesc` representing a boundary named by *btag* + on the same volume as *self*. + + An error is raised if this method is called on a non-boundary instance of + :class:`DOFDesc`. + """ + if not isinstance(self.domain_tag, BoundaryDomainTag): + raise ValueError(f"must originate on boundary, got '{self.domain_tag}'") + return replace(self, + domain_tag=replace(self.domain_tag, tag=btag)) + def as_identifier(self) -> str: """Returns a descriptive string for this :class:`DOFDesc` that is usable in Python identifiers. From bc0c4735c6c119779fd016f3be4d73c10578699a Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Tue, 13 Sep 2022 14:44:21 -0500 Subject: [PATCH 11/12] fix comment --- grudge/dof_desc.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/grudge/dof_desc.py b/grudge/dof_desc.py index cb8cd810b..bdcc5fbad 100644 --- a/grudge/dof_desc.py +++ b/grudge/dof_desc.py @@ -121,10 +121,7 @@ class ScalarDomainTag: # noqa: N801 @dataclass(frozen=True, eq=True, init=True) class VolumeDomainTag: - """A domain tag referring to a volume identified by the - volume tag :attr:`tag`. These volume identifiers are only used - when the :class:`~grudge.discretization.DiscretizationCollection` contains - more than one volume. + """A domain tag referring to a volume identified by the volume tag :attr:`tag`. .. attribute:: tag From 0dc755e85fe498c6eb49a1556f21c6badc8527d4 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Tue, 13 Sep 2022 15:39:30 -0500 Subject: [PATCH 12/12] fix deprecation timeline --- grudge/trace_pair.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 3b1f10821..1f49ae0d6 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -123,7 +123,7 @@ def __init__(self, dd: DOFDesc, *, if not isinstance(dd, DOFDesc): warn("Constructing a TracePair with a first argument that is not " "exactly a DOFDesc (but convertible to one) is deprecated. " - "This will stop working in July 2022. " + "This will stop working in December 2022. " "Pass an actual DOFDesc instead.", DeprecationWarning, stacklevel=2) dd = dof_desc.as_dofdesc(dd) @@ -217,7 +217,7 @@ def bdry_trace_pair( if not isinstance(dd, DOFDesc): warn("Calling bdry_trace_pair with a first argument that is not " "exactly a DOFDesc (but convertible to one) is deprecated. " - "This will stop working in July 2022. " + "This will stop working in December 2022. " "Pass an actual DOFDesc instead.", DeprecationWarning, stacklevel=2) dd = dof_desc.as_dofdesc(dd) @@ -251,7 +251,7 @@ def bv_trace_pair( if not isinstance(dd, DOFDesc): warn("Calling bv_trace_pair with a first argument that is not " "exactly a DOFDesc (but convertible to one) is deprecated. " - "This will stop working in July 2022. " + "This will stop working in December 2022. " "Pass an actual DOFDesc instead.", DeprecationWarning, stacklevel=2) dd = dof_desc.as_dofdesc(dd)