diff --git a/doc/conf.py b/doc/conf.py index 31bc96d9b..a71d9e55a 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -40,6 +40,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/grudge/__init__.py b/grudge/__init__.py index 04cccfb13..9510c5529 100644 --- a/grudge/__init__.py +++ b/grudge/__init__.py @@ -23,8 +23,9 @@ import grudge.symbolic as sym from grudge.execution import bind -from grudge.discretization import DiscretizationCollection +from grudge.discretization import ( + DiscretizationCollection, make_discretization_collection) __all__ = [ - "sym", "bind", "DiscretizationCollection" + "sym", "bind", "DiscretizationCollection", "make_discretization_collection" ] diff --git a/grudge/discretization.py b/grudge/discretization.py index 3ead9485f..36c85fe9e 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -1,7 +1,19 @@ """ -.. currentmodule:: grudge +.. autoclass:: DiscretizationTag + +.. currentmodule:: grudge .. autoclass:: DiscretizationCollection +.. autofunction:: make_discretization_collection + +.. currentmodule:: grudge.discretization + +.. autofunction:: relabel_partitions + +Internal things that are visble due to type annotations +------------------------------------------------------- + +.. class:: _InterPartitionConnectionPair """ __copyright__ = """ @@ -29,34 +41,89 @@ THE SOFTWARE. """ -from typing import Callable -from pytools import memoize_method +from typing import Mapping, Optional, Union, Tuple, TYPE_CHECKING + +from pytools import memoize_method, single_valued from grudge.dof_desc import ( - DD_VOLUME, - DISCR_TAG_BASE, - DISCR_TAG_MODAL, - DTAG_BOUNDARY, - DOFDesc, - DiscretizationTag, - as_dofdesc + VTAG_ALL, + BTAG_MULTIVOL_PARTITION, + DD_VOLUME_ALL, + DISCR_TAG_BASE, + DISCR_TAG_MODAL, + DTAG_VOLUME, DTAG_BOUNDARY, + DOFDesc, + VolumeTag, DomainTag, + DiscretizationTag, + as_dofdesc ) import numpy as np # noqa: F401 from arraycontext import ArrayContext -from meshmode.discretization import ElementGroupFactory +from meshmode.discretization import ElementGroupFactory, Discretization from meshmode.discretization.connection import ( FACE_RESTR_INTERIOR, FACE_RESTR_ALL, make_face_restriction, DiscretizationConnection ) -from meshmode.mesh import Mesh, BTAG_PARTITION +from meshmode.mesh import ( + InterPartitionAdjacencyGroup, Mesh, BTAG_PARTITION, BoundaryTag) 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 @@ -64,11 +131,8 @@ class DiscretizationCollection: (volume, interior facets, boundaries) and associated element groups. - .. automethod:: __init__ - .. autoattribute:: dim .. autoattribute:: ambient_dim - .. autoattribute:: mesh .. autoattribute:: real_dtype .. autoattribute:: complex_dtype @@ -88,9 +152,17 @@ 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, + inter_partition_connections: Optional[ + Mapping[Union[BTAG_MULTIVOL_PARTITION, BTAG_PARTITION], + DiscretizationConnection] + ] = None + ) -> None: """ :arg discr_tag_to_group_factory: A mapping from discretization tags (typically one of: :class:`grudge.dof_desc.DISCR_TAG_BASE`, @@ -105,45 +177,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) - ) - # NOTE: Can be removed when symbolics are completely removed # {{{ management of discretization-scoped common subexpressions @@ -178,18 +211,54 @@ def __init__(self, array_context: ArrayContext, mesh: Mesh, # }}} - self._dist_boundary_connections = _get_dist_boundary_connections_single_vol( - volume_discr=self._volume_discr, - mpi_communicator=mpi_communicator, - array_context=self._setup_actx, - get_group_factory_for_discretization_tag=( - self.group_factory_for_discretization_tag), - get_connection_to_rank_boundary=( - lambda remote_rank: - self.connection_from_dds( - DOFDesc("vol", DISCR_TAG_BASE), - DOFDesc(BTAG_PARTITION(remote_rank), - DISCR_TAG_BASE))),) + from meshmode.discretization import Discretization + + # {{{ deprecated backward compatibility yuck + + if isinstance(volume_discrs, Mesh): + 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 + + if inter_partition_connections is not None: + raise TypeError("may not pass inter_partition_connections when " + "DiscretizationCollection constructor is called in " + "legacy mode") + + self._inter_partition_connections = \ + _set_up_inter_partition_connections( + array_context=self._setup_actx, + mpi_communicator=mpi_communicator, + volume_discrs=volume_discrs, + base_group_factory=( + discr_tag_to_group_factory[DISCR_TAG_BASE])) + else: + if inter_partition_connections is None: + raise TypeError("inter_partition_connections must be passed when " + "DiscretizationCollection constructor is called in " + "'modern' mode") + + self._inter_partition_connections = inter_partition_connections + self._discr_tag_to_group_factory = discr_tag_to_group_factory + + # }}} + + self._volume_discrs = volume_discrs # }}} @@ -212,35 +281,6 @@ def is_management_rank(self): return self.mpi_communicator.Get_rank() \ == self.get_management_rank_index() - # {{{ distributed - - def distributed_boundary_swap_connection(self, dd): - """Provides a mapping from the base volume discretization - to the exterior boundary restriction on a parallel boundary - partition described by *dd*. This connection is used to - communicate across element boundaries in different parallel - partitions during distributed runs. - - :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 - associated :class:`meshmode.mesh.BTAG_PARTITION` - corresponding to a particular communication rank. - """ - if dd.discretization_tag is not DISCR_TAG_BASE: - # FIXME - raise NotImplementedError( - "Distributed communication with discretization tag " - f"{dd.discretization_tag} is not implemented." - ) - - assert isinstance(dd.domain_tag, DTAG_BOUNDARY) - assert isinstance(dd.domain_tag.tag, BTAG_PARTITION) - - return self._dist_boundary_connections[dd.domain_tag.tag.part_nr] - - # }}} - # {{{ discr_from_dd @memoize_method @@ -259,42 +299,40 @@ 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, DTAG_BOUNDARY): + 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): @@ -310,7 +348,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 \ @@ -379,12 +417,14 @@ 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, DTAG_BOUNDARY) and from_discr_tag == to_discr_tag - and to_dd.domain_tag is FACE_RESTR_ALL): + and isinstance(to_dd.domain_tag, DTAG_BOUNDARY) + and to_dd.domain_tag.tag is FACE_RESTR_ALL): faces_conn = self.connection_from_dds( - DOFDesc("vol"), - DOFDesc(from_dd.domain_tag)) + DOFDesc( + DTAG_VOLUME(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 @@ -402,7 +442,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 @@ -436,75 +476,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, DTAG_VOLUME): + if isinstance(to_dd.domain_tag, DTAG_BOUNDARY): + 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, DTAG_BOUNDARY): + 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_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. - """ 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, DTAG_VOLUME) + + 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)) @@ -552,25 +596,30 @@ 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: DTAG_BOUNDARY): 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(DTAG_VOLUME(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: DTAG_BOUNDARY) -> 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( + DTAG_VOLUME(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 @@ -579,7 +628,8 @@ def _interior_faces_connection(self): ) @memoize_method - def opposite_face_connection(self): + def opposite_face_connection( + self, domain_tag: DTAG_BOUNDARY) -> 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. @@ -587,27 +637,11 @@ 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)) # }}} @@ -616,55 +650,58 @@ def _all_faces_volume_connection(self): @property def dim(self): """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): """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): """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): """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): """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): """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() @@ -681,7 +718,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): @@ -698,45 +735,152 @@ def normal(self, dd): # }}} -# {{{ distributed setup +# {{{ distributed/multi-volume setup + +def _check_btag_against_vtag( + discr_vol_tag: VolumeTag, tag: BoundaryTag + ) -> Union[BTAG_MULTIVOL_PARTITION, BTAG_PARTITION]: + if isinstance(tag, BTAG_MULTIVOL_PARTITION): + if tag.self_volume_tag != discr_vol_tag: + raise ValueError("encountered inter-partition self-volume tag " + f"'{tag.self_volume_tag}' in a volume discretization with " + f"tag '{discr_vol_tag}'") + + return tag + + elif isinstance(tag, BTAG_PARTITION): + return tag + + else: + raise TypeError("unexpected type of inter-partition boundary tag " + f"'{type(tag)}'") + + +def _volume_discr_for_btag( + volume_discrs: Mapping[VolumeTag, Discretization], + btag: BoundaryTag) -> Discretization: + if isinstance(btag, BTAG_MULTIVOL_PARTITION): + return volume_discrs[btag.self_volume_tag] + elif isinstance(btag, BTAG_PARTITION): + return volume_discrs[VTAG_ALL] + else: + raise TypeError("unexpected type of inter-partition boundary tag " + f"'{type(btag)}'") + + +def _remote_rank_from_btag(btag: BoundaryTag) -> Optional[int]: + if isinstance(btag, BTAG_PARTITION): + return btag.part_nr + + elif isinstance(btag, BTAG_MULTIVOL_PARTITION): + return btag.other_rank -def _get_dist_boundary_connections_single_vol( - volume_discr, mpi_communicator, array_context, - get_group_factory_for_discretization_tag: Callable[ - [DiscretizationTag], ElementGroupFactory], - get_connection_to_rank_boundary: Callable[[int], DiscretizationConnection]): - boundary_connections = {} + else: + raise TypeError("unexpected type of inter-partition boundary tag " + f"'{type(btag)}'") + + +def _flip_btag( + local_rank: Optional[int], + btag: Union[BTAG_PARTITION, BTAG_MULTIVOL_PARTITION] + ) -> Union[BTAG_PARTITION, BTAG_MULTIVOL_PARTITION]: + if isinstance(btag, BTAG_PARTITION): + return BTAG_PARTITION(local_rank) + + elif isinstance(btag, BTAG_MULTIVOL_PARTITION): + return BTAG_MULTIVOL_PARTITION( + other_rank=None if btag.other_rank is None else local_rank, + self_volume_tag=btag.other_volume_tag, + other_volume_tag=btag.self_volume_tag) + else: + raise TypeError("unexpected type of inter-partition boundary tag " + f"'{type(btag)}'") + + +def _set_up_inter_partition_connections( + array_context: ArrayContext, + mpi_communicator: Optional["mpi4py.MPI.Intracomm"], + volume_discrs: Mapping[VolumeTag, Discretization], + base_group_factory: ElementGroupFactory, + ) -> Mapping[Union[BTAG_MULTIVOL_PARTITION, BTAG_PARTITION], + DiscretizationConnection]: + + from meshmode.distributed import (get_inter_partition_tags, + make_remote_group_infos, InterRankBoundaryInfo, + MPIBoundaryCommSetupHelper) + + inter_part_tags = { + _check_btag_against_vtag(discr_vol_tag, tag) + for discr_vol_tag, volume_discr in volume_discrs.items() + for tag in get_inter_partition_tags(volume_discr.mesh)} + + inter_part_conns: Mapping[Union[BTAG_MULTIVOL_PARTITION, BTAG_PARTITION], + DiscretizationConnection] = {} + + if inter_part_tags: + local_boundary_restrictions = { + btag: make_face_restriction( + array_context, _volume_discr_for_btag(volume_discrs, btag), + base_group_factory, boundary_tag=btag) + for btag in inter_part_tags} + + irbis = [] - from meshmode.distributed import get_connected_partitions - connected_parts = get_connected_partitions(volume_discr.mesh) + for btag in inter_part_tags: + other_rank = _remote_rank_from_btag(btag) + btag_restr = local_boundary_restrictions[btag] - if connected_parts: - if mpi_communicator is None: - raise RuntimeError("must supply an MPI communicator when using a " - "distributed mesh") + if other_rank is None: + # {{{ rank-local interface between multiple volumes - grp_factory = get_group_factory_for_discretization_tag(DISCR_TAG_BASE) + from meshmode.discretization.connection import \ + make_partition_connection + remote_btag = _flip_btag(None, btag) + inter_part_conns[btag] = make_partition_connection( + array_context, + local_bdry_conn=btag_restr, + remote_bdry_discr=( + local_boundary_restrictions[remote_btag].to_discr), + remote_group_infos=make_remote_group_infos( + array_context, remote_btag, btag_restr)) + + # }}} + else: + # {{{ cross-rank interface + + if mpi_communicator is None: + raise RuntimeError("must supply an MPI communicator " + "when using a distributed mesh") + + irbis.append( + InterRankBoundaryInfo( + local_btag=btag, + remote_btag=_flip_btag(mpi_communicator.rank, btag), + remote_rank=other_rank, + local_boundary_connection=btag_restr)) + + # }}} - local_boundary_connections = {} - for i_remote_part in connected_parts: - local_boundary_connections[i_remote_part] = \ - get_connection_to_rank_boundary(i_remote_part) + if irbis: + assert mpi_communicator is not None - from meshmode.distributed import MPIBoundaryCommSetupHelper - with MPIBoundaryCommSetupHelper(mpi_communicator, array_context, - local_boundary_connections, grp_factory) as bdry_setup_helper: - while True: - conns = bdry_setup_helper.complete_some() - if not conns: - break - for i_remote_part, conn in conns.items(): - boundary_connections[i_remote_part] = conn + with MPIBoundaryCommSetupHelper(mpi_communicator, array_context, + irbis, base_group_factory) as bdry_setup_helper: + while True: + conns = bdry_setup_helper.complete_some() + if not conns: + # We're done. + break - return boundary_connections + inter_part_conns.update(conns) + + return inter_part_conns # }}} +# {{{ modal group factory + def _generate_modal_group_factory(nodal_group_factory): from meshmode.discretization.poly_element import ( ModalSimplexGroupFactory, @@ -756,4 +900,132 @@ def _generate_modal_group_factory(nodal_group_factory): f"Unknown mesh element group: {mesh_group_cls}" ) +# }}} + + +# {{{ 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 single_valued, is_single_valued + + 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 + + volume_discrs = { + vtag: ( + Discretization( + array_context, mesh_or_discr, + discr_tag_to_group_factory[DISCR_TAG_BASE]) + if isinstance(mesh_or_discr, Mesh) else mesh_or_discr) + for vtag, mesh_or_discr in volumes.items()} + + return DiscretizationCollection( + array_context=array_context, + volume_discrs=volume_discrs, + discr_tag_to_group_factory=discr_tag_to_group_factory, + inter_partition_connections=_set_up_inter_partition_connections( + array_context=array_context, + mpi_communicator=getattr( + array_context, "mpi_communicator", None), + volume_discrs=volume_discrs, + base_group_factory=discr_tag_to_group_factory[DISCR_TAG_BASE], + )) + +# }}} + + +# {{{ relabel_partitions + +def relabel_partitions(mesh: Mesh, + self_rank: int, self_volume_tag: VolumeTag, + part_nr_to_rank_and_vol_tag: Mapping[int, Tuple[int, VolumeTag]]) -> Mesh: + """Given a partitioned mesh (which includes :class:`meshmode.mesh.BTAG_PARTITION` + boundary tags), relabel those boundary tags into + :class:`grudge.dof_desc.BTAG_MULTIVOL_PARTITION` tags, which map each + of the incoming partitions onto a combination of rank and volume tag, + given by *part_nr_to_rank_and_vol_tag*. + """ + + def _new_btag(btag: BoundaryTag) -> BTAG_MULTIVOL_PARTITION: + if not isinstance(btag, BTAG_PARTITION): + raise TypeError("unexpected inter-partition boundary tags of type " + f"'{type(btag)}', expected BTAG_PARTITION") + + rank, vol_tag = part_nr_to_rank_and_vol_tag[btag.part_nr] + return BTAG_MULTIVOL_PARTITION( + other_rank=(None if rank == self_rank else rank), + self_volume_tag=self_volume_tag, + other_volume_tag=vol_tag) + + assert mesh.facial_adjacency_groups is not None + + from dataclasses import replace + return mesh.copy(facial_adjacency_groups=[ + [ + replace(fagrp, + boundary_tag=_new_btag(fagrp.boundary_tag)) + if isinstance(fagrp, InterPartitionAdjacencyGroup) + else fagrp + for fagrp in grp_fagrp_list] + for grp_fagrp_list in mesh.facial_adjacency_groups]) + +# }}} + + # vim: foldmethod=marker diff --git a/grudge/dof_desc.py b/grudge/dof_desc.py index aba344198..43b80ed8c 100644 --- a/grudge/dof_desc.py +++ b/grudge/dof_desc.py @@ -1,4 +1,54 @@ -"""Degree of freedom (DOF) descriptions""" +""" +Volume tags +----------- + +.. autoclass:: VolumeTag +.. autoclass:: VTAG_ALL + +:mod:`grudge`-specific boundary tags +------------------------------------ + +.. autoclass:: BTAG_MULTIVOL_PARTITION + +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:: DTAG_VOLUME +.. autoclass:: DTAG_BOUNDARY + +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 +""" __copyright__ = """ Copyright (C) 2008 Andreas Kloeckner @@ -25,100 +75,115 @@ THE SOFTWARE. """ -from typing import Hashable -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 +# {{{ volume tags -.. autoclass:: DOFDesc -.. autofunction:: as_dofdesc +class VTAG_ALL: # noqa: N801 + pass + + +VolumeTag = Hashable + +# }}} -.. data:: DD_SCALAR -.. data:: DD_VOLUME -.. data:: DD_VOLUME_MODAL -""" +# {{{ boundary tags -# {{{ DOF description +@dataclass(init=True, eq=True, frozen=True) +class BTAG_MULTIVOL_PARTITION: # noqa: N801 + """ + .. attribute:: other_rank + + An integer, or *None*. If *None*, this marks a partition boundary + to another volume on the same rank. + + .. attribute:: self_volume_tag + .. attribute:: other_volume_tag + """ + other_rank: Optional[int] + self_volume_tag: "VolumeTag" + other_volume_tag: "VolumeTag" + +# }}} -DiscretizationTag = Hashable +# {{{ domain tag class DTAG_SCALAR: # noqa: N801 """A domain tag denoting scalar values.""" -class DTAG_VOLUME_ALL: # noqa: N801 - """ - A domain tag denoting values defined - in all cell volumes. +@dataclass(frozen=True, eq=True, init=True) +class DTAG_VOLUME: # noqa: N801 + """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 + +DTAG_VOLUME_ALL = DTAG_VOLUME(VTAG_ALL) + +@dataclass(frozen=True, eq=True, init=True) 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`. + """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[Type[DTAG_SCALAR], DTAG_VOLUME, DTAG_BOUNDARY] - def __ne__(self, other): - return not self.__eq__(other) +# }}} + + +# {{{ discretization tag - def __hash__(self): - return hash(type(self)) ^ hash(self.tag) +class _DiscretizationTag: # noqa: N801 + pass - 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 @@ -132,20 +197,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. @@ -162,194 +229,181 @@ class DOFDesc: .. automethod:: uses_quadrature + .. automethod:: with_domain_tag .. automethod:: with_discr_tag - .. automethod:: with_dtag + .. automethod:: trace .. automethod:: __eq__ .. automethod:: __ne__ .. automethod:: __hash__ """ - 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) - - 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 domain_tag 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 + def __init__(self, domain_tag: Any, + discretization_tag: Optional[type[DiscretizationTag]] = None): - # 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 + if ( + not (domain_tag is DTAG_SCALAR + or isinstance(domain_tag, (DTAG_BOUNDARY, DTAG_VOLUME))) + 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) - def is_scalar(self): + domain_tag, discretization_tag = _normalize_domain_and_discr_tag( + domain_tag, discretization_tag) + + object.__setattr__(self, "domain_tag", domain_tag) + object.__setattr__(self, "discretization_tag", discretization_tag) + + def is_scalar(self) -> bool: return self.domain_tag is DTAG_SCALAR - 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_boundary_or_partition_interface(self): - return isinstance(self.domain_tag, DTAG_BOUNDARY) + def is_volume(self) -> bool: + return isinstance(self.domain_tag, DTAG_VOLUME) - 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, DTAG_BOUNDARY) + 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, DTAG_BOUNDARY) + + 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) - - return "DOFDesc({}, {})".format( - fmt(self.domain_tag), - fmt(self.discretization_tag)) - - -DD_SCALAR = DOFDesc(DTAG_SCALAR, None) - -DD_VOLUME = DOFDesc(DTAG_VOLUME_ALL, None) + 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, DTAG_VOLUME): + raise ValueError(f"must originate on volume, got '{self.domain_tag}'") + return replace(self, + domain_tag=DTAG_BOUNDARY(btag, volume_tag=self.domain_tag.tag)) + + def with_discr_tag(self, discr_tag) -> "DOFDesc": + return replace(self, discretization_tag=discr_tag) + + +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, + ) -> Tuple[DomainTag, DiscretizationTag]: + if domain in [DTAG_SCALAR, "scalar"]: + domain = DTAG_SCALAR + elif isinstance(domain, (DTAG_BOUNDARY, DTAG_VOLUME)): + pass + elif domain == "vol": + domain = DTAG_VOLUME_ALL + elif domain in [FACE_RESTR_ALL, "all_faces"]: + domain = DTAG_BOUNDARY(FACE_RESTR_ALL) + elif domain in [FACE_RESTR_INTERIOR, "int_faces"]: + domain = DTAG_BOUNDARY(FACE_RESTR_INTERIOR) + elif isinstance(domain, BTAG_PARTITION): + domain = DTAG_BOUNDARY(domain, VTAG_ALL) + elif isinstance(domain, BTAG_MULTIVOL_PARTITION): + domain = DTAG_BOUNDARY(domain, domain.self_volume_tag) + elif domain in [BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE]: + domain = DTAG_BOUNDARY(domain, VTAG_ALL) + 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 + + return domain, discretization_tag + + +def as_dofdesc( + domain: Any, discretization_tag: Optional[DiscretizationTag] = 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) -def as_dofdesc(dd): - if isinstance(dd, DOFDesc): - return dd - return DOFDesc(dd, discretization_tag=None) + return DOFDesc(domain, discretization_tag) # }}} # {{{ Deprecated tags -_deprecated_name_to_new_name = {"QTAG_NONE": "DISCR_TAG_BASE", - "QTAG_MODAL": "DISCR_TAG_MODAL"} +_deprecated_name_to_new_name = {"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 b41c683cd..8bc9214af 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -43,12 +43,14 @@ """ +from typing import Optional, Sequence import numpy as np from arraycontext import ArrayContext, thaw, freeze, DeviceScalar from meshmode.transform_metadata import FirstAxisIsElementsTag -from grudge.dof_desc import DD_VOLUME, DOFDesc, as_dofdesc +from grudge.dof_desc import ( + DD_VOLUME_ALL, DOFDesc, as_dofdesc, DTAG_BOUNDARY, FACE_RESTR_ALL) from grudge.discretization import DiscretizationCollection import grudge.op as op @@ -59,7 +61,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 @@ -98,8 +101,8 @@ def _compute_characteristic_lengthscales(): # corresponding group non-geometric factor cng * geo_facts for cng, geo_facts in zip( - dt_non_geometric_factors(dcoll), - thaw(dt_geometric_factors(dcoll), actx) + dt_non_geometric_factors(dcoll, dd), + thaw(dt_geometric_factors(dcoll, dd), actx) ) ) ) @@ -109,7 +112,8 @@ def _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: @@ -126,7 +130,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 = [] @@ -158,7 +162,8 @@ def dt_non_geometric_factors( @memoize_on_first_arg def h_max_from_volume( - dcoll: DiscretizationCollection, dim=None, dd=None) -> "DeviceScalar": + dcoll: DiscretizationCollection, dim=None, + dd: Optional[DOFDesc] = None) -> "DeviceScalar": """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. @@ -173,7 +178,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: @@ -189,7 +194,8 @@ def h_max_from_volume( @memoize_on_first_arg def h_min_from_volume( - dcoll: DiscretizationCollection, dim=None, dd=None) -> "DeviceScalar": + dcoll: DiscretizationCollection, dim=None, + dd: Optional[DOFDesc] = None) -> "DeviceScalar": """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. @@ -204,7 +210,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: @@ -219,7 +225,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). @@ -242,7 +248,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) @@ -269,7 +275,8 @@ def dt_geometric_factors( # Inscribed "circle" radius is half the cell size return freeze(cell_vols/2) - dd_face = DOFDesc("all_faces", dd.discretization_tag) + dd_face = dd.with_domain_tag( + DTAG_BOUNDARY(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..400fee355 100644 --- a/grudge/eager.py +++ b/grudge/eager.py @@ -87,7 +87,6 @@ def nodal_max(self, dd, vec): return op.nodal_max(self, dd, vec) -connected_ranks = op.connected_ranks interior_trace_pair = op.interior_trace_pair cross_rank_trace_pairs = op.cross_rank_trace_pairs diff --git a/grudge/execution.py b/grudge/execution.py index 73b5e8948..7669ccb8a 100644 --- a/grudge/execution.py +++ b/grudge/execution.py @@ -37,6 +37,8 @@ import grudge.symbolic.mappers as mappers from grudge import sym from grudge.function_registry import base_function_registry +from grudge.dof_desc import \ + DD_VOLUME_ALL, DTAG_BOUNDARY, FACE_RESTR_INTERIOR, VTAG_ALL, DOFDesc import logging logger = logging.getLogger(__name__) @@ -323,7 +325,9 @@ def map_opposite_partition_face_swap(self, op, field_expr): return bdry_conn(remote_bdry_vec) def map_opposite_interior_face_swap(self, op, field_expr): - return self.dcoll.opposite_face_connection()(self.rec(field_expr)) + return self.dcoll.opposite_face_connection( + DTAG_BOUNDARY(FACE_RESTR_INTERIOR, VTAG_ALL) + )(self.rec(field_expr)) # }}} @@ -382,8 +386,6 @@ def prg(): return result def map_signed_face_ones(self, expr): - from grudge.dof_desc import DOFDesc, DD_VOLUME - assert expr.dd.is_trace() face_discr = self.dcoll.discr_from_dd(expr.dd) assert face_discr.dim == 0 @@ -391,7 +393,7 @@ def map_signed_face_ones(self, expr): # NOTE: ignore quadrature_tags on expr.dd, since we only care about # the face_id here all_faces_conn = self.dcoll.connection_from_dds( - DD_VOLUME, + DD_VOLUME_ALL, DOFDesc(expr.dd.domain_tag)) field = face_discr.empty(self.array_context, dtype=self.dcoll.real_dtype) @@ -682,7 +684,9 @@ def dumper(name, sym_operator): dumper("before-bind", sym_operator) sym_operator = mappers.OperatorBinder()(sym_operator) - mappers.ErrorChecker(dcoll.mesh)(sym_operator) + # the old and busted symbolic stuff will not learn to use multi-volume + mesh = dcoll.discr_from_dd(DD_VOLUME_ALL).mesh + mappers.ErrorChecker(mesh)(sym_operator) sym_operator = \ mappers.OppositeInteriorFaceSwapUniqueIDAssigner()(sym_operator) @@ -713,14 +717,14 @@ def dumper(name, sym_operator): sym_operator = post_bind_mapper(sym_operator) dumper("before-empty-flux-killer", sym_operator) - sym_operator = mappers.EmptyFluxKiller(dcoll.mesh)(sym_operator) + sym_operator = mappers.EmptyFluxKiller(mesh)(sym_operator) dumper("before-cfold", sym_operator) sym_operator = mappers.CommutativeConstantFoldingMapper()(sym_operator) dumper("before-qcheck", sym_operator) sym_operator = mappers.QuadratureCheckerAndRemover( - dcoll.discr_tag_to_group_factory)(sym_operator) + dcoll._discr_tag_to_group_factory)(sym_operator) # Work around https://github.com/numpy/numpy/issues/9438 # diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py index 492d8b7f7..56226fc3c 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 thaw, freeze, ArrayContext @@ -67,7 +68,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 pymbolic.geometric_algebra import MultiVector @@ -110,7 +111,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}}` @@ -145,7 +147,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) @@ -177,8 +179,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. @@ -202,7 +206,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. @@ -225,7 +231,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 @@ -252,7 +259,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 @@ -266,7 +273,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:: @@ -292,7 +300,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) @@ -301,7 +309,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. @@ -315,7 +324,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 @@ -331,7 +340,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: @@ -355,7 +365,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 @@ -382,7 +392,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 @@ -441,7 +452,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 @@ -463,7 +474,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: @@ -483,7 +494,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 @@ -527,7 +539,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() @@ -535,7 +547,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) ) signed_ones = dcoll.discr_from_dd(dd.with_discr_tag(DISCR_TAG_BASE)).zeros( actx, dtype=dcoll.real_dtype @@ -556,7 +568,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*. @@ -569,7 +581,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: @@ -590,8 +602,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*. @@ -603,7 +617,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, @@ -611,7 +625,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 @@ -627,7 +642,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(): @@ -646,7 +661,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 @@ -672,7 +688,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`. @@ -728,10 +744,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)) ) @@ -749,7 +765,7 @@ def _normal(): return thaw(n, actx) -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 @@ -779,8 +795,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:: @@ -798,7 +814,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) @@ -827,7 +843,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: @@ -852,7 +868,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:: @@ -869,7 +885,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 aacd2fcfd..58424e757 100644 --- a/grudge/models/advection.py +++ b/grudge/models/advection.py @@ -216,13 +216,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 @@ -236,7 +236,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, @@ -327,9 +327,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, @@ -337,7 +337,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) @@ -349,7 +349,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 70bd1caf7..3ee3f7218 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -66,6 +66,7 @@ import numpy as np import grudge.dof_desc as dof_desc +from grudge.dof_desc import DD_VOLUME_ALL from grudge.interpolation import interp # noqa: F401 from grudge.projection import project # noqa: F401 @@ -88,7 +89,6 @@ from grudge.trace_pair import ( # noqa: F401 interior_trace_pair, interior_trace_pairs, - connected_ranks, cross_rank_trace_pairs, bdry_trace_pair, bv_trace_pair @@ -162,7 +162,7 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, def _div_helper(dcoll, diff_func, *args): if len(args) == 1: vecs, = args - dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd = DD_VOLUME_ALL elif len(args) == 2: dd, vecs = args else: @@ -204,7 +204,7 @@ def _div_helper(dcoll, diff_func, *args): def _grad_helper(dcoll, scalar_grad, *args, nested): if len(args) == 1: vec, = args - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = dof_desc.DD_VOLUME_ALL elif len(args) == 2: dd_in, vec = args else: @@ -265,11 +265,11 @@ 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 dd_in == dof_desc.as_dofdesc(DD_VOLUME_ALL) from grudge.geometry import inverse_surface_metric_derivative_mat - discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + discr = dcoll.discr_from_dd(DD_VOLUME_ALL) actx = vec.array_context inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, @@ -320,7 +320,7 @@ def local_d_dx( if not isinstance(vec, DOFArray): return map_array_container(partial(local_d_dx, dcoll, xyz_axis), vec) - discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + discr = dcoll.discr_from_dd(DD_VOLUME_ALL) actx = vec.array_context from grudge.geometry import inverse_surface_metric_derivative_mat @@ -408,7 +408,7 @@ def _weak_scalar_grad(dcoll, dd_in, vec): from grudge.geometry import inverse_surface_metric_derivative_mat 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_VOLUME_ALL) actx = vec.array_context inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, dd=dd_in, @@ -478,7 +478,7 @@ def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT """ 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: @@ -493,7 +493,7 @@ def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT from grudge.geometry import inverse_surface_metric_derivative_mat 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_VOLUME_ALL) actx = vec.array_context inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, dd=dd_in, @@ -646,13 +646,13 @@ def mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: 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: raise TypeError("invalid number of arguments") - return _apply_mass_operator(dcoll, dof_desc.DD_VOLUME, dd, vec) + return _apply_mass_operator(dcoll, DD_VOLUME_ALL, dd, vec) # }}} @@ -755,7 +755,7 @@ def inverse_mass(dcoll: DiscretizationCollection, vec) -> ArrayOrContainerT: """ return _apply_inverse_mass_operator( - dcoll, dof_desc.DD_VOLUME, dof_desc.DD_VOLUME, vec + dcoll, DD_VOLUME_ALL, DD_VOLUME_ALL, vec ) # }}} @@ -857,7 +857,7 @@ def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec): from grudge.geometry import area_element - volm_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + volm_discr = dcoll.discr_from_dd(DD_VOLUME_ALL) face_discr = dcoll.discr_from_dd(dd) dtype = vec.entry_dtype actx = vec.array_context diff --git a/grudge/reductions.py b/grudge/reductions.py index a34746103..229ca529e 100644 --- a/grudge/reductions.py +++ b/grudge/reductions.py @@ -92,7 +92,7 @@ def norm(dcoll: DiscretizationCollection, vec, p, dd=None) -> "DeviceScalar": :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) @@ -126,7 +126,7 @@ def nodal_sum(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar": 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 @@ -169,7 +169,7 @@ def nodal_min(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar": if comm is None: return nodal_min_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 actx = vec.array_context @@ -213,7 +213,7 @@ def nodal_max(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar": if comm is None: return nodal_max_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 actx = vec.array_context @@ -290,7 +290,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: @@ -456,7 +456,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/symbolic/operators.py b/grudge/symbolic/operators.py index 031ec60e0..c8ae21d6d 100644 --- a/grudge/symbolic/operators.py +++ b/grudge/symbolic/operators.py @@ -25,6 +25,8 @@ import numpy as np import pymbolic.primitives +from grudge.dof_desc import DTAG_BOUNDARY + from typing import Tuple __doc__ = """ @@ -567,7 +569,9 @@ def __init__(self, dd_in=None, dd_out=None, unique_id=None): dd_out = dd_in super().__init__(dd_in, dd_out) - if self.dd_in.domain_tag is not FACE_RESTR_INTERIOR: + if not isinstance(self.dd_in.domain_tag, DTAG_BOUNDARY): + raise ValueError("dd_in must be a boundary") + if self.dd_in.domain_tag.tag is not FACE_RESTR_INTERIOR: raise ValueError("dd_in must be an interior faces domain") if self.dd_out != self.dd_in: raise ValueError("dd_out and dd_in must be identical") @@ -641,7 +645,9 @@ def __init__(self, dd_in=None, dd_out=None): if not dd_out.is_volume(): raise ValueError("dd_out must be a volume domain") - if dd_in.domain_tag is not FACE_RESTR_ALL: + if not isinstance(dd_in.domain_tag, DTAG_BOUNDARY): + raise ValueError("dd_in must be a boundary") + if dd_in.domain_tag.tag is not FACE_RESTR_ALL: raise ValueError("dd_in must be an interior faces domain") super().__init__(dd_in, dd_out) diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py index d95bc657b..4494c5bd5 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -683,12 +683,12 @@ def int_tpair(expression, qtag=None, from_dd=None): if from_dd.domain_tag == trace_dd.domain_tag: i = expression else: - i = project(from_dd, trace_dd.with_qtag(None))(expression) + i = project(from_dd, trace_dd.with_discr_tag(None))(expression) e = cse(OppositeInteriorFaceSwap()(i)) if trace_dd.uses_quadrature(): - i = cse(project(trace_dd.with_qtag(None), trace_dd)(i)) - e = cse(project(trace_dd.with_qtag(None), trace_dd)(e)) + i = cse(project(trace_dd.with_discr_tag(None), trace_dd)(i)) + e = cse(project(trace_dd.with_discr_tag(None), trace_dd)(e)) return TracePair(trace_dd, interior=i, exterior=e) diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 872028d74..1ede3efa0 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -7,18 +7,21 @@ .. autoclass:: TracePair -Boundary trace functions ------------------------- +Boundary traces +--------------- .. autofunction:: bdry_trace_pair .. autofunction:: bv_trace_pair -Interior and cross-rank trace functions ---------------------------------------- +Interior, cross-rank, and inter-volume traces +--------------------------------------------- .. autofunction:: interior_trace_pairs .. autofunction:: local_interior_trace_pair +.. autofunction:: inter_volume_trace_pairs +.. autofunction:: local_inter_volume_trace_pairs .. autofunction:: cross_rank_trace_pairs +.. autofunction:: cross_rank_inter_volume_trace_pairs """ __copyright__ = """ @@ -46,18 +49,19 @@ """ -from typing import List, Hashable, Optional, Type, Any +from typing import List, Hashable, Optional, Type, Any, Sequence, Union from pytools.persistent_dict import KeyBuilder from arraycontext import ( - ArrayContainer, - with_container_arithmetic, - dataclass_array_container, - get_container_context_recursively, - flatten, to_numpy, - unflatten, from_numpy -) + ArrayContainer, + ArrayContext, + with_container_arithmetic, + dataclass_array_container, + get_container_context_recursively, + flatten, to_numpy, + unflatten, from_numpy, + flat_size_and_dtype) from arraycontext.container import ArrayOrContainerT from dataclasses import dataclass @@ -67,13 +71,18 @@ from pytools import memoize_on_first_arg from pytools.obj_array import obj_array_vectorize -from grudge.discretization import DiscretizationCollection +from grudge.discretization import DiscretizationCollection, _remote_rank_from_btag from grudge.projection import project 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, + VolumeTag, DTAG_VOLUME, DTAG_BOUNDARY, BTAG_MULTIVOL_PARTITION, + ) # {{{ trace pair container class @@ -107,11 +116,13 @@ class TracePair: and the current interfaces, with symbolic information or concrete data. """ - dd: dof_desc.DOFDesc + dd: DOFDesc interior: ArrayContainer exterior: ArrayContainer - def __init__(self, dd, *, interior, exterior): + def __init__(self, dd: DOFDesc, *, + interior: ArrayOrContainerT, + exterior: ArrayOrContainerT): object.__setattr__(self, "dd", dof_desc.as_dofdesc(dd)) object.__setattr__(self, "interior", interior) object.__setattr__(self, "exterior", exterior) @@ -232,12 +243,14 @@ def bv_trace_pair( # {{{ 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.container.ArrayContainer` of them. @@ -250,17 +263,24 @@ 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, DTAG_VOLUME) + trace_dd = volume_dd.trace(FACE_RESTR_INTERIOR) + + interior = project(dcoll, volume_dd, trace_dd, vec) - def get_opposite_face(el): + def get_opposite_trace(el): if isinstance(el, Number): return el else: - return dcoll.opposite_face_connection()(el) + assert isinstance(trace_dd.domain_tag, DTAG_BOUNDARY) + return dcoll.opposite_face_connection(trace_dd.domain_tag)(el) - e = obj_array_vectorize(get_opposite_face, i) + e = obj_array_vectorize(get_opposite_trace, interior) - return TracePair("int_faces", interior=i, exterior=e) + return TracePair(trace_dd, interior=interior, exterior=e) def interior_trace_pair(dcoll: DiscretizationCollection, vec) -> TracePair: @@ -274,7 +294,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. @@ -302,147 +323,298 @@ 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, volume_dd=volume_dd)] + + cross_rank_trace_pairs( + dcoll, vec, comm_tag=comm_tag, volume_dd=volume_dd) + ) + +# }}} + + +# {{{ inter-volume trace pairs + +def local_inter_volume_trace_pairs( + dcoll: DiscretizationCollection, + self_volume_dd: DOFDesc, self_ary: ArrayOrContainerT, + other_volume_dd: DOFDesc, other_ary: ArrayOrContainerT, + ) -> ArrayOrContainerT: + if not isinstance(self_volume_dd.domain_tag, DTAG_VOLUME): + raise ValueError("self_volume_dd must describe a volume") + if not isinstance(other_volume_dd.domain_tag, DTAG_VOLUME): + raise ValueError("other_volume_dd must describe a volume") + if self_volume_dd.discretization_tag != DISCR_TAG_BASE: + raise TypeError( + f"expected a base-discretized self DOFDesc, got '{self_volume_dd}'") + if other_volume_dd.discretization_tag != DISCR_TAG_BASE: + raise TypeError( + f"expected a base-discretized other DOFDesc, got '{other_volume_dd}'") + + self_btag = BTAG_MULTIVOL_PARTITION( + other_rank=None, + self_volume_tag=self_volume_dd.domain_tag.tag, + other_volume_tag=other_volume_dd.domain_tag.tag) + other_btag = BTAG_MULTIVOL_PARTITION( + other_rank=None, + self_volume_tag=other_volume_dd.domain_tag.tag, + other_volume_tag=self_volume_dd.domain_tag.tag) + + self_trace_dd = self_volume_dd.trace(self_btag) + other_trace_dd = other_volume_dd.trace(other_btag) + + # FIXME: In all likelihood, these traces will be reevaluated from + # the other side, which is hard to prevent given the interface we + # have. Lazy eval will hopefully collapse those redundant evaluations... + self_trace = project( + dcoll, self_volume_dd, self_trace_dd, self_ary) + other_trace = project( + dcoll, other_volume_dd, other_trace_dd, other_ary) + + other_to_self = dcoll._inter_partition_connections[self_btag] + + def get_opposite_trace(el): + if isinstance(el, Number): + return el + else: + return other_to_self(el) + + return TracePair( + self_trace_dd, + interior=self_trace, + exterior=obj_array_vectorize(get_opposite_trace, other_trace)) + + +def inter_volume_trace_pairs(dcoll: DiscretizationCollection, + self_volume_dd: DOFDesc, self_ary: ArrayOrContainerT, + other_volume_dd: DOFDesc, other_ary: ArrayOrContainerT, + comm_tag: Hashable = None) -> List[ArrayOrContainerT]: + """ + Note that :func:`local_inter_volume_trace_pairs` provides the rank-local + contributions if those are needed in isolation. Similarly, + :func:`cross_rank_inter_volume_trace_pairs` provides only the trace pairs + defined on cross-rank boundaries. + """ + # TODO documentation + return ( - [local_interior_trace_pair(dcoll, vec)] - + cross_rank_trace_pairs(dcoll, vec, comm_tag=comm_tag) + [local_inter_volume_trace_pairs(dcoll, + self_volume_dd, self_ary, other_volume_dd, other_ary)] + + cross_rank_inter_volume_trace_pairs(dcoll, + self_volume_dd, self_ary, other_volume_dd, other_ary, + comm_tag=comm_tag) ) # }}} -# {{{ 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): - from meshmode.distributed import get_connected_partitions - return get_connected_partitions(dcoll._volume_discr.mesh) +def _remote_inter_partition_tags( + dcoll: DiscretizationCollection, + self_volume_tag: VolumeTag, + other_volume_tag: Optional[VolumeTag] = None + ) -> Sequence[Union[BTAG_MULTIVOL_PARTITION, BTAG_PARTITION]]: + if other_volume_tag is None: + other_volume_tag = self_volume_tag + + same_vol_tag = (self_volume_tag == other_volume_tag) + + result = [] + for tag in dcoll._inter_partition_connections: + if isinstance(tag, BTAG_PARTITION): + if not same_vol_tag: + raise AssertionError("encountered BTAG_PARTITION in a situation " + "with same_vol_tag != other_volume_tag") + result.append(tag) + + elif isinstance(tag, BTAG_MULTIVOL_PARTITION): + if (tag.other_rank is not None + and tag.self_volume_tag == self_volume_tag + and tag.other_volume_tag == other_volume_tag): + result.append(tag) + + else: + raise AssertionError("unexpected inter-partition tag type encountered: " + f"'{tag}'") + + return result + + +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 + 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.") + + return num_tag + +# }}} + + +# {{{ eager rank-boundary communication -class _RankBoundaryCommunication: +class _RankBoundaryCommunicationEager: base_comm_tag = 1273 def __init__(self, - dcoll: DiscretizationCollection, - array_container: ArrayOrContainerT, - remote_rank, comm_tag: Optional[int] = None): - actx = get_container_context_recursively(array_container) - btag = BTAG_PARTITION(remote_rank) + actx: ArrayContext, + dcoll: DiscretizationCollection, + btag: Union[BTAG_MULTIVOL_PARTITION, BTAG_PARTITION], + *, + local_bdry_data: ArrayOrContainerT, + send_data: ArrayOrContainerT, + comm_tag: Optional[Hashable] = None): - local_bdry_data = project(dcoll, "vol", btag, array_container) comm = dcoll.mpi_communicator + assert comm is not None + + remote_rank = _remote_rank_from_btag(btag) + assert remote_rank is not None self.dcoll = dcoll self.array_context = actx - self.remote_btag = btag + self.btag = btag self.bdry_discr = dcoll.discr_from_dd(btag) 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) - # and comm.Irecv (MPI_Irecv) respectively. These initiate non-blocking - # point-to-point communication requests and require explicit management - # via the use of wait (MPI_Wait, MPI_Waitall, MPI_Waitany, MPI_Waitsome), - # test (MPI_Test, MPI_Testall, MPI_Testany, MPI_Testsome), and cancel - # (MPI_Cancel). The rank-local data `self.local_bdry_data_np` will have its - # associated memory buffer sent across connected ranks and must not be - # modified at the Python level during this process. Completion of the - # requests is handled in :meth:`finish`. - # - # For more details on the mpi4py semantics, see: - # https://mpi4py.readthedocs.io/en/stable/overview.html#nonblocking-communications - # # NOTE: mpi4py currently (2021-11-03) holds a reference to the send # memory buffer for (i.e. `self.local_bdry_data_np`) until the send # requests is complete, however it is not clear that this is documented # behavior. We hold on to the buffer (via the instance attribute) # as well, just in case. - self.send_req = comm.Isend(self.local_bdry_data_np, - remote_rank, - tag=self.comm_tag) - self.remote_data_host_numpy = np.empty_like(self.local_bdry_data_np) - self.recv_req = comm.Irecv(self.remote_data_host_numpy, + self.send_data_np = to_numpy(flatten(send_data, actx), actx) + self.send_req = comm.Isend(self.send_data_np, remote_rank, tag=self.comm_tag) + recv_size, recv_dtype = flat_size_and_dtype(local_bdry_data) + self.recv_data_np = np.empty(recv_size, recv_dtype) + self.recv_req = comm.Irecv(self.recv_data_np, remote_rank, tag=self.comm_tag) + def finish(self): # Wait for the nonblocking receive request to complete before # accessing the data self.recv_req.Wait() - # Nonblocking receive is complete, we can now access the data and apply - # the boundary-swap connection - actx = self.array_context - remote_bdry_data_flat = from_numpy(self.remote_data_host_numpy, actx) - 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))) - swapped_remote_bdry_data = bdry_conn(remote_bdry_data) + recv_data_flat = from_numpy( + self.recv_data_np, self.array_context) + unswapped_remote_bdry_data = unflatten(self.local_bdry_data, + recv_data_flat, self.array_context) + bdry_conn = self.dcoll._inter_partition_connections[ + self.btag] + remote_bdry_data = bdry_conn(unswapped_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.btag, interior=self.local_bdry_data, - exterior=swapped_remote_bdry_data) + exterior=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: ArrayOrContainerT, - remote_rank: int, comm_tag: Hashable): + actx: ArrayContext, + dcoll: DiscretizationCollection, + btag: Union[BTAG_MULTIVOL_PARTITION, BTAG_PARTITION], + *, + local_bdry_data: ArrayOrContainerT, + send_data: ArrayOrContainerT, + comm_tag: Optional[Hashable] = None) -> None: + if comm_tag is None: - raise ValueError("lazy communication requires 'tag' to be supplied") + raise ValueError("lazy communication requires 'comm_tag' to be supplied") 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.array_context = actx + self.bdry_discr = dcoll.discr_from_dd(btag) + self.btag = btag + + remote_rank = _remote_rank_from_btag(btag) + assert remote_rank is not None + + self.local_bdry_data = local_bdry_data + + from arraycontext.container.traversal import rec_keyed_map_array_container + + key_to_send_subary = {} + + def store_send_subary(key, send_subary): + key_to_send_subary[key] = send_subary + return send_subary + rec_keyed_map_array_container(store_send_subary, send_data) - self.local_bdry_data = project( - dcoll, "vol", self.remote_btag, array_container) + from pytato import make_distributed_recv, staple_distributed_send - def communicate_single_array(key, local_bdry_ary): + def communicate_single_array(key, local_bdry_subary): ary_tag = (comm_tag, key) return staple_distributed_send( - local_bdry_ary, dest_rank=remote_rank, comm_tag=ary_tag, + key_to_send_subary[key], dest_rank=remote_rank, comm_tag=ary_tag, stapled_to=make_distributed_recv( src_rank=remote_rank, comm_tag=ary_tag, - shape=local_bdry_ary.shape, dtype=local_bdry_ary.dtype)) + shape=local_bdry_subary.shape, + dtype=local_bdry_subary.dtype)) - from arraycontext.container.traversal import rec_keyed_map_array_container self.remote_data = rec_keyed_map_array_container( communicate_single_array, self.local_bdry_data) def finish(self): - bdry_conn = self.dcoll.distributed_boundary_swap_connection( - dof_desc.as_dofdesc(dof_desc.DTAG_BOUNDARY(self.remote_btag))) + bdry_conn = self.dcoll._inter_partition_connections[self.btag] - return TracePair(self.remote_btag, + return TracePair(self.btag, 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: ArrayOrContainerT, + *, comm_tag: Hashable = None, + 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 @@ -471,6 +643,16 @@ 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, DTAG_VOLUME): + 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. " @@ -481,50 +663,128 @@ def cross_rank_trace_pairs( comm_tag = tag del tag + # }}} + + comm_btags = _remote_inter_partition_tags( + dcoll, self_volume_tag=volume_dd.domain_tag.tag) + + # This asserts that there is only one data exchange per rank, so that + # there is no risk of mismatched data reaching the wrong recipient. + # (Since we have only a single tag.) + assert len(comm_btags) == len( + {_remote_rank_from_btag(btag) for btag in comm_btags}) + 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)] + # NOTE: Assumes that the same number is passed on every rank + return [TracePair(DOFDesc(btag, DISCR_TAG_BASE), interior=ary, exterior=ary) + for btag in comm_btags] actx = get_container_context_recursively(ary) + assert actx is not None from grudge.array_context import MPIPytatoArrayContextBase if isinstance(actx, MPIPytatoArrayContextBase): rbc = _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 - - # 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) - ] - - # Complete send/receives and return communicated data + rbc = _RankBoundaryCommunicationEager + + def start_comm(btag): + if isinstance(btag, BTAG_PARTITION): + volume_dd = DD_VOLUME_ALL + elif isinstance(btag, BTAG_MULTIVOL_PARTITION): + volume_dd = DOFDesc(DTAG_VOLUME(btag.self_volume_tag), DISCR_TAG_BASE) + else: + raise AssertionError() + + local_bdry_data = project(dcoll, volume_dd, btag, ary) + + return rbc(actx, dcoll, btag, + local_bdry_data=local_bdry_data, + send_data=local_bdry_data, + comm_tag=comm_tag) + + rank_bdry_communcators = [start_comm(btag) for btag in comm_btags] + return [rc.finish() for rc in rank_bdry_communcators] + +# }}} + + +# {{{ cross_rank_inter_volume_trace_pairs + +def cross_rank_inter_volume_trace_pairs( + dcoll: DiscretizationCollection, + self_volume_dd: DOFDesc, self_ary: ArrayOrContainerT, + other_volume_dd: DOFDesc, other_ary: ArrayOrContainerT, + *, comm_tag: Hashable = None, + ) -> List[TracePair]: + # FIXME: Should this interface take in boundary data instead? + # TODO: Docs + r"""Get a :class:`list` of *ary* trace pairs for each partition boundary. + + :arg comm_tag: a hashable object used to match sent and received data + across ranks. Communication will only match if both endpoints specify + objects that compare equal. A generalization of MPI communication + tags to arbitary, potentially composite objects. + + :returns: a :class:`list` of :class:`TracePair` objects. + """ + # {{{ process arguments + + if not isinstance(self_volume_dd.domain_tag, DTAG_VOLUME): + raise ValueError("self_volume_dd must describe a volume") + if not isinstance(other_volume_dd.domain_tag, DTAG_VOLUME): + raise ValueError("other_volume_dd must describe a volume") + if self_volume_dd.discretization_tag != DISCR_TAG_BASE: + raise TypeError( + f"expected a base-discretized self DOFDesc, got '{self_volume_dd}'") + if other_volume_dd.discretization_tag != DISCR_TAG_BASE: + raise TypeError( + f"expected a base-discretized other DOFDesc, got '{other_volume_dd}'") + + # }}} + + comm_btags = _remote_inter_partition_tags( + dcoll, + self_volume_tag=self_volume_dd.domain_tag.tag, + other_volume_tag=other_volume_dd.domain_tag.tag) + + # This asserts that there is only one data exchange per rank, so that + # there is no risk of mismatched data reaching the wrong recipient. + # (Since we have only a single tag.) + assert len(comm_btags) == len( + {_remote_rank_from_btag(btag) for btag in comm_btags}) + + actx = get_container_context_recursively(self_ary) + assert actx is not None + + from grudge.array_context import MPIPytatoArrayContextBase + + if isinstance(actx, MPIPytatoArrayContextBase): + rbc = _RankBoundaryCommunicationLazy + else: + rbc = _RankBoundaryCommunicationEager + + def start_comm(btag): + assert isinstance(btag, BTAG_MULTIVOL_PARTITION) + self_volume_dd = DOFDesc( + DTAG_VOLUME(btag.self_volume_tag), DISCR_TAG_BASE) + other_volume_dd = DOFDesc( + DTAG_VOLUME(btag.other_volume_tag), DISCR_TAG_BASE) + + local_bdry_data = project(dcoll, self_volume_dd, btag, self_ary) + send_data = project(dcoll, other_volume_dd, + BTAG_MULTIVOL_PARTITION( + other_rank=btag.other_rank, + self_volume_tag=btag.other_volume_tag, + other_volume_tag=btag.self_volume_tag), other_ary) + + return rbc(actx, dcoll, btag, + local_bdry_data=local_bdry_data, + send_data=send_data, + comm_tag=comm_tag) + + rank_bdry_communcators = [start_comm(btag) for btag in comm_btags] return [rc.finish() for rc in rank_bdry_communcators] # }}} diff --git a/requirements.txt b/requirements.txt index 2107e5aeb..6d8841e9a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,7 +10,7 @@ git+https://github.com/inducer/leap.git#egg=leap git+https://github.com/inducer/meshpy.git#egg=meshpy git+https://github.com/inducer/modepy.git#egg=modepy git+https://github.com/inducer/arraycontext.git#egg=arraycontext -git+https://github.com/inducer/meshmode.git#egg=meshmode +git+https://github.com/inducer/meshmode.git@generic-part-bdry#egg=meshmode git+https://github.com/inducer/pyvisfile.git#egg=pyvisfile git+https://github.com/inducer/pymetis.git#egg=pymetis git+https://github.com/illinois-ceesd/logpyle.git#egg=logpyle diff --git a/test/test_grudge.py b/test/test_grudge.py index 06f0710cb..32ea66322 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -40,7 +40,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 @@ -343,7 +343,9 @@ def test_face_normal_surface(actx_factory, mesh_name): surf_normal = surf_normal / actx.np.sqrt(sum(surf_normal**2)) face_normal_i = thaw(dcoll.normal(df), actx) - face_normal_e = dcoll.opposite_face_connection()(face_normal_i) + face_normal_e = dcoll.opposite_face_connection( + dof_desc.DTAG_BOUNDARY(dof_desc.FACE_RESTR_INTERIOR, dof_desc.VTAG_ALL) + )(face_normal_i) if mesh.ambient_dim == 3: from grudge.geometry import pseudoscalar, area_element @@ -978,6 +980,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() @@ -1041,6 +1045,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 @@ -1060,6 +1068,39 @@ def test_empty_boundary(actx_factory): assert isinstance(component, DOFArray) assert len(component) == len(dcoll.discr_from_dd(BTAG_NONE).groups) +# }}} + + +# {{{ multi-volume + +def test_multi_volume(actx_factory): + dim = 2 + actx = actx_factory() + + mesh = mgen.generate_regular_rect_mesh( + a=(-0.5,)*dim, b=(0.5,)*dim, + nelements_per_axis=(8,)*dim, order=4) + + meg, = mesh.groups + part_per_element = ( + mesh.vertices[0, meg.vertex_indices[:, 0]] > 0).astype(np.int32) + + from meshmode.mesh.processing import partition_mesh + from grudge.discretization import relabel_partitions + parts = { + i: relabel_partitions( + partition_mesh(mesh, part_per_element, i)[0], + self_rank=0, self_volume_tag=i, + part_nr_to_rank_and_vol_tag={ + 0: (0, 0), + 1: (0, 1), + }) + for i in range(2)} + + make_discretization_collection(actx, parts, order=4) + +# }}} + # You can test individual routines by typing # $ python test_grudge.py 'test_routine()' diff --git a/test/test_grudge_sym_old.py b/test/test_grudge_sym_old.py index 557351dcb..aa3bbb332 100644 --- a/test/test_grudge_sym_old.py +++ b/test/test_grudge_sym_old.py @@ -132,7 +132,7 @@ def _get_variables_on(dd): return sym_f, sym_x, sym_ones - sym_f, sym_x, sym_ones = _get_variables_on(dof_desc.DD_VOLUME) + sym_f, sym_x, sym_ones = _get_variables_on(dof_desc.DD_VOLUME_ALL) f_volm = actx.to_numpy(flatten(bind(discr, sym.cos(sym_x[0])**2)(actx))) ones_volm = actx.to_numpy(flatten(bind(discr, sym_ones)(actx))) @@ -140,7 +140,7 @@ def _get_variables_on(dd): f_quad = bind(discr, sym.cos(sym_x[0])**2)(actx) ones_quad = bind(discr, sym_ones)(actx) - mass_op = bind(discr, sym.MassOperator(dd_quad, dof_desc.DD_VOLUME)(sym_f)) + mass_op = bind(discr, sym.MassOperator(dd_quad, dof_desc.DD_VOLUME_ALL)(sym_f)) num_integral_1 = np.dot(ones_volm, actx.to_numpy(flatten(mass_op(f=f_quad)))) err_1 = abs(num_integral_1 - true_integral) @@ -227,14 +227,14 @@ def test_mass_surface_area(actx_factory, name): for resolution in builder.resolutions: mesh = builder.get_mesh(resolution, builder.mesh_order) discr = DiscretizationCollection(actx, mesh, order=builder.order) - volume_discr = discr.discr_from_dd(dof_desc.DD_VOLUME) + volume_discr = discr.discr_from_dd(dof_desc.DD_VOLUME_ALL) logger.info("ndofs: %d", volume_discr.ndofs) logger.info("nelements: %d", volume_discr.mesh.nelements) # {{{ compute surface area - dd = dof_desc.DD_VOLUME + dd = dof_desc.DD_VOLUME_ALL sym_op = sym.NodalSum(dd)(sym.MassOperator(dd, dd)(sym.Ones(dd))) approx_surface_area = bind(discr, sym_op)(actx) @@ -285,14 +285,14 @@ def test_surface_mass_operator_inverse(actx_factory, name): for resolution in builder.resolutions: mesh = builder.get_mesh(resolution, builder.mesh_order) discr = DiscretizationCollection(actx, mesh, order=builder.order) - volume_discr = discr.discr_from_dd(dof_desc.DD_VOLUME) + volume_discr = discr.discr_from_dd(dof_desc.DD_VOLUME_ALL) logger.info("ndofs: %d", volume_discr.ndofs) logger.info("nelements: %d", volume_discr.mesh.nelements) # {{{ compute inverse mass - dd = dof_desc.DD_VOLUME + dd = dof_desc.DD_VOLUME_ALL sym_f = sym.cos(4.0 * sym.nodes(mesh.ambient_dim, dd)[0]) sym_op = sym.InverseMassOperator(dd, dd)( sym.MassOperator(dd, dd)(sym.var("f"))) @@ -342,7 +342,7 @@ def test_face_normal_surface(actx_factory, mesh_name): mesh = builder.get_mesh(builder.resolutions[0], builder.mesh_order) discr = DiscretizationCollection(actx, mesh, order=builder.order) - volume_discr = discr.discr_from_dd(dof_desc.DD_VOLUME) + volume_discr = discr.discr_from_dd(dof_desc.DD_VOLUME_ALL) logger.info("ndofs: %d", volume_discr.ndofs) logger.info("nelements: %d", volume_discr.mesh.nelements) @@ -351,7 +351,7 @@ def test_face_normal_surface(actx_factory, mesh_name): # {{{ symbolic from meshmode.discretization.connection import FACE_RESTR_INTERIOR - dv = dof_desc.DD_VOLUME + dv = dof_desc.DD_VOLUME_ALL df = dof_desc.as_dofdesc(FACE_RESTR_INTERIOR) ambient_dim = mesh.ambient_dim @@ -578,11 +578,11 @@ def f(x): } ) - volume = discr.discr_from_dd(dof_desc.DD_VOLUME) + volume = discr.discr_from_dd(dof_desc.DD_VOLUME_ALL) logger.info("ndofs: %d", volume.ndofs) logger.info("nelements: %d", volume.mesh.nelements) - dd = dof_desc.DD_VOLUME + dd = dof_desc.DD_VOLUME_ALL dq = dd.with_discr_tag("product") df = dof_desc.as_dofdesc(FACE_RESTR_ALL) ambient_dim = discr.ambient_dim @@ -657,7 +657,8 @@ def test_op_collector_order_determinism(): class TestOperator(sym.Operator): def __init__(self): - sym.Operator.__init__(self, dof_desc.DD_VOLUME, dof_desc.DD_VOLUME) + sym.Operator.__init__( + self, dof_desc.DD_VOLUME_ALL, dof_desc.DD_VOLUME_ALL) mapper_method = "map_test_operator" @@ -725,7 +726,7 @@ def double(queue, x): a=(0,) * dims, b=(1,) * dims, nelements_per_axis=(4,) * dims) discr = DiscretizationCollection(actx, mesh, order=1) - ones = sym.Ones(dof_desc.DD_VOLUME) + ones = sym.Ones(dof_desc.DD_VOLUME_ALL) op = ( ones * 3 + sym.FunctionSymbol("double")(ones)) @@ -737,7 +738,7 @@ def double(queue, x): base_function_registry, "double", implementation=double, - dd=dof_desc.DD_VOLUME) + dd=dof_desc.DD_VOLUME_ALL) bound_op = bind(discr, op, function_registry=freg) @@ -756,7 +757,7 @@ def test_function_symbol_array(actx_factory, array_type): a=(-0.5,)*dim, b=(0.5,)*dim, nelements_per_axis=(8,)*dim, order=4) discr = DiscretizationCollection(actx, mesh, order=4) - volume_discr = discr.discr_from_dd(dof_desc.DD_VOLUME) + volume_discr = discr.discr_from_dd(dof_desc.DD_VOLUME_ALL) if array_type == "scalar": sym_x = sym.var("x") @@ -897,8 +898,8 @@ def test_incorrect_assignment_aggregation(actx_factory, ambient_dim): # {{{ test with a relative norm - from grudge.dof_desc import DD_VOLUME - dd = DD_VOLUME + from grudge.dof_desc import DD_VOLUME_ALL + dd = DD_VOLUME_ALL sym_x = sym.make_sym_array("y", ambient_dim, dd=dd) sym_y = sym.make_sym_array("y", ambient_dim, dd=dd) diff --git a/test/test_mpi_communication.py b/test/test_mpi_communication.py index 24e3d6380..d53ad2f81 100644 --- a/test/test_mpi_communication.py +++ b/test/test_mpi_communication.py @@ -46,6 +46,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 _test_func_comparison_mpi_communication_entrypoint(actx): hopefully_zero = ( op.project( dcoll, "int_faces", "all_faces", - dcoll.opposite_face_connection()(int_faces_func) + dcoll.opposite_face_connection( + dof_desc.DTAG_BOUNDARY( + dof_desc.FACE_RESTR_INTERIOR, dof_desc.VTAG_ALL) + )(int_faces_func) ) + sum(op.project(dcoll, tpair.dd, "all_faces", tpair.int) for tpair in op.cross_rank_trace_pairs(dcoll, myfunc, diff --git a/test/test_reductions.py b/test/test_reductions.py index 67291e49a..6a94a8ffb 100644 --- a/test/test_reductions.py +++ b/test/test_reductions.py @@ -121,7 +121,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) @@ -244,7 +244,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)