diff --git a/.pylintrc-local.yml b/.pylintrc-local.yml index 08f36d4c9..13303e6cb 100644 --- a/.pylintrc-local.yml +++ b/.pylintrc-local.yml @@ -1,14 +1,6 @@ - arg: ignore val: - mappers - - gas_dynamics - - burgers.py - - diffusion.py - - dt_finding.py - - nd_calculus.py - - pml.py - - poisson.py - - second_order.py - arg: ignored-modules val: - sympy diff --git a/doc/conf.py b/doc/conf.py index 0cd5ba6d6..1ff3ca070 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -35,6 +35,7 @@ def get_version(): "https://documen.tician.de/arraycontext/": None, "https://documen.tician.de/meshmode/": None, "https://documen.tician.de/loopy/": None, + "https://mpi4py.readthedocs.io/en/stable": None, } # index-page demo uses pyopencl via plot_directive diff --git a/examples/hello-grudge.py b/examples/hello-grudge.py index cfb724115..017430d86 100644 --- a/examples/hello-grudge.py +++ b/examples/hello-grudge.py @@ -14,7 +14,7 @@ import grudge.op as op from meshmode.mesh.generation import generate_box_mesh from meshmode.array_context import PyOpenCLArrayContext -from grudge.dof_desc import DTAG_BOUNDARY, FACE_RESTR_INTERIOR +from grudge.dof_desc import BoundaryDomainTag, FACE_RESTR_INTERIOR ctx = cl.create_some_context() @@ -51,8 +51,8 @@ def flux(dcoll, u_tpair): vol_discr = dcoll.discr_from_dd("vol") -left_bndry = DTAG_BOUNDARY("left") -right_bndry = DTAG_BOUNDARY("right") +left_bndry = BoundaryDomainTag("left") +right_bndry = BoundaryDomainTag("right") x_vol = actx.thaw(dcoll.nodes()) x_bndry = actx.thaw(dcoll.discr_from_dd(left_bndry).nodes()) diff --git a/grudge/__init__.py b/grudge/__init__.py index aad8dbd1c..fa4a2b3b7 100644 --- a/grudge/__init__.py +++ b/grudge/__init__.py @@ -20,8 +20,9 @@ THE SOFTWARE. """ -from grudge.discretization import DiscretizationCollection +from grudge.discretization import ( + DiscretizationCollection, make_discretization_collection) __all__ = [ - "DiscretizationCollection" + "DiscretizationCollection", "make_discretization_collection" ] diff --git a/grudge/discretization.py b/grudge/discretization.py index 43bd24226..dc08cd11f 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,30 +41,91 @@ THE SOFTWARE. """ -from pytools import memoize_method +from typing import Mapping, Optional, Union, Tuple, TYPE_CHECKING, Any + +from pytools import memoize_method, single_valued from grudge.dof_desc import ( - DD_VOLUME, - DISCR_TAG_BASE, - DISCR_TAG_MODAL, - DTAG_BOUNDARY, - DOFDesc, - as_dofdesc + VTAG_ALL, + BTAG_MULTIVOL_PARTITION, + DD_VOLUME_ALL, + DISCR_TAG_BASE, + DISCR_TAG_MODAL, + VolumeDomainTag, BoundaryDomainTag, + DOFDesc, + VolumeTag, DomainTag, + DiscretizationTag, + as_dofdesc, + ConvertibleToDOFDesc ) import numpy as np # noqa: F401 from arraycontext import ArrayContext +from meshmode.discretization import ElementGroupFactory, Discretization from meshmode.discretization.connection import ( FACE_RESTR_INTERIOR, FACE_RESTR_ALL, - make_face_restriction + make_face_restriction, + DiscretizationConnection ) -from meshmode.mesh import Mesh, BTAG_PARTITION +from meshmode.mesh import ( + InterPartitionAdjacencyGroup, Mesh, BTAG_PARTITION, BoundaryTag) +from meshmode.dof_array import DOFArray from warnings import warn +if TYPE_CHECKING: + import mpi4py.MPI + + +# {{{ discr_tag_to_group_factory normalization + +def _normalize_discr_tag_to_group_factory( + dim: int, + discr_tag_to_group_factory: Optional[ + Mapping[DiscretizationTag, ElementGroupFactory]], + order: Optional[int] + ) -> Mapping[DiscretizationTag, ElementGroupFactory]: + from meshmode.discretization.poly_element import \ + default_simplex_group_factory + + if discr_tag_to_group_factory is None: + if order is None: + raise TypeError( + "one of 'order' and 'discr_tag_to_group_factory' must be given" + ) + + discr_tag_to_group_factory = { + DISCR_TAG_BASE: default_simplex_group_factory( + base_dim=dim, order=order)} + else: + discr_tag_to_group_factory = dict(discr_tag_to_group_factory) + + if order is not None: + if DISCR_TAG_BASE in discr_tag_to_group_factory: + raise ValueError( + "if 'order' is given, 'discr_tag_to_group_factory' must " + "not have a key of DISCR_TAG_BASE" + ) + + discr_tag_to_group_factory[DISCR_TAG_BASE] = \ + default_simplex_group_factory(base_dim=dim, order=order) + + assert discr_tag_to_group_factory is not None + + # Modal discr should always come from the base discretization + if DISCR_TAG_MODAL not in discr_tag_to_group_factory: + discr_tag_to_group_factory[DISCR_TAG_MODAL] = \ + _generate_modal_group_factory( + discr_tag_to_group_factory[DISCR_TAG_BASE] + ) + + return discr_tag_to_group_factory + +# }}} + class DiscretizationCollection: """A collection of discretizations, defined on the same underlying @@ -60,11 +133,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 @@ -84,11 +154,15 @@ class DiscretizationCollection: # {{{ constructor - def __init__(self, array_context: ArrayContext, mesh: Mesh, - order=None, - discr_tag_to_group_factory=None, mpi_communicator=None, - # FIXME: `quad_tag_to_group_factory` is deprecated - quad_tag_to_group_factory=None): + 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[BoundaryDomainTag, DiscretizationConnection]] = None + ) -> None: """ :arg discr_tag_to_group_factory: A mapping from discretization tags (typically one of: :class:`grudge.dof_desc.DISCR_TAG_BASE`, @@ -101,63 +175,8 @@ def __init__(self, array_context: ArrayContext, mesh: Mesh, discretization. """ - if (quad_tag_to_group_factory is not None - and discr_tag_to_group_factory is not None): - raise ValueError( - "Both `quad_tag_to_group_factory` and `discr_tag_to_group_factory` " - "are specified. Use `discr_tag_to_group_factory` instead." - ) - - # FIXME: `quad_tag_to_group_factory` is deprecated - if (quad_tag_to_group_factory is not None - and discr_tag_to_group_factory is None): - warn("`quad_tag_to_group_factory` is a deprecated kwarg and will " - "be dropped in version 2022.x. Use `discr_tag_to_group_factory` " - "instead.", - DeprecationWarning, stacklevel=2) - discr_tag_to_group_factory = quad_tag_to_group_factory - self._setup_actx = array_context.clone() - from meshmode.discretization.poly_element import \ - default_simplex_group_factory - - if discr_tag_to_group_factory is None: - if order is None: - raise TypeError( - "one of 'order' and 'discr_tag_to_group_factory' must be given" - ) - - discr_tag_to_group_factory = { - DISCR_TAG_BASE: default_simplex_group_factory( - base_dim=mesh.dim, order=order)} - else: - if order is not None: - discr_tag_to_group_factory = discr_tag_to_group_factory.copy() - if DISCR_TAG_BASE in discr_tag_to_group_factory: - raise ValueError( - "if 'order' is given, 'discr_tag_to_group_factory' must " - "not have a key of DISCR_TAG_BASE" - ) - - discr_tag_to_group_factory[DISCR_TAG_BASE] = \ - default_simplex_group_factory(base_dim=mesh.dim, order=order) - - # Modal discr should always come from the base discretization - discr_tag_to_group_factory[DISCR_TAG_MODAL] = \ - _generate_modal_group_factory( - discr_tag_to_group_factory[DISCR_TAG_BASE] - ) - - self.discr_tag_to_group_factory = discr_tag_to_group_factory - - from meshmode.discretization import Discretization - - self._volume_discr = Discretization( - array_context, mesh, - self.group_factory_for_discretization_tag(DISCR_TAG_BASE) - ) - # {{{ process mpi_communicator argument if mpi_communicator is not None: @@ -181,9 +200,56 @@ def __init__(self, array_context: ArrayContext, mesh: Mesh, # }}} - self._dist_boundary_connections = \ - self._set_up_distributed_communication( - mpi_communicator, array_context) + from meshmode.discretization import Discretization + + # {{{ 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 + + assert discr_tag_to_group_factory is not None + self._discr_tag_to_group_factory = discr_tag_to_group_factory + + # }}} + + self._volume_discrs = volume_discrs # }}} @@ -196,16 +262,6 @@ def mpi_communicator(self): return self._mpi_communicator - @property - def quad_tag_to_group_factory(self): - warn("`DiscretizationCollection.quad_tag_to_group_factory` " - "is deprecated and will go away in 2022. Use " - "`DiscretizationCollection.discr_tag_to_group_factory` " - "instead.", - DeprecationWarning, stacklevel=2) - - return self.discr_tag_to_group_factory - def get_management_rank_index(self): return 0 @@ -216,86 +272,12 @@ def is_management_rank(self): return self.mpi_communicator.Get_rank() \ == self.get_management_rank_index() - # {{{ distributed - - def _set_up_distributed_communication(self, mpi_communicator, array_context): - from_dd = DOFDesc("vol", DISCR_TAG_BASE) - - boundary_connections = {} - - from meshmode.distributed import get_connected_partitions - connected_parts = get_connected_partitions(self._volume_discr.mesh) - - if connected_parts: - if mpi_communicator is None: - raise RuntimeError("must supply an MPI communicator when using a " - "distributed mesh") - - grp_factory = \ - self.group_factory_for_discretization_tag(DISCR_TAG_BASE) - - local_boundary_connections = {} - for i_remote_part in connected_parts: - local_boundary_connections[i_remote_part] = self.connection_from_dds( - from_dd, DOFDesc(BTAG_PARTITION(i_remote_part), - DISCR_TAG_BASE)) - - from 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 - - return boundary_connections - - def get_distributed_boundary_swap_connection(self, dd): - warn("`DiscretizationCollection.get_distributed_boundary_swap_connection` " - "is deprecated and will go away in 2022. Use " - "`DiscretizationCollection.distributed_boundary_swap_connection` " - "instead.", - DeprecationWarning, stacklevel=2) - return self.distributed_boundary_swap_connection(dd) - - def distributed_boundary_swap_connection(self, dd): - """Provides a mapping from the base volume discretization - to the exterior boundary restriction on a parallel boundary - 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 - def discr_from_dd(self, dd): + def discr_from_dd(self, dd: "ConvertibleToDOFDesc") -> Discretization: """Provides a :class:`meshmode.discretization.Discretization` object from *dd*. - - :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value - convertible to one. """ dd = as_dofdesc(dd) @@ -305,45 +287,43 @@ def discr_from_dd(self, dd): return self._modal_discr(dd.domain_tag) if dd.is_volume(): - if discr_tag is not DISCR_TAG_BASE: - return self._discr_tag_volume_discr(discr_tag) - return self._volume_discr + return self._volume_discr_from_dd(dd) if discr_tag is not DISCR_TAG_BASE: - no_quad_discr = self.discr_from_dd(DOFDesc(dd.domain_tag)) + base_discr = self.discr_from_dd(dd.with_discr_tag(DISCR_TAG_BASE)) from meshmode.discretization import Discretization return Discretization( self._setup_actx, - no_quad_discr.mesh, + base_discr.mesh, self.group_factory_for_discretization_tag(discr_tag) ) assert discr_tag is DISCR_TAG_BASE - if dd.domain_tag is FACE_RESTR_ALL: - return self._all_faces_volume_connection().to_discr - elif dd.domain_tag is FACE_RESTR_INTERIOR: - return self._interior_faces_connection().to_discr - elif dd.is_boundary_or_partition_interface(): - return self._boundary_connection(dd.domain_tag.tag).to_discr + if isinstance(dd.domain_tag, BoundaryDomainTag): + if dd.domain_tag.tag in [FACE_RESTR_ALL, FACE_RESTR_INTERIOR]: + return self._faces_connection(dd.domain_tag).to_discr + else: + return self._boundary_connection(dd.domain_tag).to_discr else: - raise ValueError("DOF desc tag not understood: " + str(dd)) + raise ValueError(f"DOF desc not understood: {dd}") # }}} # {{{ _base_to_geoderiv_connection @memoize_method - def _has_affine_groups(self): + def _has_affine_groups(self, domain_tag: DomainTag) -> bool: from modepy.shapes import Simplex + discr = self.discr_from_dd(DOFDesc(domain_tag, DISCR_TAG_BASE)) return any( megrp.is_affine and issubclass(megrp._modepy_shape_cls, Simplex) - for megrp in self._volume_discr.mesh.groups) + for megrp in discr.mesh.groups) @memoize_method - def _base_to_geoderiv_connection(self, dd: DOFDesc): + def _base_to_geoderiv_connection(self, dd: DOFDesc) -> DiscretizationConnection: r"""The "geometry derivatives" discretization for a given *dd* is typically identical to the one returned by :meth:`discr_from_dd`, however for affinely-mapped simplicial elements, it will use a @@ -356,7 +336,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 \ @@ -393,7 +373,9 @@ def geo_group_factory(megrp, index): # {{{ connection_from_dds @memoize_method - def connection_from_dds(self, from_dd, to_dd): + def connection_from_dds( + self, from_dd: "ConvertibleToDOFDesc", to_dd: "ConvertibleToDOFDesc" + ) -> DiscretizationConnection: """Provides a mapping (connection) from one discretization to another, e.g. from the volume to the boundary, or from the base to the an overintegrated quadrature discretization, or from @@ -425,12 +407,15 @@ def connection_from_dds(self, from_dd, to_dd): assert (to_discr_tag is not DISCR_TAG_MODAL and from_discr_tag is not DISCR_TAG_MODAL) - if (not from_dd.is_volume() + if (isinstance(from_dd.domain_tag, BoundaryDomainTag) and from_discr_tag == to_discr_tag - and to_dd.domain_tag is FACE_RESTR_ALL): + and isinstance(to_dd.domain_tag, BoundaryDomainTag) + and to_dd.domain_tag.tag is FACE_RESTR_ALL): faces_conn = self.connection_from_dds( - DOFDesc("vol"), - DOFDesc(from_dd.domain_tag)) + DOFDesc( + VolumeDomainTag(from_dd.domain_tag.volume_tag), + DISCR_TAG_BASE), + from_dd.with_discr_tag(DISCR_TAG_BASE)) from meshmode.discretization.connection import \ make_face_to_all_faces_embedding @@ -448,7 +433,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 @@ -482,73 +467,79 @@ def connection_from_dds(self, from_dd, to_dd): # }}} if from_discr_tag is not DISCR_TAG_BASE: - raise ValueError("cannot interpolate *from* a " - "(non-interpolatory) quadrature grid") + raise ValueError("cannot get a connection *from* a " + f"(non-interpolatory) quadrature grid: '{from_dd}'") assert to_discr_tag is DISCR_TAG_BASE - if from_dd.is_volume(): - if to_dd.domain_tag is FACE_RESTR_ALL: - return self._all_faces_volume_connection() - if to_dd.domain_tag is FACE_RESTR_INTERIOR: - return self._interior_faces_connection() - elif to_dd.is_boundary_or_partition_interface(): - assert from_discr_tag is DISCR_TAG_BASE - return self._boundary_connection(to_dd.domain_tag.tag) + if isinstance(from_dd.domain_tag, VolumeDomainTag): + if isinstance(to_dd.domain_tag, BoundaryDomainTag): + if to_dd.domain_tag.volume_tag != from_dd.domain_tag.tag: + raise ValueError("cannot get a connection from one volume " + f"('{from_dd.domain_tag.tag}') " + "to the boundary of another volume " + f"('{to_dd.domain_tag.volume_tag}') ") + if to_dd.domain_tag.tag in [FACE_RESTR_ALL, FACE_RESTR_INTERIOR]: + return self._faces_connection(to_dd.domain_tag) + elif isinstance(to_dd.domain_tag, BoundaryDomainTag): + assert from_discr_tag is DISCR_TAG_BASE + return self._boundary_connection(to_dd.domain_tag) elif to_dd.is_volume(): + if to_dd.domain_tag != from_dd.domain_tag: + raise ValueError("cannot get a connection between " + "volumes of different tags: requested " + f"'{from_dd.domain_tag}' -> '{to_dd.domain_tag}'") + from meshmode.discretization.connection import \ make_same_mesh_connection - to_discr = self._discr_tag_volume_discr(to_discr_tag) - from_discr = self._volume_discr - return make_same_mesh_connection(self._setup_actx, to_discr, - from_discr) + return make_same_mesh_connection( + self._setup_actx, + self._volume_discr_from_dd(to_dd), + self._volume_discr_from_dd(from_dd)) else: - raise ValueError("cannot interpolate from volume to: " + str(to_dd)) + raise ValueError( + f"cannot get a connection from volume to: '{to_dd}'") else: - raise ValueError("cannot interpolate from: " + str(from_dd)) + raise ValueError(f"cannot get a connection from: '{from_dd}'") # }}} # {{{ group_factory_for_discretization_tag - def group_factory_for_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, VolumeDomainTag) + + try: + base_volume_discr = self._volume_discrs[dd.domain_tag.tag] + except KeyError: + raise ValueError("a volume discretization with volume tag " + f"'{dd.domain_tag.tag}' is not known") # Refuse to re-make the volume discretization - if discretization_tag is DISCR_TAG_BASE: - return self._volume_discr + if dd.discretization_tag is DISCR_TAG_BASE: + return base_volume_discr from meshmode.discretization import Discretization return Discretization( - self._setup_actx, self._volume_discr.mesh, - self.group_factory_for_discretization_tag(discretization_tag) + self._setup_actx, base_volume_discr.mesh, + self.group_factory_for_discretization_tag(dd.discretization_tag) ) @memoize_method - def _modal_discr(self, domain_tag): + def _modal_discr(self, domain_tag) -> Discretization: from meshmode.discretization import Discretization discr_base = self.discr_from_dd(DOFDesc(domain_tag, DISCR_TAG_BASE)) @@ -557,10 +548,12 @@ def _modal_discr(self, domain_tag): self.group_factory_for_discretization_tag(DISCR_TAG_MODAL) ) + # }}} + # {{{ connection factories: modal<->nodal @memoize_method - def _modal_to_nodal_connection(self, to_dd): + def _modal_to_nodal_connection(self, to_dd: DOFDesc) -> DiscretizationConnection: """ :arg to_dd: a :class:`grudge.dof_desc.DOFDesc` describing the dofs corresponding to the @@ -575,7 +568,8 @@ def _modal_to_nodal_connection(self, to_dd): ) @memoize_method - def _nodal_to_modal_connection(self, from_dd): + def _nodal_to_modal_connection( + self, from_dd: DOFDesc) -> DiscretizationConnection: """ :arg from_dd: a :class:`grudge.dof_desc.DOFDesc` describing the dofs corresponding to the @@ -594,25 +588,31 @@ def _nodal_to_modal_connection(self, from_dd): # {{{ connection factories: boundary @memoize_method - def _boundary_connection(self, boundary_tag): + def _boundary_connection( + self, domain_tag: BoundaryDomainTag) -> DiscretizationConnection: return make_face_restriction( - self._setup_actx, - self._volume_discr, - self.group_factory_for_discretization_tag(DISCR_TAG_BASE), - boundary_tag=boundary_tag - ) + self._setup_actx, + self._volume_discr_from_dd( + DOFDesc(VolumeDomainTag(domain_tag.volume_tag), DISCR_TAG_BASE)), + self.group_factory_for_discretization_tag(DISCR_TAG_BASE), + boundary_tag=domain_tag.tag) # }}} - # {{{ connection factories: interior faces + # {{{ connection factories: faces @memoize_method - def _interior_faces_connection(self): + def _faces_connection( + self, domain_tag: BoundaryDomainTag) -> DiscretizationConnection: + assert domain_tag.tag in [FACE_RESTR_INTERIOR, FACE_RESTR_ALL] + return make_face_restriction( self._setup_actx, - self._volume_discr, + self._volume_discr_from_dd( + DOFDesc( + VolumeDomainTag(domain_tag.volume_tag), DISCR_TAG_BASE)), self.group_factory_for_discretization_tag(DISCR_TAG_BASE), - FACE_RESTR_INTERIOR, + domain_tag.tag, # FIXME: This will need to change as soon as we support # pyramids or other elements with non-identical face @@ -621,7 +621,8 @@ def _interior_faces_connection(self): ) @memoize_method - def opposite_face_connection(self): + def opposite_face_connection( + self, domain_tag: BoundaryDomainTag) -> DiscretizationConnection: """Provides a mapping from the base volume discretization to the exterior boundary restriction on a neighboring element. This does not take into account parallel partitions. @@ -629,93 +630,78 @@ 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()) + self._faces_connection(domain_tag)) # }}} - # {{{ 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 - ) - - # }}} + # {{{ properties @property - def dim(self): + def dim(self) -> int: """Return the topological dimension.""" - return self._volume_discr.dim + return single_valued(discr.dim for discr in self._volume_discrs.values()) @property - def ambient_dim(self): + def ambient_dim(self) -> int: """Return the dimension of the ambient space.""" - return self._volume_discr.ambient_dim + return single_valued( + discr.ambient_dim for discr in self._volume_discrs.values()) @property - def real_dtype(self): + def real_dtype(self) -> "np.dtype[Any]": """Return the data type used for real-valued arithmetic.""" - return self._volume_discr.real_dtype + return single_valued( + discr.real_dtype for discr in self._volume_discrs.values()) @property - def complex_dtype(self): + def complex_dtype(self) -> "np.dtype[Any]": """Return the data type used for complex-valued arithmetic.""" - return self._volume_discr.complex_dtype + return single_valued( + discr.complex_dtype for discr in self._volume_discrs.values()) - @property - def mesh(self): - """Return the :class:`meshmode.mesh.Mesh` over which the discretization - collection is built. - """ - return self._volume_discr.mesh + # }}} + + # {{{ array creators - def empty(self, array_context: ArrayContext, dtype=None): + def empty(self, array_context: ArrayContext, dtype=None, + *, dd: Optional[DOFDesc] = None) -> DOFArray: """Return an empty :class:`~meshmode.dof_array.DOFArray` defined at - the volume nodes: :class:`grudge.dof_desc.DD_VOLUME`. + the volume nodes: :class:`grudge.dof_desc.DD_VOLUME_ALL`. :arg array_context: an :class:`~arraycontext.context.ArrayContext`. :arg dtype: type special value 'c' will result in a vector of dtype :attr:`complex_dtype`. If *None* (the default), a real vector will be returned. """ - return self._volume_discr.empty(array_context, dtype) + if dd is None: + dd = DD_VOLUME_ALL + return self.discr_from_dd(dd).empty(array_context, dtype) - def zeros(self, array_context: ArrayContext, dtype=None): + def zeros(self, array_context: ArrayContext, dtype=None, + *, dd: Optional[DOFDesc] = None) -> DOFArray: """Return a zero-initialized :class:`~meshmode.dof_array.DOFArray` - defined at the volume nodes, :class:`grudge.dof_desc.DD_VOLUME`. + defined at the volume nodes, :class:`grudge.dof_desc.DD_VOLUME_ALL`. :arg array_context: an :class:`~arraycontext.context.ArrayContext`. :arg dtype: type special value 'c' will result in a vector of dtype :attr:`complex_dtype`. If *None* (the default), a real vector will be returned. """ - return self._volume_discr.zeros(array_context, dtype) + if dd is None: + dd = DD_VOLUME_ALL + + return self.discr_from_dd(dd).zeros(array_context, dtype) def is_volume_where(self, where): return where is None or as_dofdesc(where).is_volume() - @property - def order(self): - warn("DiscretizationCollection.order is deprecated, " - "consider using the orders of element groups instead. " - "'order' will go away in 2021.", - DeprecationWarning, stacklevel=2) - - from pytools import single_valued - return single_valued(egrp.order for egrp in self._volume_discr.groups) + # }}} - # {{{ Discretization-specific geometric properties + # {{{ discretization-specific geometric fields def nodes(self, dd=None): r"""Return the nodes of a discretization specified by *dd*. @@ -725,7 +711,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): @@ -741,14 +727,145 @@ def normal(self, dd): # }}} -class DGDiscretizationWithBoundaries(DiscretizationCollection): - def __init__(self, *args, **kwargs): - warn("DGDiscretizationWithBoundaries is deprecated and will go away " - "in 2022. Use DiscretizationCollection instead.", - DeprecationWarning, stacklevel=2) +# {{{ distributed/multi-volume setup + +def _check_btag(tag: BoundaryTag) -> Union[BTAG_MULTIVOL_PARTITION, BTAG_PARTITION]: + if isinstance(tag, BTAG_MULTIVOL_PARTITION): + return tag + + elif isinstance(tag, BTAG_PARTITION): + return tag + + else: + raise TypeError("unexpected type of inter-partition boundary tag " + f"'{type(tag)}'") + + +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 + + else: + raise TypeError("unexpected type of inter-partition boundary tag " + f"'{type(btag)}'") + + +def _flip_dtag( + self_rank: Optional[int], + domain_tag: BoundaryDomainTag, + ) -> BoundaryDomainTag: + if isinstance(domain_tag.tag, BTAG_PARTITION): + assert self_rank is not None + return BoundaryDomainTag( + BTAG_PARTITION(self_rank), domain_tag.volume_tag) + + elif isinstance(domain_tag.tag, BTAG_MULTIVOL_PARTITION): + return BoundaryDomainTag( + BTAG_MULTIVOL_PARTITION( + other_rank=None if domain_tag.tag.other_rank is None else self_rank, + other_volume_tag=domain_tag.volume_tag), + domain_tag.tag.other_volume_tag) + + else: + raise TypeError("unexpected type of inter-partition boundary tag " + f"'{type(domain_tag.tag)}'") + + +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[ + BoundaryDomainTag, + DiscretizationConnection]: + + from meshmode.distributed import (get_inter_partition_tags, + make_remote_group_infos, InterRankBoundaryInfo, + MPIBoundaryCommSetupHelper) + + inter_part_tags = { + BoundaryDomainTag(_check_btag(btag), discr_vol_tag) + for discr_vol_tag, volume_discr in volume_discrs.items() + for btag in get_inter_partition_tags(volume_discr.mesh)} + + inter_part_conns: Mapping[ + BoundaryDomainTag, + DiscretizationConnection] = {} + + if inter_part_tags: + local_boundary_restrictions = { + domain_tag: make_face_restriction( + array_context, volume_discrs[domain_tag.volume_tag], + base_group_factory, boundary_tag=domain_tag.tag) + for domain_tag in inter_part_tags} + + irbis = [] + + for domain_tag in inter_part_tags: + assert isinstance( + domain_tag.tag, (BTAG_PARTITION, BTAG_MULTIVOL_PARTITION)) + + other_rank = _remote_rank_from_btag(domain_tag.tag) + btag_restr = local_boundary_restrictions[domain_tag] + + if other_rank is None: + # {{{ rank-local interface between multiple volumes + + assert isinstance(domain_tag.tag, BTAG_MULTIVOL_PARTITION) + + from meshmode.discretization.connection import \ + make_partition_connection + remote_dtag = _flip_dtag(None, domain_tag) + inter_part_conns[domain_tag] = make_partition_connection( + array_context, + local_bdry_conn=btag_restr, + remote_bdry_discr=( + local_boundary_restrictions[remote_dtag].to_discr), + remote_group_infos=make_remote_group_infos( + array_context, remote_dtag.tag, 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=domain_tag.tag, + local_part_id=domain_tag, + remote_part_id=_flip_dtag( + mpi_communicator.rank, domain_tag), + remote_rank=other_rank, + local_boundary_connection=btag_restr)) + + # }}} + + if irbis: + assert mpi_communicator is not None + + 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 + + inter_part_conns.update(conns) + + return inter_part_conns + +# }}} - super().__init__(*args, **kwargs) +# {{{ modal group factory def _generate_modal_group_factory(nodal_group_factory): from meshmode.discretization.poly_element import ( @@ -769,4 +886,131 @@ 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, + 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), + 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 267e4f56e..e3015d1d8 100644 --- a/grudge/dof_desc.py +++ b/grudge/dof_desc.py @@ -1,4 +1,57 @@ -"""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:: VolumeDomainTag +.. autoclass:: BoundaryDomainTag + +Discretization tags +------------------- + +A discretization tag serves as a symbolic identifier of the manner in which +meaning is assigned to degrees of freedom. + +.. autoclass:: DISCR_TAG_BASE +.. autoclass:: DISCR_TAG_QUAD +.. autoclass:: DISCR_TAG_MODAL + +DOF Descriptor +-------------- + +.. autoclass:: DOFDesc +.. autofunction:: as_dofdesc + +Shortcuts +--------- + +.. data:: DD_SCALAR +.. data:: DD_VOLUME_ALL +.. data:: DD_VOLUME_ALL_MODAL + +Internal things that are visble due to type annotations +------------------------------------------------------- + +.. class:: _DiscretizationTag +.. class:: ConvertibleToDOFDesc + + Anything that is convertible to a :class:`DOFDesc` via :func:`as_dofdesc`. +""" __copyright__ = """ Copyright (C) 2008 Andreas Kloeckner @@ -25,31 +78,18 @@ THE SOFTWARE. """ -from meshmode.discretization.connection import \ - FACE_RESTR_INTERIOR, FACE_RESTR_ALL -from meshmode.mesh import \ - BTAG_PARTITION, BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE -from warnings import warn import sys +from warnings import warn +from typing import Hashable, Union, Type, Optional, Any, Tuple +from dataclasses import dataclass, replace +from meshmode.discretization.connection import ( + FACE_RESTR_INTERIOR, FACE_RESTR_ALL) +from meshmode.mesh import ( + BTAG_PARTITION, BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE, BoundaryTag) -__doc__ = """ -.. autoclass:: DTAG_SCALAR -.. autoclass:: DTAG_VOLUME_ALL -.. autoclass:: DTAG_BOUNDARY - -.. autoclass:: DISCR_TAG_BASE -.. autoclass:: DISCR_TAG_QUAD -.. autoclass:: DISCR_TAG_MODAL - -.. autoclass:: DOFDesc -.. autofunction:: as_dofdesc - -.. data:: DD_SCALAR -.. data:: DD_VOLUME -.. data:: DD_VOLUME_MODAL -""" +# {{{ _to_identifier def _to_identifier(name: str) -> str: if not name.isidentifier(): @@ -57,71 +97,109 @@ def _to_identifier(name: str) -> str: else: return name +# }}} -# {{{ DOF description -class DTAG_SCALAR: # noqa: N801 - """A domain tag denoting scalar values.""" +# {{{ volume tags + +class VTAG_ALL: # noqa: N801 + pass + + +VolumeTag = Hashable + +# }}} -class DTAG_VOLUME_ALL: # noqa: N801 +# {{{ partition identifier + +@dataclass(init=True, eq=True, frozen=True) +class BTAG_MULTIVOL_PARTITION: # noqa: N801 """ - A domain tag denoting values defined - in all cell volumes. + .. attribute:: other_rank + + An integer, or *None*. If *None*, this marks a partition boundary + to another volume on the same rank. + + .. attribute:: other_volume_tag """ + other_rank: Optional[int] + other_volume_tag: "VolumeTag" + +# }}} -class DTAG_BOUNDARY: # noqa: N801 - """A domain tag describing the values on element - boundaries which are adjacent to elements - of another :class:`~meshmode.mesh.Mesh`. +# {{{ domain tag + +@dataclass(frozen=True, eq=True) +class ScalarDomainTag: # noqa: N801 + """A domain tag denoting scalar values.""" + + +DTAG_SCALAR = ScalarDomainTag() + + +@dataclass(frozen=True, eq=True, init=True) +class VolumeDomainTag: + """A domain tag referring to a volume identified by the + volume tag :attr:`tag`. These volume identifiers are only used + when the :class:`~grudge.discretization.DiscretizationCollection` contains + more than one volume. .. attribute:: tag .. automethod:: __init__ - .. automethod:: __eq__ - .. automethod:: __ne__ - .. automethod:: __hash__ """ + tag: VolumeTag - 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 +DTAG_VOLUME_ALL = VolumeDomainTag(VTAG_ALL) + + +@dataclass(frozen=True, eq=True, init=True) +class BoundaryDomainTag: + """A domain tag referring to a boundary identified by the + boundary tag :attr:`tag`. + + .. attribute:: tag + .. attribute:: volume_tag + + .. automethod:: __init__ + """ + tag: BoundaryTag + volume_tag: VolumeTag = VTAG_ALL + + +DomainTag = Union[ScalarDomainTag, VolumeDomainTag, BoundaryDomainTag] + +# }}} + - 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 @@ -135,20 +213,22 @@ class CustomQuadTag(DISCR_TAG_QUAD): "A custom quadrature discretization tag." dd = DOFDesc(DTAG_VOLUME_ALL, CustomQuadTag) - """ -class DISCR_TAG_MODAL: # noqa: N801 - """A discretization tag indicating the use of a - basic discretization grid with modal degrees of - freedom. This tag is used to distinguish the - modal discretization from the base (nodal) - discretization (e.g. :class:`DISCR_TAG_BASE`) or +class DISCR_TAG_MODAL(_DiscretizationTag): # noqa: N801 + """A discretization tag indicating the use of unisolvent modal degrees of + freedom. This tag is used to distinguish the modal discretization from the + base (nodal) discretization (e.g. :class:`DISCR_TAG_BASE`) or discretizations on quadrature grids (:class:`DISCR_TAG_QUAD`). """ +# }}} + + +# {{{ DOF descriptor +@dataclass(frozen=True, eq=True) class DOFDesc: """Describes the meaning of degrees of freedom. @@ -165,8 +245,9 @@ class DOFDesc: .. automethod:: uses_quadrature + .. automethod:: with_domain_tag .. automethod:: with_discr_tag - .. automethod:: with_dtag + .. automethod:: trace .. automethod:: __eq__ .. automethod:: __ne__ @@ -174,159 +255,87 @@ class DOFDesc: .. automethod:: as_identifier """ - def __init__(self, domain_tag, discretization_tag=None, - # FIXME: `quadrature_tag` is deprecated - quadrature_tag=None): - """ - :arg domain_tag: One of the following: - :class:`DTAG_SCALAR` (or the string ``"scalar"``), - :class:`DTAG_VOLUME_ALL` (or the string ``"vol"``) - for the default volume discretization, - :data:`~meshmode.discretization.connection.FACE_RESTR_ALL` - (or the string ``"all_faces"``), or - :data:`~meshmode.discretization.connection.FACE_RESTR_INTERIOR` - (or the string ``"int_faces"``), or one of - :class:`~meshmode.mesh.BTAG_ALL`, - :class:`~meshmode.mesh.BTAG_NONE`, - :class:`~meshmode.mesh.BTAG_REALLY_ALL`, - :class:`~meshmode.mesh.BTAG_PARTITION`, - or *None* to indicate that the geometry is not yet known. - - :arg discretization_tag: - *None* or :class:`DISCR_TAG_BASE` to indicate the use of the basic - discretization grid, :class:`DISCR_TAG_MODAL` to indicate a - modal discretization, or :class:`DISCR_TAG_QUAD` to indicate - the use of a quadrature grid. - """ + domain_tag: DomainTag + discretization_tag: DiscretizationTag - if domain_tag is None: - pass - elif domain_tag in [DTAG_SCALAR, "scalar"]: - domain_tag = DTAG_SCALAR - elif domain_tag in [DTAG_VOLUME_ALL, "vol"]: - domain_tag = DTAG_VOLUME_ALL - elif domain_tag in [FACE_RESTR_ALL, "all_faces"]: - domain_tag = FACE_RESTR_ALL - elif domain_tag in [FACE_RESTR_INTERIOR, "int_faces"]: - domain_tag = FACE_RESTR_INTERIOR - elif isinstance(domain_tag, BTAG_PARTITION): - domain_tag = DTAG_BOUNDARY(domain_tag) - elif domain_tag in [BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE]: - domain_tag = DTAG_BOUNDARY(domain_tag) - elif isinstance(domain_tag, DTAG_BOUNDARY): - pass - else: - raise ValueError("domain tag not understood: %s" % domain_tag) + def __init__(self, domain_tag: Any, + discretization_tag: Optional[type[DiscretizationTag]] = None): - if (quadrature_tag is not None and discretization_tag is not None): - raise ValueError( - "Both `quadrature_tag` and `discretization_tag` are specified. " - "Use `discretization_tag` instead." - ) + if ( + not (isinstance(domain_tag, + (ScalarDomainTag, BoundaryDomainTag, VolumeDomainTag))) + or discretization_tag is None + or ( + not isinstance(discretization_tag, type) + or not issubclass(discretization_tag, _DiscretizationTag))): + warn("Sloppy construction of DOFDesc is deprecated. " + "This will stop working in 2023. " + "Call as_dofdesc instead, with the same arguments. ", + DeprecationWarning, stacklevel=2) - # 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 + domain_tag, discretization_tag = _normalize_domain_and_discr_tag( + domain_tag, discretization_tag) - if domain_tag is DTAG_SCALAR and discretization_tag is not None: - raise ValueError("cannot have nontrivial discretization tag on scalar") + object.__setattr__(self, "domain_tag", domain_tag) + object.__setattr__(self, "discretization_tag", discretization_tag) - if discretization_tag is None: - discretization_tag = DISCR_TAG_BASE + def is_scalar(self) -> bool: + return isinstance(self.domain_tag, ScalarDomainTag) - # 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 - - def is_scalar(self): - 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_volume(self) -> bool: + return isinstance(self.domain_tag, VolumeDomainTag) - def is_boundary_or_partition_interface(self): - return isinstance(self.domain_tag, DTAG_BOUNDARY) - - def is_trace(self): - return (self.is_boundary_or_partition_interface() - or self.domain_tag in [ + def is_boundary_or_partition_interface(self) -> bool: + return (isinstance(self.domain_tag, BoundaryDomainTag) + and self.domain_tag.tag not in [ FACE_RESTR_ALL, FACE_RESTR_INTERIOR]) - def uses_quadrature(self): + def is_trace(self) -> bool: + return isinstance(self.domain_tag, BoundaryDomainTag) + + def uses_quadrature(self) -> bool: # FIXME: String tags are deprecated - # Check for string first, otherwise - # `issubclass` will raise an exception whenever - # its first argument is not a class. - # This can go away once support for strings is dropped - # completely. if isinstance(self.discretization_tag, str): # All strings are interpreted as quadrature-related tags return True - elif issubclass(self.discretization_tag, DISCR_TAG_QUAD): - return True - elif issubclass(self.discretization_tag, - (DISCR_TAG_BASE, DISCR_TAG_MODAL)): - return False - else: - raise ValueError( - f"Unsure how to interpret tag: {self.discretization_tag}" - ) - - def with_qtag(self, discr_tag): - warn("`DOFDesc.with_qtag` is deprecated and will be dropped " - "in version 2022.x. Use `DOFDesc.with_discr_tag` instead.", - DeprecationWarning, stacklevel=2) - return self.with_discr_tag(discr_tag) - - def with_discr_tag(self, discr_tag): - return type(self)(domain_tag=self.domain_tag, - discretization_tag=discr_tag) - - def with_dtag(self, dtag): - return type(self)(domain_tag=dtag, - discretization_tag=self.discretization_tag) - - def __eq__(self, other): - return (type(self) == type(other) - and self.domain_tag == other.domain_tag - and self.discretization_tag == other.discretization_tag) - - def __ne__(self, other): - return not self.__eq__(other) - - def __hash__(self): - return hash((type(self), self.domain_tag, self.discretization_tag)) - - def __repr__(self): - def fmt(s): - if isinstance(s, type): - return s.__name__ - else: - return repr(s) + elif isinstance(self.discretization_tag, type): + if issubclass(self.discretization_tag, DISCR_TAG_QUAD): + return True + elif issubclass(self.discretization_tag, + (DISCR_TAG_BASE, DISCR_TAG_MODAL)): + return False + + raise ValueError( + f"Invalid discretization tag: {self.discretization_tag}") + + def with_dtag(self, dtag) -> "DOFDesc": + from warnings import warn + warn("'with_dtag' is deprecated. Use 'with_domain_tag' instead. " + "This will stop working in 2023", + DeprecationWarning, stacklevel=2) + return replace(self, domain_tag=dtag) + + def with_domain_tag(self, dtag) -> "DOFDesc": + return replace(self, domain_tag=dtag) + + def trace(self, btag: BoundaryTag) -> "DOFDesc": + """Return a :class:`DOFDesc` for the restriction of the volume + descriptor *self* to the boundary named by *btag*. + + An error is raised if this method is called on a non-volume instance of + :class:`DOFDesc`. + """ + if not isinstance(self.domain_tag, VolumeDomainTag): + raise ValueError(f"must originate on volume, got '{self.domain_tag}'") + return replace(self, + domain_tag=BoundaryDomainTag(btag, volume_tag=self.domain_tag.tag)) - return "DOFDesc({}, {})".format( - fmt(self.domain_tag), - fmt(self.discretization_tag)) + def with_discr_tag(self, discr_tag) -> "DOFDesc": + return replace(self, discretization_tag=discr_tag) def as_identifier(self) -> str: """Returns a descriptive string for this :class:`DOFDesc` that is usable @@ -341,7 +350,16 @@ def as_identifier(self) -> str: dom_id = "f_all" elif self.domain_tag is FACE_RESTR_INTERIOR: dom_id = "f_int" - elif isinstance(self.domain_tag, DTAG_BOUNDARY): + elif isinstance(self.domain_tag, VolumeDomainTag): + vtag = self.domain_tag.tag + if isinstance(vtag, type): + vtag = vtag.__name__.replace("VTAG_", "").lower() + elif isinstance(vtag, str): + vtag = _to_identifier(vtag) + else: + vtag = _to_identifier(str(vtag)) + dom_id = f"v_{vtag}" + elif isinstance(self.domain_tag, BoundaryDomainTag): btag = self.domain_tag.tag if isinstance(btag, type): btag = btag.__name__.replace("BTAG_", "").lower() @@ -369,31 +387,101 @@ def as_identifier(self) -> str: return f"{dom_id}{discr_id}" -DD_SCALAR = DOFDesc(DTAG_SCALAR, None) +DD_SCALAR = DOFDesc(DTAG_SCALAR, DISCR_TAG_BASE) +DD_VOLUME_ALL = DOFDesc(DTAG_VOLUME_ALL, DISCR_TAG_BASE) +DD_VOLUME_ALL_MODAL = DOFDesc(DTAG_VOLUME_ALL, DISCR_TAG_MODAL) + + +def _normalize_domain_and_discr_tag( + domain: Any, + discretization_tag: Optional[DiscretizationTag] = None, + *, _contextual_volume_tag: Optional[VolumeTag] = None + ) -> Tuple[DomainTag, DiscretizationTag]: + + if _contextual_volume_tag is None: + _contextual_volume_tag = VTAG_ALL + + if domain in [DTAG_SCALAR, "scalar"]: + domain = DTAG_SCALAR + elif isinstance(domain, (BoundaryDomainTag, VolumeDomainTag)): + pass + elif domain == "vol": + domain = DTAG_VOLUME_ALL + elif domain in [FACE_RESTR_ALL, "all_faces"]: + domain = BoundaryDomainTag(FACE_RESTR_ALL) + elif domain in [FACE_RESTR_INTERIOR, "int_faces"]: + domain = BoundaryDomainTag(FACE_RESTR_INTERIOR) + elif isinstance(domain, (BTAG_PARTITION, BTAG_MULTIVOL_PARTITION)): + domain = BoundaryDomainTag(domain, _contextual_volume_tag) + elif domain in [BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE]: + domain = BoundaryDomainTag(domain, _contextual_volume_tag) + else: + raise ValueError("domain tag not understood: %s" % domain) + + if domain is DTAG_SCALAR and discretization_tag is not None: + raise ValueError("cannot have nontrivial discretization tag on scalar") + + if discretization_tag is None: + discretization_tag = DISCR_TAG_BASE + + return domain, discretization_tag + + +ConvertibleToDOFDesc = Any -DD_VOLUME = DOFDesc(DTAG_VOLUME_ALL, None) -DD_VOLUME_MODAL = DOFDesc(DTAG_VOLUME_ALL, DISCR_TAG_MODAL) +def as_dofdesc( + domain: "ConvertibleToDOFDesc", + discretization_tag: Optional[DiscretizationTag] = None, + *, _contextual_volume_tag: Optional[VolumeTag] = None) -> DOFDesc: + """ + :arg domain_tag: One of the following: + :class:`DTAG_SCALAR` (or the string ``"scalar"``), + :class:`DTAG_VOLUME_ALL` (or the string ``"vol"``) + for the default volume discretization, + :data:`~meshmode.discretization.connection.FACE_RESTR_ALL` + (or the string ``"all_faces"``), or + :data:`~meshmode.discretization.connection.FACE_RESTR_INTERIOR` + (or the string ``"int_faces"``), or one of + :class:`~meshmode.mesh.BTAG_ALL`, + :class:`~meshmode.mesh.BTAG_NONE`, + :class:`~meshmode.mesh.BTAG_REALLY_ALL`, + :class:`~meshmode.mesh.BTAG_PARTITION`, + or *None* to indicate that the geometry is not yet known. + + :arg discretization_tag: + *None* or :class:`DISCR_TAG_BASE` to indicate the use of the basic + discretization grid, :class:`DISCR_TAG_MODAL` to indicate a + modal discretization, or :class:`DISCR_TAG_QUAD` to indicate + the use of a quadrature grid. + """ + + if isinstance(domain, DOFDesc): + return domain + domain, discretization_tag = _normalize_domain_and_discr_tag( + domain, discretization_tag, + _contextual_volume_tag=_contextual_volume_tag) -def as_dofdesc(dd): - if isinstance(dd, DOFDesc): - return dd - return DOFDesc(dd, discretization_tag=None) + return DOFDesc(domain, discretization_tag) # }}} -# {{{ Deprecated tags +# {{{ deprecations -_deprecated_name_to_new_name = {"QTAG_NONE": "DISCR_TAG_BASE", - "QTAG_MODAL": "DISCR_TAG_MODAL"} +_deprecated_name_to_new_name = { + "DTAG_VOLUME": "VolumeDomainTag", + "DTAG_BOUNDARY": "BoundaryDomainTag", + "DD_VOLUME": "DD_VOLUME_ALL", + "DD_VOLUME_MODAL": "DD_VOLUME_ALL_MODAL" + } def __getattr__(name): if name in _deprecated_name_to_new_name: warn(f"'{name}' is deprecated and will be dropped " - f"in version 2022.x. Use '{_deprecated_name_to_new_name[name]}' " + f"in version 2023.x. Use '{_deprecated_name_to_new_name[name]}' " "instead.", DeprecationWarning, stacklevel=2) return globals()[_deprecated_name_to_new_name[name]] diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 6c846cc49..a108cddbd 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -43,6 +43,7 @@ """ +from typing import Optional, Sequence import numpy as np from arraycontext import ArrayContext, Scalar, tag_axes @@ -52,7 +53,8 @@ DiscretizationFaceAxisTag, DiscretizationElementAxisTag) -from grudge.dof_desc import DD_VOLUME, DOFDesc, as_dofdesc +from grudge.dof_desc import ( + DD_VOLUME_ALL, DOFDesc, as_dofdesc, BoundaryDomainTag, FACE_RESTR_ALL) from grudge.discretization import DiscretizationCollection import grudge.op as op @@ -63,7 +65,8 @@ def characteristic_lengthscales( - actx: ArrayContext, dcoll: DiscretizationCollection) -> DOFArray: + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None) -> DOFArray: r"""Computes the characteristic length scale :math:`h_{\text{loc}}` at each node. The characteristic length scale is mainly useful for estimating the stable time step size. E.g. for a hyperbolic system, an estimate of the @@ -103,15 +106,16 @@ def _compute_characteristic_lengthscales(): # corresponding group non-geometric factor cng * geo_facts for cng, geo_facts in zip( - dt_non_geometric_factors(dcoll), - actx.thaw(dt_geometric_factors(dcoll))))))) + dt_non_geometric_factors(dcoll, dd), + actx.thaw(dt_geometric_factors(dcoll, dd))))))) return actx.thaw(_compute_characteristic_lengthscales()) @memoize_on_first_arg def dt_non_geometric_factors( - dcoll: DiscretizationCollection, dd=None) -> list: + dcoll: DiscretizationCollection, dd: Optional[DOFDesc] = None + ) -> Sequence[float]: r"""Computes the non-geometric scale factors following [Hesthaven_2008]_, section 6.4, for each element group in the *dd* discretization: @@ -128,7 +132,7 @@ def dt_non_geometric_factors( node distance on the reference element for each group. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL discr = dcoll.discr_from_dd(dd) min_delta_rs = [] @@ -160,7 +164,8 @@ def dt_non_geometric_factors( @memoize_on_first_arg def h_max_from_volume( - dcoll: DiscretizationCollection, dim=None, dd=None) -> Scalar: + dcoll: DiscretizationCollection, dim=None, + dd: Optional[DOFDesc] = None) -> Scalar: """Returns a (maximum) characteristic length based on the volume of the elements. This length may not be representative if the elements have very high aspect ratios. @@ -175,7 +180,7 @@ def h_max_from_volume( from grudge.reductions import nodal_max, elementwise_sum if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dd = as_dofdesc(dd) if dim is None: @@ -191,7 +196,8 @@ def h_max_from_volume( @memoize_on_first_arg def h_min_from_volume( - dcoll: DiscretizationCollection, dim=None, dd=None) -> Scalar: + dcoll: DiscretizationCollection, dim=None, + dd: Optional[DOFDesc] = None) -> Scalar: """Returns a (minimum) characteristic length based on the volume of the elements. This length may not be representative if the elements have very high aspect ratios. @@ -206,7 +212,7 @@ def h_min_from_volume( from grudge.reductions import nodal_min, elementwise_sum if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dd = as_dofdesc(dd) if dim is None: @@ -221,7 +227,7 @@ def h_min_from_volume( def dt_geometric_factors( - dcoll: DiscretizationCollection, dd=None) -> DOFArray: + dcoll: DiscretizationCollection, dd: Optional[DOFDesc] = None) -> DOFArray: r"""Computes a geometric scaling factor for each cell following [Hesthaven_2008]_, section 6.4, defined as the inradius (radius of an inscribed circle/sphere). @@ -245,7 +251,7 @@ def dt_geometric_factors( from meshmode.discretization.poly_element import SimplexElementGroupBase if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL actx = dcoll._setup_actx volm_discr = dcoll.discr_from_dd(dd) @@ -272,7 +278,8 @@ def dt_geometric_factors( # Inscribed "circle" radius is half the cell size return actx.freeze(cell_vols/2) - dd_face = DOFDesc("all_faces", dd.discretization_tag) + dd_face = dd.with_domain_tag( + BoundaryDomainTag(FACE_RESTR_ALL, dd.domain_tag.tag)) face_discr = dcoll.discr_from_dd(dd_face) # Compute areas of each face diff --git a/grudge/eager.py b/grudge/eager.py index 1886cfd04..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/geometry/metrics.py b/grudge/geometry/metrics.py index 89e1f1f2c..f5c622cd5 100644 --- a/grudge/geometry/metrics.py +++ b/grudge/geometry/metrics.py @@ -58,6 +58,7 @@ """ +from typing import Optional, Tuple, Union import numpy as np from arraycontext import ArrayContext, tag_axes @@ -68,7 +69,7 @@ import grudge.dof_desc as dof_desc from grudge.dof_desc import ( - DD_VOLUME, DOFDesc, DISCR_TAG_BASE + DD_VOLUME_ALL, DOFDesc, DISCR_TAG_BASE ) from meshmode.transform_metadata import (DiscretizationAmbientDimAxisTag, @@ -115,7 +116,8 @@ def to_quad(vec): def forward_metric_nth_derivative( actx: ArrayContext, dcoll: DiscretizationCollection, - xyz_axis, ref_axes, dd=None, + xyz_axis: int, ref_axes: Union[int, Tuple[Tuple[int, int], ...]], + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> DOFArray: r"""Pointwise metric derivatives representing repeated derivatives of the physical coordinate enumerated by *xyz_axis*: :math:`x_{\mathrm{xyz\_axis}}` @@ -150,7 +152,7 @@ def forward_metric_nth_derivative( metric derivative at each nodal coordinate. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL inner_dd = dd.with_discr_tag(DISCR_TAG_BASE) @@ -182,8 +184,10 @@ def forward_metric_nth_derivative( def forward_metric_derivative_vector( - actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, dd=None, - *, _use_geoderiv_connection=False) -> np.ndarray: + actx: ArrayContext, dcoll: DiscretizationCollection, + rst_axis: Union[int, Tuple[Tuple[int, int], ...]], + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False + ) -> np.ndarray: r"""Computes an object array containing the forward metric derivatives of each physical coordinate. @@ -207,7 +211,9 @@ def forward_metric_derivative_vector( def forward_metric_derivative_mv( - actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + rst_axis: Union[int, Tuple[Tuple[int, int], ...]], + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> MultiVector: r"""Computes a :class:`pymbolic.geometric_algebra.MultiVector` containing the forward metric derivatives of each physical coordinate. @@ -230,7 +236,8 @@ def forward_metric_derivative_mv( def forward_metric_derivative_mat( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> np.ndarray: r"""Computes the forward metric derivative matrix, also commonly called the Jacobian matrix, with entries defined as the @@ -257,7 +264,7 @@ def forward_metric_derivative_mat( ambient_dim = dcoll.ambient_dim if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim @@ -271,7 +278,8 @@ def forward_metric_derivative_mat( def first_fundamental_form(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None, *, _use_geoderiv_connection=False) -> np.ndarray: + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False + ) -> np.ndarray: r"""Computes the first fundamental form using the Jacobian matrix: .. math:: @@ -297,7 +305,7 @@ def first_fundamental_form(actx: ArrayContext, dcoll: DiscretizationCollection, form. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL mder = forward_metric_derivative_mat( actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection) @@ -306,7 +314,8 @@ def first_fundamental_form(actx: ArrayContext, dcoll: DiscretizationCollection, def inverse_metric_derivative_mat( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> np.ndarray: r"""Computes the inverse metric derivative matrix, which is the inverse of the Jacobian (forward metric derivative) matrix. @@ -320,7 +329,7 @@ def inverse_metric_derivative_mat( ambient_dim = dcoll.ambient_dim if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim @@ -336,7 +345,8 @@ def inverse_metric_derivative_mat( def inverse_first_fundamental_form( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> np.ndarray: r"""Computes the inverse of the first fundamental form: @@ -360,7 +370,7 @@ def inverse_first_fundamental_form( first fundamental form. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim @@ -387,7 +397,8 @@ def inverse_first_fundamental_form( def inverse_metric_derivative( - actx: ArrayContext, dcoll: DiscretizationCollection, rst_axis, xyz_axis, dd, + actx: ArrayContext, dcoll: DiscretizationCollection, + rst_axis: int, xyz_axis: int, dd: DOFDesc, *, _use_geoderiv_connection=False ) -> DOFArray: r"""Computes the inverse metric derivative of the physical @@ -446,7 +457,7 @@ def outprod_with_unit(i, at): def inverse_surface_metric_derivative( actx: ArrayContext, dcoll: DiscretizationCollection, - rst_axis, xyz_axis, dd=None, + rst_axis, xyz_axis, dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False): r"""Computes the inverse surface metric derivative of the physical coordinate enumerated by *xyz_axis* with respect to the @@ -468,7 +479,7 @@ def inverse_surface_metric_derivative( ambient_dim = dcoll.ambient_dim if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dd = dof_desc.as_dofdesc(dd) if ambient_dim == dim: @@ -488,7 +499,8 @@ def inverse_surface_metric_derivative( def inverse_surface_metric_derivative_mat( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, times_area_element=False, _use_geoderiv_connection=False): r"""Computes the matrix of inverse surface metric derivatives, indexed by ``(xyz_axis, rst_axis)``. It returns all values of @@ -509,7 +521,7 @@ def inverse_surface_metric_derivative_mat( """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dd = dof_desc.as_dofdesc(dd) @memoize_in(dcoll, (inverse_surface_metric_derivative_mat, dd, @@ -542,7 +554,7 @@ def _inv_surf_metric_deriv(): def _signed_face_ones( - actx: ArrayContext, dcoll: DiscretizationCollection, dd + actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc ) -> DOFArray: assert dd.is_trace() @@ -550,7 +562,7 @@ def _signed_face_ones( # NOTE: ignore quadrature_tags on dd, since we only care about # the face_id here all_faces_conn = dcoll.connection_from_dds( - DD_VOLUME, DOFDesc(dd.domain_tag) + DD_VOLUME_ALL, DOFDesc(dd.domain_tag, DISCR_TAG_BASE) ) signed_ones = dcoll.discr_from_dd(dd.with_discr_tag(DISCR_TAG_BASE)).zeros( actx, dtype=dcoll.real_dtype @@ -571,7 +583,7 @@ def _signed_face_ones( def parametrization_derivative( - actx: ArrayContext, dcoll: DiscretizationCollection, dd, + actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc, *, _use_geoderiv_connection=False) -> MultiVector: r"""Computes the product of forward metric derivatives spanning the tangent space with topological dimension *dim*. @@ -584,7 +596,7 @@ def parametrization_derivative( the product of metric derivatives. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim if dim == 0: @@ -605,8 +617,10 @@ def parametrization_derivative( ) -def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None, *, _use_geoderiv_connection=False) -> MultiVector: +def pseudoscalar( + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False + ) -> MultiVector: r"""Computes the field of pseudoscalars for the domain/discretization identified by *dd*. @@ -618,7 +632,7 @@ def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection, :class:`~meshmode.dof_array.DOFArray`\ s. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL return parametrization_derivative( actx, dcoll, dd, @@ -626,7 +640,8 @@ def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection, def area_element( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False ) -> DOFArray: r"""Computes the scale factor used to transform integrals from reference @@ -642,7 +657,7 @@ def area_element( volumes for each element. """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL @memoize_in(dcoll, (area_element, dd, _use_geoderiv_connection)) def _area_elements(): @@ -662,7 +677,8 @@ def _area_elements(): # {{{ surface normal vectors def rel_mv_normal( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None, + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None, *, _use_geoderiv_connection=False) -> MultiVector: r"""Computes surface normals at each nodal location as a :class:`~pymbolic.geometric_algebra.MultiVector` relative to the @@ -688,7 +704,7 @@ def rel_mv_normal( def mv_normal( - actx: ArrayContext, dcoll: DiscretizationCollection, dd, + actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc, *, _use_geoderiv_connection=False ) -> MultiVector: r"""Exterior unit normal as a :class:`~pymbolic.geometric_algebra.MultiVector`. @@ -744,10 +760,10 @@ def _normal(): from grudge.op import project volm_normal = MultiVector( - project(dcoll, dof_desc.DD_VOLUME, dd, + project(dcoll, DD_VOLUME_ALL, dd, rel_mv_normal( actx, dcoll, - dd=dof_desc.DD_VOLUME, + dd=DD_VOLUME_ALL, _use_geoderiv_connection=_use_geoderiv_connection ).as_vector(dtype=object)) ) @@ -768,7 +784,7 @@ def _normal(): return actx.thaw(_normal()) -def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd, +def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc, *, _use_geoderiv_connection=None): """Get the unit normal to the specified surface discretization, *dd*. This supports both volume discretizations @@ -798,8 +814,8 @@ def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd, # {{{ Curvature computations def second_fundamental_form( - actx: ArrayContext, dcoll: DiscretizationCollection, dd=None - ) -> np.ndarray: + actx: ArrayContext, dcoll: DiscretizationCollection, + dd: Optional[DOFDesc] = None) -> np.ndarray: r"""Computes the second fundamental form: .. math:: @@ -817,7 +833,7 @@ def second_fundamental_form( """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.discr_from_dd(dd).dim normal = rel_mv_normal(actx, dcoll, dd=dd).as_vector(dtype=object) @@ -846,7 +862,7 @@ def second_fundamental_form( def shape_operator(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None) -> np.ndarray: + dd: Optional[DOFDesc] = None) -> np.ndarray: r"""Computes the shape operator (also called the curvature tensor) containing second order derivatives: @@ -871,7 +887,7 @@ def shape_operator(actx: ArrayContext, dcoll: DiscretizationCollection, def summed_curvature(actx: ArrayContext, dcoll: DiscretizationCollection, - dd=None) -> DOFArray: + dd: Optional[DOFDesc] = None) -> DOFArray: r"""Computes the sum of the principal curvatures: .. math:: @@ -888,7 +904,7 @@ def summed_curvature(actx: ArrayContext, dcoll: DiscretizationCollection, """ if dd is None: - dd = DD_VOLUME + dd = DD_VOLUME_ALL dim = dcoll.ambient_dim - 1 diff --git a/grudge/models/advection.py b/grudge/models/advection.py index cfe1a4920..c5d35e2d1 100644 --- a/grudge/models/advection.py +++ b/grudge/models/advection.py @@ -214,13 +214,13 @@ def __init__(self, dcoll, v, inflow_u, flux_type="central", quad_tag=None): self.quad_tag = quad_tag def flux(self, u_tpair): - from grudge.dof_desc import DD_VOLUME + from grudge.dof_desc import DD_VOLUME_ALL - surf_v = op.project(self.dcoll, DD_VOLUME, u_tpair.dd, self.v) + surf_v = op.project(self.dcoll, DD_VOLUME_ALL, u_tpair.dd, self.v) return advection_weak_flux(self.dcoll, self.flux_type, u_tpair, surf_v) def operator(self, t, u): - from grudge.dof_desc import DOFDesc, DD_VOLUME, DTAG_VOLUME_ALL + from grudge.dof_desc import DOFDesc, DD_VOLUME_ALL, DTAG_VOLUME_ALL from meshmode.mesh import BTAG_ALL from meshmode.discretization.connection import FACE_RESTR_ALL @@ -234,7 +234,7 @@ def flux(tpair): return op.project(dcoll, tpair.dd, face_dd, self.flux(tpair)) def to_quad(arg): - return op.project(dcoll, DD_VOLUME, quad_dd, arg) + return op.project(dcoll, DD_VOLUME_ALL, quad_dd, arg) if self.inflow_u is not None: inflow_flux = flux(op.bv_trace_pair(dcoll, @@ -279,7 +279,7 @@ def to_quad(arg): # {{{ closed surface advection def v_dot_n_tpair(actx, dcoll, velocity, trace_dd): - from grudge.dof_desc import DTAG_BOUNDARY + from grudge.dof_desc import BoundaryDomainTag from grudge.trace_pair import TracePair from meshmode.discretization.connection import FACE_RESTR_INTERIOR @@ -287,10 +287,9 @@ def v_dot_n_tpair(actx, dcoll, velocity, trace_dd): v_dot_n = velocity.dot(normal) i = op.project(dcoll, trace_dd.with_discr_tag(None), trace_dd, v_dot_n) - if trace_dd.domain_tag is FACE_RESTR_INTERIOR: - e = dcoll.opposite_face_connection()(i) - elif isinstance(trace_dd.domain_tag, DTAG_BOUNDARY): - e = dcoll.distributed_boundary_swap_connection(trace_dd)(i) + assert isinstance(trace_dd.domain_tag, BoundaryDomainTag) + if trace_dd.domain_tag.tag is FACE_RESTR_INTERIOR: + e = dcoll.opposite_face_connection(trace_dd.domain_tag)(i) else: raise ValueError("Unrecognized domain tag: %s" % trace_dd.domain_tag) @@ -325,9 +324,9 @@ def __init__(self, dcoll, v, flux_type="central", quad_tag=None): self.quad_tag = quad_tag def flux(self, u_tpair): - from grudge.dof_desc import DD_VOLUME + from grudge.dof_desc import DD_VOLUME_ALL - surf_v = op.project(self.dcoll, DD_VOLUME, + surf_v = op.project(self.dcoll, DD_VOLUME_ALL, u_tpair.dd.with_discr_tag(None), self.v) return surface_advection_weak_flux(self.dcoll, self.flux_type, @@ -335,7 +334,7 @@ def flux(self, u_tpair): surf_v) def operator(self, t, u): - from grudge.dof_desc import DOFDesc, DD_VOLUME, DTAG_VOLUME_ALL + from grudge.dof_desc import DOFDesc, DD_VOLUME_ALL, DTAG_VOLUME_ALL from meshmode.discretization.connection import FACE_RESTR_ALL face_dd = DOFDesc(FACE_RESTR_ALL, self.quad_tag) @@ -347,7 +346,7 @@ def flux(tpair): return op.project(dcoll, tpair.dd, face_dd, self.flux(tpair)) def to_quad(arg): - return op.project(dcoll, DD_VOLUME, quad_dd, arg) + return op.project(dcoll, DD_VOLUME_ALL, quad_dd, arg) quad_v = to_quad(self.v) quad_u = to_quad(u) diff --git a/grudge/op.py b/grudge/op.py index 015e0718b..967a66ac3 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -88,6 +88,7 @@ import numpy as np import grudge.dof_desc as dof_desc +from grudge.dof_desc import DD_VOLUME_ALL, FACE_RESTR_ALL from grudge.interpolation import interp from grudge.projection import project @@ -113,8 +114,10 @@ interior_trace_pair, interior_trace_pairs, local_interior_trace_pair, - connected_ranks, + inter_volume_trace_pairs, + local_inter_volume_trace_pairs, cross_rank_trace_pairs, + cross_rank_inter_volume_trace_pairs, bdry_trace_pair, bv_trace_pair ) @@ -142,8 +145,10 @@ "interior_trace_pair", "interior_trace_pairs", "local_interior_trace_pair", - "connected_ranks", + "inter_volume_trace_pairs", + "local_inter_volume_trace_pairs", "cross_rank_trace_pairs", + "cross_rank_inter_volume_trace_pairs", "bdry_trace_pair", "bv_trace_pair", @@ -244,11 +249,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, @@ -276,7 +281,7 @@ def local_grad( :class:`~meshmode.dof_array.DOFArray`\ s or :class:`~arraycontext.ArrayContainer` of object arrays. """ - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = DD_VOLUME_ALL from grudge.tools import rec_map_subarrays return rec_map_subarrays( partial(_strong_scalar_grad, dcoll, dd_in), @@ -303,7 +308,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 @@ -392,7 +397,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, @@ -429,7 +434,7 @@ def weak_local_grad( """ if len(args) == 1: vecs, = args - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = DD_VOLUME_ALL elif len(args) == 2: dd_in, vecs = args else: @@ -474,7 +479,7 @@ def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: """ if len(args) == 2: xyz_axis, vec = args - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = dof_desc.DD_VOLUME_ALL elif len(args) == 3: dd_in, xyz_axis, vec = args else: @@ -489,7 +494,7 @@ def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: 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, @@ -533,7 +538,7 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: """ if len(args) == 1: vecs, = args - dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd_in = DD_VOLUME_ALL elif len(args) == 2: dd_in, vecs = args else: @@ -650,13 +655,13 @@ def mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: 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) # }}} @@ -751,7 +756,7 @@ def inverse_mass(dcoll: DiscretizationCollection, vec) -> ArrayOrContainer: """ return _apply_inverse_mass_operator( - dcoll, dof_desc.DD_VOLUME, dof_desc.DD_VOLUME, vec + dcoll, DD_VOLUME_ALL, DD_VOLUME_ALL, vec ) # }}} @@ -858,7 +863,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 @@ -932,7 +937,7 @@ def face_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: if len(args) == 1: vec, = args - dd = dof_desc.DOFDesc("all_faces", dof_desc.DISCR_TAG_BASE) + dd = DD_VOLUME_ALL.trace(FACE_RESTR_ALL) elif len(args) == 2: dd, vec = args else: diff --git a/grudge/projection.py b/grudge/projection.py index 425239591..ba8d4bc3c 100644 --- a/grudge/projection.py +++ b/grudge/projection.py @@ -37,13 +37,15 @@ from arraycontext import ArrayOrContainer from grudge.discretization import DiscretizationCollection -from grudge.dof_desc import as_dofdesc +from grudge.dof_desc import as_dofdesc, VolumeDomainTag, ConvertibleToDOFDesc from numbers import Number def project( - dcoll: DiscretizationCollection, src, tgt, vec) -> ArrayOrContainer: + dcoll: DiscretizationCollection, + src: "ConvertibleToDOFDesc", + tgt: "ConvertibleToDOFDesc", vec) -> ArrayOrContainer: """Project from one discretization to another, e.g. from the volume to the boundary, or from the base to the an overintegrated quadrature discretization. @@ -55,10 +57,22 @@ def project( :returns: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` like *vec*. """ - src = as_dofdesc(src) - tgt = as_dofdesc(tgt) + # {{{ process dofdesc arguments - if isinstance(vec, Number) or src == tgt: + src_dofdesc = as_dofdesc(src) + + contextual_volume_tag = None + if isinstance(src_dofdesc.domain_tag, VolumeDomainTag): + contextual_volume_tag = src_dofdesc.domain_tag.tag + + tgt_dofdesc = as_dofdesc(tgt, _contextual_volume_tag=contextual_volume_tag) + + del src + del tgt + + # }}} + + if isinstance(vec, Number) or src_dofdesc == tgt_dofdesc: return vec - return dcoll.connection_from_dds(src, tgt)(vec) + return dcoll.connection_from_dds(src_dofdesc, tgt_dofdesc)(vec) diff --git a/grudge/reductions.py b/grudge/reductions.py index 95ed44726..ab106c8c4 100644 --- a/grudge/reductions.py +++ b/grudge/reductions.py @@ -94,7 +94,7 @@ def norm(dcoll: DiscretizationCollection, vec, p, dd=None) -> Scalar: :returns: a nonegative scalar denoting the norm. """ if dd is None: - dd = dof_desc.DD_VOLUME + dd = dof_desc.DD_VOLUME_ALL from arraycontext import get_container_context_recursively actx = get_container_context_recursively(vec) @@ -128,7 +128,7 @@ def nodal_sum(dcoll: DiscretizationCollection, dd, vec) -> Scalar: if comm is None: return nodal_sum_loc(dcoll, dd, vec) - # NOTE: Don't move this + # NOTE: Do not move, we do not want to import mpi4py in single-rank computations from mpi4py import MPI from arraycontext import get_container_context_recursively @@ -174,7 +174,7 @@ def nodal_min(dcoll: DiscretizationCollection, dd, vec, *, initial=None) -> Scal if comm is None: return nodal_min_loc(dcoll, dd, vec, initial=initial) - # NOTE: Don't move this + # NOTE: Do not move, we do not want to import mpi4py in single-rank computations from mpi4py import MPI actx = vec.array_context @@ -231,7 +231,7 @@ def nodal_max(dcoll: DiscretizationCollection, dd, vec, *, initial=None) -> Scal if comm is None: return nodal_max_loc(dcoll, dd, vec, initial=initial) - # NOTE: Don't move this + # NOTE: Do not move, we do not want to import mpi4py in single-rank computations from mpi4py import MPI actx = vec.array_context @@ -320,7 +320,7 @@ def _apply_elementwise_reduction( """ if len(args) == 1: vec, = args - dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd = dof_desc.DD_VOLUME_ALL elif len(args) == 2: dd, vec = args else: @@ -485,7 +485,7 @@ def elementwise_integral( """ if len(args) == 1: vec, = args - dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + dd = dof_desc.DD_VOLUME_ALL elif len(args) == 2: dd, vec = args else: diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index cb2de38f6..3db3517a6 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -18,12 +18,15 @@ .. 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__ = """ @@ -51,17 +54,20 @@ """ -from typing import List, Hashable, Optional, Type, Any +from warnings import warn +from typing import List, Hashable, Optional, Type, Any, Sequence from pytools.persistent_dict import KeyBuilder from arraycontext import ( ArrayContainer, + ArrayContext, with_container_arithmetic, dataclass_array_container, get_container_context_recursively, flatten, to_numpy, unflatten, from_numpy, + flat_size_and_dtype, ArrayOrContainer ) @@ -72,13 +78,19 @@ 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, VolumeDomainTag, BoundaryDomainTag, BTAG_MULTIVOL_PARTITION, + ConvertibleToDOFDesc, + ) # {{{ trace pair container class @@ -107,12 +119,22 @@ class TracePair: .. automethod:: __len__ """ - dd: dof_desc.DOFDesc + dd: DOFDesc interior: ArrayContainer exterior: ArrayContainer - def __init__(self, dd, *, interior, exterior): - object.__setattr__(self, "dd", dof_desc.as_dofdesc(dd)) + def __init__(self, dd: DOFDesc, *, + interior: ArrayOrContainer, + exterior: ArrayOrContainer): + if not isinstance(dd, DOFDesc): + warn("Constructing a TracePair with a first argument that is not " + "exactly a DOFDesc (but convertible to one) is deprecated. " + "This will stop working in July 2022. " + "Pass an actual DOFDesc instead.", + DeprecationWarning, stacklevel=2) + dd = dof_desc.as_dofdesc(dd) + + object.__setattr__(self, "dd", dd) object.__setattr__(self, "interior", interior) object.__setattr__(self, "exterior", exterior) @@ -178,7 +200,8 @@ def diff(self): # {{{ boundary trace pairs def bdry_trace_pair( - dcoll: DiscretizationCollection, dd, interior, exterior) -> TracePair: + dcoll: DiscretizationCollection, dd: "ConvertibleToDOFDesc", + interior, exterior) -> TracePair: """Returns a trace pair defined on the exterior boundary. Input arguments are assumed to already be defined on the boundary denoted by *dd*. If the input arguments *interior* and *exterior* are @@ -197,11 +220,19 @@ def bdry_trace_pair( be used for the flux. :returns: a :class:`TracePair` on the boundary. """ + if not isinstance(dd, DOFDesc): + warn("Calling bdry_trace_pair with a first argument that is not " + "exactly a DOFDesc (but convertible to one) is deprecated. " + "This will stop working in July 2022. " + "Pass an actual DOFDesc instead.", + DeprecationWarning, stacklevel=2) + dd = dof_desc.as_dofdesc(dd) return TracePair(dd, interior=interior, exterior=exterior) def bv_trace_pair( - dcoll: DiscretizationCollection, dd, interior, exterior) -> TracePair: + dcoll: DiscretizationCollection, dd: "ConvertibleToDOFDesc", + interior, exterior) -> TracePair: """Returns a trace pair defined on the exterior boundary. The interior argument is assumed to be defined on the volume discretization, and will therefore be restricted to the boundary *dd* prior to creating a @@ -223,21 +254,29 @@ def bv_trace_pair( be used for the flux. :returns: a :class:`TracePair` on the boundary. """ + if not isinstance(dd, DOFDesc): + warn("Calling bv_trace_pair with a first argument that is not " + "exactly a DOFDesc (but convertible to one) is deprecated. " + "This will stop working in July 2022. " + "Pass an actual DOFDesc instead.", + DeprecationWarning, stacklevel=2) + dd = dof_desc.as_dofdesc(dd) return bdry_trace_pair( - dcoll, dd, project(dcoll, "vol", dd, interior), exterior - ) + dcoll, dd, project(dcoll, "vol", dd, interior), exterior) # }}} # {{{ interior trace pairs -def local_interior_trace_pair(dcoll: DiscretizationCollection, vec) -> TracePair: +def local_interior_trace_pair( + dcoll: DiscretizationCollection, vec, *, + volume_dd: Optional[DOFDesc] = None, + ) -> TracePair: r"""Return a :class:`TracePair` for the interior faces of *dcoll* with a discretization tag specified by *discr_tag*. This does not include interior faces on different MPI ranks. - :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.ArrayContainer` of them. @@ -250,21 +289,27 @@ 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 - def get_opposite_face(el): + assert isinstance(volume_dd.domain_tag, VolumeDomainTag) + trace_dd = volume_dd.trace(FACE_RESTR_INTERIOR) + + interior = project(dcoll, volume_dd, trace_dd, vec) + + def get_opposite_trace(el): if isinstance(el, Number): return el else: - return dcoll.opposite_face_connection()(el) + assert isinstance(trace_dd.domain_tag, BoundaryDomainTag) + 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: - from warnings import warn warn("`grudge.op.interior_trace_pair` is deprecated and will be dropped " "in version 2022.x. Use `local_interior_trace_pair` " "instead, or `interior_trace_pairs` which also includes contributions " @@ -274,7 +319,8 @@ def interior_trace_pair(dcoll: DiscretizationCollection, vec) -> TracePair: def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *, - comm_tag: Hashable = None, tag: Hashable = None) -> List[TracePair]: + comm_tag: Hashable = None, tag: Hashable = None, + volume_dd: Optional[DOFDesc] = None) -> List[TracePair]: r"""Return a :class:`list` of :class:`TracePair` objects defined on the interior faces of *dcoll* and any faces connected to a parallel boundary. @@ -293,7 +339,6 @@ def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *, """ if tag is not None: - from warnings import warn warn("Specifying 'tag' is deprecated and will stop working in July of 2022. " "Specify 'comm_tag' instead.", DeprecationWarning, stacklevel=2) if comm_tag is not None: @@ -302,148 +347,293 @@ def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *, comm_tag = tag del tag + if volume_dd is None: + volume_dd = DD_VOLUME_ALL + return ( - [local_interior_trace_pair(dcoll, vec)] - + cross_rank_trace_pairs(dcoll, vec, comm_tag=comm_tag) + [local_interior_trace_pair( + dcoll, vec, volume_dd=volume_dd)] + + cross_rank_trace_pairs( + dcoll, vec, comm_tag=comm_tag, volume_dd=volume_dd) ) # }}} -# {{{ distributed-memory functionality +# {{{ inter-volume trace pairs + +def local_inter_volume_trace_pairs( + dcoll: DiscretizationCollection, + self_volume_dd: DOFDesc, self_ary: ArrayOrContainer, + other_volume_dd: DOFDesc, other_ary: ArrayOrContainer, + ) -> ArrayOrContainer: + if not isinstance(self_volume_dd.domain_tag, VolumeDomainTag): + raise ValueError("self_volume_dd must describe a volume") + if not isinstance(other_volume_dd.domain_tag, VolumeDomainTag): + 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, + other_volume_tag=other_volume_dd.domain_tag.tag) + other_btag = BTAG_MULTIVOL_PARTITION( + other_rank=None, + 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[ + BoundaryDomainTag(self_btag, self_volume_dd.domain_tag.tag)] + + 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: ArrayOrContainer, + other_volume_dd: DOFDesc, other_ary: ArrayOrContainer, + comm_tag: Hashable = None) -> List[ArrayOrContainer]: + """ + 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_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: 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[BoundaryDomainTag]: + if other_volume_tag is None: + other_volume_tag = self_volume_tag + + result: List[BoundaryDomainTag] = [] + for domain_tag in dcoll._inter_partition_connections: + if isinstance(domain_tag.tag, BTAG_PARTITION): + if domain_tag.volume_tag == self_volume_tag: + result.append(domain_tag) + + elif isinstance(domain_tag.tag, BTAG_MULTIVOL_PARTITION): + if (domain_tag.tag.other_rank is not None + and domain_tag.volume_tag == self_volume_tag + and domain_tag.tag.other_volume_tag == other_volume_tag): + result.append(domain_tag) + + else: + raise AssertionError("unexpected inter-partition tag type encountered: " + f"'{domain_tag.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 -class _RankBoundaryCommunication: + 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 _RankBoundaryCommunicationEager: base_comm_tag = 1273 def __init__(self, - dcoll: DiscretizationCollection, - array_container: ArrayOrContainer, - remote_rank, comm_tag: Optional[int] = None): - actx = get_container_context_recursively(array_container) - btag = BTAG_PARTITION(remote_rank) + actx: ArrayContext, + dcoll: DiscretizationCollection, + domain_tag: BoundaryDomainTag, + *, local_bdry_data: ArrayOrContainer, + send_data: ArrayOrContainer, + 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(domain_tag.tag) + assert remote_rank is not None self.dcoll = dcoll self.array_context = actx - self.remote_btag = btag - self.bdry_discr = dcoll.discr_from_dd(btag) + self.domain_tag = domain_tag + self.bdry_discr = dcoll.discr_from_dd(domain_tag) 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.domain_tag] + 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, - interior=self.local_bdry_data, - exterior=swapped_remote_bdry_data) + return TracePair( + DOFDesc(self.domain_tag, DISCR_TAG_BASE), + interior=self.local_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: ArrayOrContainer, - remote_rank: int, comm_tag: Hashable): + actx: ArrayContext, + dcoll: DiscretizationCollection, + domain_tag: BoundaryDomainTag, + *, + local_bdry_data: ArrayOrContainer, + send_data: ArrayOrContainer, + 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(domain_tag) + self.domain_tag = domain_tag + + remote_rank = _remote_rank_from_btag(domain_tag.tag) + 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, - axes=local_bdry_ary.axes)) + shape=local_bdry_subary.shape, + dtype=local_bdry_subary.dtype, + axes=local_bdry_subary.axes)) - 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.domain_tag] - return TracePair(self.remote_btag, - interior=self.local_bdry_data, - exterior=bdry_conn(self.remote_data)) + return TracePair( + DOFDesc(self.domain_tag, DISCR_TAG_BASE), + interior=self.local_bdry_data, + exterior=bdry_conn(self.remote_data)) +# }}} -class _TagKeyBuilder(KeyBuilder): - def update_for_type(self, key_hash, key: Type[Any]): - self.rec(key_hash, (key.__module__, key.__name__, key.__name__,)) +# {{{ cross_rank_trace_pairs def cross_rank_trace_pairs( - dcoll: DiscretizationCollection, ary, - comm_tag: Hashable = None, - tag: Hashable = None) -> List[TracePair]: + dcoll: DiscretizationCollection, ary: ArrayOrContainer, + tag: Hashable = None, + *, comm_tag: Hashable = None, + volume_dd: Optional[DOFDesc] = None) -> List[TracePair]: r"""Get a :class:`list` of *ary* trace pairs for each partition boundary. For each partition boundary, the field data values in *ary* are @@ -472,60 +662,144 @@ def cross_rank_trace_pairs( :returns: a :class:`list` of :class:`TracePair` objects. """ + # {{{ process arguments + + if volume_dd is None: + volume_dd = DD_VOLUME_ALL + + if not isinstance(volume_dd.domain_tag, VolumeDomainTag): + raise TypeError(f"expected a volume DOFDesc, got '{volume_dd}'") + if volume_dd.discretization_tag != DISCR_TAG_BASE: + raise TypeError(f"expected a base-discretized DOFDesc, got '{volume_dd}'") + if tag is not None: - from warnings import warn warn("Specifying 'tag' is deprecated and will stop working in July of 2022. " - "Specify 'comm_tag' instead.", DeprecationWarning, stacklevel=2) + "Specify 'comm_tag' (keyword-only) instead.", + DeprecationWarning, stacklevel=2) if comm_tag is not None: raise TypeError("may only specify one of 'tag' and 'comm_tag'") else: comm_tag = tag del tag + # }}} + + comm_bdtags = _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_bdtags) == len( + {_remote_rank_from_btag(bdtag.tag) for bdtag in comm_bdtags}) + 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(bdtag, DISCR_TAG_BASE), interior=ary, exterior=ary) + for bdtag in comm_bdtags] 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(bdtag): + local_bdry_data = project( + dcoll, + DOFDesc(VolumeDomainTag(bdtag.volume_tag), DISCR_TAG_BASE), + DOFDesc(bdtag, DISCR_TAG_BASE), + ary) + + return rbc(actx, dcoll, bdtag, + local_bdry_data=local_bdry_data, + send_data=local_bdry_data, + comm_tag=comm_tag) + + rank_bdry_communcators = [start_comm(bdtag) for bdtag in comm_bdtags] + 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: ArrayOrContainer, + other_volume_dd: DOFDesc, other_ary: ArrayOrContainer, + *, 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, VolumeDomainTag): + raise ValueError("self_volume_dd must describe a volume") + if not isinstance(other_volume_dd.domain_tag, VolumeDomainTag): + 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_bdtags = _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_bdtags) == len( + {_remote_rank_from_btag(bdtag.tag) for bdtag in comm_bdtags}) + + 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(bdtag): + assert isinstance(bdtag.tag, BTAG_MULTIVOL_PARTITION) + self_volume_dd = DOFDesc( + VolumeDomainTag(bdtag.volume_tag), DISCR_TAG_BASE) + other_volume_dd = DOFDesc( + VolumeDomainTag(bdtag.tag.other_volume_tag), DISCR_TAG_BASE) + + local_bdry_data = project(dcoll, self_volume_dd, bdtag, self_ary) + send_data = project(dcoll, other_volume_dd, + BTAG_MULTIVOL_PARTITION( + other_rank=bdtag.tag.other_rank, + other_volume_tag=bdtag.volume_tag), other_ary) + + return rbc(actx, dcoll, bdtag, + local_bdry_data=local_bdry_data, + send_data=send_data, + comm_tag=comm_tag) + + rank_bdry_communcators = [start_comm(bdtag) for bdtag in comm_bdtags] 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 ce0b199b8..a0bb9ac18 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -38,7 +38,7 @@ from pytools.obj_array import flat_obj_array -from grudge import DiscretizationCollection +from grudge import DiscretizationCollection, make_discretization_collection import grudge.dof_desc as dof_desc import grudge.op as op @@ -341,7 +341,10 @@ def test_face_normal_surface(actx_factory, mesh_name): surf_normal = surf_normal / actx.np.sqrt(sum(surf_normal**2)) face_normal_i = actx.thaw(dcoll.normal(df)) - face_normal_e = dcoll.opposite_face_connection()(face_normal_i) + face_normal_e = dcoll.opposite_face_connection( + dof_desc.BoundaryDomainTag( + dof_desc.FACE_RESTR_INTERIOR, dof_desc.VTAG_ALL) + )(face_normal_i) if mesh.ambient_dim == 3: from grudge.geometry import pseudoscalar, area_element @@ -979,6 +982,8 @@ def bessel_j(actx, n, r): # }}} +# {{{ test norms + @pytest.mark.parametrize("p", [2, np.inf]) def test_norm_real(actx_factory, p): actx = actx_factory() @@ -1042,6 +1047,10 @@ def test_norm_obj_array(actx_factory, p): logger.info("norm: %.5e %.5e", norm, ref_norm) assert abs(norm-ref_norm) / abs(ref_norm) < 1e-14 +# }}} + + +# {{{ empty boundaries def test_empty_boundary(actx_factory): # https://github.com/inducer/grudge/issues/54 @@ -1061,6 +1070,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, + 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_mpi_communication.py b/test/test_mpi_communication.py index 47100b415..424b78199 100644 --- a/test/test_mpi_communication.py +++ b/test/test_mpi_communication.py @@ -45,6 +45,7 @@ from pytools.obj_array import flat_obj_array import grudge.op as op +import grudge.dof_desc as dof_desc class SimpleTag: @@ -153,7 +154,10 @@ def hopefully_zero(): return ( op.project( dcoll, "int_faces", "all_faces", - dcoll.opposite_face_connection()(int_faces_func) + dcoll.opposite_face_connection( + dof_desc.BoundaryDomainTag( + dof_desc.FACE_RESTR_INTERIOR, dof_desc.VTAG_ALL) + )(int_faces_func) ) + sum(op.project(dcoll, tpair.dd, "all_faces", tpair.ext) for tpair in op.cross_rank_trace_pairs(dcoll, myfunc, diff --git a/test/test_reductions.py b/test/test_reductions.py index 9e7387bad..e0e424fe4 100644 --- a/test/test_reductions.py +++ b/test/test_reductions.py @@ -149,7 +149,7 @@ def f(x): min_res = np.empty(grp_f.shape) max_res = np.empty(grp_f.shape) sum_res = np.empty(grp_f.shape) - for eidx in range(dcoll.mesh.nelements): + for eidx in range(mesh.nelements): element_data = actx.to_numpy(grp_f[eidx]) min_res[eidx, :] = np.min(element_data) max_res[eidx, :] = np.max(element_data) @@ -272,7 +272,7 @@ def _get_ref_data(field): min_res = np.empty(grp_f.shape) max_res = np.empty(grp_f.shape) sum_res = np.empty(grp_f.shape) - for eidx in range(dcoll.mesh.nelements): + for eidx in range(mesh.nelements): element_data = actx.to_numpy(grp_f[eidx]) min_res[eidx, :] = np.min(element_data) max_res[eidx, :] = np.max(element_data)