From ef79f3576679a344f6c2ae341c66fee754b59cf8 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 10 Jan 2023 12:36:56 -0600 Subject: [PATCH 01/97] Support new find_distributed_partition https://github.com/inducer/pytato/pull/393 changes the function signature. --- grudge/array_context.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index 541c83b34..5a27bc9fd 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -239,7 +239,15 @@ def _dag_to_compiled_func(self, dict_of_named_arrays, self.actx._compile_trace_callback(self.f, "pre_find_distributed_partition", dict_of_named_arrays) - distributed_partition = pt.find_distributed_partition(dict_of_named_arrays) + try: + distributed_partition = pt.find_distributed_partition( + dict_of_named_arrays) + except TypeError: + # https://github.com/inducer/pytato/pull/393 changes the + # function signature + distributed_partition = pt.find_distributed_partition( + self.actx.mpi_communicator, + dict_of_named_arrays) if __debug__: # pylint-ignore-reason: From 439a41f9acb6f9995254177f47e6f6fbcf3ff6e8 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 10 Jan 2023 13:00:13 -0600 Subject: [PATCH 02/97] pylint --- grudge/array_context.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/grudge/array_context.py b/grudge/array_context.py index 5a27bc9fd..5c7437f19 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -245,7 +245,11 @@ def _dag_to_compiled_func(self, dict_of_named_arrays, except TypeError: # https://github.com/inducer/pytato/pull/393 changes the # function signature + # pylint: disable=too-many-function-args distributed_partition = pt.find_distributed_partition( + # pylint-ignore-reason: + # '_BasePytatoArrayContext' has no 'mpi_communicator' member + # pylint: disable=no-member self.actx.mpi_communicator, dict_of_named_arrays) From 6f64d77afb7f600b1eeebf6b96fa9fc24a05b772 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 10 Jan 2023 14:04:35 -0600 Subject: [PATCH 03/97] flake8 --- grudge/array_context.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index 5c7437f19..5b11418d6 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -247,8 +247,9 @@ def _dag_to_compiled_func(self, dict_of_named_arrays, # function signature # pylint: disable=too-many-function-args distributed_partition = pt.find_distributed_partition( - # pylint-ignore-reason: - # '_BasePytatoArrayContext' has no 'mpi_communicator' member + # pylint-ignore-reason: + # '_BasePytatoArrayContext' has no + # 'mpi_communicator' member # pylint: disable=no-member self.actx.mpi_communicator, dict_of_named_arrays) From 6f52ae50f53bcd5e7bb6dad39f7d0c80f72d16de Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 19 Jan 2023 15:08:18 -0600 Subject: [PATCH 04/97] only catch one type of TypeError Co-authored-by: Matt Smith --- grudge/array_context.py | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index 5b11418d6..daef3df10 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -239,20 +239,21 @@ def _dag_to_compiled_func(self, dict_of_named_arrays, self.actx._compile_trace_callback(self.f, "pre_find_distributed_partition", dict_of_named_arrays) - try: - distributed_partition = pt.find_distributed_partition( - dict_of_named_arrays) - except TypeError: - # https://github.com/inducer/pytato/pull/393 changes the - # function signature - # pylint: disable=too-many-function-args - distributed_partition = pt.find_distributed_partition( - # pylint-ignore-reason: - # '_BasePytatoArrayContext' has no - # 'mpi_communicator' member - # pylint: disable=no-member - self.actx.mpi_communicator, - dict_of_named_arrays) + # https://github.com/inducer/pytato/pull/393 changes the function signature + try: + # pylint: disable=too-many-function-args + distributed_partition = pt.find_distributed_partition( + # pylint-ignore-reason: + # '_BasePytatoArrayContext' has no + # 'mpi_communicator' member + # pylint: disable=no-member + self.actx.mpi_communicator, dict_of_named_arrays) + except TypeError as e: + if "find_distributed_partition() takes 1 positional" in str(e): + distributed_partition = pt.find_distributed_partition( + dict_of_named_arrays) + else: + raise if __debug__: # pylint-ignore-reason: From bc2d41b2647ba55b60dd6aa02a8a3cedcd478afa Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 19 Jan 2023 15:10:04 -0600 Subject: [PATCH 05/97] flake8 --- grudge/array_context.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index daef3df10..06e2cc493 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -239,21 +239,21 @@ def _dag_to_compiled_func(self, dict_of_named_arrays, self.actx._compile_trace_callback(self.f, "pre_find_distributed_partition", dict_of_named_arrays) - # https://github.com/inducer/pytato/pull/393 changes the function signature - try: - # pylint: disable=too-many-function-args - distributed_partition = pt.find_distributed_partition( - # pylint-ignore-reason: - # '_BasePytatoArrayContext' has no - # 'mpi_communicator' member - # pylint: disable=no-member - self.actx.mpi_communicator, dict_of_named_arrays) - except TypeError as e: - if "find_distributed_partition() takes 1 positional" in str(e): - distributed_partition = pt.find_distributed_partition( - dict_of_named_arrays) - else: - raise + # https://github.com/inducer/pytato/pull/393 changes the function signature + try: + # pylint: disable=too-many-function-args + distributed_partition = pt.find_distributed_partition( + # pylint-ignore-reason: + # '_BasePytatoArrayContext' has no + # 'mpi_communicator' member + # pylint: disable=no-member + self.actx.mpi_communicator, dict_of_named_arrays) + except TypeError as e: + if "find_distributed_partition() takes 1 positional" in str(e): + distributed_partition = pt.find_distributed_partition( + dict_of_named_arrays) + else: + raise if __debug__: # pylint-ignore-reason: From dec4fce52fdbcc637aad90a39bab49e866fa6967 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 21 Sep 2022 09:49:52 -0500 Subject: [PATCH 06/97] set up connections between volumes --- grudge/discretization.py | 281 +++++++++++++++++++++++++++++---------- 1 file changed, 208 insertions(+), 73 deletions(-) diff --git a/grudge/discretization.py b/grudge/discretization.py index 8e57ca503..fd4c39728 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -7,6 +7,7 @@ .. autofunction:: make_discretization_collection .. currentmodule:: grudge.discretization +.. autoclass:: PartID """ __copyright__ = """ @@ -34,10 +35,12 @@ THE SOFTWARE. """ -from typing import Mapping, Optional, Union, TYPE_CHECKING, Any +from typing import Sequence, Mapping, Optional, Union, Tuple, TYPE_CHECKING, Any from pytools import memoize_method, single_valued +from dataclasses import dataclass, replace + from grudge.dof_desc import ( VTAG_ALL, DD_VOLUME_ALL, @@ -71,6 +74,75 @@ import mpi4py.MPI +@dataclass(frozen=True) +class PartID: + """Unique identifier for a piece of a partitioned mesh. + + .. attribute:: volume_tag + + The volume of the part. + + .. attribute:: rank + + The (optional) MPI rank of the part. + + """ + volume_tag: VolumeTag + rank: Optional[int] = None + + +# {{{ part ID normalization + +def _normalize_mesh_part_ids( + mesh: Mesh, + self_volume_tag: VolumeTag, + all_volume_tags: Sequence[VolumeTag], + mpi_communicator: Optional["mpi4py.MPI.Intracomm"] = None): + """Convert a mesh's configuration-dependent "part ID" into a fixed type.""" + from numbers import Integral + if mpi_communicator is not None: + # Accept PartID or rank (assume intra-volume for the latter) + def as_part_id(mesh_part_id): + if isinstance(mesh_part_id, PartID): + return mesh_part_id + elif isinstance(mesh_part_id, Integral): + return PartID(self_volume_tag, int(mesh_part_id)) + else: + raise TypeError(f"Unable to convert {mesh_part_id} to PartID.") + else: + # Accept PartID or volume tag + def as_part_id(mesh_part_id): + if isinstance(mesh_part_id, PartID): + return mesh_part_id + elif mesh_part_id in all_volume_tags: + return PartID(mesh_part_id) + else: + raise TypeError(f"Unable to convert {mesh_part_id} to PartID.") + + facial_adjacency_groups = mesh.facial_adjacency_groups + + new_facial_adjacency_groups = [] + + from meshmode.mesh import InterPartAdjacencyGroup + for grp_list in facial_adjacency_groups: + new_grp_list = [] + for fagrp in grp_list: + if isinstance(fagrp, InterPartAdjacencyGroup): + part_id = as_part_id(fagrp.part_id) + new_fagrp = replace( + fagrp, + boundary_tag=BTAG_PARTITION(part_id), + part_id=part_id) + else: + new_fagrp = fagrp + new_grp_list.append(new_fagrp) + new_facial_adjacency_groups.append(new_grp_list) + + return mesh.copy(facial_adjacency_groups=new_facial_adjacency_groups) + +# }}} + + # {{{ discr_tag_to_group_factory normalization def _normalize_discr_tag_to_group_factory( @@ -156,6 +228,9 @@ def __init__(self, array_context: ArrayContext, discr_tag_to_group_factory: Optional[ Mapping[DiscretizationTag, ElementGroupFactory]] = None, mpi_communicator: Optional["mpi4py.MPI.Intracomm"] = None, + inter_part_connections: Optional[ + Mapping[Tuple[PartID, PartID], + DiscretizationConnection]] = None, ) -> None: """ :arg discr_tag_to_group_factory: A mapping from discretization tags @@ -206,6 +281,9 @@ def __init__(self, array_context: ArrayContext, mesh = volume_discrs + mesh = _normalize_mesh_part_ids( + mesh, VTAG_ALL, [VTAG_ALL], mpi_communicator=mpi_communicator) + discr_tag_to_group_factory = _normalize_discr_tag_to_group_factory( dim=mesh.dim, discr_tag_to_group_factory=discr_tag_to_group_factory, @@ -219,17 +297,32 @@ def __init__(self, array_context: ArrayContext, del mesh + if inter_part_connections is not None: + raise TypeError("may not pass inter_part_connections when " + "DiscretizationCollection constructor is called in " + "legacy mode") + + self._inter_part_connections = \ + _set_up_inter_part_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: 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 + if inter_part_connections is None: + raise TypeError("inter_part_connections must be passed when " + "DiscretizationCollection constructor is called in " + "'modern' mode") + + self._inter_part_connections = inter_part_connections - self._dist_boundary_connections = { - vtag: self._set_up_distributed_communication( - vtag, mpi_communicator, array_context) - for vtag in self._volume_discrs.keys()} + self._volume_discrs = volume_discrs # }}} @@ -252,71 +345,6 @@ def is_management_rank(self): return self.mpi_communicator.Get_rank() \ == self.get_management_rank_index() - # {{{ distributed - - def _set_up_distributed_communication( - self, vtag, mpi_communicator, array_context): - from_dd = DOFDesc(VolumeDomainTag(vtag), DISCR_TAG_BASE) - - boundary_connections = {} - - from meshmode.distributed import get_connected_partitions - connected_parts = get_connected_partitions(self._volume_discrs[vtag].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, from_dd.trace(BTAG_PARTITION(i_remote_part))) - - 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 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.BoundaryDomainTag` 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, BoundaryDomainTag) - assert isinstance(dd.domain_tag.tag, BTAG_PARTITION) - - vtag = dd.domain_tag.volume_tag - - return self._dist_boundary_connections[vtag][dd.domain_tag.tag.part_nr] - - # }}} - # {{{ discr_from_dd @memoize_method @@ -772,6 +800,105 @@ def normal(self, dd): # }}} +# {{{ distributed/multi-volume setup + +def _set_up_inter_part_connections( + array_context: ArrayContext, + mpi_communicator: Optional["mpi4py.MPI.Intracomm"], + volume_discrs: Mapping[VolumeTag, Discretization], + base_group_factory: ElementGroupFactory, + ) -> Mapping[ + Tuple[PartID, PartID], + DiscretizationConnection]: + + from meshmode.distributed import (get_connected_parts, + make_remote_group_infos, InterRankBoundaryInfo, + MPIBoundaryCommSetupHelper) + + rank = mpi_communicator.Get_rank() if mpi_communicator is not None else None + + # Save boundary restrictions as they're created to avoid potentially creating + # them twice in the loop below + cached_part_bdry_restrictions: Mapping[ + Tuple[PartID, PartID], + DiscretizationConnection] = {} + + def get_part_bdry_restriction(self_part_id, other_part_id): + cached_result = cached_part_bdry_restrictions.get( + (self_part_id, other_part_id), None) + if cached_result is not None: + return cached_result + return cached_part_bdry_restrictions.setdefault( + (self_part_id, other_part_id), + make_face_restriction( + array_context, volume_discrs[self_part_id.volume_tag], + base_group_factory, + boundary_tag=BTAG_PARTITION(other_part_id))) + + inter_part_conns: Mapping[ + Tuple[PartID, PartID], + DiscretizationConnection] = {} + + irbis = [] + + for vtag, volume_discr in volume_discrs.items(): + part_id = PartID(vtag, rank) + connected_part_ids = get_connected_parts(volume_discr.mesh) + for connected_part_id in connected_part_ids: + bdry_restr = get_part_bdry_restriction( + self_part_id=part_id, other_part_id=connected_part_id) + + if connected_part_id.rank == rank: + # {{{ rank-local interface between multiple volumes + + connected_bdry_restr = get_part_bdry_restriction( + self_part_id=connected_part_id, other_part_id=part_id) + + from meshmode.discretization.connection import \ + make_partition_connection + inter_part_conns[connected_part_id, part_id] = \ + make_partition_connection( + array_context, + local_bdry_conn=bdry_restr, + remote_bdry_discr=connected_bdry_restr.to_discr, + remote_group_infos=make_remote_group_infos( + array_context, part_id, connected_bdry_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_part_id=part_id, + remote_part_id=connected_part_id, + remote_rank=connected_part_id.rank, + local_boundary_connection=bdry_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 + +# }}} + + # {{{ modal group factory def _generate_modal_group_factory(nodal_group_factory): @@ -860,6 +987,8 @@ def make_discretization_collection( del order + mpi_communicator = getattr(array_context, "mpi_communicator", None) + if any( isinstance(mesh_or_discr, Discretization) for mesh_or_discr in volumes.values()): @@ -868,14 +997,20 @@ def make_discretization_collection( volume_discrs = { vtag: Discretization( array_context, - mesh, + _normalize_mesh_part_ids( + mesh, vtag, volumes.keys(), mpi_communicator=mpi_communicator), discr_tag_to_group_factory[DISCR_TAG_BASE]) for vtag, mesh in volumes.items()} return DiscretizationCollection( array_context=array_context, volume_discrs=volume_discrs, - discr_tag_to_group_factory=discr_tag_to_group_factory) + discr_tag_to_group_factory=discr_tag_to_group_factory, + inter_part_connections=_set_up_inter_part_connections( + array_context=array_context, + mpi_communicator=mpi_communicator, + volume_discrs=volume_discrs, + base_group_factory=discr_tag_to_group_factory[DISCR_TAG_BASE])) # }}} From a9690f9ade97b0730499ce1793f7afb9f3b29898 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 21 Sep 2022 09:51:54 -0500 Subject: [PATCH 07/97] add inter-volume communication --- grudge/eager.py | 3 +- grudge/op.py | 10 +- grudge/trace_pair.py | 627 +++++++++++++++++++++++++++++++++++-------- 3 files changed, 522 insertions(+), 118 deletions(-) diff --git a/grudge/eager.py b/grudge/eager.py index 626e15592..08cf08f2a 100644 --- a/grudge/eager.py +++ b/grudge/eager.py @@ -87,7 +87,8 @@ def nodal_max(self, dd, vec): return op.nodal_max(self, dd, vec) -connected_ranks = op.connected_ranks +# FIXME: Deprecate connected_ranks instead of removing +connected_parts = op.connected_parts interior_trace_pair = op.interior_trace_pair cross_rank_trace_pairs = op.cross_rank_trace_pairs diff --git a/grudge/op.py b/grudge/op.py index f5781f4be..a6cef8ffa 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -118,8 +118,11 @@ interior_trace_pair, interior_trace_pairs, local_interior_trace_pair, - connected_ranks, + connected_parts, + 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 ) @@ -147,8 +150,11 @@ "interior_trace_pair", "interior_trace_pairs", "local_interior_trace_pair", - "connected_ranks", + "connected_parts", + "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", diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 1f49ae0d6..0b0400f12 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__ = """ @@ -52,17 +55,18 @@ from warnings import warn -from typing import List, Hashable, Optional, Type, Any +from typing import List, Hashable, Optional, Tuple, Type, Any, Sequence, Mapping 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, + get_container_context_recursively_opt, + to_numpy, + from_numpy, ArrayOrContainer ) @@ -72,7 +76,7 @@ from pytools import memoize_on_first_arg -from grudge.discretization import DiscretizationCollection +from grudge.discretization import DiscretizationCollection, PartID from grudge.projection import project from meshmode.mesh import BTAG_PARTITION @@ -82,7 +86,7 @@ import grudge.dof_desc as dof_desc from grudge.dof_desc import ( DOFDesc, DD_VOLUME_ALL, FACE_RESTR_INTERIOR, DISCR_TAG_BASE, - VolumeDomainTag, + VolumeTag, VolumeDomainTag, BoundaryDomainTag, ConvertibleToDOFDesc, ) @@ -360,6 +364,124 @@ def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *, # }}} +# {{{ inter-volume trace pairs + +def local_inter_volume_trace_pairs( + dcoll: DiscretizationCollection, + pairwise_volume_data: Mapping[ + Tuple[DOFDesc, DOFDesc], + Tuple[ArrayOrContainer, ArrayOrContainer]] + ) -> Mapping[Tuple[DOFDesc, DOFDesc], TracePair]: + for vol_dd_pair in pairwise_volume_data.keys(): + for vol_dd in vol_dd_pair: + if not isinstance(vol_dd.domain_tag, VolumeDomainTag): + raise ValueError( + "pairwise_volume_data keys must describe volumes, " + f"got '{vol_dd}'") + if vol_dd.discretization_tag != DISCR_TAG_BASE: + raise ValueError( + "expected base-discretized DOFDesc in pairwise_volume_data, " + f"got '{vol_dd}'") + + rank = ( + dcoll.mpi_communicator.Get_rank() + if dcoll.mpi_communicator is not None + else None) + + result: Mapping[Tuple[DOFDesc, DOFDesc], TracePair] = {} + + for vol_dd_pair, vol_data_pair in pairwise_volume_data.items(): + from meshmode.mesh import mesh_has_boundary + if not mesh_has_boundary( + dcoll.discr_from_dd(vol_dd_pair[0]).mesh, + BTAG_PARTITION(PartID(vol_dd_pair[1].domain_tag.tag, rank))): + continue + + directional_vol_dd_pairs = [ + (vol_dd_pair[1], vol_dd_pair[0]), + (vol_dd_pair[0], vol_dd_pair[1])] + + trace_dd_pair = tuple( + self_vol_dd.trace( + BTAG_PARTITION( + PartID(other_vol_dd.domain_tag.tag, rank))) + for other_vol_dd, self_vol_dd in directional_vol_dd_pairs) + + # Pre-compute the projections out here to avoid doing it twice inside + # the loop below + trace_data = { + trace_dd: project(dcoll, vol_dd, trace_dd, vol_data) + for vol_dd, trace_dd, vol_data in zip( + vol_dd_pair, trace_dd_pair, vol_data_pair)} + + for other_vol_dd, self_vol_dd in directional_vol_dd_pairs: + self_part_id = PartID(self_vol_dd.domain_tag.tag, rank) + other_part_id = PartID(other_vol_dd.domain_tag.tag, rank) + + self_trace_dd = self_vol_dd.trace(BTAG_PARTITION(other_part_id)) + other_trace_dd = other_vol_dd.trace(BTAG_PARTITION(self_part_id)) + + self_trace_data = trace_data[self_trace_dd] + unswapped_other_trace_data = trace_data[other_trace_dd] + + other_to_self = dcoll._inter_part_connections[ + other_part_id, self_part_id] + + def get_opposite_trace(ary): + if isinstance(ary, Number): + return ary + else: + return other_to_self(ary) # noqa: B023 + + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + other_trace_data = rec_map_array_container( + get_opposite_trace, + unswapped_other_trace_data, + leaf_class=DOFArray) + + result[other_vol_dd, self_vol_dd] = TracePair( + self_trace_dd, + interior=self_trace_data, + exterior=other_trace_data) + + return result + + +def inter_volume_trace_pairs(dcoll: DiscretizationCollection, + pairwise_volume_data: Mapping[ + Tuple[DOFDesc, DOFDesc], + Tuple[ArrayOrContainer, ArrayOrContainer]], + comm_tag: Hashable = None) -> Mapping[ + Tuple[DOFDesc, DOFDesc], + List[TracePair]]: + """ + 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 + + result: Mapping[ + Tuple[DOFDesc, DOFDesc], + List[TracePair]] = {} + + local_tpairs = local_inter_volume_trace_pairs(dcoll, pairwise_volume_data) + cross_rank_tpairs = cross_rank_inter_volume_trace_pairs( + dcoll, pairwise_volume_data, comm_tag=comm_tag) + + for directional_vol_dd_pair, tpair in local_tpairs.items(): + result[directional_vol_dd_pair] = [tpair] + + for directional_vol_dd_pair, tpairs in cross_rank_tpairs.items(): + result.setdefault(directional_vol_dd_pair, []).extend(tpairs) + + return result + +# }}} + + # {{{ distributed: helper functions class _TagKeyBuilder(KeyBuilder): @@ -367,16 +489,21 @@ def update_for_type(self, key_hash, key: Type[Any]): self.rec(key_hash, (key.__module__, key.__name__, key.__name__,)) +# FIXME: Deprecate connected_ranks instead of removing @memoize_on_first_arg -def connected_ranks( +def connected_parts( dcoll: DiscretizationCollection, - volume_dd: Optional[DOFDesc] = None): - if volume_dd is None: - volume_dd = DD_VOLUME_ALL + self_volume_tag: VolumeTag, + other_volume_tag: VolumeTag + ) -> Sequence[PartID]: + result: List[PartID] = [ + connected_part_id + for connected_part_id, part_id in dcoll._inter_part_connections.keys() + if ( + part_id.volume_tag == self_volume_tag + and connected_part_id.volume_tag == other_volume_tag)] - from meshmode.distributed import get_connected_partitions - return get_connected_partitions( - dcoll._volume_discrs[volume_dd.domain_tag.tag].mesh) + return result def _sym_tag_to_num_tag(comm_tag: Optional[Hashable]) -> Optional[int]: @@ -414,24 +541,33 @@ class _RankBoundaryCommunicationEager: base_comm_tag = 1273 def __init__(self, - dcoll: DiscretizationCollection, - array_container: ArrayOrContainer, - remote_rank, comm_tag: Optional[int] = None, - volume_dd=DD_VOLUME_ALL): - actx = get_container_context_recursively(array_container) - bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_rank)) - - local_bdry_data = project(dcoll, volume_dd, bdry_dd, array_container) + actx: ArrayContext, + dcoll: DiscretizationCollection, + *, + local_part_id: PartID, + remote_part_id: PartID, + local_bdry_data: ArrayOrContainer, + remote_bdry_data_template: ArrayOrContainer, + comm_tag: Optional[Hashable] = None): + comm = dcoll.mpi_communicator assert comm is not None + remote_rank = remote_part_id.rank + assert remote_rank is not None + self.dcoll = dcoll self.array_context = actx - self.remote_bdry_dd = bdry_dd - self.bdry_discr = dcoll.discr_from_dd(bdry_dd) + self.local_part_id = local_part_id + self.remote_part_id = remote_part_id + self.local_bdry_dd = DOFDesc( + BoundaryDomainTag( + BTAG_PARTITION(remote_part_id), + volume_tag=local_part_id.volume_tag), + DISCR_TAG_BASE) + self.bdry_discr = dcoll.discr_from_dd(self.local_bdry_dd) self.local_bdry_data = local_bdry_data - self.local_bdry_data_np = \ - to_numpy(flatten(self.local_bdry_data, actx), actx) + self.remote_bdry_data_template = remote_bdry_data_template self.comm_tag = self.base_comm_tag comm_tag = _sym_tag_to_num_tag(comm_tag) @@ -439,55 +575,80 @@ def __init__(self, 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, - remote_rank, - tag=self.comm_tag) + self.send_reqs = [] + self.send_data = [] + + def send_single_array(key, local_subary): + if not isinstance(local_subary, Number): + local_subary_np = to_numpy(local_subary, actx) + self.send_reqs.append( + comm.Isend(local_subary_np, remote_rank, tag=self.comm_tag)) + self.send_data.append(local_subary_np) + return local_subary + + self.recv_reqs = [] + self.recv_data = {} + + def recv_single_array(key, remote_subary_template): + if not isinstance(remote_subary_template, Number): + remote_subary_np = np.empty( + remote_subary_template.shape, + remote_subary_template.dtype) + self.recv_reqs.append( + comm.Irecv(remote_subary_np, remote_rank, tag=self.comm_tag)) + self.recv_data[key] = remote_subary_np + return remote_subary_template + + from arraycontext.container.traversal import rec_keyed_map_array_container + rec_keyed_map_array_container(send_single_array, local_bdry_data) + rec_keyed_map_array_container(recv_single_array, remote_bdry_data_template) def finish(self): - # Wait for the nonblocking receive request to complete before + from mpi4py import MPI + + # Wait for the nonblocking receive requests 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( - self.remote_bdry_dd) - swapped_remote_bdry_data = bdry_conn(remote_bdry_data) - - # Complete the nonblocking send request associated with communicating - # `self.local_bdry_data_np` - self.send_req.Wait() - - return TracePair(self.remote_bdry_dd, - interior=self.local_bdry_data, - exterior=swapped_remote_bdry_data) + MPI.Request.waitall(self.recv_reqs) + + def finish_single_array(key, remote_subary_template): + if isinstance(remote_subary_template, Number): + # NOTE: Assumes that the same number is passed on every rank + return remote_subary_template + else: + return from_numpy(self.recv_data[key], self.array_context) + + from arraycontext.container.traversal import rec_keyed_map_array_container + unswapped_remote_bdry_data = rec_keyed_map_array_container( + finish_single_array, self.remote_bdry_data_template) + + remote_to_local = self.dcoll._inter_part_connections[ + self.remote_part_id, self.local_part_id] + + def get_opposite_trace(ary): + if isinstance(ary, Number): + return ary + else: + return remote_to_local(ary) + + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + remote_bdry_data = rec_map_array_container( + get_opposite_trace, + unswapped_remote_bdry_data, + leaf_class=DOFArray) + + # Complete the nonblocking send requests + MPI.Request.waitall(self.send_reqs) + + return TracePair( + self.local_bdry_dd, + interior=self.local_bdry_data, + exterior=remote_bdry_data) # }}} @@ -496,51 +657,112 @@ def finish(self): class _RankBoundaryCommunicationLazy: def __init__(self, - dcoll: DiscretizationCollection, - array_container: ArrayOrContainer, - remote_rank: int, comm_tag: Hashable, - volume_dd=DD_VOLUME_ALL): + actx: ArrayContext, + dcoll: DiscretizationCollection, + *, + local_part_id: PartID, + remote_part_id: PartID, + local_bdry_data: ArrayOrContainer, + remote_bdry_data_template: 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") - bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_rank)) + remote_rank = remote_part_id.rank + assert remote_rank is not None self.dcoll = dcoll - self.array_context = get_container_context_recursively(array_container) - self.remote_bdry_dd = bdry_dd - self.bdry_discr = dcoll.discr_from_dd(self.remote_bdry_dd) - - self.local_bdry_data = project( - dcoll, volume_dd, bdry_dd, array_container) - - from pytato import make_distributed_recv, staple_distributed_send - - def communicate_single_array(key, local_bdry_ary): - ary_tag = (comm_tag, key) - return staple_distributed_send( - local_bdry_ary, dest_rank=remote_rank, comm_tag=ary_tag, - stapled_to=make_distributed_recv( + self.array_context = actx + self.local_bdry_dd = DOFDesc( + BoundaryDomainTag( + BTAG_PARTITION(remote_part_id), + volume_tag=local_part_id.volume_tag), + DISCR_TAG_BASE) + self.bdry_discr = dcoll.discr_from_dd(self.local_bdry_dd) + self.local_part_id = local_part_id + self.remote_part_id = remote_part_id + + from pytato import ( + make_distributed_recv, + make_distributed_send, + DistributedSendRefHolder) + + # TODO: This currently assumes that local_bdry_data and + # remote_bdry_data_template have the same structure. This is not true + # in general. Find a way to staple the sends appropriately when the number + # of recvs is not equal to the number of sends + # FIXME: Overly restrictive (just needs to be the same structure) + assert type(local_bdry_data) == type(remote_bdry_data_template) + + sends = {} + + def send_single_array(key, local_subary): + if isinstance(local_subary, Number): + return + else: + ary_tag = (comm_tag, key) + sends[key] = make_distributed_send( + local_subary, dest_rank=remote_rank, comm_tag=ary_tag) + + def recv_single_array(key, remote_subary_template): + if isinstance(remote_subary_template, Number): + # NOTE: Assumes that the same number is passed on every rank + return remote_subary_template + else: + ary_tag = (comm_tag, key) + return DistributedSendRefHolder( + sends[key], + 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=remote_subary_template.shape, + dtype=remote_subary_template.dtype, + axes=remote_subary_template.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( - self.remote_bdry_dd) + rec_keyed_map_array_container(send_single_array, local_bdry_data) + self.local_bdry_data = local_bdry_data - return TracePair(self.remote_bdry_dd, - interior=self.local_bdry_data, - exterior=bdry_conn(self.remote_data)) + self.unswapped_remote_bdry_data = rec_keyed_map_array_container( + recv_single_array, remote_bdry_data_template) + + def finish(self): + remote_to_local = self.dcoll._inter_part_connections[ + self.remote_part_id, self.local_part_id] + + def get_opposite_trace(ary): + if isinstance(ary, Number): + return ary + else: + return remote_to_local(ary) + + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + remote_bdry_data = rec_map_array_container( + get_opposite_trace, + self.unswapped_remote_bdry_data, + leaf_class=DOFArray) + + return TracePair( + self.local_bdry_dd, + interior=self.local_bdry_data, + exterior=remote_bdry_data) # }}} # {{{ cross_rank_trace_pairs +def _replace_dof_arrays(array_container, dof_array): + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + return rec_map_array_container( + lambda x: dof_array if isinstance(x, DOFArray) else x, + array_container, + leaf_class=DOFArray) + + def cross_rank_trace_pairs( dcoll: DiscretizationCollection, ary: ArrayOrContainer, tag: Hashable = None, @@ -549,9 +771,9 @@ def cross_rank_trace_pairs( r"""Get a :class:`list` of *ary* trace pairs for each partition boundary. For each partition boundary, the field data values in *ary* are - communicated to/from the neighboring partition. Presumably, this - communication is MPI (but strictly speaking, may not be, and this - routine is agnostic to the underlying communication). + communicated to/from the neighboring part. Presumably, this communication + is MPI (but strictly speaking, may not be, and this routine is agnostic to + the underlying communication). For each face on each partition boundary, a :class:`TracePair` is created with the locally, and @@ -596,14 +818,36 @@ def cross_rank_trace_pairs( # }}} - if isinstance(ary, Number): - # NOTE: Assumed that the same number is passed on every rank - return [TracePair( - volume_dd.trace(BTAG_PARTITION(remote_rank)), - interior=ary, exterior=ary) - for remote_rank in connected_ranks(dcoll, volume_dd=volume_dd)] + if dcoll.mpi_communicator is None: + return [] + + rank = dcoll.mpi_communicator.Get_rank() + + local_part_id = PartID(volume_dd.domain_tag.tag, rank) + + connected_part_ids = connected_parts( + dcoll, self_volume_tag=volume_dd.domain_tag.tag, + other_volume_tag=volume_dd.domain_tag.tag) + + remote_part_ids = [ + part_id + for part_id in connected_part_ids + if part_id.rank != rank] + + # 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(remote_part_ids) == len({part_id.rank for part_id in remote_part_ids}) - actx = get_container_context_recursively(ary) + actx = get_container_context_recursively_opt(ary) + + if actx is None: + # NOTE: Assumes that the same number is passed on every rank + return [ + TracePair( + volume_dd.trace(BTAG_PARTITION(remote_part_id)), + interior=ary, exterior=ary) + for remote_part_id in remote_part_ids] from grudge.array_context import MPIPytatoArrayContextBase @@ -612,14 +856,167 @@ def cross_rank_trace_pairs( else: rbc_class = _RankBoundaryCommunicationEager - # Initialize and post all sends/receives - rank_bdry_communcators = [ - rbc_class(dcoll, ary, remote_rank, comm_tag=comm_tag, volume_dd=volume_dd) - for remote_rank in connected_ranks(dcoll, volume_dd=volume_dd) - ] + rank_bdry_communicators = [] + + for remote_part_id in remote_part_ids: + bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_part_id)) + + local_bdry_data = project(dcoll, volume_dd, bdry_dd, ary) + + from arraycontext import tag_axes + from meshmode.transform_metadata import ( + DiscretizationElementAxisTag, + DiscretizationDOFAxisTag) + remote_bdry_zeros = tag_axes( + actx, { + 0: DiscretizationElementAxisTag(), + 1: DiscretizationDOFAxisTag()}, + dcoll._inter_part_connections[ + remote_part_id, local_part_id].from_discr.zeros(actx)) + + remote_bdry_data_template = _replace_dof_arrays( + local_bdry_data, remote_bdry_zeros) + + rank_bdry_communicators.append( + rbc_class(actx, dcoll, + local_part_id=local_part_id, + remote_part_id=remote_part_id, + local_bdry_data=local_bdry_data, + remote_bdry_data_template=remote_bdry_data_template, + comm_tag=comm_tag)) + + return [rbc.finish() for rbc in rank_bdry_communicators] + +# }}} + + +# {{{ cross_rank_inter_volume_trace_pairs + +def cross_rank_inter_volume_trace_pairs( + dcoll: DiscretizationCollection, + pairwise_volume_data: Mapping[ + Tuple[DOFDesc, DOFDesc], + Tuple[ArrayOrContainer, ArrayOrContainer]], + *, comm_tag: Hashable = None, + ) -> Mapping[ + Tuple[DOFDesc, DOFDesc], + 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 + + for vol_dd_pair in pairwise_volume_data.keys(): + for vol_dd in vol_dd_pair: + if not isinstance(vol_dd.domain_tag, VolumeDomainTag): + raise ValueError( + "pairwise_volume_data keys must describe volumes, " + f"got '{vol_dd}'") + if vol_dd.discretization_tag != DISCR_TAG_BASE: + raise ValueError( + "expected base-discretized DOFDesc in pairwise_volume_data, " + f"got '{vol_dd}'") + + # }}} + + if dcoll.mpi_communicator is None: + return {} + + rank = dcoll.mpi_communicator.Get_rank() + + for vol_data_pair in pairwise_volume_data.values(): + for vol_data in vol_data_pair: + actx = get_container_context_recursively_opt(vol_data) + if actx is not None: + break + if actx is not None: + break + + def get_remote_connected_parts(local_vol_dd, remote_vol_dd): + connected_part_ids = connected_parts( + dcoll, self_volume_tag=local_vol_dd.domain_tag.tag, + other_volume_tag=remote_vol_dd.domain_tag.tag) + return [ + part_id + for part_id in connected_part_ids + if part_id.rank != rank] + + if actx is None: + # NOTE: Assumes that the same number is passed on every rank for a + # given volume + return { + (remote_vol_dd, local_vol_dd): [ + TracePair( + local_vol_dd.trace(BTAG_PARTITION(remote_part_id)), + interior=local_vol_ary, exterior=remote_vol_ary) + for remote_part_id in get_remote_connected_parts( + local_vol_dd, remote_vol_dd)] + for (remote_vol_dd, local_vol_dd), (remote_vol_ary, local_vol_ary) + in pairwise_volume_data.items()} + + from grudge.array_context import MPIPytatoArrayContextBase + + if isinstance(actx, MPIPytatoArrayContextBase): + rbc_class = _RankBoundaryCommunicationLazy + else: + rbc_class = _RankBoundaryCommunicationEager - # Complete send/receives and return communicated data - return [rc.finish() for rc in rank_bdry_communcators] + rank_bdry_communicators = {} + + for vol_dd_pair, vol_data_pair in pairwise_volume_data.items(): + directional_volume_data = { + (vol_dd_pair[0], vol_dd_pair[1]): (vol_data_pair[0], vol_data_pair[1]), + (vol_dd_pair[1], vol_dd_pair[0]): (vol_data_pair[1], vol_data_pair[0])} + + for dd_pair, data_pair in directional_volume_data.items(): + other_vol_dd, self_vol_dd = dd_pair + other_vol_data, self_vol_data = data_pair + + self_part_id = PartID(self_vol_dd.domain_tag.tag, rank) + other_part_ids = get_remote_connected_parts(self_vol_dd, other_vol_dd) + + rbcs = [] + + for other_part_id in other_part_ids: + self_bdry_dd = self_vol_dd.trace(BTAG_PARTITION(other_part_id)) + self_bdry_data = project( + dcoll, self_vol_dd, self_bdry_dd, self_vol_data) + + from arraycontext import tag_axes + from meshmode.transform_metadata import ( + DiscretizationElementAxisTag, + DiscretizationDOFAxisTag) + other_bdry_zeros = tag_axes( + actx, { + 0: DiscretizationElementAxisTag(), + 1: DiscretizationDOFAxisTag()}, + dcoll._inter_part_connections[ + other_part_id, self_part_id].from_discr.zeros(actx)) + + other_bdry_data_template = _replace_dof_arrays( + other_vol_data, other_bdry_zeros) + + rbcs.append( + rbc_class(actx, dcoll, + local_part_id=self_part_id, + remote_part_id=other_part_id, + local_bdry_data=self_bdry_data, + remote_bdry_data_template=other_bdry_data_template, + comm_tag=comm_tag)) + + rank_bdry_communicators[other_vol_dd, self_vol_dd] = rbcs + + return { + directional_vol_dd_pair: [rbc.finish() for rbc in rbcs] + for directional_vol_dd_pair, rbcs in rank_bdry_communicators.items()} # }}} From c536905462a8153b66f12dea28395fab272ac3de Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 3 Nov 2022 09:20:20 -0700 Subject: [PATCH 08/97] add fixme --- grudge/trace_pair.py | 1 + 1 file changed, 1 insertion(+) diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 0b0400f12..84dedf386 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -525,6 +525,7 @@ def _sym_tag_to_num_tag(comm_tag: Optional[Hashable]) -> Optional[int]: num_tag = sum(ord(ch) << i for i, ch in enumerate(digest)) % tag_ub + # FIXME: This prints the wrong numerical tag because of base_comm_tag below warn("Encountered unknown symbolic tag " f"'{comm_tag}', assigning a value of '{num_tag}'. " "This is a temporary workaround, please ensure that " From d5576fbfba127cfc68fa1fe2a5a8352e4f1bcffa Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 3 Nov 2022 10:07:09 -0700 Subject: [PATCH 09/97] check for heterogeneous inter-volume data --- grudge/trace_pair.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 84dedf386..7358e5af8 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -915,7 +915,7 @@ def cross_rank_inter_volume_trace_pairs( """ # {{{ process arguments - for vol_dd_pair in pairwise_volume_data.keys(): + for vol_dd_pair, vol_data_pair in pairwise_volume_data.items(): for vol_dd in vol_dd_pair: if not isinstance(vol_dd.domain_tag, VolumeDomainTag): raise ValueError( @@ -925,6 +925,9 @@ def cross_rank_inter_volume_trace_pairs( raise ValueError( "expected base-discretized DOFDesc in pairwise_volume_data, " f"got '{vol_dd}'") + # FIXME: This check could probably be made more robust + if type(vol_data_pair[0]) != type(vol_data_pair[1]): # noqa: E721 + raise ValueError("heterogeneous inter-volume data not supported.") # }}} From 1a38e2df87a657ad964de674aa1b6acf874005b0 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 3 Nov 2022 10:07:26 -0700 Subject: [PATCH 10/97] tag communication by destination volume --- grudge/trace_pair.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 7358e5af8..acc086505 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -570,10 +570,17 @@ def __init__(self, self.local_bdry_data = local_bdry_data self.remote_bdry_data_template = remote_bdry_data_template - 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 + def _generate_num_comm_tag(sym_comm_tag): + result = self.base_comm_tag + num_comm_tag = _sym_tag_to_num_tag(sym_comm_tag) + if num_comm_tag is not None: + result += num_comm_tag + return result + + send_sym_comm_tag = (remote_part_id.volume_tag, comm_tag) + recv_sym_comm_tag = (local_part_id.volume_tag, comm_tag) + self.send_comm_tag = _generate_num_comm_tag(send_sym_comm_tag) + self.recv_comm_tag = _generate_num_comm_tag(recv_sym_comm_tag) del comm_tag # NOTE: mpi4py currently (2021-11-03) holds a reference to the send @@ -588,7 +595,7 @@ def send_single_array(key, local_subary): if not isinstance(local_subary, Number): local_subary_np = to_numpy(local_subary, actx) self.send_reqs.append( - comm.Isend(local_subary_np, remote_rank, tag=self.comm_tag)) + comm.Isend(local_subary_np, remote_rank, tag=self.send_comm_tag)) self.send_data.append(local_subary_np) return local_subary @@ -601,7 +608,8 @@ def recv_single_array(key, remote_subary_template): remote_subary_template.shape, remote_subary_template.dtype) self.recv_reqs.append( - comm.Irecv(remote_subary_np, remote_rank, tag=self.comm_tag)) + comm.Irecv(remote_subary_np, remote_rank, + tag=self.recv_comm_tag)) self.recv_data[key] = remote_subary_np return remote_subary_template @@ -702,7 +710,7 @@ def send_single_array(key, local_subary): if isinstance(local_subary, Number): return else: - ary_tag = (comm_tag, key) + ary_tag = (remote_part_id.volume_tag, comm_tag, key) sends[key] = make_distributed_send( local_subary, dest_rank=remote_rank, comm_tag=ary_tag) @@ -711,7 +719,7 @@ def recv_single_array(key, remote_subary_template): # NOTE: Assumes that the same number is passed on every rank return remote_subary_template else: - ary_tag = (comm_tag, key) + ary_tag = (local_part_id.volume_tag, comm_tag, key) return DistributedSendRefHolder( sends[key], make_distributed_recv( From 50e2d42ed72e9033a7db8c0ac124d6ac5471c0c3 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 3 Apr 2023 15:28:38 -0500 Subject: [PATCH 11/97] add filter_part_boundaries eases setting up boundaries when calling operators on only one volume (i.e., uncoupled) --- grudge/discretization.py | 44 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 43 insertions(+), 1 deletion(-) diff --git a/grudge/discretization.py b/grudge/discretization.py index fd4c39728..25d5fa797 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -8,6 +8,7 @@ .. currentmodule:: grudge.discretization .. autoclass:: PartID +.. autofunction:: filter_part_boundaries """ __copyright__ = """ @@ -35,7 +36,8 @@ THE SOFTWARE. """ -from typing import Sequence, Mapping, Optional, Union, Tuple, TYPE_CHECKING, Any +from typing import ( + Sequence, Mapping, Optional, Union, List, Tuple, TYPE_CHECKING, Any) from pytools import memoize_method, single_valued @@ -1015,4 +1017,44 @@ def make_discretization_collection( # }}} +# {{{ filter_part_boundaries + +def filter_part_boundaries( + dcoll: DiscretizationCollection, + *, + volume_dd: DOFDesc = DD_VOLUME_ALL, + neighbor_volume_dd: Optional[DOFDesc] = None, + neighbor_rank: Optional[int] = None) -> List[DOFDesc]: + """ + Retrieve tags of part boundaries that match *neighbor_volume_dd* and/or + *neighbor_rank*. + """ + vol_mesh = dcoll.discr_from_dd(volume_dd).mesh + + from meshmode.mesh import InterPartAdjacencyGroup + filtered_part_bdry_dds = [ + volume_dd.trace(fagrp.boundary_tag) + for fagrp_list in vol_mesh.facial_adjacency_groups + for fagrp in fagrp_list + if isinstance(fagrp, InterPartAdjacencyGroup)] + + if neighbor_volume_dd is not None: + filtered_part_bdry_dds = [ + bdry_dd + for bdry_dd in filtered_part_bdry_dds + if ( + bdry_dd.domain_tag.tag.part_id.volume_tag + == neighbor_volume_dd.domain_tag.tag)] + + if neighbor_rank is not None: + filtered_part_bdry_dds = [ + bdry_dd + for bdry_dd in filtered_part_bdry_dds + if bdry_dd.domain_tag.tag.part_id.rank == neighbor_rank] + + return filtered_part_bdry_dds + +# }}} + + # vim: foldmethod=marker From dc8d8dc0fd21d9dc27d7d2be48934199ec08c914 Mon Sep 17 00:00:00 2001 From: "Thomas H. Gibson" Date: Thu, 13 Apr 2023 20:44:25 -0500 Subject: [PATCH 12/97] Entropy stable DG and flux-differencing This is https://github.com/inducer/grudge/pull/214 squashed and rebased. Co-authored-by: Andreas Kloeckner --- doc/operators.rst | 1 + examples/euler/acoustic_pulse.py | 40 ++- examples/euler/sod.py | 227 ++++++++++++++++ examples/euler/vortex.py | 29 ++- grudge/flux_differencing.py | 268 +++++++++++++++++++ grudge/interpolation.py | 141 ++++++++++ grudge/models/euler.py | 428 ++++++++++++++++++++++++++----- grudge/op.py | 3 +- grudge/projection.py | 70 ++++- test/test_euler_model.py | 78 +++++- test/test_sbp_ops.py | 171 ++++++++++++ 11 files changed, 1370 insertions(+), 86 deletions(-) create mode 100644 examples/euler/sod.py create mode 100644 grudge/flux_differencing.py create mode 100644 test/test_sbp_ops.py diff --git a/doc/operators.rst b/doc/operators.rst index 550a78023..881afd945 100644 --- a/doc/operators.rst +++ b/doc/operators.rst @@ -3,6 +3,7 @@ Discontinuous Galerkin operators .. automodule:: grudge.op .. automodule:: grudge.trace_pair +.. automodule:: grudge.flux_differencing Transfering data between discretizations diff --git a/examples/euler/acoustic_pulse.py b/examples/euler/acoustic_pulse.py index 779062910..d5017de22 100644 --- a/examples/euler/acoustic_pulse.py +++ b/examples/euler/acoustic_pulse.py @@ -35,6 +35,7 @@ from grudge.models.euler import ( ConservedEulerField, EulerOperator, + EntropyStableEulerOperator, InviscidWallBC ) from grudge.shortcuts import rk4_step @@ -111,9 +112,24 @@ def run_acoustic_pulse(actx, order=3, final_time=1, resolution=16, + esdg=False, overintegration=False, visualize=False): + logger.info( + """ + Acoustic pulse parameters:\n + order: %s\n + final_time: %s\n + resolution: %s\n + entropy stable: %s\n + overintegration: %s\n + visualize: %s + """, + order, final_time, resolution, esdg, + overintegration, visualize + ) + # eos-related parameters gamma = 1.4 @@ -135,7 +151,15 @@ def run_acoustic_pulse(actx, (default_simplex_group_factory, QuadratureSimplexGroupFactory) - exp_name = f"fld-acoustic-pulse-N{order}-K{resolution}" + if esdg: + case = "esdg-pulse" + operator_cls = EntropyStableEulerOperator + else: + case = "pulse" + operator_cls = EulerOperator + + exp_name = f"fld-{case}-N{order}-K{resolution}" + if overintegration: exp_name += "-overintegrated" quad_tag = DISCR_TAG_QUAD @@ -155,7 +179,7 @@ def run_acoustic_pulse(actx, # {{{ Euler operator - euler_operator = EulerOperator( + euler_operator = operator_cls( dcoll, bdry_conditions={BTAG_ALL: InviscidWallBC()}, flux_type="lf", @@ -212,7 +236,7 @@ def rhs(t, q): def main(ctx_factory, order=3, final_time=1, resolution=16, - overintegration=False, visualize=False, lazy=False): + esdg=False, overintegration=False, visualize=False, lazy=False): cl_ctx = ctx_factory() queue = cl.CommandQueue(cl_ctx) @@ -228,10 +252,17 @@ def main(ctx_factory, order=3, final_time=1, resolution=16, force_device_scalars=True, ) + if not actx.supports_nonscalar_broadcasting and esdg is True: + raise RuntimeError( + "Cannot use ESDG with an array context that cannot perform " + "nonscalar broadcasting. Run with --lazy instead." + ) + run_acoustic_pulse( actx, order=order, resolution=resolution, + esdg=esdg, overintegration=overintegration, final_time=final_time, visualize=visualize @@ -245,6 +276,8 @@ def main(ctx_factory, order=3, final_time=1, resolution=16, parser.add_argument("--order", default=3, type=int) parser.add_argument("--tfinal", default=0.1, type=float) parser.add_argument("--resolution", default=16, type=int) + parser.add_argument("--esdg", action="store_true", + help="use entropy stable dg") parser.add_argument("--oi", action="store_true", help="use overintegration") parser.add_argument("--visualize", action="store_true", @@ -258,6 +291,7 @@ def main(ctx_factory, order=3, final_time=1, resolution=16, order=args.order, final_time=args.tfinal, resolution=args.resolution, + esdg=args.esdg, overintegration=args.oi, visualize=args.visualize, lazy=args.lazy) diff --git a/examples/euler/sod.py b/examples/euler/sod.py new file mode 100644 index 000000000..abf138f41 --- /dev/null +++ b/examples/euler/sod.py @@ -0,0 +1,227 @@ +__copyright__ = """ +Copyright (C) 2021 University of Illinois Board of Trustees +""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +from grudge.dof_desc import BoundaryDomainTag +import pyopencl as cl +import pyopencl.tools as cl_tools + +from arraycontext import thaw, freeze +from grudge.array_context import PytatoPyOpenCLArrayContext +from grudge.models.euler import ( + EntropyStableEulerOperator, + ConservedEulerField, + PrescribedBC, + conservative_to_primitive_vars, +) +from grudge.shortcuts import rk4_step + +from pytools.obj_array import make_obj_array + +import grudge.op as op + +import logging +logger = logging.getLogger(__name__) + + +def sod_shock_initial_condition(nodes, t=0): + gamma = 1.4 + dim = len(nodes) + gmn1 = 1.0 / (gamma - 1.0) + x = nodes[0] + actx = x.array_context + zeros = 0*x + + _x0 = 0.5 + _rhoin = 1.0 + _rhoout = 0.125 + _pin = 1.0 + _pout = 0.1 + rhoin = zeros + _rhoin + rhoout = zeros + _rhoout + energyin = zeros + gmn1 * _pin + energyout = zeros + gmn1 * _pout + + x0 = zeros + _x0 + sigma = 1e-13 + weight = 0.5 * (1.0 - actx.np.tanh(1.0/sigma * (x - x0))) + + mass = rhoout + (rhoin - rhoout)*weight + energy = energyout + (energyin - energyout)*weight + momentum = make_obj_array([zeros for _ in range(dim)]) + + return ConservedEulerField(mass=mass, energy=energy, momentum=momentum) + + +def run_sod_shock_tube( + actx, order=4, resolution=32, final_time=0.2, visualize=False): + + logger.info( + """ + Sod 1-D parameters:\n + order: %s\n + final_time: %s\n + resolution: %s\n + visualize: %s + """, + order, final_time, resolution, visualize + ) + + # eos-related parameters + gamma = 1.4 + + # {{{ discretization + + from meshmode.mesh.generation import generate_regular_rect_mesh + + dim = 1 + box_ll = 0.0 + box_ur = 1.0 + mesh = generate_regular_rect_mesh( + a=(box_ll,)*dim, + b=(box_ur,)*dim, + nelements_per_axis=(resolution,)*dim, + boundary_tag_to_face={ + "prescribed": ["+x", "-x"], + } + ) + + from grudge import DiscretizationCollection + from grudge.dof_desc import \ + DISCR_TAG_BASE, DISCR_TAG_QUAD + from meshmode.discretization.poly_element import \ + (default_simplex_group_factory, + QuadratureSimplexGroupFactory) + + exp_name = f"fld-sod-1d-N{order}-K{resolution}" + quad_tag = DISCR_TAG_QUAD + + dcoll = DiscretizationCollection( + actx, mesh, + discr_tag_to_group_factory={ + DISCR_TAG_BASE: default_simplex_group_factory(dim, order), + DISCR_TAG_QUAD: QuadratureSimplexGroupFactory(order + 2) + } + ) + + # }}} + + # {{{ Euler operator + + dd_prescribed = BoundaryDomainTag("prescribed") + bcs = { + dd_prescribed: PrescribedBC(prescribed_state=sod_shock_initial_condition) + } + + euler_operator = EntropyStableEulerOperator( + dcoll, + bdry_conditions=bcs, + flux_type="lf", + gamma=gamma, + quadrature_tag=quad_tag + ) + + def rhs(t, q): + return euler_operator.operator(t, q) + + compiled_rhs = actx.compile(rhs) + + fields = sod_shock_initial_condition(thaw(dcoll.nodes(), actx)) + + from grudge.dt_utils import h_min_from_volume + + cfl = 0.01 + cn = 0.5*(order + 1)**2 + dt = cfl * actx.to_numpy(h_min_from_volume(dcoll)) / cn + + logger.info("Timestep size: %g", dt) + + # }}} + + from grudge.shortcuts import make_visualizer + + vis = make_visualizer(dcoll) + + # {{{ time stepping + + step = 0 + t = 0.0 + while t < final_time: + if step % 10 == 0: + norm_q = actx.to_numpy(op.norm(dcoll, fields, 2)) + logger.info("[%04d] t = %.5f |q| = %.5e", step, t, norm_q) + if visualize: + rho, velocity, pressure = \ + conservative_to_primitive_vars(fields, gamma=gamma) + vis.write_vtk_file( + f"{exp_name}-{step:04d}.vtu", + [ + ("rho", rho), + ("energy", fields.energy), + ("momentum", fields.momentum), + ("velocity", velocity), + ("pressure", pressure) + ] + ) + assert norm_q < 10000 + + fields = thaw(freeze(fields, actx), actx) + fields = rk4_step(fields, t, dt, compiled_rhs) + t += dt + step += 1 + + # }}} + + +def main(ctx_factory, order=4, final_time=0.2, resolution=32, visualize=False): + cl_ctx = ctx_factory() + queue = cl.CommandQueue(cl_ctx) + actx = PytatoPyOpenCLArrayContext( + queue, + allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)), + ) + + run_sod_shock_tube( + actx, order=order, + resolution=resolution, + final_time=final_time, + visualize=visualize) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--order", default=4, type=int) + parser.add_argument("--tfinal", default=0.2, type=float) + parser.add_argument("--resolution", default=32, type=int) + parser.add_argument("--visualize", action="store_true") + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + main(cl.create_some_context, + order=args.order, + final_time=args.tfinal, + resolution=args.resolution, + visualize=args.visualize) diff --git a/examples/euler/vortex.py b/examples/euler/vortex.py index 9f00743e5..ab94529db 100644 --- a/examples/euler/vortex.py +++ b/examples/euler/vortex.py @@ -29,7 +29,8 @@ from grudge.array_context import PytatoPyOpenCLArrayContext, PyOpenCLArrayContext from grudge.models.euler import ( vortex_initial_condition, - EulerOperator + EulerOperator, + EntropyStableEulerOperator ) from grudge.shortcuts import rk4_step @@ -40,6 +41,7 @@ def run_vortex(actx, order=3, resolution=8, final_time=5, + esdg=False, overintegration=False, flux_type="central", visualize=False): @@ -50,11 +52,12 @@ def run_vortex(actx, order=3, resolution=8, final_time=5, order: %s\n final_time: %s\n resolution: %s\n + entropy stable: %s\n overintegration: %s\n flux_type: %s\n visualize: %s """, - order, final_time, resolution, + order, final_time, resolution, esdg, overintegration, flux_type, visualize ) @@ -76,7 +79,14 @@ def run_vortex(actx, order=3, resolution=8, final_time=5, from meshmode.discretization.poly_element import \ default_simplex_group_factory, QuadratureSimplexGroupFactory - exp_name = f"fld-vortex-N{order}-K{resolution}-{flux_type}" + if esdg: + case = "esdg-vortex" + operator_cls = EntropyStableEulerOperator + else: + case = "vortex" + operator_cls = EulerOperator + + exp_name = f"fld-{case}-N{order}-K{resolution}-{flux_type}" if overintegration: exp_name += "-overintegrated" @@ -97,7 +107,7 @@ def run_vortex(actx, order=3, resolution=8, final_time=5, # {{{ Euler operator - euler_operator = EulerOperator( + euler_operator = operator_cls( dcoll, flux_type=flux_type, gamma=gamma, @@ -154,6 +164,7 @@ def rhs(t, q): def main(ctx_factory, order=3, final_time=5, resolution=8, + esdg=False, overintegration=False, lf_stabilization=False, visualize=False, @@ -173,6 +184,12 @@ def main(ctx_factory, order=3, final_time=5, resolution=8, force_device_scalars=True, ) + if not actx.supports_nonscalar_broadcasting and esdg is True: + raise RuntimeError( + "Cannot use ESDG with an array context that cannot perform " + "nonscalar broadcasting. Run with --lazy instead." + ) + if lf_stabilization: flux_type = "lf" else: @@ -183,6 +200,7 @@ def main(ctx_factory, order=3, final_time=5, resolution=8, order=order, resolution=resolution, final_time=final_time, + esdg=esdg, overintegration=overintegration, flux_type=flux_type, visualize=visualize @@ -196,6 +214,8 @@ def main(ctx_factory, order=3, final_time=5, resolution=8, parser.add_argument("--order", default=3, type=int) parser.add_argument("--tfinal", default=0.015, type=float) parser.add_argument("--resolution", default=8, type=int) + parser.add_argument("--esdg", action="store_true", + help="use entropy stable dg") parser.add_argument("--oi", action="store_true", help="use overintegration") parser.add_argument("--lf", action="store_true", @@ -211,6 +231,7 @@ def main(ctx_factory, order=3, final_time=5, resolution=8, order=args.order, final_time=args.tfinal, resolution=args.resolution, + esdg=args.esdg, overintegration=args.oi, lf_stabilization=args.lf, visualize=args.visualize, diff --git a/grudge/flux_differencing.py b/grudge/flux_differencing.py new file mode 100644 index 000000000..782f5c9dc --- /dev/null +++ b/grudge/flux_differencing.py @@ -0,0 +1,268 @@ +"""Grudge module for flux-differencing in entropy-stable DG methods + +Flux-differencing +----------------- + +.. autofunction:: volume_flux_differencing +""" + +__copyright__ = """ +Copyright (C) 2021 University of Illinois Board of Trustees +""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +from arraycontext import ( + ArrayContext, + map_array_container, + freeze +) +from arraycontext import ArrayOrContainer + +from functools import partial + +from meshmode.transform_metadata import FirstAxisIsElementsTag +from meshmode.dof_array import DOFArray + +from grudge.discretization import DiscretizationCollection +from grudge.dof_desc import DOFDesc + +from pytools import memoize_in, keyed_memoize_in + +import numpy as np + + +def _reference_skew_symmetric_hybridized_sbp_operators( + actx: ArrayContext, + base_element_group, + vol_quad_element_group, + face_quad_element_group, dtype): + @keyed_memoize_in( + actx, _reference_skew_symmetric_hybridized_sbp_operators, + lambda base_grp, quad_vol_grp, face_quad_grp: ( + base_grp.discretization_key(), + quad_vol_grp.discretization_key(), + face_quad_grp.discretization_key())) + def get_reference_skew_symetric_hybridized_diff_mats( + base_grp, quad_vol_grp, face_quad_grp): + from meshmode.discretization.poly_element import diff_matrices + from modepy import faces_for_shape, face_normal + from grudge.interpolation import ( + volume_quadrature_interpolation_matrix, + surface_quadrature_interpolation_matrix + ) + from grudge.op import reference_inverse_mass_matrix + + # {{{ Volume operators + + weights = quad_vol_grp.quadrature_rule().weights + vdm_q = actx.to_numpy( + volume_quadrature_interpolation_matrix(actx, base_grp, quad_vol_grp)) + inv_mass_mat = actx.to_numpy( + reference_inverse_mass_matrix(actx, base_grp)) + p_mat = inv_mass_mat @ (vdm_q.T * weights) + + # }}} + + # {{{ Surface operators + + faces = faces_for_shape(base_grp.shape) + nfaces = len(faces) + # NOTE: assumes same quadrature rule on all faces + face_weights = np.tile(face_quad_grp.quadrature_rule().weights, nfaces) + face_normals = [face_normal(face) for face in faces] + nnods_per_face = face_quad_grp.nunit_dofs + e = np.ones(shape=(nnods_per_face,)) + nrstj = [ + # nsrtJ = nhat * Jhatf, where nhat is the reference normal + # and Jhatf is the Jacobian det. of the transformation from + # the face of the reference element to the reference face. + np.concatenate([np.sign(nhat[idx])*e for nhat in face_normals]) + for idx in range(base_grp.dim) + ] + b_mats = [np.diag(face_weights*nrstj[d]) for d in range(base_grp.dim)] + vf_mat = actx.to_numpy( + surface_quadrature_interpolation_matrix( + actx, + base_element_group=base_grp, + face_quad_element_group=face_quad_grp)) + zero_mat = np.zeros((nfaces*nnods_per_face, nfaces*nnods_per_face), + dtype=dtype) + + # }}} + + # {{{ Hybridized (volume + surface) operators + + q_mats = [p_mat.T @ (weights * vdm_q.T @ vdm_q) @ diff_mat @ p_mat + for diff_mat in diff_matrices(base_grp)] + e_mat = vf_mat @ p_mat + q_skew_hybridized = np.asarray( + [ + np.block( + [[q_mats[d] - q_mats[d].T, e_mat.T @ b_mats[d]], + [-b_mats[d] @ e_mat, zero_mat]] + ) for d in range(base_grp.dim) + ], + order="C" + ) + + # }}} + + return actx.freeze(actx.from_numpy(q_skew_hybridized)) + + return get_reference_skew_symetric_hybridized_diff_mats( + base_element_group, + vol_quad_element_group, + face_quad_element_group + ) + + +def _single_axis_hybridized_derivative_kernel( + dcoll, dd_quad, dd_face_quad, xyz_axis, flux_matrix): + if not dcoll._has_affine_groups(dd_quad.domain_tag): + raise NotImplementedError("Not implemented for non-affine elements yet.") + + if not isinstance(flux_matrix, DOFArray): + return map_array_container( + partial(_single_axis_hybridized_derivative_kernel, + dcoll, dd_quad, dd_face_quad, xyz_axis), + flux_matrix + ) + + from grudge.geometry import \ + area_element, inverse_surface_metric_derivative + from grudge.interpolation import ( + volume_and_surface_interpolation_matrix, + volume_and_surface_quadrature_interpolation + ) + + actx = flux_matrix.array_context + + # FIXME: This is kinda meh + def inverse_jac_matrix(): + @memoize_in( + dcoll, + (_single_axis_hybridized_derivative_kernel, dd_quad, dd_face_quad)) + def _inv_surf_metric_deriv(): + return freeze( + actx.np.stack( + [ + actx.np.stack( + [ + volume_and_surface_quadrature_interpolation( + dcoll, dd_quad, dd_face_quad, + area_element(actx, dcoll) + * inverse_surface_metric_derivative( + actx, dcoll, + rst_ax, xyz_axis + ) + ) for rst_ax in range(dcoll.dim) + ] + ) for xyz_axis in range(dcoll.ambient_dim) + ] + ), + actx + ) + return _inv_surf_metric_deriv() + + return DOFArray( + actx, + data=tuple( + # r for rst axis + actx.einsum("ik,rej,rij,eij->ek", + volume_and_surface_interpolation_matrix( + actx, + base_element_group=bgrp, + vol_quad_element_group=qvgrp, + face_quad_element_group=qafgrp + ), + ijm_i[xyz_axis], + _reference_skew_symmetric_hybridized_sbp_operators( + actx, + bgrp, + qvgrp, + qafgrp, + fmat_i.dtype + ), + fmat_i, + arg_names=("Vh_mat_t", "inv_jac_t", "Q_mat", "F_mat"), + tagged=(FirstAxisIsElementsTag(),)) + + for bgrp, qvgrp, qafgrp, fmat_i, ijm_i in zip( + dcoll.discr_from_dd("vol").groups, + dcoll.discr_from_dd(dd_quad).groups, + dcoll.discr_from_dd(dd_face_quad).groups, + flux_matrix, + inverse_jac_matrix() + ) + ) + ) + + +def volume_flux_differencing( + dcoll: DiscretizationCollection, + dd_quad: DOFDesc, + dd_face_quad: DOFDesc, + flux_matrices: ArrayOrContainer) -> ArrayOrContainer: + r"""Computes the volume contribution of the DG divergence operator using + flux-differencing: + + .. math:: + + \mathrm{VOL} = \sum_{i=1}^{d} + \begin{bmatrix} + \mathbf{V}_q \\ \mathbf{V}_f + \end{bmatrix}^T + \left( + \left( \mathbf{Q}_{i} - \mathbf{Q}^T_{i} \right) + \circ \mathbf{F}_{i} + \right)\mathbf{1} + + where :math:`\circ` denotes the + `Hadamard product `__, + :math:`\mathbf{F}_{i}` are matrices whose entries are computed + as the evaluation of an entropy-conserving two-point flux function + (e.g. :func:`grudge.models.euler.divergence_flux_chandrashekar`) + and :math:`\mathbf{Q}_{i} - \mathbf{Q}^T_{i}` are the skew-symmetric + hybridized differentiation operators defined in (15) of + `this paper `__. + + :arg flux_matrices: a :class:`~meshmode.dof_array.DOFArray` or an + :class:`~arraycontext.container.ArrayContainer` of them containing + evaluations of two-point flux. + :returns: a :class:`~meshmode.dof_array.DOFArray` or an + :class:`~arraycontext.container.ArrayContainer` of them. + """ + + def _hybridized_div(fmats): + return sum(_single_axis_hybridized_derivative_kernel( + dcoll, dd_quad, dd_face_quad, i, fmat_i) + for i, fmat_i in enumerate(fmats)) + + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + _hybridized_div, + (dcoll.ambient_dim,), (), + flux_matrices, scalar_cls=DOFArray) + + +# vim: foldmethod=marker diff --git a/grudge/interpolation.py b/grudge/interpolation.py index 61bdf1a13..8976ae79b 100644 --- a/grudge/interpolation.py +++ b/grudge/interpolation.py @@ -32,7 +32,24 @@ """ +import numpy as np + +from arraycontext import ( + ArrayContext, + map_array_container +) +from arraycontext import ArrayOrContainerT + +from functools import partial + +from meshmode.transform_metadata import FirstAxisIsElementsTag + from grudge.discretization import DiscretizationCollection +from grudge.dof_desc import DOFDesc + +from meshmode.dof_array import DOFArray + +from pytools import keyed_memoize_in # FIXME: Should revamp interp and make clear distinctions @@ -46,3 +63,127 @@ def interp(dcoll: DiscretizationCollection, src, tgt, vec): from grudge.projection import project return project(dcoll, src, tgt, vec) + + +# {{{ Interpolation matrices + +def volume_quadrature_interpolation_matrix( + actx: ArrayContext, base_element_group, vol_quad_element_group): + @keyed_memoize_in( + actx, volume_quadrature_interpolation_matrix, + lambda base_grp, vol_quad_grp: (base_grp.discretization_key(), + vol_quad_grp.discretization_key())) + def get_volume_vand(base_grp, vol_quad_grp): + from modepy import vandermonde + + basis = base_grp.basis_obj() + vdm_inv = np.linalg.inv(vandermonde(basis.functions, + base_grp.unit_nodes)) + vdm_q = vandermonde(basis.functions, vol_quad_grp.unit_nodes) @ vdm_inv + return actx.freeze(actx.from_numpy(vdm_q)) + + return get_volume_vand(base_element_group, vol_quad_element_group) + + +def surface_quadrature_interpolation_matrix( + actx: ArrayContext, base_element_group, face_quad_element_group): + @keyed_memoize_in( + actx, surface_quadrature_interpolation_matrix, + lambda base_grp, face_quad_grp: (base_grp.discretization_key(), + face_quad_grp.discretization_key())) + def get_surface_vand(base_grp, face_quad_grp): + nfaces = base_grp.mesh_el_group.nfaces + assert face_quad_grp.nelements == nfaces * base_grp.nelements + + from modepy import vandermonde, faces_for_shape + + basis = base_grp.basis_obj() + vdm_inv = np.linalg.inv(vandermonde(basis.functions, + base_grp.unit_nodes)) + faces = faces_for_shape(base_grp.shape) + # NOTE: Assumes same quadrature rule on each face + face_quadrature = face_quad_grp.quadrature_rule() + + surface_nodes = faces[0].map_to_volume(face_quadrature.nodes) + for fidx in range(1, nfaces): + surface_nodes = np.append( + surface_nodes, + faces[fidx].map_to_volume(face_quadrature.nodes), + axis=1 + ) + vdm_f = vandermonde(basis.functions, surface_nodes) @ vdm_inv + return actx.freeze(actx.from_numpy(vdm_f)) + + return get_surface_vand(base_element_group, face_quad_element_group) + + +def volume_and_surface_interpolation_matrix( + actx: ArrayContext, + base_element_group, vol_quad_element_group, face_quad_element_group): + @keyed_memoize_in( + actx, volume_and_surface_interpolation_matrix, + lambda base_grp, vol_quad_grp, face_quad_grp: ( + base_grp.discretization_key(), + vol_quad_grp.discretization_key(), + face_quad_grp.discretization_key())) + def get_vol_surf_interpolation_matrix(base_grp, vol_quad_grp, face_quad_grp): + vq_mat = actx.to_numpy( + volume_quadrature_interpolation_matrix( + actx, + base_element_group=base_grp, + vol_quad_element_group=vol_quad_grp)) + vf_mat = actx.to_numpy( + surface_quadrature_interpolation_matrix( + actx, + base_element_group=base_grp, + face_quad_element_group=face_quad_grp)) + return actx.freeze(actx.from_numpy(np.block([[vq_mat], [vf_mat]]))) + + return get_vol_surf_interpolation_matrix( + base_element_group, vol_quad_element_group, face_quad_element_group + ) + +# }}} + + +def volume_and_surface_quadrature_interpolation( + dcoll: DiscretizationCollection, + dd_quad: DOFDesc, + dd_face_quad: DOFDesc, + vec: ArrayOrContainerT) -> ArrayOrContainerT: + """todo. + """ + if not isinstance(vec, DOFArray): + return map_array_container( + partial(volume_and_surface_quadrature_interpolation, + dcoll, dd_quad, dd_face_quad), vec + ) + + actx = vec.array_context + discr = dcoll.discr_from_dd("vol") + quad_volm_discr = dcoll.discr_from_dd(dd_quad) + quad_face_discr = dcoll.discr_from_dd(dd_face_quad) + + return DOFArray( + actx, + data=tuple( + actx.einsum("ij,ej->ei", + volume_and_surface_interpolation_matrix( + actx, + base_element_group=bgrp, + vol_quad_element_group=qvgrp, + face_quad_element_group=qfgrp + ), + vec_i, + arg_names=("Vh_mat", "vec"), + tagged=(FirstAxisIsElementsTag(),)) + + for bgrp, qvgrp, qfgrp, vec_i in zip( + discr.groups, + quad_volm_discr.groups, + quad_face_discr.groups, vec) + ) + ) + + +# vim: foldmethod=marker diff --git a/grudge/models/euler.py b/grudge/models/euler.py index f4d6f8f4c..70c765e31 100644 --- a/grudge/models/euler.py +++ b/grudge/models/euler.py @@ -4,6 +4,7 @@ ----------------- .. autoclass:: EulerOperator +.. autoclass:: EntropyStableEulerOperator Predefined initial conditions ----------------------------- @@ -20,6 +21,9 @@ .. autofunction:: euler_volume_flux .. autofunction:: euler_numerical_flux + +.. autofunction:: divergence_flux_chandrashekar +.. autofunction:: entropy_stable_numerical_flux_chandrashekar """ __copyright__ = """ @@ -49,9 +53,12 @@ from abc import ABCMeta, abstractmethod from dataclasses import dataclass + from arraycontext import ( dataclass_array_container, - with_container_arithmetic + with_container_arithmetic, + map_array_container, thaw, + outer ) from meshmode.dof_array import DOFArray @@ -124,14 +131,13 @@ def vortex_initial_condition( # {{{ Variable transformation and helper routines -def conservative_to_primitive_vars(cv_state: ConservedEulerField, gamma=1.4): +def conservative_to_primitive_vars(cv_state: ConservedEulerField, gamma: float): """Converts from conserved variables (density, momentum, total energy) into primitive variables (density, velocity, pressure). :arg cv_state: A :class:`ConservedEulerField` containing the conserved variables. - :arg gamma: The isentropic expansion factor for a single-species gas - (default set to 1.4). + :arg gamma: The isentropic expansion factor. :returns: A :class:`Tuple` containing the primitive variables: (density, velocity, pressure). """ @@ -141,20 +147,76 @@ def conservative_to_primitive_vars(cv_state: ConservedEulerField, gamma=1.4): u = rho_u / rho p = (gamma - 1) * (rho_e - 0.5 * sum(rho_u * u)) - return rho, u, p + return (rho, u, p) -def compute_wavespeed(cv_state: ConservedEulerField, gamma=1.4): - """Computes the total translational wavespeed. +def conservative_to_entropy_vars(cv_state: ConservedEulerField, gamma: float): + """Converts from conserved variables (density, momentum, total energy) + into entropy variables. :arg cv_state: A :class:`ConservedEulerField` containing the conserved variables. :arg gamma: The isentropic expansion factor for a single-species gas (default set to 1.4). + :returns: A :class:`ConservedEulerField` containing the entropy variables. + """ + actx = cv_state.array_context + rho, u, p = conservative_to_primitive_vars(cv_state, gamma) + + u_square = sum(v ** 2 for v in u) + s = actx.np.log(p) - gamma*actx.np.log(rho) + rho_p = rho / p + + return ConservedEulerField( + mass=((gamma - s)/(gamma - 1)) - 0.5 * rho_p * u_square, + energy=-rho_p, + momentum=rho_p * u + ) + + +def entropy_to_conservative_vars(ev_state: ConservedEulerField, gamma: float): + """Converts from entropy variables into conserved variables + (density, momentum, total energy). + + :arg ev_state: A :class:`ConservedEulerField` containing the entropy + variables. + :arg gamma: The isentropic expansion factor. + :returns: A :class:`ConservedEulerField` containing the conserved variables. + """ + actx = ev_state.array_context + # See Hughes, Franca, Mallet (1986) A new finite element + # formulation for CFD: (DOI: 10.1016/0045-7825(86)90127-1) + inv_gamma_minus_one = 1/(gamma - 1) + + # Convert to entropy `-rho * s` used by Hughes, France, Mallet (1986) + ev_state = ev_state * (gamma - 1) + v1 = ev_state.mass + v2t4 = ev_state.momentum + v5 = ev_state.energy + + v_square = sum(v**2 for v in v2t4) + s = gamma - v1 + v_square/(2*v5) + rho_iota = ( + ((gamma - 1) / (-v5)**gamma)**(inv_gamma_minus_one) + ) * actx.np.exp(-s * inv_gamma_minus_one) + + return ConservedEulerField( + mass=-rho_iota * v5, + energy=rho_iota * (1 - v_square/(2*v5)), + momentum=rho_iota * v2t4 + ) + + +def compute_wavespeed(cv_state: ConservedEulerField, gamma: float): + """Computes the total translational wavespeed. + + :arg cv_state: A :class:`ConservedEulerField` containing the conserved + variables. + :arg gamma: The isentropic expansion factor. :returns: A :class:`~meshmode.dof_array.DOFArray` containing local wavespeeds. """ actx = cv_state.array_context - rho, u, p = conservative_to_primitive_vars(cv_state, gamma=gamma) + rho, u, p = conservative_to_primitive_vars(cv_state, gamma) return actx.np.sqrt(np.dot(u, u)) + actx.np.sqrt(gamma * (p / rho)) @@ -173,7 +235,7 @@ def boundary_tpair( self, dcoll: DiscretizationCollection, dd_bc: DOFDesc, - state: ConservedEulerField, t=0): + restricted_state: ConservedEulerField, t=0): pass @@ -183,14 +245,13 @@ def boundary_tpair( self, dcoll: DiscretizationCollection, dd_bc: DOFDesc, - state: ConservedEulerField, t=0): - actx = state.array_context - dd_base = as_dofdesc("vol").with_discr_tag(DISCR_TAG_BASE) + restricted_state: ConservedEulerField, t=0): + actx = restricted_state.array_context return TracePair( dd_bc, - interior=op.project(dcoll, dd_base, dd_bc, state), - exterior=self.prescribed_state(actx.thaw(dcoll.nodes(dd_bc)), t=t) + interior=restricted_state, + exterior=self.prescribed_state(thaw(dcoll.nodes(dd_bc), actx), t=t) ) @@ -200,11 +261,10 @@ def boundary_tpair( self, dcoll: DiscretizationCollection, dd_bc: DOFDesc, - state: ConservedEulerField, t=0): - actx = state.array_context - dd_base = as_dofdesc("vol").with_discr_tag(DISCR_TAG_BASE) - nhat = actx.thaw(dcoll.normal(dd_bc)) - interior = op.project(dcoll, dd_base, dd_bc, state) + restricted_state: ConservedEulerField, t=0): + actx = restricted_state.array_context + nhat = thaw(dcoll.normal(dd_bc), actx) + interior = restricted_state return TracePair( dd_bc, @@ -224,19 +284,17 @@ def boundary_tpair( # {{{ Euler operator def euler_volume_flux( - dcoll: DiscretizationCollection, cv_state: ConservedEulerField, gamma=1.4): + dcoll: DiscretizationCollection, + cv_state: ConservedEulerField, gamma: float): """Computes the (non-linear) volume flux for the Euler operator. :arg cv_state: A :class:`ConservedEulerField` containing the conserved variables. - :arg gamma: The isentropic expansion factor for a single-species gas - (default set to 1.4). + :arg gamma: The isentropic expansion factor. :returns: A :class:`ConservedEulerField` containing the volume fluxes. """ - from arraycontext import outer - - rho, u, p = conservative_to_primitive_vars(cv_state, gamma=gamma) + rho, u, p = conservative_to_primitive_vars(cv_state, gamma) return ConservedEulerField( mass=cv_state.momentum, @@ -247,40 +305,35 @@ def euler_volume_flux( def euler_numerical_flux( dcoll: DiscretizationCollection, tpair: TracePair, - gamma=1.4, lf_stabilization=False): + gamma: float, dissipation=False): """Computes the interface numerical flux for the Euler operator. :arg tpair: A :class:`grudge.trace_pair.TracePair` containing the conserved variables on the interior and exterior sides of element facets. - :arg gamma: The isentropic expansion factor for a single-species gas - (default set to 1.4). - :arg lf_stabilization: A boolean denoting whether to apply Lax-Friedrichs + :arg gamma: The isentropic expansion factor. + :arg dissipation: A boolean denoting whether to apply Lax-Friedrichs dissipation. :returns: A :class:`ConservedEulerField` containing the interface fluxes. """ - dd_intfaces = tpair.dd - dd_allfaces = dd_intfaces.with_dtag("all_faces") q_ll = tpair.int q_rr = tpair.ext actx = q_ll.array_context flux_tpair = TracePair( tpair.dd, - interior=euler_volume_flux(dcoll, q_ll, gamma=gamma), - exterior=euler_volume_flux(dcoll, q_rr, gamma=gamma) + interior=euler_volume_flux(dcoll, q_ll, gamma), + exterior=euler_volume_flux(dcoll, q_rr, gamma) ) num_flux = flux_tpair.avg - normal = actx.thaw(dcoll.normal(dd_intfaces)) - - if lf_stabilization: - from arraycontext import outer + normal = thaw(dcoll.normal(tpair.dd), actx) + if dissipation: # Compute jump penalization parameter - lam = actx.np.maximum(compute_wavespeed(q_ll, gamma=gamma), - compute_wavespeed(q_rr, gamma=gamma)) + lam = actx.np.maximum(compute_wavespeed(q_ll, gamma), + compute_wavespeed(q_rr, gamma)) num_flux -= lam*outer(tpair.diff, normal)/2 - return op.project(dcoll, dd_intfaces, dd_allfaces, num_flux @ normal) + return num_flux @ normal class EulerOperator(HyperbolicOperator): @@ -309,59 +362,298 @@ def __init__(self, dcoll: DiscretizationCollection, def max_characteristic_velocity(self, actx, **kwargs): state = kwargs["state"] - return compute_wavespeed(state, gamma=self.gamma) + return compute_wavespeed(state, self.gamma) def operator(self, t, q): dcoll = self.dcoll gamma = self.gamma qtag = self.qtag - dq = DOFDesc("vol", qtag) - df = DOFDesc("all_faces", qtag) - def interp_to_quad(u): - return op.project(dcoll, "vol", dq, u) + dissipation = self.lf_stabilization + + dd_base = as_dofdesc("vol", DISCR_TAG_BASE) + dd_vol_quad = as_dofdesc("vol", qtag) + dd_face_quad = as_dofdesc("all_faces", qtag) + + def interp_to_quad_surf(tpair): + dd = tpair.dd + dd_quad = dd.with_discr_tag(qtag) + return TracePair( + dd_quad, + interior=op.project(dcoll, dd, dd_quad, tpair.int), + exterior=op.project(dcoll, dd, dd_quad, tpair.ext) + ) + + interior_trace_pairs = [ + interp_to_quad_surf(tpair) + for tpair in op.interior_trace_pairs(dcoll, q) + ] - # Compute volume fluxes - volume_fluxes = op.weak_local_div( - dcoll, dq, - interp_to_quad(euler_volume_flux(dcoll, q, gamma=gamma)) + # Compute volume derivatives + volume_derivs = op.weak_local_div( + dcoll, dd_vol_quad, + euler_volume_flux( + dcoll, op.project(dcoll, dd_base, dd_vol_quad, q), gamma) ) # Compute interior interface fluxes interface_fluxes = ( sum( - euler_numerical_flux( - dcoll, - op.tracepair_with_discr_tag(dcoll, qtag, tpair), - gamma=gamma, - lf_stabilization=self.lf_stabilization - ) for tpair in op.interior_trace_pairs(dcoll, q) + op.project(dcoll, qtpair.dd, dd_face_quad, + euler_numerical_flux(dcoll, qtpair, gamma, + dissipation=dissipation)) + for qtpair in interior_trace_pairs ) ) # Compute boundary fluxes if self.bdry_conditions is not None: - bc_fluxes = sum( - euler_numerical_flux( + for btag in self.bdry_conditions: + boundary_condition = self.bdry_conditions[btag] + dd_bc = as_dofdesc(btag).with_discr_tag(qtag) + bc_flux = op.project( dcoll, - self.bdry_conditions[btag].boundary_tpair( + dd_bc, + dd_face_quad, + euler_numerical_flux( dcoll, - as_dofdesc(btag).with_discr_tag(qtag), - q, - t=t - ), - gamma=gamma, - lf_stabilization=self.lf_stabilization - ) for btag in self.bdry_conditions - ) - interface_fluxes = interface_fluxes + bc_fluxes + boundary_condition.boundary_tpair( + dcoll=dcoll, + dd_bc=dd_bc, + restricted_state=op.project(dcoll, dd_base, dd_bc, q), + t=t + ), + gamma, + dissipation=dissipation + ) + ) + interface_fluxes = interface_fluxes + bc_flux return op.inverse_mass( dcoll, - volume_fluxes - op.face_mass(dcoll, df, interface_fluxes) + volume_derivs - op.face_mass(dcoll, dd_face_quad, interface_fluxes) ) # }}} +# {{{ Entropy stable Euler operator + +def divergence_flux_chandrashekar( + dcoll: DiscretizationCollection, + q_left: ConservedEulerField, + q_right: ConservedEulerField, gamma: float): + """Two-point volume flux based on the entropy conserving + and kinetic energy preserving two-point flux in: + + Chandrashekar (2013) Kinetic Energy Preserving and Entropy Stable Finite + Volume Schemes for Compressible Euler and Navier-Stokes Equations: + `DOI `__. + + :args q_left: A :class:`ConservedEulerField` containing the "left" state. + :args q_right: A :class:`ConservedEulerField` containing the "right" state. + :arg gamma: The isentropic expansion factor. + """ + dim = dcoll.dim + actx = q_left.array_context + + def ln_mean(x: DOFArray, y: DOFArray, epsilon=1e-4): + f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) + return actx.np.where( + actx.np.less(f2, epsilon), + (x + y) / (2 + f2*2/3 + f2*f2*2/5 + f2*f2*f2*2/7), + (y - x) / actx.np.log(y / x) + ) + + rho_left, u_left, p_left = conservative_to_primitive_vars(q_left, gamma) + rho_right, u_right, p_right = conservative_to_primitive_vars(q_right, gamma) + + beta_left = 0.5 * rho_left / p_left + beta_right = 0.5 * rho_right / p_right + specific_kin_left = 0.5 * sum(v**2 for v in u_left) + specific_kin_right = 0.5 * sum(v**2 for v in u_right) + + rho_avg = 0.5 * (rho_left + rho_right) + rho_mean = ln_mean(rho_left, rho_right) + beta_mean = ln_mean(beta_left, beta_right) + beta_avg = 0.5 * (beta_left + beta_right) + u_avg = 0.5 * (u_left + u_right) + p_mean = 0.5 * rho_avg / beta_avg + + velocity_square_avg = specific_kin_left + specific_kin_right + + mass_flux = rho_mean * u_avg + momentum_flux = outer(mass_flux, u_avg) + np.eye(dim) * p_mean + energy_flux = ( + mass_flux * 0.5 * (1/(gamma - 1)/beta_mean - velocity_square_avg) + + np.dot(momentum_flux, u_avg) + ) + + return ConservedEulerField(mass=mass_flux, + energy=energy_flux, + momentum=momentum_flux) + + +def entropy_stable_numerical_flux_chandrashekar( + dcoll: DiscretizationCollection, tpair: TracePair, + gamma: float, dissipation=False): + """Entropy stable numerical flux based on the entropy conserving + and kinetic energy preserving two-point flux in: + + Chandrashekar (2013) Kinetic Energy Preserving and Entropy Stable Finite + Volume Schemes for Compressible Euler and Navier-Stokes Equations + `DOI `__. + + :arg tpair: A :class:`grudge.trace_pair.TracePair` containing the conserved + variables on the interior and exterior sides of element facets. + :arg gamma: The isentropic expansion factor. + :arg dissipation: A boolean denoting whether to apply Lax-Friedrichs + dissipation. + :returns: A :class:`ConservedEulerField` containing the interface fluxes. + """ + q_int = tpair.int + q_ext = tpair.ext + actx = q_int.array_context + + num_flux = divergence_flux_chandrashekar( + dcoll, q_left=q_int, q_right=q_ext, gamma=gamma) + normal = thaw(dcoll.normal(tpair.dd), actx) + + if dissipation: + # Compute jump penalization parameter + lam = actx.np.maximum(compute_wavespeed(q_int, gamma), + compute_wavespeed(q_ext, gamma)) + num_flux -= lam*outer(tpair.diff, normal)/2 + + return num_flux @ normal + + +class EntropyStableEulerOperator(EulerOperator): + """Discretizes the Euler equations using an entropy-stable + discontinuous Galerkin discretization as outlined in (15) + of `this paper `__. + """ + + def operator(self, t, q): + from grudge.projection import volume_quadrature_project + from grudge.interpolation import \ + volume_and_surface_quadrature_interpolation + + dcoll = self.dcoll + gamma = self.gamma + qtag = self.qtag + dissipation = self.lf_stabilization + + dd_base = DOFDesc("vol", DISCR_TAG_BASE) + dd_vol_quad = DOFDesc("vol", qtag) + dd_face_quad = DOFDesc("all_faces", qtag) + + # Convert to projected entropy variables: v_q = P_q v(u_q) + proj_entropy_vars = \ + volume_quadrature_project( + dcoll, dd_vol_quad, + conservative_to_entropy_vars( + # Interpolate state to vol quad grid: u_q = V_q u + op.project(dcoll, dd_base, dd_vol_quad, q), gamma)) + + def modified_conserved_vars_tpair(tpair): + dd = tpair.dd + dd_quad = dd.with_discr_tag(qtag) + # Interpolate entropy variables to the surface quadrature grid + ev_tpair = op.project(dcoll, dd, dd_quad, tpair) + return TracePair( + dd_quad, + # Convert interior and exterior states to conserved variables + interior=entropy_to_conservative_vars(ev_tpair.int, gamma), + exterior=entropy_to_conservative_vars(ev_tpair.ext, gamma) + ) + + # Compute interior trace pairs containing the modified conserved + # variables (in terms of projected entropy variables) + interior_trace_pairs = [ + modified_conserved_vars_tpair(tpair) + for tpair in op.interior_trace_pairs(dcoll, proj_entropy_vars) + ] + + from functools import partial + from grudge.flux_differencing import volume_flux_differencing + + def _reshape(shape, ary): + if not isinstance(ary, DOFArray): + return map_array_container(partial(_reshape, shape), ary) + + return DOFArray(ary.array_context, data=tuple( + subary.reshape(grp.nelements, *shape) + # Just need group for determining the number of elements + for grp, subary in zip(dcoll.discr_from_dd(dd_base).groups, ary))) + + # Compute the (modified) conserved state in terms of the projected + # entropy variables on both the volume and surface nodes + qtilde_vol_and_surf = \ + entropy_to_conservative_vars( + # Interpolate projected entropy variables to + # volume + surface quadrature grids + volume_and_surface_quadrature_interpolation( + dcoll, dd_vol_quad, dd_face_quad, proj_entropy_vars), gamma) + + # FIXME: These matrices are actually symmetric. Could make use + # of that to avoid redundant computation. + flux_matrices = divergence_flux_chandrashekar( + dcoll, + _reshape((1, -1), qtilde_vol_and_surf), + _reshape((-1, 1), qtilde_vol_and_surf), + gamma + ) + + # Compute volume derivatives using flux differencing + volume_derivs = -volume_flux_differencing( + dcoll, dd_vol_quad, dd_face_quad, flux_matrices) + + # Computing interface numerical fluxes + interface_fluxes = ( + sum( + op.project(dcoll, qtpair.dd, dd_face_quad, + entropy_stable_numerical_flux_chandrashekar( + dcoll, qtpair, gamma, dissipation=dissipation)) + for qtpair in interior_trace_pairs + ) + ) + + # Compute boundary fluxes + if self.bdry_conditions is not None: + for btag in self.bdry_conditions: + boundary_condition = self.bdry_conditions[btag] + dd_bc = as_dofdesc(btag).with_discr_tag(qtag) + bc_flux = op.project( + dcoll, + dd_bc, + dd_face_quad, + entropy_stable_numerical_flux_chandrashekar( + dcoll, + boundary_condition.boundary_tpair( + dcoll=dcoll, + dd_bc=dd_bc, + # Pass modified conserved state to be used as + # the "interior" state for computing the boundary + # trace pair + restricted_state=entropy_to_conservative_vars( + op.project( + dcoll, dd_base, dd_bc, proj_entropy_vars), + gamma + ), + t=t + ), + gamma, + dissipation=dissipation + ) + ) + interface_fluxes = interface_fluxes + bc_flux + + return op.inverse_mass( + dcoll, + volume_derivs - op.face_mass(dcoll, dd_face_quad, interface_fluxes) + ) + +# }}} + # vim: foldmethod=marker diff --git a/grudge/op.py b/grudge/op.py index f5781f4be..61df22b14 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -95,7 +95,7 @@ ) from grudge.interpolation import interp -from grudge.projection import project +from grudge.projection import project, volume_quadrature_project from grudge.reductions import ( norm, @@ -127,6 +127,7 @@ __all__ = ( "project", + "volume_quadrature_project", "interp", "norm", diff --git a/grudge/projection.py b/grudge/projection.py index e21e02295..38b4f0123 100644 --- a/grudge/projection.py +++ b/grudge/projection.py @@ -5,6 +5,7 @@ ----------- .. autofunction:: project +.. autofunction:: volume_quadrature_project """ from __future__ import annotations @@ -33,8 +34,12 @@ THE SOFTWARE. """ +from functools import partial +from numbers import Number + +import numpy as np -from arraycontext import ArrayOrContainer +from arraycontext import ArrayOrContainer, map_array_container from grudge.discretization import DiscretizationCollection from grudge.dof_desc import ( @@ -43,7 +48,10 @@ BoundaryDomainTag, ConvertibleToDOFDesc) -from numbers import Number +from meshmode.dof_array import DOFArray +from meshmode.transform_metadata import FirstAxisIsElementsTag + +from pytools import keyed_memoize_in def project( @@ -82,3 +90,61 @@ def project( return vec return dcoll.connection_from_dds(src_dofdesc, tgt_dofdesc)(vec) + + +def volume_quadrature_project( + dcoll: DiscretizationCollection, dd_q, vec) -> ArrayOrContainer: + """Projects a field on the quadrature discreization, described by *dd_q*, + into the polynomial space described by the volume discretization. + + :arg dd_q: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. + :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an + :class:`~arraycontext.container.ArrayContainer` of them. + :returns: a :class:`~meshmode.dof_array.DOFArray` or an + :class:`~arraycontext.container.ArrayContainer` like *vec*. + """ + if not isinstance(vec, DOFArray): + return map_array_container( + partial(volume_quadrature_project, dcoll, dd_q), vec + ) + + from grudge.geometry import area_element + from grudge.interpolation import volume_quadrature_interpolation_matrix + from grudge.op import inverse_mass + + actx = vec.array_context + discr = dcoll.discr_from_dd("vol") + quad_discr = dcoll.discr_from_dd(dd_q) + jacobians = area_element( + actx, dcoll, dd=dd_q, + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + + @keyed_memoize_in( + actx, volume_quadrature_project, + lambda base_grp, vol_quad_grp: (base_grp.discretization_key(), + vol_quad_grp.discretization_key())) + def get_mat(base_grp, vol_quad_grp): + vdm_q = actx.to_numpy( + volume_quadrature_interpolation_matrix( + actx, base_grp, vol_quad_grp + ) + ) + weights = np.diag(vol_quad_grp.quadrature_rule().weights) + return actx.freeze(actx.from_numpy(vdm_q.T @ weights)) + + return inverse_mass( + dcoll, + DOFArray( + actx, + data=tuple( + actx.einsum("ij,ej,ej->ei", + get_mat(bgrp, qgrp), + jac_i, + vec_i, + arg_names=("vqw_t", "jac", "vec"), + tagged=(FirstAxisIsElementsTag(),)) + for bgrp, qgrp, vec_i, jac_i in zip( + discr.groups, quad_discr.groups, vec, jacobians) + ) + ) + ) diff --git a/test/test_euler_model.py b/test/test_euler_model.py index 13a479d48..5fec819ba 100644 --- a/test/test_euler_model.py +++ b/test/test_euler_model.py @@ -24,22 +24,26 @@ import pytest -from grudge.array_context import PytestPyOpenCLArrayContextFactory +from grudge.array_context import \ + PytestPyOpenCLArrayContextFactory, PytestPytatoPyOpenCLArrayContextFactory from arraycontext import ( - pytest_generate_tests_for_array_contexts, + pytest_generate_tests_for_array_contexts, thaw, ) pytest_generate_tests = pytest_generate_tests_for_array_contexts( - [PytestPyOpenCLArrayContextFactory]) + [PytestPyOpenCLArrayContextFactory, + PytestPytatoPyOpenCLArrayContextFactory]) import grudge.op as op +import numpy as np + import logging logger = logging.getLogger(__name__) @pytest.mark.parametrize("order", [1, 2, 3]) -def test_euler_vortex_convergence(actx_factory, order): - +@pytest.mark.parametrize("esdg", [False, True]) +def test_euler_vortex_convergence(actx_factory, order, esdg): from meshmode.mesh.generation import generate_regular_rect_mesh from grudge import DiscretizationCollection @@ -47,7 +51,8 @@ def test_euler_vortex_convergence(actx_factory, order): from grudge.dt_utils import h_max_from_volume from grudge.models.euler import ( vortex_initial_condition, - EulerOperator + EulerOperator, + EntropyStableEulerOperator ) from grudge.shortcuts import rk4_step @@ -60,6 +65,16 @@ def test_euler_vortex_convergence(actx_factory, order): actx = actx_factory() eoc_rec = EOCRecorder() quad_tag = DISCR_TAG_QUAD + if esdg: + operator_cls = EntropyStableEulerOperator + else: + operator_cls = EulerOperator + + if esdg and not actx.supports_nonscalar_broadcasting: + pytest.xfail( + "Flux-differencing computations requires an array context " + "that supports non-scalar broadcasting" + ) for resolution in [8, 16, 32]: @@ -85,7 +100,7 @@ def test_euler_vortex_convergence(actx_factory, order): # }}} - euler_operator = EulerOperator( + euler_operator = operator_cls( dcoll, flux_type="lf", gamma=1.4, @@ -135,8 +150,55 @@ def rhs(t, q): logger.info("\n%s", eoc_rec.pretty_print(abscissa_label="h", error_label="L2 Error")) + assert eoc_rec.order_estimate() >= order + 0.5 + + +def test_entropy_variable_roundtrip(actx_factory): + from grudge.models.euler import ( + entropy_to_conservative_vars, + conservative_to_entropy_vars, + vortex_initial_condition + ) + + actx = actx_factory() + gamma = 1.4 # Adiabatic expansion factor for single-gas Euler model + + from meshmode.mesh.generation import generate_regular_rect_mesh + + dim = 2 + res = 5 + mesh = generate_regular_rect_mesh( + a=(0, -5), + b=(20, 5), + nelements_per_axis=(2*res, res), + periodic=(True, True)) + + from meshmode.discretization.poly_element import \ + default_simplex_group_factory + from grudge import DiscretizationCollection + from grudge.dof_desc import DISCR_TAG_BASE + + order = 3 + dcoll = DiscretizationCollection( + actx, mesh, + discr_tag_to_group_factory={ + DISCR_TAG_BASE: default_simplex_group_factory(dim, order) + } + ) + + # Fields in conserved variables + fields = vortex_initial_condition(thaw(dcoll.nodes(), actx)) + + # Map back and forth between entropy and conserved vars + fields_ev = conservative_to_entropy_vars(fields, gamma) + ev_fields_to_cons = entropy_to_conservative_vars(fields_ev, gamma) + residual = ev_fields_to_cons - fields + + assert actx.to_numpy(op.norm(dcoll, residual.mass, np.inf)) < 1e-13 + assert actx.to_numpy(op.norm(dcoll, residual.energy, np.inf)) < 1e-13 assert ( - eoc_rec.order_estimate() >= order + 0.5 + actx.to_numpy(op.norm(dcoll, residual.momentum[i], np.inf)) < 1e-13 + for i in range(dim) ) diff --git a/test/test_sbp_ops.py b/test/test_sbp_ops.py new file mode 100644 index 000000000..41bf0d6d1 --- /dev/null +++ b/test/test_sbp_ops.py @@ -0,0 +1,171 @@ +__copyright__ = "Copyright (C) 2021 University of Illinois Board of Trustees" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +import numpy as np + +from grudge import DiscretizationCollection +from grudge.dof_desc import DOFDesc, DISCR_TAG_BASE, DISCR_TAG_QUAD + +import pytest + +from grudge.array_context import PytestPyOpenCLArrayContextFactory +from arraycontext import pytest_generate_tests_for_array_contexts +pytest_generate_tests = pytest_generate_tests_for_array_contexts( + [PytestPyOpenCLArrayContextFactory]) + +import logging + +logger = logging.getLogger(__name__) + + +@pytest.mark.parametrize("dim", [1, 2, 3]) +@pytest.mark.parametrize("order", [1, 2, 3, 4]) +def test_reference_element_sbp_operators(actx_factory, dim, order): + actx = actx_factory() + + from meshmode.mesh.generation import generate_regular_rect_mesh + + nel_1d = 5 + box_ll = -5.0 + box_ur = 5.0 + mesh = generate_regular_rect_mesh( + a=(box_ll,)*dim, + b=(box_ur,)*dim, + nelements_per_axis=(nel_1d,)*dim) + + from meshmode.discretization.poly_element import \ + default_simplex_group_factory, QuadratureSimplexGroupFactory + + dcoll = DiscretizationCollection( + actx, mesh, + discr_tag_to_group_factory={ + DISCR_TAG_BASE: default_simplex_group_factory(dim, order), + DISCR_TAG_QUAD: QuadratureSimplexGroupFactory(2*order) + } + ) + + dd_q = DOFDesc("vol", DISCR_TAG_QUAD) + dd_f = DOFDesc("all_faces", DISCR_TAG_QUAD) + + volm_discr = dcoll.discr_from_dd("vol") + quad_discr = dcoll.discr_from_dd(dd_q) + quad_face_discr = dcoll.discr_from_dd(dd_f) + + from meshmode.discretization.poly_element import diff_matrices + from modepy import faces_for_shape, face_normal + from grudge.interpolation import ( + volume_quadrature_interpolation_matrix, + surface_quadrature_interpolation_matrix + ) + from grudge.op import reference_inverse_mass_matrix + + for vgrp, qgrp, qfgrp in zip(volm_discr.groups, + quad_discr.groups, + quad_face_discr.groups): + nq_vol = qgrp.nunit_dofs + nq_per_face = qfgrp.nunit_dofs + nfaces = vgrp.shape.nfaces + nq_faces = nfaces * nq_per_face + nq_total = nq_vol + nq_faces + + # {{{ Volume operators + + weights = qgrp.quadrature_rule().weights + vdm_q = actx.to_numpy( + volume_quadrature_interpolation_matrix(actx, vgrp, qgrp)) + inv_mass_mat = actx.to_numpy( + reference_inverse_mass_matrix(actx, vgrp)) + p_mat = inv_mass_mat @ (vdm_q.T * weights) + + # }}} + + # Checks Pq @ Vq = Minv @ Vq.T @ W @ Vq = I + assert np.allclose(p_mat @ vdm_q, + np.identity(len(inv_mass_mat)), rtol=1e-15) + + # {{{ Surface operators + + faces = faces_for_shape(vgrp.shape) + # NOTE: assumes same quadrature rule on all faces + face_weights = np.tile(qfgrp.quadrature_rule().weights, nfaces) + face_normals = [face_normal(face) for face in faces] + e = np.ones(shape=(nq_per_face,)) + nrstj = [np.concatenate([np.sign(nhat[idx])*e + for nhat in face_normals]) + for idx in range(vgrp.dim)] + b_mats = [np.diag(face_weights*nrstj[d]) for d in range(vgrp.dim)] + vf_mat = actx.to_numpy( + surface_quadrature_interpolation_matrix( + actx, + base_element_group=vgrp, + face_quad_element_group=qfgrp + ) + ) + + # }}} + + # {{{ Hybridized (volume + surface) operators + + q_mats = [p_mat.T @ (weights * vdm_q.T @ vdm_q) @ diff_mat @ p_mat + for diff_mat in diff_matrices(vgrp)] + e_mat = vf_mat @ p_mat + qtilde_mats = 0.5 * np.asarray( + [np.block([[q_mats[d] - q_mats[d].T, e_mat.T @ b_mats[d]], + [-b_mats[d] @ e_mat, b_mats[d]]]) + for d in range(dcoll.dim)] + ) + + # }}} + + ones = np.ones(shape=(nq_total,)) + zeros = np.zeros(shape=(nq_total,)) + for idx in range(dim): + # Checks the generalized SBP property: + # Qi + Qi.T = E.T @ Bi @ E + # c.f. Lemma 1. https://arxiv.org/pdf/1708.01243.pdf + assert np.allclose(q_mats[idx] + q_mats[idx].T, + e_mat.T @ b_mats[idx] @ e_mat, rtol=1e-15) + + # Checks the SBP-like property for the skew hybridized operator + # Qitilde + Qitilde.T = [0 0; 0 Bi] + # c.f. Theorem 1 and Lemma 1. https://arxiv.org/pdf/1902.01828.pdf + residual = qtilde_mats[idx] + qtilde_mats[idx].T + residual[nq_vol:nq_vol+nq_faces, nq_vol:nq_vol+nq_faces] -= b_mats[idx] + assert np.allclose(residual, np.zeros(residual.shape), rtol=1e-15) + + # Checks quadrature condition for: Qiskew @ ones = zeros + # Qiskew + Qiskew.T = [0 0; 0 Bi] + # c.f. Lemma 2. https://arxiv.org/pdf/1902.01828.pdf + assert np.allclose(np.dot(qtilde_mats[idx], ones), + zeros, rtol=1e-15) + + +# You can test individual routines by typing +# $ python test_grudge.py 'test_routine()' + +if __name__ == "__main__": + import sys + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + pytest.main([__file__]) From 7c532c6a2113101acc0c406f57752a0848c3a644 Mon Sep 17 00:00:00 2001 From: "Michael T. Campbell" Date: Wed, 26 Apr 2023 08:32:24 -0700 Subject: [PATCH 13/97] Update for multi-volume cases. --- grudge/flux_differencing.py | 13 ++++++++----- grudge/interpolation.py | 4 +++- grudge/projection.py | 6 +++++- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/grudge/flux_differencing.py b/grudge/flux_differencing.py index 782f5c9dc..61ef1c1bb 100644 --- a/grudge/flux_differencing.py +++ b/grudge/flux_differencing.py @@ -148,6 +148,9 @@ def _single_axis_hybridized_derivative_kernel( flux_matrix ) + from grudge.dof_desc import DISCR_TAG_BASE + dd_vol = dd_quad.with_discr_tag(DISCR_TAG_BASE) + from grudge.geometry import \ area_element, inverse_surface_metric_derivative from grudge.interpolation import ( @@ -208,11 +211,11 @@ def _inv_surf_metric_deriv(): tagged=(FirstAxisIsElementsTag(),)) for bgrp, qvgrp, qafgrp, fmat_i, ijm_i in zip( - dcoll.discr_from_dd("vol").groups, - dcoll.discr_from_dd(dd_quad).groups, - dcoll.discr_from_dd(dd_face_quad).groups, - flux_matrix, - inverse_jac_matrix() + dcoll.discr_from_dd(dd_vol).groups, + dcoll.discr_from_dd(dd_quad).groups, + dcoll.discr_from_dd(dd_face_quad).groups, + flux_matrix, + inverse_jac_matrix() ) ) ) diff --git a/grudge/interpolation.py b/grudge/interpolation.py index 8976ae79b..a2b9d1f77 100644 --- a/grudge/interpolation.py +++ b/grudge/interpolation.py @@ -159,8 +159,10 @@ def volume_and_surface_quadrature_interpolation( dcoll, dd_quad, dd_face_quad), vec ) + from grudge.dof_desc import DISCR_TAG_BASE + dd_vol = dd_quad.with_discr_tag(DISCR_TAG_BASE) actx = vec.array_context - discr = dcoll.discr_from_dd("vol") + discr = dcoll.discr_from_dd(dd_vol) quad_volm_discr = dcoll.discr_from_dd(dd_quad) quad_face_discr = dcoll.discr_from_dd(dd_face_quad) diff --git a/grudge/projection.py b/grudge/projection.py index 38b4f0123..093919673 100644 --- a/grudge/projection.py +++ b/grudge/projection.py @@ -112,8 +112,11 @@ def volume_quadrature_project( from grudge.interpolation import volume_quadrature_interpolation_matrix from grudge.op import inverse_mass + from grudge.dof_desc import DISCR_TAG_BASE + dd_vol = dd_q.with_discr_tag(DISCR_TAG_BASE) + actx = vec.array_context - discr = dcoll.discr_from_dd("vol") + discr = dcoll.discr_from_dd(dd_vol) quad_discr = dcoll.discr_from_dd(dd_q) jacobians = area_element( actx, dcoll, dd=dd_q, @@ -134,6 +137,7 @@ def get_mat(base_grp, vol_quad_grp): return inverse_mass( dcoll, + dd_vol, DOFArray( actx, data=tuple( From ffa5592be5a525c7a93ddbbccbdba33dfb19ed19 Mon Sep 17 00:00:00 2001 From: Michael Campbell Date: Wed, 26 Apr 2023 13:55:21 -0500 Subject: [PATCH 14/97] Update for multivol --- grudge/flux_differencing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/grudge/flux_differencing.py b/grudge/flux_differencing.py index 61ef1c1bb..816138df9 100644 --- a/grudge/flux_differencing.py +++ b/grudge/flux_differencing.py @@ -173,10 +173,10 @@ def _inv_surf_metric_deriv(): [ volume_and_surface_quadrature_interpolation( dcoll, dd_quad, dd_face_quad, - area_element(actx, dcoll) + area_element(actx, dcoll, dd=dd_vol) * inverse_surface_metric_derivative( actx, dcoll, - rst_ax, xyz_axis + rst_ax, xyz_axis, dd=dd_vol ) ) for rst_ax in range(dcoll.dim) ] From 248043c6a7fe0d78667d965215eeb5c603a3773a Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 24 May 2023 19:07:18 -0700 Subject: [PATCH 15/97] tag axes in reshape in _apply_elementwise_reduction --- grudge/reductions.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/grudge/reductions.py b/grudge/reductions.py index 6087b5725..9d68aff9e 100644 --- a/grudge/reductions.py +++ b/grudge/reductions.py @@ -65,6 +65,7 @@ make_loopy_program, map_array_container, serialize_container, + tag_axes, Scalar, ArrayOrContainer ) @@ -73,7 +74,9 @@ from pytools import memoize_in from meshmode.dof_array import DOFArray -from meshmode.transform_metadata import DiscretizationDOFAxisTag +from meshmode.transform_metadata import ( + DiscretizationElementAxisTag, + DiscretizationDOFAxisTag) import numpy as np import grudge.dof_desc as dof_desc @@ -339,7 +342,10 @@ def _apply_elementwise_reduction( return DOFArray( actx, data=tuple( - getattr(actx.np, op_name)(vec_i, axis=1).reshape(-1, 1) + tag_axes(actx, { + 0: DiscretizationElementAxisTag(), + 1: DiscretizationDOFAxisTag()}, + getattr(actx.np, op_name)(vec_i, axis=1).reshape(-1, 1)) for vec_i in vec ) ) From 889444542abadd1d7e97e2cbebb0182f4136ec27 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 21 Sep 2022 09:49:52 -0500 Subject: [PATCH 16/97] set up connections between volumes --- grudge/discretization.py | 281 +++++++++++++++++++++++++++++---------- 1 file changed, 208 insertions(+), 73 deletions(-) diff --git a/grudge/discretization.py b/grudge/discretization.py index 8e57ca503..fd4c39728 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -7,6 +7,7 @@ .. autofunction:: make_discretization_collection .. currentmodule:: grudge.discretization +.. autoclass:: PartID """ __copyright__ = """ @@ -34,10 +35,12 @@ THE SOFTWARE. """ -from typing import Mapping, Optional, Union, TYPE_CHECKING, Any +from typing import Sequence, Mapping, Optional, Union, Tuple, TYPE_CHECKING, Any from pytools import memoize_method, single_valued +from dataclasses import dataclass, replace + from grudge.dof_desc import ( VTAG_ALL, DD_VOLUME_ALL, @@ -71,6 +74,75 @@ import mpi4py.MPI +@dataclass(frozen=True) +class PartID: + """Unique identifier for a piece of a partitioned mesh. + + .. attribute:: volume_tag + + The volume of the part. + + .. attribute:: rank + + The (optional) MPI rank of the part. + + """ + volume_tag: VolumeTag + rank: Optional[int] = None + + +# {{{ part ID normalization + +def _normalize_mesh_part_ids( + mesh: Mesh, + self_volume_tag: VolumeTag, + all_volume_tags: Sequence[VolumeTag], + mpi_communicator: Optional["mpi4py.MPI.Intracomm"] = None): + """Convert a mesh's configuration-dependent "part ID" into a fixed type.""" + from numbers import Integral + if mpi_communicator is not None: + # Accept PartID or rank (assume intra-volume for the latter) + def as_part_id(mesh_part_id): + if isinstance(mesh_part_id, PartID): + return mesh_part_id + elif isinstance(mesh_part_id, Integral): + return PartID(self_volume_tag, int(mesh_part_id)) + else: + raise TypeError(f"Unable to convert {mesh_part_id} to PartID.") + else: + # Accept PartID or volume tag + def as_part_id(mesh_part_id): + if isinstance(mesh_part_id, PartID): + return mesh_part_id + elif mesh_part_id in all_volume_tags: + return PartID(mesh_part_id) + else: + raise TypeError(f"Unable to convert {mesh_part_id} to PartID.") + + facial_adjacency_groups = mesh.facial_adjacency_groups + + new_facial_adjacency_groups = [] + + from meshmode.mesh import InterPartAdjacencyGroup + for grp_list in facial_adjacency_groups: + new_grp_list = [] + for fagrp in grp_list: + if isinstance(fagrp, InterPartAdjacencyGroup): + part_id = as_part_id(fagrp.part_id) + new_fagrp = replace( + fagrp, + boundary_tag=BTAG_PARTITION(part_id), + part_id=part_id) + else: + new_fagrp = fagrp + new_grp_list.append(new_fagrp) + new_facial_adjacency_groups.append(new_grp_list) + + return mesh.copy(facial_adjacency_groups=new_facial_adjacency_groups) + +# }}} + + # {{{ discr_tag_to_group_factory normalization def _normalize_discr_tag_to_group_factory( @@ -156,6 +228,9 @@ def __init__(self, array_context: ArrayContext, discr_tag_to_group_factory: Optional[ Mapping[DiscretizationTag, ElementGroupFactory]] = None, mpi_communicator: Optional["mpi4py.MPI.Intracomm"] = None, + inter_part_connections: Optional[ + Mapping[Tuple[PartID, PartID], + DiscretizationConnection]] = None, ) -> None: """ :arg discr_tag_to_group_factory: A mapping from discretization tags @@ -206,6 +281,9 @@ def __init__(self, array_context: ArrayContext, mesh = volume_discrs + mesh = _normalize_mesh_part_ids( + mesh, VTAG_ALL, [VTAG_ALL], mpi_communicator=mpi_communicator) + discr_tag_to_group_factory = _normalize_discr_tag_to_group_factory( dim=mesh.dim, discr_tag_to_group_factory=discr_tag_to_group_factory, @@ -219,17 +297,32 @@ def __init__(self, array_context: ArrayContext, del mesh + if inter_part_connections is not None: + raise TypeError("may not pass inter_part_connections when " + "DiscretizationCollection constructor is called in " + "legacy mode") + + self._inter_part_connections = \ + _set_up_inter_part_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: 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 + if inter_part_connections is None: + raise TypeError("inter_part_connections must be passed when " + "DiscretizationCollection constructor is called in " + "'modern' mode") + + self._inter_part_connections = inter_part_connections - self._dist_boundary_connections = { - vtag: self._set_up_distributed_communication( - vtag, mpi_communicator, array_context) - for vtag in self._volume_discrs.keys()} + self._volume_discrs = volume_discrs # }}} @@ -252,71 +345,6 @@ def is_management_rank(self): return self.mpi_communicator.Get_rank() \ == self.get_management_rank_index() - # {{{ distributed - - def _set_up_distributed_communication( - self, vtag, mpi_communicator, array_context): - from_dd = DOFDesc(VolumeDomainTag(vtag), DISCR_TAG_BASE) - - boundary_connections = {} - - from meshmode.distributed import get_connected_partitions - connected_parts = get_connected_partitions(self._volume_discrs[vtag].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, from_dd.trace(BTAG_PARTITION(i_remote_part))) - - 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 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.BoundaryDomainTag` 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, BoundaryDomainTag) - assert isinstance(dd.domain_tag.tag, BTAG_PARTITION) - - vtag = dd.domain_tag.volume_tag - - return self._dist_boundary_connections[vtag][dd.domain_tag.tag.part_nr] - - # }}} - # {{{ discr_from_dd @memoize_method @@ -772,6 +800,105 @@ def normal(self, dd): # }}} +# {{{ distributed/multi-volume setup + +def _set_up_inter_part_connections( + array_context: ArrayContext, + mpi_communicator: Optional["mpi4py.MPI.Intracomm"], + volume_discrs: Mapping[VolumeTag, Discretization], + base_group_factory: ElementGroupFactory, + ) -> Mapping[ + Tuple[PartID, PartID], + DiscretizationConnection]: + + from meshmode.distributed import (get_connected_parts, + make_remote_group_infos, InterRankBoundaryInfo, + MPIBoundaryCommSetupHelper) + + rank = mpi_communicator.Get_rank() if mpi_communicator is not None else None + + # Save boundary restrictions as they're created to avoid potentially creating + # them twice in the loop below + cached_part_bdry_restrictions: Mapping[ + Tuple[PartID, PartID], + DiscretizationConnection] = {} + + def get_part_bdry_restriction(self_part_id, other_part_id): + cached_result = cached_part_bdry_restrictions.get( + (self_part_id, other_part_id), None) + if cached_result is not None: + return cached_result + return cached_part_bdry_restrictions.setdefault( + (self_part_id, other_part_id), + make_face_restriction( + array_context, volume_discrs[self_part_id.volume_tag], + base_group_factory, + boundary_tag=BTAG_PARTITION(other_part_id))) + + inter_part_conns: Mapping[ + Tuple[PartID, PartID], + DiscretizationConnection] = {} + + irbis = [] + + for vtag, volume_discr in volume_discrs.items(): + part_id = PartID(vtag, rank) + connected_part_ids = get_connected_parts(volume_discr.mesh) + for connected_part_id in connected_part_ids: + bdry_restr = get_part_bdry_restriction( + self_part_id=part_id, other_part_id=connected_part_id) + + if connected_part_id.rank == rank: + # {{{ rank-local interface between multiple volumes + + connected_bdry_restr = get_part_bdry_restriction( + self_part_id=connected_part_id, other_part_id=part_id) + + from meshmode.discretization.connection import \ + make_partition_connection + inter_part_conns[connected_part_id, part_id] = \ + make_partition_connection( + array_context, + local_bdry_conn=bdry_restr, + remote_bdry_discr=connected_bdry_restr.to_discr, + remote_group_infos=make_remote_group_infos( + array_context, part_id, connected_bdry_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_part_id=part_id, + remote_part_id=connected_part_id, + remote_rank=connected_part_id.rank, + local_boundary_connection=bdry_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 + +# }}} + + # {{{ modal group factory def _generate_modal_group_factory(nodal_group_factory): @@ -860,6 +987,8 @@ def make_discretization_collection( del order + mpi_communicator = getattr(array_context, "mpi_communicator", None) + if any( isinstance(mesh_or_discr, Discretization) for mesh_or_discr in volumes.values()): @@ -868,14 +997,20 @@ def make_discretization_collection( volume_discrs = { vtag: Discretization( array_context, - mesh, + _normalize_mesh_part_ids( + mesh, vtag, volumes.keys(), mpi_communicator=mpi_communicator), discr_tag_to_group_factory[DISCR_TAG_BASE]) for vtag, mesh in volumes.items()} return DiscretizationCollection( array_context=array_context, volume_discrs=volume_discrs, - discr_tag_to_group_factory=discr_tag_to_group_factory) + discr_tag_to_group_factory=discr_tag_to_group_factory, + inter_part_connections=_set_up_inter_part_connections( + array_context=array_context, + mpi_communicator=mpi_communicator, + volume_discrs=volume_discrs, + base_group_factory=discr_tag_to_group_factory[DISCR_TAG_BASE])) # }}} From deb2e46e006767d980318cda15a98e13fbd2b784 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 21 Sep 2022 09:51:54 -0500 Subject: [PATCH 17/97] add inter-volume communication --- grudge/eager.py | 3 +- grudge/op.py | 10 +- grudge/trace_pair.py | 627 +++++++++++++++++++++++++++++++++++-------- 3 files changed, 522 insertions(+), 118 deletions(-) diff --git a/grudge/eager.py b/grudge/eager.py index 626e15592..08cf08f2a 100644 --- a/grudge/eager.py +++ b/grudge/eager.py @@ -87,7 +87,8 @@ def nodal_max(self, dd, vec): return op.nodal_max(self, dd, vec) -connected_ranks = op.connected_ranks +# FIXME: Deprecate connected_ranks instead of removing +connected_parts = op.connected_parts interior_trace_pair = op.interior_trace_pair cross_rank_trace_pairs = op.cross_rank_trace_pairs diff --git a/grudge/op.py b/grudge/op.py index f5781f4be..a6cef8ffa 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -118,8 +118,11 @@ interior_trace_pair, interior_trace_pairs, local_interior_trace_pair, - connected_ranks, + connected_parts, + 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 ) @@ -147,8 +150,11 @@ "interior_trace_pair", "interior_trace_pairs", "local_interior_trace_pair", - "connected_ranks", + "connected_parts", + "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", diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 1f49ae0d6..0b0400f12 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__ = """ @@ -52,17 +55,18 @@ from warnings import warn -from typing import List, Hashable, Optional, Type, Any +from typing import List, Hashable, Optional, Tuple, Type, Any, Sequence, Mapping 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, + get_container_context_recursively_opt, + to_numpy, + from_numpy, ArrayOrContainer ) @@ -72,7 +76,7 @@ from pytools import memoize_on_first_arg -from grudge.discretization import DiscretizationCollection +from grudge.discretization import DiscretizationCollection, PartID from grudge.projection import project from meshmode.mesh import BTAG_PARTITION @@ -82,7 +86,7 @@ import grudge.dof_desc as dof_desc from grudge.dof_desc import ( DOFDesc, DD_VOLUME_ALL, FACE_RESTR_INTERIOR, DISCR_TAG_BASE, - VolumeDomainTag, + VolumeTag, VolumeDomainTag, BoundaryDomainTag, ConvertibleToDOFDesc, ) @@ -360,6 +364,124 @@ def interior_trace_pairs(dcoll: DiscretizationCollection, vec, *, # }}} +# {{{ inter-volume trace pairs + +def local_inter_volume_trace_pairs( + dcoll: DiscretizationCollection, + pairwise_volume_data: Mapping[ + Tuple[DOFDesc, DOFDesc], + Tuple[ArrayOrContainer, ArrayOrContainer]] + ) -> Mapping[Tuple[DOFDesc, DOFDesc], TracePair]: + for vol_dd_pair in pairwise_volume_data.keys(): + for vol_dd in vol_dd_pair: + if not isinstance(vol_dd.domain_tag, VolumeDomainTag): + raise ValueError( + "pairwise_volume_data keys must describe volumes, " + f"got '{vol_dd}'") + if vol_dd.discretization_tag != DISCR_TAG_BASE: + raise ValueError( + "expected base-discretized DOFDesc in pairwise_volume_data, " + f"got '{vol_dd}'") + + rank = ( + dcoll.mpi_communicator.Get_rank() + if dcoll.mpi_communicator is not None + else None) + + result: Mapping[Tuple[DOFDesc, DOFDesc], TracePair] = {} + + for vol_dd_pair, vol_data_pair in pairwise_volume_data.items(): + from meshmode.mesh import mesh_has_boundary + if not mesh_has_boundary( + dcoll.discr_from_dd(vol_dd_pair[0]).mesh, + BTAG_PARTITION(PartID(vol_dd_pair[1].domain_tag.tag, rank))): + continue + + directional_vol_dd_pairs = [ + (vol_dd_pair[1], vol_dd_pair[0]), + (vol_dd_pair[0], vol_dd_pair[1])] + + trace_dd_pair = tuple( + self_vol_dd.trace( + BTAG_PARTITION( + PartID(other_vol_dd.domain_tag.tag, rank))) + for other_vol_dd, self_vol_dd in directional_vol_dd_pairs) + + # Pre-compute the projections out here to avoid doing it twice inside + # the loop below + trace_data = { + trace_dd: project(dcoll, vol_dd, trace_dd, vol_data) + for vol_dd, trace_dd, vol_data in zip( + vol_dd_pair, trace_dd_pair, vol_data_pair)} + + for other_vol_dd, self_vol_dd in directional_vol_dd_pairs: + self_part_id = PartID(self_vol_dd.domain_tag.tag, rank) + other_part_id = PartID(other_vol_dd.domain_tag.tag, rank) + + self_trace_dd = self_vol_dd.trace(BTAG_PARTITION(other_part_id)) + other_trace_dd = other_vol_dd.trace(BTAG_PARTITION(self_part_id)) + + self_trace_data = trace_data[self_trace_dd] + unswapped_other_trace_data = trace_data[other_trace_dd] + + other_to_self = dcoll._inter_part_connections[ + other_part_id, self_part_id] + + def get_opposite_trace(ary): + if isinstance(ary, Number): + return ary + else: + return other_to_self(ary) # noqa: B023 + + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + other_trace_data = rec_map_array_container( + get_opposite_trace, + unswapped_other_trace_data, + leaf_class=DOFArray) + + result[other_vol_dd, self_vol_dd] = TracePair( + self_trace_dd, + interior=self_trace_data, + exterior=other_trace_data) + + return result + + +def inter_volume_trace_pairs(dcoll: DiscretizationCollection, + pairwise_volume_data: Mapping[ + Tuple[DOFDesc, DOFDesc], + Tuple[ArrayOrContainer, ArrayOrContainer]], + comm_tag: Hashable = None) -> Mapping[ + Tuple[DOFDesc, DOFDesc], + List[TracePair]]: + """ + 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 + + result: Mapping[ + Tuple[DOFDesc, DOFDesc], + List[TracePair]] = {} + + local_tpairs = local_inter_volume_trace_pairs(dcoll, pairwise_volume_data) + cross_rank_tpairs = cross_rank_inter_volume_trace_pairs( + dcoll, pairwise_volume_data, comm_tag=comm_tag) + + for directional_vol_dd_pair, tpair in local_tpairs.items(): + result[directional_vol_dd_pair] = [tpair] + + for directional_vol_dd_pair, tpairs in cross_rank_tpairs.items(): + result.setdefault(directional_vol_dd_pair, []).extend(tpairs) + + return result + +# }}} + + # {{{ distributed: helper functions class _TagKeyBuilder(KeyBuilder): @@ -367,16 +489,21 @@ def update_for_type(self, key_hash, key: Type[Any]): self.rec(key_hash, (key.__module__, key.__name__, key.__name__,)) +# FIXME: Deprecate connected_ranks instead of removing @memoize_on_first_arg -def connected_ranks( +def connected_parts( dcoll: DiscretizationCollection, - volume_dd: Optional[DOFDesc] = None): - if volume_dd is None: - volume_dd = DD_VOLUME_ALL + self_volume_tag: VolumeTag, + other_volume_tag: VolumeTag + ) -> Sequence[PartID]: + result: List[PartID] = [ + connected_part_id + for connected_part_id, part_id in dcoll._inter_part_connections.keys() + if ( + part_id.volume_tag == self_volume_tag + and connected_part_id.volume_tag == other_volume_tag)] - from meshmode.distributed import get_connected_partitions - return get_connected_partitions( - dcoll._volume_discrs[volume_dd.domain_tag.tag].mesh) + return result def _sym_tag_to_num_tag(comm_tag: Optional[Hashable]) -> Optional[int]: @@ -414,24 +541,33 @@ class _RankBoundaryCommunicationEager: base_comm_tag = 1273 def __init__(self, - dcoll: DiscretizationCollection, - array_container: ArrayOrContainer, - remote_rank, comm_tag: Optional[int] = None, - volume_dd=DD_VOLUME_ALL): - actx = get_container_context_recursively(array_container) - bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_rank)) - - local_bdry_data = project(dcoll, volume_dd, bdry_dd, array_container) + actx: ArrayContext, + dcoll: DiscretizationCollection, + *, + local_part_id: PartID, + remote_part_id: PartID, + local_bdry_data: ArrayOrContainer, + remote_bdry_data_template: ArrayOrContainer, + comm_tag: Optional[Hashable] = None): + comm = dcoll.mpi_communicator assert comm is not None + remote_rank = remote_part_id.rank + assert remote_rank is not None + self.dcoll = dcoll self.array_context = actx - self.remote_bdry_dd = bdry_dd - self.bdry_discr = dcoll.discr_from_dd(bdry_dd) + self.local_part_id = local_part_id + self.remote_part_id = remote_part_id + self.local_bdry_dd = DOFDesc( + BoundaryDomainTag( + BTAG_PARTITION(remote_part_id), + volume_tag=local_part_id.volume_tag), + DISCR_TAG_BASE) + self.bdry_discr = dcoll.discr_from_dd(self.local_bdry_dd) self.local_bdry_data = local_bdry_data - self.local_bdry_data_np = \ - to_numpy(flatten(self.local_bdry_data, actx), actx) + self.remote_bdry_data_template = remote_bdry_data_template self.comm_tag = self.base_comm_tag comm_tag = _sym_tag_to_num_tag(comm_tag) @@ -439,55 +575,80 @@ def __init__(self, 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, - remote_rank, - tag=self.comm_tag) + self.send_reqs = [] + self.send_data = [] + + def send_single_array(key, local_subary): + if not isinstance(local_subary, Number): + local_subary_np = to_numpy(local_subary, actx) + self.send_reqs.append( + comm.Isend(local_subary_np, remote_rank, tag=self.comm_tag)) + self.send_data.append(local_subary_np) + return local_subary + + self.recv_reqs = [] + self.recv_data = {} + + def recv_single_array(key, remote_subary_template): + if not isinstance(remote_subary_template, Number): + remote_subary_np = np.empty( + remote_subary_template.shape, + remote_subary_template.dtype) + self.recv_reqs.append( + comm.Irecv(remote_subary_np, remote_rank, tag=self.comm_tag)) + self.recv_data[key] = remote_subary_np + return remote_subary_template + + from arraycontext.container.traversal import rec_keyed_map_array_container + rec_keyed_map_array_container(send_single_array, local_bdry_data) + rec_keyed_map_array_container(recv_single_array, remote_bdry_data_template) def finish(self): - # Wait for the nonblocking receive request to complete before + from mpi4py import MPI + + # Wait for the nonblocking receive requests 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( - self.remote_bdry_dd) - swapped_remote_bdry_data = bdry_conn(remote_bdry_data) - - # Complete the nonblocking send request associated with communicating - # `self.local_bdry_data_np` - self.send_req.Wait() - - return TracePair(self.remote_bdry_dd, - interior=self.local_bdry_data, - exterior=swapped_remote_bdry_data) + MPI.Request.waitall(self.recv_reqs) + + def finish_single_array(key, remote_subary_template): + if isinstance(remote_subary_template, Number): + # NOTE: Assumes that the same number is passed on every rank + return remote_subary_template + else: + return from_numpy(self.recv_data[key], self.array_context) + + from arraycontext.container.traversal import rec_keyed_map_array_container + unswapped_remote_bdry_data = rec_keyed_map_array_container( + finish_single_array, self.remote_bdry_data_template) + + remote_to_local = self.dcoll._inter_part_connections[ + self.remote_part_id, self.local_part_id] + + def get_opposite_trace(ary): + if isinstance(ary, Number): + return ary + else: + return remote_to_local(ary) + + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + remote_bdry_data = rec_map_array_container( + get_opposite_trace, + unswapped_remote_bdry_data, + leaf_class=DOFArray) + + # Complete the nonblocking send requests + MPI.Request.waitall(self.send_reqs) + + return TracePair( + self.local_bdry_dd, + interior=self.local_bdry_data, + exterior=remote_bdry_data) # }}} @@ -496,51 +657,112 @@ def finish(self): class _RankBoundaryCommunicationLazy: def __init__(self, - dcoll: DiscretizationCollection, - array_container: ArrayOrContainer, - remote_rank: int, comm_tag: Hashable, - volume_dd=DD_VOLUME_ALL): + actx: ArrayContext, + dcoll: DiscretizationCollection, + *, + local_part_id: PartID, + remote_part_id: PartID, + local_bdry_data: ArrayOrContainer, + remote_bdry_data_template: 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") - bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_rank)) + remote_rank = remote_part_id.rank + assert remote_rank is not None self.dcoll = dcoll - self.array_context = get_container_context_recursively(array_container) - self.remote_bdry_dd = bdry_dd - self.bdry_discr = dcoll.discr_from_dd(self.remote_bdry_dd) - - self.local_bdry_data = project( - dcoll, volume_dd, bdry_dd, array_container) - - from pytato import make_distributed_recv, staple_distributed_send - - def communicate_single_array(key, local_bdry_ary): - ary_tag = (comm_tag, key) - return staple_distributed_send( - local_bdry_ary, dest_rank=remote_rank, comm_tag=ary_tag, - stapled_to=make_distributed_recv( + self.array_context = actx + self.local_bdry_dd = DOFDesc( + BoundaryDomainTag( + BTAG_PARTITION(remote_part_id), + volume_tag=local_part_id.volume_tag), + DISCR_TAG_BASE) + self.bdry_discr = dcoll.discr_from_dd(self.local_bdry_dd) + self.local_part_id = local_part_id + self.remote_part_id = remote_part_id + + from pytato import ( + make_distributed_recv, + make_distributed_send, + DistributedSendRefHolder) + + # TODO: This currently assumes that local_bdry_data and + # remote_bdry_data_template have the same structure. This is not true + # in general. Find a way to staple the sends appropriately when the number + # of recvs is not equal to the number of sends + # FIXME: Overly restrictive (just needs to be the same structure) + assert type(local_bdry_data) == type(remote_bdry_data_template) + + sends = {} + + def send_single_array(key, local_subary): + if isinstance(local_subary, Number): + return + else: + ary_tag = (comm_tag, key) + sends[key] = make_distributed_send( + local_subary, dest_rank=remote_rank, comm_tag=ary_tag) + + def recv_single_array(key, remote_subary_template): + if isinstance(remote_subary_template, Number): + # NOTE: Assumes that the same number is passed on every rank + return remote_subary_template + else: + ary_tag = (comm_tag, key) + return DistributedSendRefHolder( + sends[key], + 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=remote_subary_template.shape, + dtype=remote_subary_template.dtype, + axes=remote_subary_template.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( - self.remote_bdry_dd) + rec_keyed_map_array_container(send_single_array, local_bdry_data) + self.local_bdry_data = local_bdry_data - return TracePair(self.remote_bdry_dd, - interior=self.local_bdry_data, - exterior=bdry_conn(self.remote_data)) + self.unswapped_remote_bdry_data = rec_keyed_map_array_container( + recv_single_array, remote_bdry_data_template) + + def finish(self): + remote_to_local = self.dcoll._inter_part_connections[ + self.remote_part_id, self.local_part_id] + + def get_opposite_trace(ary): + if isinstance(ary, Number): + return ary + else: + return remote_to_local(ary) + + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + remote_bdry_data = rec_map_array_container( + get_opposite_trace, + self.unswapped_remote_bdry_data, + leaf_class=DOFArray) + + return TracePair( + self.local_bdry_dd, + interior=self.local_bdry_data, + exterior=remote_bdry_data) # }}} # {{{ cross_rank_trace_pairs +def _replace_dof_arrays(array_container, dof_array): + from arraycontext import rec_map_array_container + from meshmode.dof_array import DOFArray + return rec_map_array_container( + lambda x: dof_array if isinstance(x, DOFArray) else x, + array_container, + leaf_class=DOFArray) + + def cross_rank_trace_pairs( dcoll: DiscretizationCollection, ary: ArrayOrContainer, tag: Hashable = None, @@ -549,9 +771,9 @@ def cross_rank_trace_pairs( r"""Get a :class:`list` of *ary* trace pairs for each partition boundary. For each partition boundary, the field data values in *ary* are - communicated to/from the neighboring partition. Presumably, this - communication is MPI (but strictly speaking, may not be, and this - routine is agnostic to the underlying communication). + communicated to/from the neighboring part. Presumably, this communication + is MPI (but strictly speaking, may not be, and this routine is agnostic to + the underlying communication). For each face on each partition boundary, a :class:`TracePair` is created with the locally, and @@ -596,14 +818,36 @@ def cross_rank_trace_pairs( # }}} - if isinstance(ary, Number): - # NOTE: Assumed that the same number is passed on every rank - return [TracePair( - volume_dd.trace(BTAG_PARTITION(remote_rank)), - interior=ary, exterior=ary) - for remote_rank in connected_ranks(dcoll, volume_dd=volume_dd)] + if dcoll.mpi_communicator is None: + return [] + + rank = dcoll.mpi_communicator.Get_rank() + + local_part_id = PartID(volume_dd.domain_tag.tag, rank) + + connected_part_ids = connected_parts( + dcoll, self_volume_tag=volume_dd.domain_tag.tag, + other_volume_tag=volume_dd.domain_tag.tag) + + remote_part_ids = [ + part_id + for part_id in connected_part_ids + if part_id.rank != rank] + + # 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(remote_part_ids) == len({part_id.rank for part_id in remote_part_ids}) - actx = get_container_context_recursively(ary) + actx = get_container_context_recursively_opt(ary) + + if actx is None: + # NOTE: Assumes that the same number is passed on every rank + return [ + TracePair( + volume_dd.trace(BTAG_PARTITION(remote_part_id)), + interior=ary, exterior=ary) + for remote_part_id in remote_part_ids] from grudge.array_context import MPIPytatoArrayContextBase @@ -612,14 +856,167 @@ def cross_rank_trace_pairs( else: rbc_class = _RankBoundaryCommunicationEager - # Initialize and post all sends/receives - rank_bdry_communcators = [ - rbc_class(dcoll, ary, remote_rank, comm_tag=comm_tag, volume_dd=volume_dd) - for remote_rank in connected_ranks(dcoll, volume_dd=volume_dd) - ] + rank_bdry_communicators = [] + + for remote_part_id in remote_part_ids: + bdry_dd = volume_dd.trace(BTAG_PARTITION(remote_part_id)) + + local_bdry_data = project(dcoll, volume_dd, bdry_dd, ary) + + from arraycontext import tag_axes + from meshmode.transform_metadata import ( + DiscretizationElementAxisTag, + DiscretizationDOFAxisTag) + remote_bdry_zeros = tag_axes( + actx, { + 0: DiscretizationElementAxisTag(), + 1: DiscretizationDOFAxisTag()}, + dcoll._inter_part_connections[ + remote_part_id, local_part_id].from_discr.zeros(actx)) + + remote_bdry_data_template = _replace_dof_arrays( + local_bdry_data, remote_bdry_zeros) + + rank_bdry_communicators.append( + rbc_class(actx, dcoll, + local_part_id=local_part_id, + remote_part_id=remote_part_id, + local_bdry_data=local_bdry_data, + remote_bdry_data_template=remote_bdry_data_template, + comm_tag=comm_tag)) + + return [rbc.finish() for rbc in rank_bdry_communicators] + +# }}} + + +# {{{ cross_rank_inter_volume_trace_pairs + +def cross_rank_inter_volume_trace_pairs( + dcoll: DiscretizationCollection, + pairwise_volume_data: Mapping[ + Tuple[DOFDesc, DOFDesc], + Tuple[ArrayOrContainer, ArrayOrContainer]], + *, comm_tag: Hashable = None, + ) -> Mapping[ + Tuple[DOFDesc, DOFDesc], + 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 + + for vol_dd_pair in pairwise_volume_data.keys(): + for vol_dd in vol_dd_pair: + if not isinstance(vol_dd.domain_tag, VolumeDomainTag): + raise ValueError( + "pairwise_volume_data keys must describe volumes, " + f"got '{vol_dd}'") + if vol_dd.discretization_tag != DISCR_TAG_BASE: + raise ValueError( + "expected base-discretized DOFDesc in pairwise_volume_data, " + f"got '{vol_dd}'") + + # }}} + + if dcoll.mpi_communicator is None: + return {} + + rank = dcoll.mpi_communicator.Get_rank() + + for vol_data_pair in pairwise_volume_data.values(): + for vol_data in vol_data_pair: + actx = get_container_context_recursively_opt(vol_data) + if actx is not None: + break + if actx is not None: + break + + def get_remote_connected_parts(local_vol_dd, remote_vol_dd): + connected_part_ids = connected_parts( + dcoll, self_volume_tag=local_vol_dd.domain_tag.tag, + other_volume_tag=remote_vol_dd.domain_tag.tag) + return [ + part_id + for part_id in connected_part_ids + if part_id.rank != rank] + + if actx is None: + # NOTE: Assumes that the same number is passed on every rank for a + # given volume + return { + (remote_vol_dd, local_vol_dd): [ + TracePair( + local_vol_dd.trace(BTAG_PARTITION(remote_part_id)), + interior=local_vol_ary, exterior=remote_vol_ary) + for remote_part_id in get_remote_connected_parts( + local_vol_dd, remote_vol_dd)] + for (remote_vol_dd, local_vol_dd), (remote_vol_ary, local_vol_ary) + in pairwise_volume_data.items()} + + from grudge.array_context import MPIPytatoArrayContextBase + + if isinstance(actx, MPIPytatoArrayContextBase): + rbc_class = _RankBoundaryCommunicationLazy + else: + rbc_class = _RankBoundaryCommunicationEager - # Complete send/receives and return communicated data - return [rc.finish() for rc in rank_bdry_communcators] + rank_bdry_communicators = {} + + for vol_dd_pair, vol_data_pair in pairwise_volume_data.items(): + directional_volume_data = { + (vol_dd_pair[0], vol_dd_pair[1]): (vol_data_pair[0], vol_data_pair[1]), + (vol_dd_pair[1], vol_dd_pair[0]): (vol_data_pair[1], vol_data_pair[0])} + + for dd_pair, data_pair in directional_volume_data.items(): + other_vol_dd, self_vol_dd = dd_pair + other_vol_data, self_vol_data = data_pair + + self_part_id = PartID(self_vol_dd.domain_tag.tag, rank) + other_part_ids = get_remote_connected_parts(self_vol_dd, other_vol_dd) + + rbcs = [] + + for other_part_id in other_part_ids: + self_bdry_dd = self_vol_dd.trace(BTAG_PARTITION(other_part_id)) + self_bdry_data = project( + dcoll, self_vol_dd, self_bdry_dd, self_vol_data) + + from arraycontext import tag_axes + from meshmode.transform_metadata import ( + DiscretizationElementAxisTag, + DiscretizationDOFAxisTag) + other_bdry_zeros = tag_axes( + actx, { + 0: DiscretizationElementAxisTag(), + 1: DiscretizationDOFAxisTag()}, + dcoll._inter_part_connections[ + other_part_id, self_part_id].from_discr.zeros(actx)) + + other_bdry_data_template = _replace_dof_arrays( + other_vol_data, other_bdry_zeros) + + rbcs.append( + rbc_class(actx, dcoll, + local_part_id=self_part_id, + remote_part_id=other_part_id, + local_bdry_data=self_bdry_data, + remote_bdry_data_template=other_bdry_data_template, + comm_tag=comm_tag)) + + rank_bdry_communicators[other_vol_dd, self_vol_dd] = rbcs + + return { + directional_vol_dd_pair: [rbc.finish() for rbc in rbcs] + for directional_vol_dd_pair, rbcs in rank_bdry_communicators.items()} # }}} From d16767acb711c967b54d73b25e4c9e9b31469ae4 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 3 Nov 2022 09:20:20 -0700 Subject: [PATCH 18/97] add fixme --- grudge/trace_pair.py | 1 + 1 file changed, 1 insertion(+) diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 0b0400f12..84dedf386 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -525,6 +525,7 @@ def _sym_tag_to_num_tag(comm_tag: Optional[Hashable]) -> Optional[int]: num_tag = sum(ord(ch) << i for i, ch in enumerate(digest)) % tag_ub + # FIXME: This prints the wrong numerical tag because of base_comm_tag below warn("Encountered unknown symbolic tag " f"'{comm_tag}', assigning a value of '{num_tag}'. " "This is a temporary workaround, please ensure that " From a3810cec76f01b3ac0377855650bdba6573b73ae Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 3 Nov 2022 10:07:09 -0700 Subject: [PATCH 19/97] check for heterogeneous inter-volume data --- grudge/trace_pair.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 84dedf386..7358e5af8 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -915,7 +915,7 @@ def cross_rank_inter_volume_trace_pairs( """ # {{{ process arguments - for vol_dd_pair in pairwise_volume_data.keys(): + for vol_dd_pair, vol_data_pair in pairwise_volume_data.items(): for vol_dd in vol_dd_pair: if not isinstance(vol_dd.domain_tag, VolumeDomainTag): raise ValueError( @@ -925,6 +925,9 @@ def cross_rank_inter_volume_trace_pairs( raise ValueError( "expected base-discretized DOFDesc in pairwise_volume_data, " f"got '{vol_dd}'") + # FIXME: This check could probably be made more robust + if type(vol_data_pair[0]) != type(vol_data_pair[1]): # noqa: E721 + raise ValueError("heterogeneous inter-volume data not supported.") # }}} From bfad1f77dc2aaf0b8340190e4d39b21a3ad7b297 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 3 Nov 2022 10:07:26 -0700 Subject: [PATCH 20/97] tag communication by destination volume --- grudge/trace_pair.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 7358e5af8..acc086505 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -570,10 +570,17 @@ def __init__(self, self.local_bdry_data = local_bdry_data self.remote_bdry_data_template = remote_bdry_data_template - 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 + def _generate_num_comm_tag(sym_comm_tag): + result = self.base_comm_tag + num_comm_tag = _sym_tag_to_num_tag(sym_comm_tag) + if num_comm_tag is not None: + result += num_comm_tag + return result + + send_sym_comm_tag = (remote_part_id.volume_tag, comm_tag) + recv_sym_comm_tag = (local_part_id.volume_tag, comm_tag) + self.send_comm_tag = _generate_num_comm_tag(send_sym_comm_tag) + self.recv_comm_tag = _generate_num_comm_tag(recv_sym_comm_tag) del comm_tag # NOTE: mpi4py currently (2021-11-03) holds a reference to the send @@ -588,7 +595,7 @@ def send_single_array(key, local_subary): if not isinstance(local_subary, Number): local_subary_np = to_numpy(local_subary, actx) self.send_reqs.append( - comm.Isend(local_subary_np, remote_rank, tag=self.comm_tag)) + comm.Isend(local_subary_np, remote_rank, tag=self.send_comm_tag)) self.send_data.append(local_subary_np) return local_subary @@ -601,7 +608,8 @@ def recv_single_array(key, remote_subary_template): remote_subary_template.shape, remote_subary_template.dtype) self.recv_reqs.append( - comm.Irecv(remote_subary_np, remote_rank, tag=self.comm_tag)) + comm.Irecv(remote_subary_np, remote_rank, + tag=self.recv_comm_tag)) self.recv_data[key] = remote_subary_np return remote_subary_template @@ -702,7 +710,7 @@ def send_single_array(key, local_subary): if isinstance(local_subary, Number): return else: - ary_tag = (comm_tag, key) + ary_tag = (remote_part_id.volume_tag, comm_tag, key) sends[key] = make_distributed_send( local_subary, dest_rank=remote_rank, comm_tag=ary_tag) @@ -711,7 +719,7 @@ def recv_single_array(key, remote_subary_template): # NOTE: Assumes that the same number is passed on every rank return remote_subary_template else: - ary_tag = (comm_tag, key) + ary_tag = (local_part_id.volume_tag, comm_tag, key) return DistributedSendRefHolder( sends[key], make_distributed_recv( From c112288a0cc2058060a6cf72bbb5e2efc987356c Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 3 Apr 2023 15:28:38 -0500 Subject: [PATCH 21/97] add filter_part_boundaries eases setting up boundaries when calling operators on only one volume (i.e., uncoupled) --- grudge/discretization.py | 44 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 43 insertions(+), 1 deletion(-) diff --git a/grudge/discretization.py b/grudge/discretization.py index fd4c39728..25d5fa797 100644 --- a/grudge/discretization.py +++ b/grudge/discretization.py @@ -8,6 +8,7 @@ .. currentmodule:: grudge.discretization .. autoclass:: PartID +.. autofunction:: filter_part_boundaries """ __copyright__ = """ @@ -35,7 +36,8 @@ THE SOFTWARE. """ -from typing import Sequence, Mapping, Optional, Union, Tuple, TYPE_CHECKING, Any +from typing import ( + Sequence, Mapping, Optional, Union, List, Tuple, TYPE_CHECKING, Any) from pytools import memoize_method, single_valued @@ -1015,4 +1017,44 @@ def make_discretization_collection( # }}} +# {{{ filter_part_boundaries + +def filter_part_boundaries( + dcoll: DiscretizationCollection, + *, + volume_dd: DOFDesc = DD_VOLUME_ALL, + neighbor_volume_dd: Optional[DOFDesc] = None, + neighbor_rank: Optional[int] = None) -> List[DOFDesc]: + """ + Retrieve tags of part boundaries that match *neighbor_volume_dd* and/or + *neighbor_rank*. + """ + vol_mesh = dcoll.discr_from_dd(volume_dd).mesh + + from meshmode.mesh import InterPartAdjacencyGroup + filtered_part_bdry_dds = [ + volume_dd.trace(fagrp.boundary_tag) + for fagrp_list in vol_mesh.facial_adjacency_groups + for fagrp in fagrp_list + if isinstance(fagrp, InterPartAdjacencyGroup)] + + if neighbor_volume_dd is not None: + filtered_part_bdry_dds = [ + bdry_dd + for bdry_dd in filtered_part_bdry_dds + if ( + bdry_dd.domain_tag.tag.part_id.volume_tag + == neighbor_volume_dd.domain_tag.tag)] + + if neighbor_rank is not None: + filtered_part_bdry_dds = [ + bdry_dd + for bdry_dd in filtered_part_bdry_dds + if bdry_dd.domain_tag.tag.part_id.rank == neighbor_rank] + + return filtered_part_bdry_dds + +# }}} + + # vim: foldmethod=marker From bfe2d4b5db703a4ad9e2447445bbaa8c02caec4b Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Tue, 27 Jun 2023 09:37:28 -0500 Subject: [PATCH 22/97] Add Numpy actx --- grudge/array_context.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/grudge/array_context.py b/grudge/array_context.py index d744d5fb4..bf57eb63c 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -3,6 +3,7 @@ .. autoclass:: PytatoPyOpenCLArrayContext .. autoclass:: MPIBasedArrayContext .. autoclass:: MPIPyOpenCLArrayContext +.. autoclass:: MPINumpyArrayContext .. class:: MPIPytatoArrayContext .. autofunction:: get_reasonable_array_context_class """ @@ -607,4 +608,27 @@ def get_reasonable_array_context_class( # }}} +# {{{ distributed + numpy +try: + from arraycontext import NumpyArrayContext + + class MPINumpyArrayContext(NumpyArrayContext, MPIBasedArrayContext): + """An array context for using distributed computation with :mod:`numpy` + eager evaluation. + .. autofunction:: __init__ + """ + + def __init__(self, mpi_communicator) -> None: + super().__init__() + self.mpi_communicator = mpi_communicator + + def clone(self): + return type(self)(self.mpi_communicator) + +except ImportError: + print("Failed to import numpy array context.") + pass +# }}} + + # vim: foldmethod=marker From b1be58a5a463aa03d4b87a7034899c5fdbb3b94c Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Mon, 31 Jul 2023 20:59:55 -0500 Subject: [PATCH 23/97] Update _reference_derivative_matrices to recognize TensorProductElementGroup --- grudge/op.py | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index f5781f4be..09e17884b 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -259,16 +259,35 @@ def _reference_derivative_matrices(actx: ArrayContext, # _reference_stiffness_transpose_matrices. assert out_element_group is in_element_group + from meshmode.mesh import TensorProductElementGroup + @keyed_memoize_in( actx, _reference_derivative_matrices, lambda grp: grp.discretization_key()) def get_ref_derivative_mats(grp): - from meshmode.discretization.poly_element import diff_matrices - return actx.freeze( - actx.tag_axis( - 1, DiscretizationDOFAxisTag(), - actx.from_numpy( - np.asarray(diff_matrices(grp))))) + + if isinstance(grp, TensorProductElementGroup): + import modepy as mp + import numpy.linalg as la + + space1d = grp.space.bases[0] + shape1d = grp.shape.bases[0] + + nodes1d = mp.edge_clustered_nodes_for_space(space1d, shape1d) + basis1d = mp.basis_for_space(space1d, shape1d) + + vdm1d = mp.vandermonde(basis1d.functions, nodes1d) + vdm_p1d = mp.vandermonde(basis1d.gradients, nodes1d)[0] + + return actx.freeze(actx.from_numpy(vdm_p1d @ la.inv(vdm1d))) + + else: + from meshmode.discretization.poly_element import diff_matrices + return actx.freeze( + actx.tag_axis( + 1, DiscretizationDOFAxisTag(), + actx.from_numpy( + np.asarray(diff_matrices(grp))))) return get_ref_derivative_mats(out_element_group) From 8fa3321ff01bea49ba4f81b159a580d7e77e7956 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Mon, 31 Jul 2023 21:10:08 -0500 Subject: [PATCH 24/97] Stub in tensor product gradient computation in _gradient_kernel --- grudge/op.py | 34 ++++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 09e17884b..cfd40eed5 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -202,19 +202,33 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, *, metric_in_matvec): # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. + from meshmode.mesh import TensorProductElementGroup + + def compute_tensor_product_grad(actx, diff_mat, vec): + """Exploits tensor product structure to differentiate each coordinate + axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) + """ + pass + per_group_grads = [ + + compute_tensor_product_grad(actx, get_diff_mat, vec_i) + if isinstance(in_grp, TensorProductElementGroup) + # r for rst axis # x for xyz axis - actx.einsum("xrej,rij,ej->xei" if metric_in_matvec else "xrei,rij,ej->xei", - ijm_i, - get_diff_mat( - actx, - out_element_group=out_grp, - in_element_group=in_grp - ), - vec_i, - arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), - tagged=(FirstAxisIsElementsTag(),)) + else actx.einsum( + "xrej,rij,ej->xei" if metric_in_matvec else "xrei,rij,ej->xei", + ijm_i, + get_diff_mat( + actx, + out_element_group=out_grp, + in_element_group=in_grp + ), + vec_i, + arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), + tagged=(FirstAxisIsElementsTag(),)) + for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, inv_jac_mat)] From 3778edfcfc872e058abf8f8655367b510e425632 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Mon, 31 Jul 2023 21:59:06 -0500 Subject: [PATCH 25/97] First version of grad routine --- grudge/op.py | 61 +++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 58 insertions(+), 3 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index cfd40eed5..3a3cd9919 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -204,15 +204,70 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # (both strong and weak derivative) and their differences. from meshmode.mesh import TensorProductElementGroup - def compute_tensor_product_grad(actx, diff_mat, vec): + def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): """Exploits tensor product structure to differentiate each coordinate axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ - pass + + from modepy.tools import ( + reshape_array_for_tensor_product_space, + unreshape_array_for_tensor_product_space) + + # reshape u to expose tensor product structure + vec = make_obj_array([ + reshape_array_for_tensor_product_space(grp.space, vec[i]) + for i in range(vec.shape[0]) + ]) + + # apply differentiation matrix to vec + if vec.shape[0] == 2: + specs = ["il,elj->eij", + "jl,eil->eij"] + elif vec.shape[1] == 3: + specs = ["il,eljk->eijk", + "jl,eilk->eijk", + "kl,eijl->eijk"] + else: + specs = None + assert specs is not None + + grad = make_obj_array([ + make_obj_array([ + actx.einsum( + spec, + diff_mat, + vec[i], + arg_names=("diff_mat", "vec"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + for i in range(vec.shape[0]) + ]) + for spec in specs + ]) + + # unreshape grad to apply geometric factors + # NOTE: In a future version, do not reshape before application of + # geometric factors. Can possibly "chain" the einsum as it is below + grad = make_obj_array([ + unreshape_array_for_tensor_product_space(grp.space, grad[i][0]) + for i in range(grad.shape[0]) + ]) + + # apply geometric factors to current grad + grad = make_obj_array([ + actx.einsum( + "rei,ei->ei", + ijm[i], + grad[i], + tagged=(FirstAxisIsElementsTag(),)) + for i in range(grad.shape[0]) + ]) + + return grad per_group_grads = [ - compute_tensor_product_grad(actx, get_diff_mat, vec_i) + compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i) if isinstance(in_grp, TensorProductElementGroup) # r for rst axis From 45e859e5a788a68dbba2fbadaadd53edc231b523 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Tue, 1 Aug 2023 12:39:48 -0500 Subject: [PATCH 26/97] Initial working version of tensor product gradient operator application --- grudge/op.py | 73 +++++++++++++++++++++++++++++++++++++------------ test/test_op.py | 68 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 17 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 3a3cd9919..c6e6ce36e 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -166,6 +166,36 @@ ) +# {{{ Temporary tools for tensor product operators +from pytools.tag import Tag +class OutputIsTensorProductDOFArrayOrdered(Tag): + pass + + +from grudge.array_context import PyOpenCLArrayContext +class TensorProductArrayContext(PyOpenCLArrayContext): + def transform_loopy_program(self, t_unit): + if len(t_unit.callables_table) == 1: + knl = t_unit.default_entrypoint + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) + + return super().transform_loopy_program(t_unit) +# }}} + + # {{{ common derivative "kernels" def _single_axis_derivative_kernel( @@ -202,28 +232,30 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, *, metric_in_matvec): # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. - from meshmode.mesh import TensorProductElementGroup def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): """Exploits tensor product structure to differentiate each coordinate axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ + actx_tp = TensorProductArrayContext( + actx.queue, + allocator=actx.allocator, + force_device_scalars=actx._force_device_scalars) + from modepy.tools import ( reshape_array_for_tensor_product_space, unreshape_array_for_tensor_product_space) # reshape u to expose tensor product structure - vec = make_obj_array([ - reshape_array_for_tensor_product_space(grp.space, vec[i]) - for i in range(vec.shape[0]) - ]) + vec = reshape_array_for_tensor_product_space(grp.space, vec) # apply differentiation matrix to vec - if vec.shape[0] == 2: + # check len(vec.shape) since shape is expected to be (nelements, ndofs) + if len(vec.shape) == 3: specs = ["il,elj->eij", "jl,eil->eij"] - elif vec.shape[1] == 3: + elif len(vec.shape) == 4: specs = ["il,eljk->eijk", "jl,eilk->eijk", "kl,eijl->eijk"] @@ -231,31 +263,34 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): specs = None assert specs is not None + diff_mat = get_diff_mat(actx, grp, grp) grad = make_obj_array([ - make_obj_array([ - actx.einsum( + actx_tp.einsum( spec, diff_mat, - vec[i], + vec, arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered())) - for i in range(vec.shape[0]) - ]) for spec in specs ]) # unreshape grad to apply geometric factors # NOTE: In a future version, do not reshape before application of - # geometric factors. Can possibly "chain" the einsum as it is below + # geometric factors. Can possibly "chain" the einsum. For example, the + # simplicial case below has einsum with spec + # ("xrei,rij,ei->ei") + # for the strong local gradient case grad = make_obj_array([ - unreshape_array_for_tensor_product_space(grp.space, grad[i][0]) + unreshape_array_for_tensor_product_space(grp.space, grad[i]) for i in range(grad.shape[0]) ]) # apply geometric factors to current grad + # FIXME: using einsum spec ("xrei,xei->xei") throws error: + # "Loopy does not directly support object arrays" grad = make_obj_array([ - actx.einsum( + actx_tp.einsum( "rei,ei->ei", ijm[i], grad[i], @@ -265,10 +300,12 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): return grad + from meshmode.discretization.poly_element import \ + TensorProductElementGroupBase per_group_grads = [ compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i) - if isinstance(in_grp, TensorProductElementGroup) + if isinstance(in_grp, TensorProductElementGroupBase) # r for rst axis # x for xyz axis @@ -335,7 +372,9 @@ def _reference_derivative_matrices(actx: ArrayContext, lambda grp: grp.discretization_key()) def get_ref_derivative_mats(grp): - if isinstance(grp, TensorProductElementGroup): + from meshmode.discretization.poly_element import \ + TensorProductElementGroupBase + if isinstance(grp, TensorProductElementGroupBase): import modepy as mp import numpy.linalg as la diff --git a/test/test_op.py b/test/test_op.py index fa7ee0bbd..8b20b2bb8 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -32,6 +32,7 @@ import pytest +from grudge.discretization import make_discretization_collection from grudge.array_context import PytestPyOpenCLArrayContextFactory from arraycontext import pytest_generate_tests_for_array_contexts pytest_generate_tests = pytest_generate_tests_for_array_contexts( @@ -159,6 +160,73 @@ def get_flux(u_tpair): assert (eoc_rec.order_estimate() >= order - 0.5 or eoc_rec.max_error() < 1e-11) + +@pytest.mark.parametrize("form", ["strong"]) +@pytest.mark.parametrize("dim", [2]) +@pytest.mark.parametrize("order", [2]) +@pytest.mark.parametrize(("vectorize", "nested"), [ + (False, False) + ]) +def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize, + nested, visualize=False): + + actx = actx_factory() + from pytools.convergence import EOCRecorder + eoc_rec = EOCRecorder() + + from meshmode.mesh import TensorProductElementGroup + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + for n in [4, 6, 8]: + mesh = mgen.generate_regular_rect_mesh( + a=(-1,)*dim, + b=(1,)*dim, + nelements_per_axis=(n,)*dim, + group_cls=TensorProductElementGroup) + + import grudge.dof_desc as dd + dcoll = make_discretization_collection( + actx, + mesh, + discr_tag_to_group_factory={ + dd.DISCR_TAG_BASE: LGL(order)}) + + + def f(x): + ret = actx.np.cos(np.pi*x[0]) + actx.np.sin(np.pi*x[1]) + + if dim == 3: + ret = ret + actx.np.sin(np.pi*x[2]) + + return ret + + + def grad_f(x): + ret = make_obj_array([dcoll.zeros(actx) for _ in range(dim)]) + + ret[0] = -np.pi*actx.np.sin(np.pi*x[0]) + ret[1] = np.pi*actx.np.cos(np.pi*x[1]) + + if dim == 3: + ret[2] = np.pi*actx.np.cos(np.pi*x[2]) + + return ret + + + x = actx.thaw(dcoll.nodes()) + u = f(x) + ref_grad = grad_f(x) + grad = op.local_grad(dcoll, u) + + rel_linf_error = actx.to_numpy(op.norm(dcoll, ref_grad - grad, np.inf) / + op.norm(dcoll, ref_grad, np.inf)) + eoc_rec.add_data_point(1./n, rel_linf_error) + + print("L^inf error:") + print(eoc_rec) + assert (eoc_rec.order_estimate() >= order - 0.5 or + eoc_rec.max_error() < 1e-11) + # }}} From 22ffacdf969f15aa4f3018d1f287c02e1455adf5 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Tue, 1 Aug 2023 12:48:08 -0500 Subject: [PATCH 27/97] Add 3 dimensional test and order 3 test for 2D and 3D --- test/test_op.py | 30 +++++++++++++++++++----------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/test/test_op.py b/test/test_op.py index 8b20b2bb8..1e45a4556 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -162,14 +162,16 @@ def get_flux(u_tpair): @pytest.mark.parametrize("form", ["strong"]) -@pytest.mark.parametrize("dim", [2]) -@pytest.mark.parametrize("order", [2]) +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize(("vectorize", "nested"), [ (False, False) ]) def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize, nested, visualize=False): - + """A "one-dimensional tensor product element" does not make sense, so the + one-dimensional case is excluded from this test. + """ actx = actx_factory() from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() @@ -193,10 +195,14 @@ def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize, def f(x): - ret = actx.np.cos(np.pi*x[0]) + actx.np.sin(np.pi*x[1]) - - if dim == 3: - ret = ret + actx.np.sin(np.pi*x[2]) + if dim == 2: + ret = actx.np.cos(np.pi*x[0]) + actx.np.sin(np.pi*x[1]) + elif dim == 3: + ret = actx.np.cos(np.pi*x[0]) + actx.np.sin(np.pi*x[1]) \ + + actx.np.sin(np.pi*x[2]) + else: + ret = None + assert ret is not None return ret @@ -204,10 +210,12 @@ def f(x): def grad_f(x): ret = make_obj_array([dcoll.zeros(actx) for _ in range(dim)]) - ret[0] = -np.pi*actx.np.sin(np.pi*x[0]) - ret[1] = np.pi*actx.np.cos(np.pi*x[1]) - - if dim == 3: + if dim == 2: + ret[0] = -np.pi*actx.np.sin(np.pi*x[0]) + ret[1] = np.pi*actx.np.cos(np.pi*x[1]) + elif dim == 3: + ret[0] = -np.pi*actx.np.sin(np.pi*x[0]) + ret[1] = np.pi*actx.np.cos(np.pi*x[1]) ret[2] = np.pi*actx.np.cos(np.pi*x[2]) return ret From d0bd17e6bafce8bc98888bb9309bf4089ae295a0 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Tue, 1 Aug 2023 13:15:07 -0500 Subject: [PATCH 28/97] Add arg names to geometric factor application, refine some comments --- grudge/op.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index c6e6ce36e..c2f73d0e9 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -167,6 +167,8 @@ # {{{ Temporary tools for tensor product operators +# NOTE: Will possibly be removed in a future version of tensor product operator +# development since (I think) it is not entirely necessary from pytools.tag import Tag class OutputIsTensorProductDOFArrayOrdered(Tag): pass @@ -251,7 +253,8 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): vec = reshape_array_for_tensor_product_space(grp.space, vec) # apply differentiation matrix to vec - # check len(vec.shape) since shape is expected to be (nelements, ndofs) + # check len(vec.shape) since shape is expected to be + # (nelements, nnodes1d, nnodes1d) if len(vec.shape) == 3: specs = ["il,elj->eij", "jl,eil->eij"] @@ -294,7 +297,8 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): "rei,ei->ei", ijm[i], grad[i], - tagged=(FirstAxisIsElementsTag(),)) + tagged=(FirstAxisIsElementsTag(),)), + arg_names=("inv_jac_t", "vec") for i in range(grad.shape[0]) ]) From ef667bbcd6bc77b83cdeb85eb93952b21bfa3c95 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Wed, 2 Aug 2023 00:42:50 -0500 Subject: [PATCH 29/97] Divergence operator version 0.0 --- grudge/op.py | 96 +++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 84 insertions(+), 12 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index c2f73d0e9..449710ed8 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -297,8 +297,8 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): "rei,ei->ei", ijm[i], grad[i], - tagged=(FirstAxisIsElementsTag(),)), - arg_names=("inv_jac_t", "vec") + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", "vec")) for i in range(grad.shape[0]) ]) @@ -339,19 +339,91 @@ def _divergence_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec *, metric_in_matvec): # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. + + + def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): + """Exploits tensor product structure to differentiate each coordinate + axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) + """ + + actx_tp = TensorProductArrayContext( + actx.queue, + allocator=actx.allocator, + force_device_scalars=actx._force_device_scalars) + + from modepy.tools import ( + reshape_array_for_tensor_product_space, + unreshape_array_for_tensor_product_space) + + # reshape u to expose tensor product structure + vec = reshape_array_for_tensor_product_space(grp.space, vec) + + # define specs to extract dr, ds, dt + if len(vec.shape) == 3: + specs = ["il,elj->eij", + "jl,eil->eij"] + elif len(vec.shape) == 4: + specs = ["il,eljk->eijk", + "jl,eilk->eijk", + "kl,eijl->eijk"] + else: + specs = None + assert specs is not None + + diff_mat = get_diff_mat(actx, grp, grp) + drdsdt = make_obj_array([ + actx_tp.einsum( + spec, + diff_mat, + vec, + arg_names=("diff_mat", "vec"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + for spec in specs + ]) + + pu.db + if len(vec) == 3: + div = drdsdt[0] + drdsdt[1] + elif len(vec) == 4: + div = drdsdt[0] + drdsdt[1] + drdsdt[2] + else: + div = None + assert div is not None + + # see compute_tensor_product_grad for note on reshape before applying + # geometric factors + div = unreshape_array_for_tensor_product_space(grp.space, div) + + div = actx.einsum("xrei,ej->ej", + ijm, + div, + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", "vec")) + + return div + + + from meshmode.discretization.poly_element import \ + TensorProductElementGroupBase per_group_divs = [ + + compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) + if isinstance(in_grp, TensorProductElementGroupBase) # r for rst axis # x for xyz axis - actx.einsum("xrej,rij,xej->ei" if metric_in_matvec else "xrei,rij,xej->ei", - ijm_i, - get_diff_mat( - actx, - out_element_group=out_grp, - in_element_group=in_grp - ), - vec_i, - arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), - tagged=(FirstAxisIsElementsTag(),)) + else actx.einsum( + "xrej,rij,xej->ei" if metric_in_matvec else "xrei,rij,xej->ei", + ijm_i, + get_diff_mat( + actx, + out_element_group=out_grp, + in_element_group=in_grp + ), + vec_i, + arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), + tagged=(FirstAxisIsElementsTag(),)) + for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, inv_jac_mat)] From eed2516ee1c504768b78767c71728f5badde11ce Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Fri, 4 Aug 2023 15:26:33 -0500 Subject: [PATCH 30/97] Prototype of divergence kernel. Needs work, but it passes currently included convergence tests --- grudge/op.py | 106 ++++++++++++++++++++++++++---------------------- test/test_op.py | 78 +++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 48 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 449710ed8..ee5e536bd 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -42,6 +42,7 @@ """ from __future__ import annotations +from re import I __copyright__ = """ Copyright (C) 2021 Andreas Kloeckner @@ -79,6 +80,7 @@ DiscretizationDOFAxisTag, DiscretizationElementAxisTag, DiscretizationFaceAxisTag) +from meshmode.discretization.poly_element import TensorProductElementGroupBase from grudge.discretization import DiscretizationCollection from grudge.dof_desc import as_dofdesc @@ -235,6 +237,7 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. + def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): """Exploits tensor product structure to differentiate each coordinate axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) @@ -263,8 +266,9 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): "jl,eilk->eijk", "kl,eijl->eijk"] else: - specs = None - assert specs is not None + raise Exception("found dimension = {len(vec.shape)-1}. Special-case" + " tensor product operations are only valid for " + " 2 <= dimension <= 3.") diff_mat = get_diff_mat(actx, grp, grp) grad = make_obj_array([ @@ -274,7 +278,7 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): vec, arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) + OutputIsTensorProductDOFArrayOrdered())) for spec in specs ]) @@ -304,8 +308,7 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): return grad - from meshmode.discretization.poly_element import \ - TensorProductElementGroupBase + per_group_grads = [ compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i) @@ -341,7 +344,7 @@ def _divergence_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec # (both strong and weak derivative) and their differences. - def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): + def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): """Exploits tensor product structure to differentiate each coordinate axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ @@ -358,58 +361,73 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): # reshape u to expose tensor product structure vec = reshape_array_for_tensor_product_space(grp.space, vec) - # define specs to extract dr, ds, dt - if len(vec.shape) == 3: - specs = ["il,elj->eij", - "jl,eil->eij"] - elif len(vec.shape) == 4: - specs = ["il,eljk->eijk", - "jl,eilk->eijk", - "kl,eijl->eijk"] + # apply differentiation matrix to vec + # check len(vec.shape) since shape is expected to be + # (nelements, nnodes1d, nnodes1d) + # FIXME: make this "dimension independent" + if len(vec.shape) == 4: + specs = ["il,xelj->eij", + "jl,xeil->eij"] + elif len(vec.shape) == 5: + specs = ["il,xeljk->eijk", + "jl,xeilk->eijk", + "kl,xeijl->eijk"] else: - specs = None - assert specs is not None + raise Exception("found dimension = {len(vec.shape)-2}. Special-case" + " tensor product operations are only valid for " + " 2 <= dimension <= 3.") diff_mat = get_diff_mat(actx, grp, grp) - drdsdt = make_obj_array([ - actx_tp.einsum( + + # get partial derivatives for each ref. coord. axis + partials = make_obj_array([ + actx_tp.einsum( spec, diff_mat, vec, arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - for spec in specs - ]) - - pu.db - if len(vec) == 3: - div = drdsdt[0] + drdsdt[1] - elif len(vec) == 4: - div = drdsdt[0] + drdsdt[1] + drdsdt[2] - else: - div = None - assert div is not None + OutputIsTensorProductDOFArrayOrdered())) + for spec in specs + ]) - # see compute_tensor_product_grad for note on reshape before applying - # geometric factors - div = unreshape_array_for_tensor_product_space(grp.space, div) + # unreshape partials to apply geometric factors + # NOTE: In a future version, do not reshape before application of + # geometric factors. Can possibly "chain" the einsum. For example, the + # simplicial case below has einsum with spec + # ("xrei,rij,xej->ei") + # for the strong local divergence case + partials = make_obj_array([ + unreshape_array_for_tensor_product_space(grp.space, partials[i]) + for i in range(partials.shape[0]) + ]) - div = actx.einsum("xrei,ej->ej", - ijm, - div, + # apply geometric factors to partial derivatives + # FIXME: using einsum spec ("xrei,xei->xei") throws error: + # "Loopy does not directly support object arrays" + partials = make_obj_array([ + actx_tp.einsum( + "rei,ei->ei", + ijm[i], + partials[i], tagged=(FirstAxisIsElementsTag(),), arg_names=("inv_jac_t", "vec")) + for i in range(partials.shape[0]) + ]) + + if partials.shape[0] == 2: + div = partials[0] + partials[1] + else: + div = partials[0] + partials[1] + partials[2] return div - from meshmode.discretization.poly_element import \ - TensorProductElementGroupBase per_group_divs = [ compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) if isinstance(in_grp, TensorProductElementGroupBase) + # r for rst axis # x for xyz axis else actx.einsum( @@ -441,24 +459,16 @@ def _reference_derivative_matrices(actx: ArrayContext, # _reference_stiffness_transpose_matrices. assert out_element_group is in_element_group - from meshmode.mesh import TensorProductElementGroup - @keyed_memoize_in( actx, _reference_derivative_matrices, lambda grp: grp.discretization_key()) def get_ref_derivative_mats(grp): - - from meshmode.discretization.poly_element import \ - TensorProductElementGroupBase if isinstance(grp, TensorProductElementGroupBase): import modepy as mp import numpy.linalg as la - space1d = grp.space.bases[0] - shape1d = grp.shape.bases[0] - - nodes1d = mp.edge_clustered_nodes_for_space(space1d, shape1d) - basis1d = mp.basis_for_space(space1d, shape1d) + nodes1d = grp.unit_nodes_1d + basis1d = grp.basis_1d_obj() vdm1d = mp.vandermonde(basis1d.functions, nodes1d) vdm_p1d = mp.vandermonde(basis1d.gradients, nodes1d)[0] diff --git a/test/test_op.py b/test/test_op.py index 1e45a4556..90860faae 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -351,6 +351,84 @@ def get_flux(u_tpair): assert (eoc_rec.order_estimate() >= order - 0.5 or eoc_rec.max_error() < 1e-11) + +@pytest.mark.parametrize("form", ["strong"]) +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("order", [2, 3]) +@pytest.mark.parametrize(("vectorize", "nested"), [ + (False, False) + ]) +def test_tensor_product_divergence(actx_factory, form, dim, order, vectorize, + nested, visualize=False): + """A "one-dimensional tensor product element" does not make sense, so the + one-dimensional case is excluded from this test. + """ + actx = actx_factory() + from pytools.convergence import EOCRecorder + eoc_rec = EOCRecorder() + + from meshmode.mesh import TensorProductElementGroup + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + for n in [4, 6, 8]: + mesh = mgen.generate_regular_rect_mesh( + a=(-1,)*dim, + b=(1,)*dim, + nelements_per_axis=(n,)*dim, + group_cls=TensorProductElementGroup) + + import grudge.dof_desc as dd + dcoll = make_discretization_collection( + actx, + mesh, + discr_tag_to_group_factory={ + dd.DISCR_TAG_BASE: LGL(order)}) + + + def f(x): + if dim == 2: + ret = make_obj_array([dcoll.empty(actx) for _ in range(dim)]) + ret[0] = actx.np.cos(np.pi*x[0]) + ret[1] = actx.np.sin(np.pi*x[1]) + + return ret + elif dim == 3: + ret = make_obj_array([dcoll.empty(actx) for _ in range(dim)]) + ret[0] = actx.np.cos(np.pi*x[0]) + ret[1] = actx.np.sin(np.pi*x[1]) + ret[2] = actx.np.sin(np.pi*x[2]) + + return ret + + + def div_f(x): + + if dim == 2: + ret = -np.pi*actx.np.sin(np.pi*x[0]) + \ + np.pi*actx.np.cos(np.pi*x[1]) + return ret + elif dim == 3: + ret = -np.pi*actx.np.sin(np.pi*x[0]) + \ + np.pi*actx.np.cos(np.pi*x[1]) + \ + np.pi*actx.np.cos(np.pi*x[2]) + + return ret + + + x = actx.thaw(dcoll.nodes()) + u = f(x) + ref_div = div_f(x) + div = op.local_div(dcoll, u) + + rel_linf_error = actx.to_numpy(op.norm(dcoll, ref_div - div, np.inf) / + op.norm(dcoll, ref_div, np.inf)) + eoc_rec.add_data_point(1./n, rel_linf_error) + + print("L^inf error:") + print(eoc_rec) + assert (eoc_rec.order_estimate() >= order - 0.5 or + eoc_rec.max_error() < 1e-11) + # }}} From 1e40a1199782d62543ba59f030382da10e5fe313 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Fri, 4 Aug 2023 15:28:19 -0500 Subject: [PATCH 31/97] Remove random import included by CoC autocomplete --- grudge/op.py | 1 - 1 file changed, 1 deletion(-) diff --git a/grudge/op.py b/grudge/op.py index ee5e536bd..39c99fd08 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -42,7 +42,6 @@ """ from __future__ import annotations -from re import I __copyright__ = """ Copyright (C) 2021 Andreas Kloeckner From f2b0275a3e7e6b242b691136f0c96e5c20b65c7a Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sat, 5 Aug 2023 16:11:00 -0500 Subject: [PATCH 32/97] Generate einsum specification dynamically instead of using if-else --- grudge/op.py | 81 ++++++++++++++++++++++++++++------------------------ 1 file changed, 44 insertions(+), 37 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 39c99fd08..3447487c7 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -254,31 +254,36 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): # reshape u to expose tensor product structure vec = reshape_array_for_tensor_product_space(grp.space, vec) - # apply differentiation matrix to vec - # check len(vec.shape) since shape is expected to be - # (nelements, nnodes1d, nnodes1d) - if len(vec.shape) == 3: - specs = ["il,elj->eij", - "jl,eil->eij"] - elif len(vec.shape) == 4: - specs = ["il,eljk->eijk", - "jl,eilk->eijk", - "kl,eijl->eijk"] - else: - raise Exception("found dimension = {len(vec.shape)-1}. Special-case" - " tensor product operations are only valid for " - " 2 <= dimension <= 3.") + # apply differentiation matrix to function data + def pre_dims(axis): + return "ijk"[0:axis] + + + def post_dims(axis): + return "ijk"[axis+1:grp.dim] + + + def out_dims(): + return "ijk"[:grp.dim] + + + def axis(i): + return "ijk"[i] + diff_mat = get_diff_mat(actx, grp, grp) + # einsum specs will look something like: + # "il,eljk->eijk" (3D first coordinate partial) + # "jl,eil->eij" (2D second coordinate partial) grad = make_obj_array([ actx_tp.einsum( - spec, + f"{axis(i)}l,e{pre_dims(i)}l{post_dims(i)}->e{out_dims()}", diff_mat, vec, arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered())) - for spec in specs + for i in range(grp.dim) ]) # unreshape grad to apply geometric factors @@ -289,7 +294,7 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): # for the strong local gradient case grad = make_obj_array([ unreshape_array_for_tensor_product_space(grp.space, grad[i]) - for i in range(grad.shape[0]) + for i in range(grp.dim) ]) # apply geometric factors to current grad @@ -360,34 +365,36 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): # reshape u to expose tensor product structure vec = reshape_array_for_tensor_product_space(grp.space, vec) - # apply differentiation matrix to vec - # check len(vec.shape) since shape is expected to be - # (nelements, nnodes1d, nnodes1d) - # FIXME: make this "dimension independent" - if len(vec.shape) == 4: - specs = ["il,xelj->eij", - "jl,xeil->eij"] - elif len(vec.shape) == 5: - specs = ["il,xeljk->eijk", - "jl,xeilk->eijk", - "kl,xeijl->eijk"] - else: - raise Exception("found dimension = {len(vec.shape)-2}. Special-case" - " tensor product operations are only valid for " - " 2 <= dimension <= 3.") + # apply differentiation matrix to function data + def pre_dims(axis): + return "ijk"[0:axis] + + + def post_dims(axis): + return "ijk"[axis+1:grp.dim] + + + def out_dims(): + return "ijk"[:grp.dim] + + + def axis(i): + return "ijk"[i] - diff_mat = get_diff_mat(actx, grp, grp) # get partial derivatives for each ref. coord. axis + diff_mat = get_diff_mat(actx, grp, grp) + + # see comment on einsum spec in `compute_tensor_product_grad` partials = make_obj_array([ actx_tp.einsum( - spec, + f"{axis(i)}l,xe{pre_dims(i)}l{post_dims(i)}->e{out_dims()}", diff_mat, vec, arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered())) - for spec in specs + for i in range(grp.dim) ]) # unreshape partials to apply geometric factors @@ -469,8 +476,8 @@ def get_ref_derivative_mats(grp): nodes1d = grp.unit_nodes_1d basis1d = grp.basis_1d_obj() - vdm1d = mp.vandermonde(basis1d.functions, nodes1d) - vdm_p1d = mp.vandermonde(basis1d.gradients, nodes1d)[0] + vdm_1d = mp.vandermonde(basis1d.functions, nodes1d) + vdm_p_1d = mp.vandermonde(basis1d.gradients, nodes1d)[0] return actx.freeze(actx.from_numpy(vdm_p1d @ la.inv(vdm1d))) From 036681c01e8732aff8da627d2fc841ba7cb10298 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sat, 5 Aug 2023 16:16:15 -0500 Subject: [PATCH 33/97] Rename vandermonde and vandermonde derivative matrices --- grudge/op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/grudge/op.py b/grudge/op.py index 3447487c7..1e1c378a1 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -479,7 +479,7 @@ def get_ref_derivative_mats(grp): vdm_1d = mp.vandermonde(basis1d.functions, nodes1d) vdm_p_1d = mp.vandermonde(basis1d.gradients, nodes1d)[0] - return actx.freeze(actx.from_numpy(vdm_p1d @ la.inv(vdm1d))) + return actx.freeze(actx.from_numpy(vdm_p_1d @ la.inv(vdm_1d))) else: from meshmode.discretization.poly_element import diff_matrices From 7ad9017a92be718649bda3884c940d0abde18be3 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Mon, 7 Aug 2023 10:47:21 -0500 Subject: [PATCH 34/97] Give einsums a single source of truth. Still only valid for dim <= 3 --- grudge/op.py | 50 +++++++------------------------------------------- 1 file changed, 7 insertions(+), 43 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 1e1c378a1..784e60667 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -254,36 +254,18 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): # reshape u to expose tensor product structure vec = reshape_array_for_tensor_product_space(grp.space, vec) - # apply differentiation matrix to function data - def pre_dims(axis): - return "ijk"[0:axis] - - - def post_dims(axis): - return "ijk"[axis+1:grp.dim] - - - def out_dims(): - return "ijk"[:grp.dim] - - - def axis(i): - return "ijk"[i] - - + # apply operators to function data + dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) - # einsum specs will look something like: - # "il,eljk->eijk" (3D first coordinate partial) - # "jl,eil->eij" (2D second coordinate partial) grad = make_obj_array([ actx_tp.einsum( - f"{axis(i)}l,e{pre_dims(i)}l{post_dims(i)}->e{out_dims()}", + f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", diff_mat, vec, arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered())) - for i in range(grp.dim) + for i in range(dim) ]) # unreshape grad to apply geometric factors @@ -366,35 +348,17 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): vec = reshape_array_for_tensor_product_space(grp.space, vec) # apply differentiation matrix to function data - def pre_dims(axis): - return "ijk"[0:axis] - - - def post_dims(axis): - return "ijk"[axis+1:grp.dim] - - - def out_dims(): - return "ijk"[:grp.dim] - - - def axis(i): - return "ijk"[i] - - - # get partial derivatives for each ref. coord. axis + dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) - - # see comment on einsum spec in `compute_tensor_product_grad` partials = make_obj_array([ actx_tp.einsum( - f"{axis(i)}l,xe{pre_dims(i)}l{post_dims(i)}->e{out_dims()}", + f"ij,xe{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", diff_mat, vec, arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered())) - for i in range(grp.dim) + for i in range(dim) ]) # unreshape partials to apply geometric factors From c645dbe43f1620d50b6e89cfb0f1ffd02f0c8654 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 10 Aug 2023 15:17:24 -0500 Subject: [PATCH 35/97] Move TP array context to array_context.py, other minor changes --- grudge/array_context.py | 39 ++++++++++++++++++++++++++++++++++++ grudge/op.py | 44 +++++++++-------------------------------- 2 files changed, 48 insertions(+), 35 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index 2e82519e2..f29e8ef5d 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -127,6 +127,45 @@ def __init__(self, queue: "pyopencl.CommandQueue", # }}} +# {{{ Tensor product array context + +class OutputIsTensorProductDOFArrayOrdered(Tag): + """Signify that the strides will not be of order "C" or "F". See + :class:`grudge.array_context.TensorProductArrayContext` for more details. + """ + pass + + +class TensorProductArrayContext(_PyOpenCLArrayContextBase): + """Specialized array context for use with tensor product elements. + + The strides for the arrays containing tensor product element data are of the + form (slow, fastest, faster, fast). These strides are not "C" or "F" order. + Hence, this specialized array context takes care of specifying the + particular strides required. + """ + + def transform_loopy_program(self, t_unit): + if len(t_unit.callables_table) == 1: + knl = t_unit.default_entrypoint + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) + + return super().transform_loopy_program(t_unit) +# }}} + # {{{ pytato diff --git a/grudge/op.py b/grudge/op.py index 784e60667..08eeacbd2 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -83,6 +83,9 @@ from grudge.discretization import DiscretizationCollection from grudge.dof_desc import as_dofdesc +from grudge.array_context import ( + TensorProductArrayContext, + OutputIsTensorProductDOFArrayOrdered) from pytools import keyed_memoize_in from pytools.obj_array import make_obj_array @@ -167,38 +170,6 @@ ) -# {{{ Temporary tools for tensor product operators -# NOTE: Will possibly be removed in a future version of tensor product operator -# development since (I think) it is not entirely necessary -from pytools.tag import Tag -class OutputIsTensorProductDOFArrayOrdered(Tag): - pass - - -from grudge.array_context import PyOpenCLArrayContext -class TensorProductArrayContext(PyOpenCLArrayContext): - def transform_loopy_program(self, t_unit): - if len(t_unit.callables_table) == 1: - knl = t_unit.default_entrypoint - if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): - new_args = [] - for arg in knl.args: - if arg.is_output: - arg = arg.copy(dim_tags=( - f"N{len(arg.shape)-1}," - + ",".join(f"N{i}" - for i in range(len(arg.shape)-1)) - )) - - new_args.append(arg) - - knl = knl.copy(args=new_args) - t_unit = t_unit.with_kernel(knl) - - return super().transform_loopy_program(t_unit) -# }}} - - # {{{ common derivative "kernels" def _single_axis_derivative_kernel( @@ -437,11 +408,14 @@ def get_ref_derivative_mats(grp): import modepy as mp import numpy.linalg as la + # not functional in current state nodes1d = grp.unit_nodes_1d - basis1d = grp.basis_1d_obj() + bases_1d = grp.bases_1d() - vdm_1d = mp.vandermonde(basis1d.functions, nodes1d) - vdm_p_1d = mp.vandermonde(basis1d.gradients, nodes1d)[0] + diff_mats = [] + for i in range(len(bases_1d)): + vdm_1d = mp.vandermonde(bases_1d.functions, nodes1d) + vdm_p_1d = mp.vandermonde(bases_1d.gradients, nodes1d)[0] return actx.freeze(actx.from_numpy(vdm_p_1d @ la.inv(vdm_1d))) From dcd7ca06b58bc4fd7d9deb2ebb738fc938eb565d Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 10 Aug 2023 18:01:40 -0500 Subject: [PATCH 36/97] Update tensor product grad test to match the other grad test case --- grudge/op.py | 7 +-- test/test_op.py | 124 ++++++++++++++++++++++++++++++++++-------------- 2 files changed, 91 insertions(+), 40 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 08eeacbd2..5a64cf624 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -408,14 +408,11 @@ def get_ref_derivative_mats(grp): import modepy as mp import numpy.linalg as la - # not functional in current state nodes1d = grp.unit_nodes_1d bases_1d = grp.bases_1d() - diff_mats = [] - for i in range(len(bases_1d)): - vdm_1d = mp.vandermonde(bases_1d.functions, nodes1d) - vdm_p_1d = mp.vandermonde(bases_1d.gradients, nodes1d)[0] + vdm_1d = mp.vandermonde(bases_1d.functions, nodes1d) + vdm_p_1d = mp.vandermonde(bases_1d.gradients, nodes1d)[0] return actx.freeze(actx.from_numpy(vdm_p_1d @ la.inv(vdm_1d))) diff --git a/test/test_op.py b/test/test_op.py index 90860faae..ec96209c4 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -165,75 +165,129 @@ def get_flux(u_tpair): @pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize(("vectorize", "nested"), [ - (False, False) + (False, False), + (True, False), + (True, True) ]) -def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize, +def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize, nested, visualize=False): """A "one-dimensional tensor product element" does not make sense, so the one-dimensional case is excluded from this test. """ + actx = actx_factory() + from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() from meshmode.mesh import TensorProductElementGroup from meshmode.discretization.poly_element import \ - LegendreGaussLobattoTensorProductGroupFactory as LGL + LegendreGaussLobattoTensorProductGroupFactory as LGL for n in [4, 6, 8]: mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, - b=(1,)*dim, + a=(-1,)*dim, b=(1,)*dim, nelements_per_axis=(n,)*dim, group_cls=TensorProductElementGroup) import grudge.dof_desc as dd - dcoll = make_discretization_collection( + dcoll = DiscretizationCollection( actx, mesh, discr_tag_to_group_factory={ dd.DISCR_TAG_BASE: LGL(order)}) - def f(x): - if dim == 2: - ret = actx.np.cos(np.pi*x[0]) + actx.np.sin(np.pi*x[1]) - elif dim == 3: - ret = actx.np.cos(np.pi*x[0]) + actx.np.sin(np.pi*x[1]) \ - + actx.np.sin(np.pi*x[2]) - else: - ret = None - assert ret is not None + result = dcoll.zeros(actx) + 1 + for i in range(dim-1): + result = result * actx.np.sin(np.pi*x[i]) + result = result * actx.np.cos(np.pi/2*x[dim-1]) + return result - return ret + def grad_f(x): + result = make_obj_array([dcoll.zeros(actx) + 1 for _ in range(dim)]) + for i in range(dim-1): + for j in range(i): + result[i] = result[i] * actx.np.sin(np.pi*x[j]) + result[i] = result[i] * np.pi*actx.np.cos(np.pi*x[i]) + for j in range(i+1, dim-1): + result[i] = result[i] * actx.np.sin(np.pi*x[j]) + result[i] = result[i] * actx.np.cos(np.pi/2*x[dim-1]) + for j in range(dim-1): + result[dim-1] = result[dim-1] * actx.np.sin(np.pi*x[j]) + result[dim-1] = result[dim-1] * (-np.pi/2*actx.np.sin(np.pi/2*x[dim-1])) + return result + x = actx.thaw(dcoll.nodes()) - def grad_f(x): - ret = make_obj_array([dcoll.zeros(actx) for _ in range(dim)]) + if vectorize: + u = make_obj_array([(i+1)*f(x) for i in range(dim)]) + else: + u = f(x) - if dim == 2: - ret[0] = -np.pi*actx.np.sin(np.pi*x[0]) - ret[1] = np.pi*actx.np.cos(np.pi*x[1]) - elif dim == 3: - ret[0] = -np.pi*actx.np.sin(np.pi*x[0]) - ret[1] = np.pi*actx.np.cos(np.pi*x[1]) - ret[2] = np.pi*actx.np.cos(np.pi*x[2]) + def get_flux(u_tpair): + dd = u_tpair.dd + dd_allfaces = dd.with_dtag("all_faces") + normal = actx.thaw(dcoll.normal(dd)) + u_avg = u_tpair.avg + if vectorize: + if nested: + flux = make_obj_array([u_avg_i * normal for u_avg_i in u_avg]) + else: + flux = np.outer(u_avg, normal) + else: + flux = u_avg * normal + return op.project(dcoll, dd, dd_allfaces, flux) - return ret + dd_allfaces = DOFDesc("all_faces") + if form == "strong": + grad_u = ( + op.local_grad(dcoll, u, nested=nested) + # No flux terms because u doesn't have inter-el jumps + ) + elif form == "weak": + grad_u = op.inverse_mass(dcoll, + -op.weak_local_grad(dcoll, u, nested=nested) # pylint: disable=E1130 + + # noqa: W504 + op.face_mass(dcoll, + dd_allfaces, + # Note: no boundary flux terms here because u_ext == u_int == 0 + sum(get_flux(utpair) + for utpair in op.interior_trace_pairs(dcoll, u)) + ) + ) + else: + raise ValueError("Invalid form argument.") - x = actx.thaw(dcoll.nodes()) - u = f(x) - ref_grad = grad_f(x) - grad = op.local_grad(dcoll, u) + if vectorize: + expected_grad_u = make_obj_array( + [(i+1)*grad_f(x) for i in range(dim)]) + if not nested: + expected_grad_u = np.stack(expected_grad_u, axis=0) + else: + expected_grad_u = grad_f(x) - rel_linf_error = actx.to_numpy(op.norm(dcoll, ref_grad - grad, np.inf) / - op.norm(dcoll, ref_grad, np.inf)) - eoc_rec.add_data_point(1./n, rel_linf_error) + if visualize: + from grudge.shortcuts import make_visualizer + vis = make_visualizer(dcoll, vis_order=order if dim == 3 else dim+3) + + filename = (f"test_gradient_{form}_{dim}_{order}" + f"{'_vec' if vectorize else ''}{'_nested' if nested else ''}.vtu") + vis.write_vtk_file(filename, [ + ("u", u), + ("grad_u", grad_u), + ("expected_grad_u", expected_grad_u), + ], overwrite=True) + + rel_linf_err = actx.to_numpy( + op.norm(dcoll, grad_u - expected_grad_u, np.inf) + / op.norm(dcoll, expected_grad_u, np.inf)) + eoc_rec.add_data_point(1./n, rel_linf_err) print("L^inf error:") print(eoc_rec) - assert (eoc_rec.order_estimate() >= order - 0.5 or - eoc_rec.max_error() < 1e-11) + assert (eoc_rec.order_estimate() >= order - 0.5 + or eoc_rec.max_error() < 1e-11) # }}} From 683cdd839df7b48781eed59f1bdfc1fc6f0c5f3b Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 10 Aug 2023 18:14:22 -0500 Subject: [PATCH 37/97] Update tensor product divergence test to match original test case. --- test/test_op.py | 121 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 83 insertions(+), 38 deletions(-) diff --git a/test/test_op.py b/test/test_op.py index ec96209c4..92309b3a2 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -169,7 +169,7 @@ def get_flux(u_tpair): (True, False), (True, True) ]) -def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize, +def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize, nested, visualize=False): """A "one-dimensional tensor product element" does not make sense, so the one-dimensional case is excluded from this test. @@ -410,14 +410,17 @@ def get_flux(u_tpair): @pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize(("vectorize", "nested"), [ - (False, False) + (False, False), + (True, False), + (True, True) ]) def test_tensor_product_divergence(actx_factory, form, dim, order, vectorize, - nested, visualize=False): + nested, visualize=False): """A "one-dimensional tensor product element" does not make sense, so the one-dimensional case is excluded from this test. """ actx = actx_factory() + from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() @@ -426,63 +429,105 @@ def test_tensor_product_divergence(actx_factory, form, dim, order, vectorize, LegendreGaussLobattoTensorProductGroupFactory as LGL for n in [4, 6, 8]: mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, - b=(1,)*dim, + a=(-1,)*dim, b=(1,)*dim, nelements_per_axis=(n,)*dim, group_cls=TensorProductElementGroup) import grudge.dof_desc as dd - dcoll = make_discretization_collection( + dcoll = DiscretizationCollection( actx, mesh, discr_tag_to_group_factory={ dd.DISCR_TAG_BASE: LGL(order)}) - def f(x): - if dim == 2: - ret = make_obj_array([dcoll.empty(actx) for _ in range(dim)]) - ret[0] = actx.np.cos(np.pi*x[0]) - ret[1] = actx.np.sin(np.pi*x[1]) + result = make_obj_array([dcoll.zeros(actx) + (i+1) for i in range(dim)]) + for i in range(dim-1): + result = result * actx.np.sin(np.pi*x[i]) + result = result * actx.np.cos(np.pi/2*x[dim-1]) + return result - return ret - elif dim == 3: - ret = make_obj_array([dcoll.empty(actx) for _ in range(dim)]) - ret[0] = actx.np.cos(np.pi*x[0]) - ret[1] = actx.np.sin(np.pi*x[1]) - ret[2] = actx.np.sin(np.pi*x[2]) + def div_f(x): + result = dcoll.zeros(actx) + for i in range(dim-1): + deriv = dcoll.zeros(actx) + (i+1) + for j in range(i): + deriv = deriv * actx.np.sin(np.pi*x[j]) + deriv = deriv * np.pi*actx.np.cos(np.pi*x[i]) + for j in range(i+1, dim-1): + deriv = deriv * actx.np.sin(np.pi*x[j]) + deriv = deriv * actx.np.cos(np.pi/2*x[dim-1]) + result = result + deriv + deriv = dcoll.zeros(actx) + dim + for j in range(dim-1): + deriv = deriv * actx.np.sin(np.pi*x[j]) + deriv = deriv * (-np.pi/2*actx.np.sin(np.pi/2*x[dim-1])) + result = result + deriv + return result - return ret + x = actx.thaw(dcoll.nodes()) + if vectorize: + u = make_obj_array([(i+1)*f(x) for i in range(dim)]) + if not nested: + u = np.stack(u, axis=0) + else: + u = f(x) - def div_f(x): + def get_flux(u_tpair): + dd = u_tpair.dd + dd_allfaces = dd.with_dtag("all_faces") + normal = actx.thaw(dcoll.normal(dd)) + flux = u_tpair.avg @ normal + return op.project(dcoll, dd, dd_allfaces, flux) - if dim == 2: - ret = -np.pi*actx.np.sin(np.pi*x[0]) + \ - np.pi*actx.np.cos(np.pi*x[1]) - return ret - elif dim == 3: - ret = -np.pi*actx.np.sin(np.pi*x[0]) + \ - np.pi*actx.np.cos(np.pi*x[1]) + \ - np.pi*actx.np.cos(np.pi*x[2]) + dd_allfaces = DOFDesc("all_faces") - return ret + if form == "strong": + div_u = ( + op.local_div(dcoll, u) + # No flux terms because u doesn't have inter-el jumps + ) + elif form == "weak": + div_u = op.inverse_mass(dcoll, + -op.weak_local_div(dcoll, u) + + # noqa: W504 + op.face_mass(dcoll, + dd_allfaces, + # Note: no boundary flux terms here because u_ext == u_int == 0 + sum(get_flux(utpair) + for utpair in op.interior_trace_pairs(dcoll, u)) + ) + ) + else: + raise ValueError("Invalid form argument.") + if vectorize: + expected_div_u = make_obj_array([(i+1)*div_f(x) for i in range(dim)]) + else: + expected_div_u = div_f(x) - x = actx.thaw(dcoll.nodes()) - u = f(x) - ref_div = div_f(x) - div = op.local_div(dcoll, u) + if visualize: + from grudge.shortcuts import make_visualizer + vis = make_visualizer(dcoll, vis_order=order if dim == 3 else dim+3) - rel_linf_error = actx.to_numpy(op.norm(dcoll, ref_div - div, np.inf) / - op.norm(dcoll, ref_div, np.inf)) - eoc_rec.add_data_point(1./n, rel_linf_error) + filename = (f"test_divergence_{form}_{dim}_{order}" + f"{'_vec' if vectorize else ''}{'_nested' if nested else ''}.vtu") + vis.write_vtk_file(filename, [ + ("u", u), + ("div_u", div_u), + ("expected_div_u", expected_div_u), + ], overwrite=True) + + rel_linf_err = actx.to_numpy( + op.norm(dcoll, div_u - expected_div_u, np.inf) + / op.norm(dcoll, expected_div_u, np.inf)) + eoc_rec.add_data_point(1./n, rel_linf_err) print("L^inf error:") print(eoc_rec) - assert (eoc_rec.order_estimate() >= order - 0.5 or - eoc_rec.max_error() < 1e-11) - + assert (eoc_rec.order_estimate() >= order - 0.5 + or eoc_rec.max_error() < 1e-11) # }}} From 15639503e3e1e097bcfe938d5ded17605e8ade66 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Mon, 14 Aug 2023 13:05:09 -0500 Subject: [PATCH 38/97] Divergence kernel functioning again --- grudge/op.py | 50 ++++++++++++++++++------------------------------- test/test_op.py | 5 +++-- 2 files changed, 21 insertions(+), 34 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 5a64cf624..9431ac901 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -250,18 +250,14 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): for i in range(grp.dim) ]) - # apply geometric factors to current grad - # FIXME: using einsum spec ("xrei,xei->xei") throws error: - # "Loopy does not directly support object arrays" - grad = make_obj_array([ - actx_tp.einsum( - "rei,ei->ei", - ijm[i], - grad[i], - tagged=(FirstAxisIsElementsTag(),), - arg_names=("inv_jac_t", "vec")) - for i in range(grad.shape[0]) - ]) + # apply geometric factors + grad = actx.np.stack([grad[i] for i in range(dim)]) + grad = actx.einsum( + "xrei,xei->xei", + ijm, + grad, + arg_names=("inv_jac_t", "vec"), + tagged=(FirstAxisIsElementsTag(),)) return grad @@ -305,7 +301,6 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): """Exploits tensor product structure to differentiate each coordinate axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ - actx_tp = TensorProductArrayContext( actx.queue, allocator=actx.allocator, @@ -323,9 +318,9 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): diff_mat = get_diff_mat(actx, grp, grp) partials = make_obj_array([ actx_tp.einsum( - f"ij,xe{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", + f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", diff_mat, - vec, + vec[i], arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered())) @@ -343,23 +338,14 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): for i in range(partials.shape[0]) ]) - # apply geometric factors to partial derivatives - # FIXME: using einsum spec ("xrei,xei->xei") throws error: - # "Loopy does not directly support object arrays" - partials = make_obj_array([ - actx_tp.einsum( - "rei,ei->ei", - ijm[i], - partials[i], - tagged=(FirstAxisIsElementsTag(),), - arg_names=("inv_jac_t", "vec")) - for i in range(partials.shape[0]) - ]) - - if partials.shape[0] == 2: - div = partials[0] + partials[1] - else: - div = partials[0] + partials[1] + partials[2] + # apply geometric factors + partials = actx.np.stack([partials[i] for i in range(dim)]) + div = actx.einsum( + "xrei,xei->ei", + ijm, + partials, + arg_names=("inv_jac_t", "vec"), + tagged=(FirstAxisIsElementsTag(),)) return div diff --git a/test/test_op.py b/test/test_op.py index 92309b3a2..e836a63e1 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -429,12 +429,13 @@ def test_tensor_product_divergence(actx_factory, form, dim, order, vectorize, LegendreGaussLobattoTensorProductGroupFactory as LGL for n in [4, 6, 8]: mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, b=(1,)*dim, + a=(-1,)*dim, + b=(1,)*dim, nelements_per_axis=(n,)*dim, group_cls=TensorProductElementGroup) import grudge.dof_desc as dd - dcoll = DiscretizationCollection( + dcoll = make_discretization_collection( actx, mesh, discr_tag_to_group_factory={ From e1380fed4fb48d8a8cca4e44cb4aa5fbd72b6d54 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sun, 20 Aug 2023 20:42:36 -0500 Subject: [PATCH 39/97] Update some comments, begin weak form matrices work --- grudge/op.py | 33 +++++++++++++++++++++++---------- requirements.txt | 2 +- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 9431ac901..951249f11 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -240,17 +240,13 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): ]) # unreshape grad to apply geometric factors - # NOTE: In a future version, do not reshape before application of - # geometric factors. Can possibly "chain" the einsum. For example, the - # simplicial case below has einsum with spec - # ("xrei,rij,ei->ei") - # for the strong local gradient case grad = make_obj_array([ unreshape_array_for_tensor_product_space(grp.space, grad[i]) for i in range(grp.dim) ]) # apply geometric factors + # TODO: chain the einsum above with the einsum below grad = actx.np.stack([grad[i] for i in range(dim)]) grad = actx.einsum( "xrei,xei->xei", @@ -328,11 +324,7 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): ]) # unreshape partials to apply geometric factors - # NOTE: In a future version, do not reshape before application of - # geometric factors. Can possibly "chain" the einsum. For example, the - # simplicial case below has einsum with spec - # ("xrei,rij,xej->ei") - # for the strong local divergence case + # TODO: chain the einsum above with the einsum below partials = make_obj_array([ unreshape_array_for_tensor_product_space(grp.space, partials[i]) for i in range(partials.shape[0]) @@ -579,7 +571,28 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): from meshmode.discretization.poly_element import \ mass_matrix, diff_matrices + if isinstance(out_grp, TensorProductElementGroupBase): + import modepy as mp + import numpy.linalg as la + + basis_1d = out_grp.bases_1d() + nodes_1d = out_grp.unit_nodes_1d + + vdm = mp.vandermonde(basis_1d.functions, nodes_1d) + vdm_p = mp.vandermonde(basis_1d.gradients, nodes_1d)[0] + + # NOTE: possibly work special-case matrices like differentiation + # matrix, mass matrix, into modepy + mmat = la.inv(vdm @ vdm.T) + diff_mat = vdm_p @ la.inv(vdm) + return actx.freeze( + actx.tag_axis(1, DiscretizationDOFAxisTag(), + actx.from_numpy( + np.asarray( + diff_mat.T @ mmat.T)))) + mmat = mass_matrix(out_grp) + return actx.freeze( actx.tag_axis(1, DiscretizationDOFAxisTag(), actx.from_numpy( diff --git a/requirements.txt b/requirements.txt index 2107e5aeb..f56f10888 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/a-alveyblanc/meshmode.git@tensor-product-1d-nodes-and-1d-basis#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 From 33a54e40c7966a5cfc24108cd8782ede7ddc0035 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Wed, 6 Sep 2023 11:31:04 -0500 Subject: [PATCH 40/97] TMP: Use outside actx in TP grad --- grudge/op.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 951249f11..212200600 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -213,10 +213,10 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ - actx_tp = TensorProductArrayContext( - actx.queue, - allocator=actx.allocator, - force_device_scalars=actx._force_device_scalars) + # actx_tp = TensorProductArrayContext( + # actx.queue, + # allocator=actx.allocator, + # force_device_scalars=actx._force_device_scalars) from modepy.tools import ( reshape_array_for_tensor_product_space, @@ -229,7 +229,7 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) grad = make_obj_array([ - actx_tp.einsum( + actx.einsum( f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", diff_mat, vec, From e446bb76abc548014e54412fab45a45c8d7c01dc Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Wed, 6 Sep 2023 11:31:21 -0500 Subject: [PATCH 41/97] Add TP transform cartoon --- examples/tp-transform-cartoon.py | 56 ++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 examples/tp-transform-cartoon.py diff --git a/examples/tp-transform-cartoon.py b/examples/tp-transform-cartoon.py new file mode 100644 index 000000000..7b2472076 --- /dev/null +++ b/examples/tp-transform-cartoon.py @@ -0,0 +1,56 @@ +import numpy as np +import pyopencl as cl +from meshmode.array_context import PytatoPyOpenCLArrayContext +import meshmode.mesh.generation as mgen +from grudge import op, DiscretizationCollection +from pytools.obj_array import make_obj_array + + +class MyArrayContext(PytatoPyOpenCLArrayContext): + pass + + +def main(): + order = 4 + + ctx = cl.create_some_context() + queue = cl.CommandQueue(ctx) + actx = MyArrayContext(queue) + + dim = 3 + n = 5 + + from meshmode.mesh import TensorProductElementGroup + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + + mesh = mgen.generate_regular_rect_mesh( + a=(-1,)*dim, b=(1,)*dim, + nelements_per_axis=(n,)*dim, + group_cls=TensorProductElementGroup) + + import grudge.dof_desc as dd + dcoll = DiscretizationCollection( + actx, + mesh, + discr_tag_to_group_factory={ + dd.DISCR_TAG_BASE: LGL(order)}) + + def f(x): + result = dcoll.zeros(actx) + 1 + for i in range(dim-1): + result = result * actx.np.sin(np.pi*x[i]) + result = result * actx.np.cos(np.pi/2*x[dim-1]) + return result + + + x = actx.thaw(dcoll.nodes()) + + u = f(x) + + op.local_grad(dcoll, u) + + +if __name__ == "__main__": + main() + From 264192c1d4dbec10e2b394b80caebb394b3859a4 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Tue, 12 Sep 2023 10:54:57 -0500 Subject: [PATCH 42/97] Temporary changes to get tensor product gradient working again --- examples/tp-transform-cartoon.py | 2 +- grudge/op.py | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/examples/tp-transform-cartoon.py b/examples/tp-transform-cartoon.py index 7b2472076..cbc23c267 100644 --- a/examples/tp-transform-cartoon.py +++ b/examples/tp-transform-cartoon.py @@ -48,7 +48,7 @@ def f(x): u = f(x) - op.local_grad(dcoll, u) + grad_u = op.local_grad(dcoll, u) if __name__ == "__main__": diff --git a/grudge/op.py b/grudge/op.py index 212200600..0b2e95a7d 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -213,10 +213,10 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ - # actx_tp = TensorProductArrayContext( - # actx.queue, - # allocator=actx.allocator, - # force_device_scalars=actx._force_device_scalars) + actx_tp = TensorProductArrayContext( + actx.queue, + allocator=actx.allocator, + force_device_scalars=actx._force_device_scalars) from modepy.tools import ( reshape_array_for_tensor_product_space, @@ -229,7 +229,7 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) grad = make_obj_array([ - actx.einsum( + actx_tp.einsum( f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", diff_mat, vec, @@ -247,8 +247,8 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): # apply geometric factors # TODO: chain the einsum above with the einsum below - grad = actx.np.stack([grad[i] for i in range(dim)]) - grad = actx.einsum( + grad = actx_tp.np.stack([grad[i] for i in range(dim)]) + grad = actx_tp.einsum( "xrei,xei->xei", ijm, grad, From ba03b3fab772fc1b78fc05d2d712ae1943c81a06 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Wed, 13 Sep 2023 18:42:07 -0500 Subject: [PATCH 43/97] Tensor product array context related changes --- examples/tp-transform-cartoon.py | 28 +++++++++++++++++++++++++--- grudge/op.py | 21 ++++++--------------- test/test_op.py | 22 +++++++++++++++------- 3 files changed, 46 insertions(+), 25 deletions(-) diff --git a/examples/tp-transform-cartoon.py b/examples/tp-transform-cartoon.py index cbc23c267..4ebade58b 100644 --- a/examples/tp-transform-cartoon.py +++ b/examples/tp-transform-cartoon.py @@ -3,11 +3,32 @@ from meshmode.array_context import PytatoPyOpenCLArrayContext import meshmode.mesh.generation as mgen from grudge import op, DiscretizationCollection +from grudge.array_context import OutputIsTensorProductDOFArrayOrdered from pytools.obj_array import make_obj_array -class MyArrayContext(PytatoPyOpenCLArrayContext): - pass +class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext): + def transform_loopy_program(self, t_unit): + + if len(t_unit.callables_table) == 1: + knl = t_unit.default_entrypoint + + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) + + return super().transform_loopy_program(t_unit) def main(): @@ -15,7 +36,7 @@ def main(): ctx = cl.create_some_context() queue = cl.CommandQueue(ctx) - actx = MyArrayContext(queue) + actx = PytatoTensorProductArrayContext(queue) dim = 3 n = 5 @@ -50,6 +71,7 @@ def f(x): grad_u = op.local_grad(dcoll, u) + pu.db if __name__ == "__main__": main() diff --git a/grudge/op.py b/grudge/op.py index 0b2e95a7d..b60ab62d7 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -213,11 +213,6 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ - actx_tp = TensorProductArrayContext( - actx.queue, - allocator=actx.allocator, - force_device_scalars=actx._force_device_scalars) - from modepy.tools import ( reshape_array_for_tensor_product_space, unreshape_array_for_tensor_product_space) @@ -229,7 +224,7 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) grad = make_obj_array([ - actx_tp.einsum( + actx.einsum( f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", diff_mat, vec, @@ -247,8 +242,8 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): # apply geometric factors # TODO: chain the einsum above with the einsum below - grad = actx_tp.np.stack([grad[i] for i in range(dim)]) - grad = actx_tp.einsum( + grad = actx.np.stack([grad[i] for i in range(dim)]) + grad = actx.einsum( "xrei,xei->xei", ijm, grad, @@ -297,10 +292,6 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): """Exploits tensor product structure to differentiate each coordinate axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ - actx_tp = TensorProductArrayContext( - actx.queue, - allocator=actx.allocator, - force_device_scalars=actx._force_device_scalars) from modepy.tools import ( reshape_array_for_tensor_product_space, @@ -313,7 +304,7 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) partials = make_obj_array([ - actx_tp.einsum( + actx.einsum( f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", diff_mat, vec[i], @@ -321,14 +312,14 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): tagged=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered())) for i in range(dim) - ]) + ]) # unreshape partials to apply geometric factors # TODO: chain the einsum above with the einsum below partials = make_obj_array([ unreshape_array_for_tensor_product_space(grp.space, partials[i]) for i in range(partials.shape[0]) - ]) + ]) # apply geometric factors partials = actx.np.stack([partials[i] for i in range(dim)]) diff --git a/test/test_op.py b/test/test_op.py index e836a63e1..05f6c1f99 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -169,13 +169,18 @@ def get_flux(u_tpair): (True, False), (True, True) ]) -def test_tensor_product_gradient(actx_factory, form, dim, order, vectorize, +def test_tensor_product_gradient(form, dim, order, vectorize, nested, visualize=False): """A "one-dimensional tensor product element" does not make sense, so the one-dimensional case is excluded from this test. """ - actx = actx_factory() + import pyopencl as cl + from grudge.array_context import TensorProductArrayContext + + ctx = cl.create_some_context() + queue = cl.CommandQueue(ctx) + actx = TensorProductArrayContext(queue) from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() @@ -303,9 +308,7 @@ def get_flux(u_tpair): (True, True) ]) def test_divergence(actx_factory, form, dim, order, vectorize, nested, - visualize=False): - actx = actx_factory() - + visualize=False): from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() @@ -414,12 +417,17 @@ def get_flux(u_tpair): (True, False), (True, True) ]) -def test_tensor_product_divergence(actx_factory, form, dim, order, vectorize, +def test_tensor_product_divergence(form, dim, order, vectorize, nested, visualize=False): """A "one-dimensional tensor product element" does not make sense, so the one-dimensional case is excluded from this test. """ - actx = actx_factory() + import pyopencl as cl + from grudge.array_context import TensorProductArrayContext + + ctx = cl.create_some_context() + queue = cl.CommandQueue(ctx) + actx = TensorProductArrayContext(queue) from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() From 8263333024437c6b4e0bab7fe6dcf8d2d60ff9f4 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 14 Sep 2023 09:28:02 -0500 Subject: [PATCH 44/97] Update example --- examples/tp-transform-cartoon.py | 33 +++++++++++++++----------------- 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/examples/tp-transform-cartoon.py b/examples/tp-transform-cartoon.py index 4ebade58b..a3f54e210 100644 --- a/examples/tp-transform-cartoon.py +++ b/examples/tp-transform-cartoon.py @@ -4,29 +4,26 @@ import meshmode.mesh.generation as mgen from grudge import op, DiscretizationCollection from grudge.array_context import OutputIsTensorProductDOFArrayOrdered -from pytools.obj_array import make_obj_array class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext): def transform_loopy_program(self, t_unit): - if len(t_unit.callables_table) == 1: - knl = t_unit.default_entrypoint + knl = t_unit.default_entrypoint + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) - if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): - new_args = [] - for arg in knl.args: - if arg.is_output: - arg = arg.copy(dim_tags=( - f"N{len(arg.shape)-1}," - + ",".join(f"N{i}" - for i in range(len(arg.shape)-1)) - )) + new_args.append(arg) - new_args.append(arg) - - knl = knl.copy(args=new_args) - t_unit = t_unit.with_kernel(knl) + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) return super().transform_loopy_program(t_unit) @@ -39,7 +36,7 @@ def main(): actx = PytatoTensorProductArrayContext(queue) dim = 3 - n = 5 + res = 5 from meshmode.mesh import TensorProductElementGroup from meshmode.discretization.poly_element import \ @@ -47,7 +44,7 @@ def main(): mesh = mgen.generate_regular_rect_mesh( a=(-1,)*dim, b=(1,)*dim, - nelements_per_axis=(n,)*dim, + nelements_per_axis=(res,)*dim, group_cls=TensorProductElementGroup) import grudge.dof_desc as dd From 6b5002846026dea998a52d686a00d094c8a22722 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 14 Sep 2023 13:10:02 -0500 Subject: [PATCH 45/97] Add code for printing generated differentiation code --- examples/tp-transform-cartoon.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/examples/tp-transform-cartoon.py b/examples/tp-transform-cartoon.py index a3f54e210..76ce6d889 100644 --- a/examples/tp-transform-cartoon.py +++ b/examples/tp-transform-cartoon.py @@ -1,5 +1,7 @@ import numpy as np import pyopencl as cl +import pytato as pt +import loopy as lp from meshmode.array_context import PytatoPyOpenCLArrayContext import meshmode.mesh.generation as mgen from grudge import op, DiscretizationCollection @@ -24,7 +26,6 @@ def transform_loopy_program(self, t_unit): knl = knl.copy(args=new_args) t_unit = t_unit.with_kernel(knl) - return super().transform_loopy_program(t_unit) @@ -67,7 +68,12 @@ def f(x): u = f(x) grad_u = op.local_grad(dcoll, u) + grad_u = actx.np.stack(grad_u)[0] + + prg = pt.generate_loopy(grad_u).program + code = lp.generate_code_v2(prg).device_code() + print(code) pu.db if __name__ == "__main__": From 5d36bfb916f9405e8d7728feda8a1c5ad6b703c1 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Mon, 18 Sep 2023 09:42:41 -0500 Subject: [PATCH 46/97] Update strong tp diff example --- examples/tp-transform-cartoon.py | 51 ++++++++++++++++++++++---------- grudge/op.py | 13 ++++++-- 2 files changed, 47 insertions(+), 17 deletions(-) diff --git a/examples/tp-transform-cartoon.py b/examples/tp-transform-cartoon.py index 76ce6d889..b0bb13cfd 100644 --- a/examples/tp-transform-cartoon.py +++ b/examples/tp-transform-cartoon.py @@ -1,17 +1,31 @@ +import loopy as lp + +import meshmode.mesh.generation as mgen + import numpy as np import pyopencl as cl import pytato as pt -import loopy as lp -from meshmode.array_context import PytatoPyOpenCLArrayContext -import meshmode.mesh.generation as mgen -from grudge import op, DiscretizationCollection + +from grudge import op from grudge.array_context import OutputIsTensorProductDOFArrayOrdered +from grudge.discretization import make_discretization_collection + +from meshmode.array_context import PytatoPyOpenCLArrayContext class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext): - def transform_loopy_program(self, t_unit): + def transform_dag(self, dag): + if "dag_dots" not in dir(self): + self.dag_dots = [] + self.dag_dots.append(pt.get_dot_graph(dag)) + + return super().transform_dag(dag) + + def transform_loopy_program(self, t_unit): knl = t_unit.default_entrypoint + + # {{{ adjust strides according to tensor product structure if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): new_args = [] for arg in knl.args: @@ -25,19 +39,31 @@ def transform_loopy_program(self, t_unit): new_args.append(arg) knl = knl.copy(args=new_args) - t_unit = t_unit.with_kernel(knl) + # }}} + + # {{{ prefetch + # }}} + + # {{{ tile + # }}} + + # FIXME: remove this (eventually) + knl = lp.set_options(knl, insert_gbarriers=True) + t_unit = t_unit.with_kernel(knl) + self.dev_code = lp.generate_code_v2(t_unit).device_code() + return super().transform_loopy_program(t_unit) def main(): - order = 4 + order = 1 ctx = cl.create_some_context() queue = cl.CommandQueue(ctx) actx = PytatoTensorProductArrayContext(queue) dim = 3 - res = 5 + res = 2 from meshmode.mesh import TensorProductElementGroup from meshmode.discretization.poly_element import \ @@ -49,7 +75,7 @@ def main(): group_cls=TensorProductElementGroup) import grudge.dof_desc as dd - dcoll = DiscretizationCollection( + dcoll = make_discretization_collection( actx, mesh, discr_tag_to_group_factory={ @@ -69,12 +95,7 @@ def f(x): grad_u = op.local_grad(dcoll, u) grad_u = actx.np.stack(grad_u)[0] - - prg = pt.generate_loopy(grad_u).program - code = lp.generate_code_v2(prg).device_code() - - print(code) - pu.db + pt.show_dot_graph(grad_u) if __name__ == "__main__": main() diff --git a/grudge/op.py b/grudge/op.py index b60ab62d7..38df02418 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -242,13 +242,15 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): # apply geometric factors # TODO: chain the einsum above with the einsum below + from arraycontext.metadata import NameHint grad = actx.np.stack([grad[i] for i in range(dim)]) grad = actx.einsum( "xrei,xei->xei", ijm, grad, arg_names=("inv_jac_t", "vec"), - tagged=(FirstAxisIsElementsTag(),)) + tagged=(FirstAxisIsElementsTag(), + NameHint("tp_gradient"),)) return grad @@ -383,7 +385,14 @@ def get_ref_derivative_mats(grp): vdm_1d = mp.vandermonde(bases_1d.functions, nodes1d) vdm_p_1d = mp.vandermonde(bases_1d.gradients, nodes1d)[0] - return actx.freeze(actx.from_numpy(vdm_p_1d @ la.inv(vdm_1d))) + diff_mat = actx.from_numpy(vdm_p_1d @ la.inv(vdm_1d)) + + from arraycontext.metadata import NameHint + return actx.freeze( + actx.tag(NameHint("tp_diff_mat_1d"), + tag_axes(actx, { + 1: DiscretizationDOFAxisTag()}, + diff_mat))) else: from meshmode.discretization.poly_element import diff_matrices From 92991a3d0524021333d2d4089a50e53df39340e9 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sat, 30 Sep 2023 13:28:43 -0500 Subject: [PATCH 47/97] Version 0.1 of weak gradient computation --- grudge/op.py | 158 +++++++++++++++++++++++++++++++++++++----------- test/test_op.py | 4 +- 2 files changed, 126 insertions(+), 36 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 38df02418..062c88101 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -207,12 +207,11 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. - - def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): + def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, + metric_in_matvec): """Exploits tensor product structure to differentiate each coordinate axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) """ - from modepy.tools import ( reshape_array_for_tensor_product_space, unreshape_array_for_tensor_product_space) @@ -223,41 +222,122 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm): # apply operators to function data dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) - grad = make_obj_array([ - actx.einsum( - f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", - diff_mat, - vec, - arg_names=("diff_mat", "vec"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - for i in range(dim) + + # weak form case: + # 3D weak_x: einsum("estu,ps,qt,ru->epqr", + # f, stiff_1D, mass_1D, mass_1D) + if metric_in_matvec: + stiff_1D, mass_1D = diff_mat + + if dim == 3: + weak_x = actx.einsum( + "estu,ps,qt,ru->epqr", + vec, + stiff_1D, + mass_1D, + mass_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + weak_y = actx.einsum( + "estu,ps,qt,ru->epqr", + vec, + mass_1D, + stiff_1D, + mass_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + weak_z = actx.einsum( + "estu,ps,qt,ru->epqr", + vec, + mass_1D, + mass_1D, + stiff_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + grad = make_obj_array([ + weak_x, + weak_y, + weak_z + ]) + + elif dim == 2: + weak_x = actx.einsum( + "est,ps,qt->epq", + vec, + stiff_1D, + mass_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + weak_y = actx.einsum( + "est,ps,qt->epq", + vec, + mass_1D, + stiff_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + grad = make_obj_array([ + weak_x, + weak_y + ]) + + # strong form case: + # x partial: einsum("il,eljk->eijk", D, f) + else: + grad = make_obj_array([ + actx.einsum( + f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", + diff_mat, + vec, + arg_names=("diff_mat", "vec"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + for i in range(dim) ]) # unreshape grad to apply geometric factors grad = make_obj_array([ unreshape_array_for_tensor_product_space(grp.space, grad[i]) for i in range(grp.dim) - ]) + ]) - # apply geometric factors - # TODO: chain the einsum above with the einsum below + # apply geometric factors in strong case from arraycontext.metadata import NameHint - grad = actx.np.stack([grad[i] for i in range(dim)]) - grad = actx.einsum( - "xrei,xei->xei", - ijm, - grad, - arg_names=("inv_jac_t", "vec"), - tagged=(FirstAxisIsElementsTag(), - NameHint("tp_gradient"),)) + if metric_in_matvec: + grad = make_obj_array([ + actx.einsum( + "rei,ei->ei", + ijm[i], + grad[i], + arg_names=("inv_jac_t", "vec"), + tagged=FirstAxisIsElementsTag()) + for i in range(dim) + ]) + else: + grad = actx.np.stack([grad[i] for i in range(dim)]) + grad = actx.einsum( + "xrei,xei->xei", + ijm, + grad, + arg_names=("inv_jac_t", "vec"), + tagged=(FirstAxisIsElementsTag(), + NameHint("tp_gradient"),)) return grad - per_group_grads = [ - compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i) + compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i, + metric_in_matvec) if isinstance(in_grp, TensorProductElementGroupBase) # r for rst axis @@ -289,7 +369,6 @@ def _divergence_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. - def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): """Exploits tensor product structure to differentiate each coordinate axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) @@ -571,6 +650,8 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): from meshmode.discretization.poly_element import \ mass_matrix, diff_matrices + # {{{ tensor product case + if isinstance(out_grp, TensorProductElementGroupBase): import modepy as mp import numpy.linalg as la @@ -581,15 +662,24 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): vdm = mp.vandermonde(basis_1d.functions, nodes_1d) vdm_p = mp.vandermonde(basis_1d.gradients, nodes_1d)[0] - # NOTE: possibly work special-case matrices like differentiation - # matrix, mass matrix, into modepy - mmat = la.inv(vdm @ vdm.T) - diff_mat = vdm_p @ la.inv(vdm) - return actx.freeze( + mass_1D = la.inv(vdm @ vdm.T) + diff_mat = la.solve(vdm.T, vdm_p.T).T + + stiff_1D = actx.freeze( actx.tag_axis(1, DiscretizationDOFAxisTag(), - actx.from_numpy( - np.asarray( - diff_mat.T @ mmat.T)))) + actx.from_numpy( + np.asarray( + diff_mat.T @ mass_1D.T)))) + + mass_1D = actx.freeze( + actx.tag_axis(1, DiscretizationDOFAxisTag(), + actx.from_numpy( + np.asarray( + mass_1D)))) + + return (stiff_1D, mass_1D) + + # }}} mmat = mass_matrix(out_grp) diff --git a/test/test_op.py b/test/test_op.py index 05f6c1f99..7d26280da 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -161,7 +161,7 @@ def get_flux(u_tpair): or eoc_rec.max_error() < 1e-11) -@pytest.mark.parametrize("form", ["strong"]) +@pytest.mark.parametrize("form", ["weak"]) @pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize(("vectorize", "nested"), [ @@ -308,7 +308,7 @@ def get_flux(u_tpair): (True, True) ]) def test_divergence(actx_factory, form, dim, order, vectorize, nested, - visualize=False): + visualize=False): from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() From 49116ab388a4245857066fc5e9d078bc1f164b96 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sat, 30 Sep 2023 15:27:44 -0500 Subject: [PATCH 48/97] Weak form divergence version 0.1 --- grudge/op.py | 141 +++++++++++++++++++++++++++++++++++++++--------- test/test_op.py | 6 ++- 2 files changed, 120 insertions(+), 27 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 062c88101..8539f305a 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -381,37 +381,128 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): # reshape u to expose tensor product structure vec = reshape_array_for_tensor_product_space(grp.space, vec) - # apply differentiation matrix to function data dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) - partials = make_obj_array([ - actx.einsum( - f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", - diff_mat, - vec[i], - arg_names=("diff_mat", "vec"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) + + # weak form + if metric_in_matvec: + stiff_1D, mass_1D = diff_mat + + if dim == 3: + weak_x = actx.einsum( + "estu,ps,qt,ru->epqr", + vec[0], + stiff_1D, + mass_1D, + mass_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + weak_y = actx.einsum( + "estu,ps,qt,ru->epqr", + vec[1], + mass_1D, + stiff_1D, + mass_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + weak_z = actx.einsum( + "estu,ps,qt,ru->epqr", + vec[2], + mass_1D, + mass_1D, + stiff_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + partials = make_obj_array([ + weak_x, weak_y, weak_z + ]) + + elif dim == 2: + weak_x = actx.einsum( + "est,ps,qt->epq", + vec[0], + stiff_1D, + mass_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + weak_y = actx.einsum( + "est,ps,qt->epq", + vec[1], + mass_1D, + stiff_1D, + arg_names=("vec", "stiff_1D_r", "mass_1D_s"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + + partials = make_obj_array([ + weak_x, weak_y + ]) + + else: + raise Exception("Dimensions of 2 and 3 are supported by " + "tensor product elements. Found dim = {dim}") + + + partials = make_obj_array([ + unreshape_array_for_tensor_product_space(grp.space, partials[i]) for i in range(dim) - ]) + ]) - # unreshape partials to apply geometric factors - # TODO: chain the einsum above with the einsum below - partials = make_obj_array([ - unreshape_array_for_tensor_product_space(grp.space, partials[i]) - for i in range(partials.shape[0]) - ]) + partials = actx.np.stack(partials) - # apply geometric factors - partials = actx.np.stack([partials[i] for i in range(dim)]) - div = actx.einsum( - "xrei,xei->ei", - ijm, - partials, - arg_names=("inv_jac_t", "vec"), - tagged=(FirstAxisIsElementsTag(),)) + div = make_obj_array([ + actx.einsum("rei,ei->ei", + ijm[i], + partials[i], + arg_names=("inv_jac_t", "vec"), + tagged=(FirstAxisIsElementsTag(),)) + for i in range(dim) + ]) + + ret = 0 + for i in range(dim): + ret += div[i] + return ret + + # strong form + else: + partials = make_obj_array([ + actx.einsum( + f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", + diff_mat, + vec[i], + arg_names=("diff_mat", "vec"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + for i in range(dim) + ]) + + # unreshape partials to apply geometric factors + # TODO: chain the einsum above with the einsum below + partials = make_obj_array([ + unreshape_array_for_tensor_product_space(grp.space, partials[i]) + for i in range(partials.shape[0]) + ]) + + # apply geometric factors + partials = actx.np.stack([partials[i] for i in range(dim)]) + + div = actx.einsum( + "xrei,xei->ei", + ijm, + partials, + arg_names=("inv_jac_t", "vec"), + tagged=(FirstAxisIsElementsTag(),)) - return div + return div per_group_divs = [ diff --git a/test/test_op.py b/test/test_op.py index 7d26280da..a5d731533 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -161,7 +161,7 @@ def get_flux(u_tpair): or eoc_rec.max_error() < 1e-11) -@pytest.mark.parametrize("form", ["weak"]) +@pytest.mark.parametrize("form", ["strong", "weak"]) @pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize(("vectorize", "nested"), [ @@ -309,6 +309,8 @@ def get_flux(u_tpair): ]) def test_divergence(actx_factory, form, dim, order, vectorize, nested, visualize=False): + actx = actx_factory() + from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() @@ -409,7 +411,7 @@ def get_flux(u_tpair): or eoc_rec.max_error() < 1e-11) -@pytest.mark.parametrize("form", ["strong"]) +@pytest.mark.parametrize("form", ["strong", "weak"]) @pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize(("vectorize", "nested"), [ From d87c19fee3af3e3d5958486bb27a8d8d0ee21f68 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sun, 1 Oct 2023 18:28:55 -0500 Subject: [PATCH 49/97] Move TP array contexts. Add acoustic pulse TP example --- .../tensor-product-examples/acoustic_pulse.py | 264 ++++++++++++++++++ examples/tp-transform-cartoon.py | 2 +- grudge/array_context.py | 121 +++++--- 3 files changed, 347 insertions(+), 40 deletions(-) create mode 100644 examples/tensor-product-examples/acoustic_pulse.py diff --git a/examples/tensor-product-examples/acoustic_pulse.py b/examples/tensor-product-examples/acoustic_pulse.py new file mode 100644 index 000000000..13c2194cf --- /dev/null +++ b/examples/tensor-product-examples/acoustic_pulse.py @@ -0,0 +1,264 @@ +__copyright__ = """ +Copyright (C) 2021 University of Illinois Board of Trustees +""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +from meshmode.mesh import TensorProductElementGroup +import numpy as np + +import pyopencl as cl +import pyopencl.tools as cl_tools + +from grudge.array_context import ( + PyOpenCLArrayContext, + PytatoPyOpenCLArrayContext +) +from grudge.models.euler import ( + ConservedEulerField, + EulerOperator, + InviscidWallBC +) +from grudge.shortcuts import rk4_step + +from meshmode.mesh import BTAG_ALL + +from pytools.obj_array import make_obj_array + +import grudge.op as op + +import logging +logger = logging.getLogger(__name__) + + +def gaussian_profile( + x_vec, t=0, rho0=1.0, rhoamp=1.0, p0=1.0, gamma=1.4, + center=None, velocity=None): + + dim = len(x_vec) + if center is None: + center = np.zeros(shape=(dim,)) + if velocity is None: + velocity = np.zeros(shape=(dim,)) + + lump_loc = center + t * velocity + + # coordinates relative to lump center + rel_center = make_obj_array( + [x_vec[i] - lump_loc[i] for i in range(dim)] + ) + actx = x_vec[0].array_context + r = actx.np.sqrt(np.dot(rel_center, rel_center)) + expterm = rhoamp * actx.np.exp(1 - r ** 2) + + mass = expterm + rho0 + mom = velocity * mass + energy = (p0 / (gamma - 1.0)) + np.dot(mom, mom) / (2.0 * mass) + + return ConservedEulerField(mass=mass, energy=energy, momentum=mom) + + +def make_pulse(amplitude, r0, w, r): + dim = len(r) + r_0 = np.zeros(dim) + r_0 = r_0 + r0 + rel_center = make_obj_array( + [r[i] - r_0[i] for i in range(dim)] + ) + actx = r[0].array_context + rms2 = w * w + r2 = np.dot(rel_center, rel_center) / rms2 + return amplitude * actx.np.exp(-.5 * r2) + + +def acoustic_pulse_condition(x_vec, t=0): + dim = len(x_vec) + vel = np.zeros(shape=(dim,)) + orig = np.zeros(shape=(dim,)) + uniform_gaussian = gaussian_profile( + x_vec, t=t, center=orig, velocity=vel, rhoamp=0.0) + + amplitude = 1.0 + width = 0.1 + pulse = make_pulse(amplitude, orig, width, x_vec) + + return ConservedEulerField( + mass=uniform_gaussian.mass, + energy=uniform_gaussian.energy + pulse, + momentum=uniform_gaussian.momentum + ) + + +def run_acoustic_pulse(actx, + order=3, + final_time=1, + resolution=4, + overintegration=False, + visualize=False): + + # eos-related parameters + gamma = 1.4 + + # {{{ discretization + + from meshmode.mesh.generation import generate_regular_rect_mesh + + dim = 3 + box_ll = -0.5 + box_ur = 0.5 + mesh = generate_regular_rect_mesh( + a=(box_ll,)*dim, + b=(box_ur,)*dim, + nelements_per_axis=(resolution,)*dim, + group_cls=TensorProductElementGroup) + + from grudge import DiscretizationCollection + from grudge.dof_desc import DISCR_TAG_BASE, DISCR_TAG_QUAD + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + + exp_name = f"fld-acoustic-pulse-N{order}-K{resolution}" + if overintegration: + exp_name += "-overintegrated" + quad_tag = DISCR_TAG_QUAD + else: + quad_tag = None + + dcoll = DiscretizationCollection( + actx, mesh, + discr_tag_to_group_factory={ + DISCR_TAG_BASE: LGL(order) + } + ) + + # }}} + + # {{{ Euler operator + + euler_operator = EulerOperator( + dcoll, + bdry_conditions={BTAG_ALL: InviscidWallBC()}, + flux_type="lf", + gamma=gamma, + quadrature_tag=quad_tag + ) + + def rhs(t, q): + return euler_operator.operator(t, q) + + compiled_rhs = actx.compile(rhs) + + from grudge.dt_utils import h_min_from_volume + + cfl = 0.125 + cn = 0.5*(order + 1)**2 + dt = cfl * actx.to_numpy(h_min_from_volume(dcoll)) / cn + + fields = acoustic_pulse_condition(actx.thaw(dcoll.nodes())) + + logger.info("Timestep size: %g", dt) + + # }}} + + from grudge.shortcuts import make_visualizer + + vis = make_visualizer(dcoll) + + # {{{ time stepping + + step = 0 + t = 0.0 + while t < final_time: + if step % 10 == 0: + norm_q = actx.to_numpy(op.norm(dcoll, fields, 2)) + logger.info("[%04d] t = %.5f |q| = %.5e", step, t, norm_q) + if visualize: + vis.write_vtk_file( + f"{exp_name}-{step:04d}.vtu", + [ + ("rho", fields.mass), + ("energy", fields.energy), + ("momentum", fields.momentum) + ] + ) + assert norm_q < 5 + + fields = actx.thaw(actx.freeze(fields)) + fields = rk4_step(fields, t, dt, compiled_rhs) + t += dt + step += 1 + + # }}} + + +def main(ctx_factory, order=3, final_time=1, resolution=16, + overintegration=False, visualize=False, lazy=False): + cl_ctx = ctx_factory() + queue = cl.CommandQueue(cl_ctx) + + if lazy: + from grudge.array_context import PytatoTensorProductArrayContext + actx = PytatoTensorProductArrayContext( + queue, + allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)), + ) + else: + from grudge.array_context import TensorProductArrayContext + actx = TensorProductArrayContext( + queue, + allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)), + force_device_scalars=True, + ) + + run_acoustic_pulse( + actx, + order=order, + resolution=resolution, + overintegration=overintegration, + final_time=final_time, + visualize=visualize + ) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--order", default=3, type=int) + parser.add_argument("--tfinal", default=0.1, type=float) + parser.add_argument("--resolution", default=16, type=int) + parser.add_argument("--oi", action="store_true", + help="use overintegration") + parser.add_argument("--visualize", action="store_true", + help="write out vtk output") + parser.add_argument("--lazy", action="store_true", + help="switch to a lazy computation mode") + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + main(cl.create_some_context, + order=args.order, + final_time=args.tfinal, + resolution=args.resolution, + overintegration=args.oi, + visualize=args.visualize, + lazy=args.lazy) diff --git a/examples/tp-transform-cartoon.py b/examples/tp-transform-cartoon.py index b0bb13cfd..f26a58737 100644 --- a/examples/tp-transform-cartoon.py +++ b/examples/tp-transform-cartoon.py @@ -62,7 +62,7 @@ def main(): queue = cl.CommandQueue(ctx) actx = PytatoTensorProductArrayContext(queue) - dim = 3 + dim = 2 res = 2 from meshmode.mesh import TensorProductElementGroup diff --git a/grudge/array_context.py b/grudge/array_context.py index f29e8ef5d..4df7df0e2 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -127,45 +127,6 @@ def __init__(self, queue: "pyopencl.CommandQueue", # }}} -# {{{ Tensor product array context - -class OutputIsTensorProductDOFArrayOrdered(Tag): - """Signify that the strides will not be of order "C" or "F". See - :class:`grudge.array_context.TensorProductArrayContext` for more details. - """ - pass - - -class TensorProductArrayContext(_PyOpenCLArrayContextBase): - """Specialized array context for use with tensor product elements. - - The strides for the arrays containing tensor product element data are of the - form (slow, fastest, faster, fast). These strides are not "C" or "F" order. - Hence, this specialized array context takes care of specifying the - particular strides required. - """ - - def transform_loopy_program(self, t_unit): - if len(t_unit.callables_table) == 1: - knl = t_unit.default_entrypoint - if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): - new_args = [] - for arg in knl.args: - if arg.is_output: - arg = arg.copy(dim_tags=( - f"N{len(arg.shape)-1}," - + ",".join(f"N{i}" - for i in range(len(arg.shape)-1)) - )) - - new_args.append(arg) - - knl = knl.copy(args=new_args) - t_unit = t_unit.with_kernel(knl) - - return super().transform_loopy_program(t_unit) -# }}} - # {{{ pytato @@ -631,4 +592,86 @@ def get_reasonable_array_context_class( # }}} +# {{{ Tensor product array context + +# {{{ Relevant tags +class OutputIsTensorProductDOFArrayOrdered(Tag): + """Signify that the strides will not be of order "C" or "F". See + :class:`grudge.array_context.TensorProductArrayContext` for more details. + """ + pass +# }}} + +# {{{ Eager TP array context +class TensorProductArrayContext(_PyOpenCLArrayContextBase): + """Specialized array context for use with tensor product elements. + + The strides for the arrays containing tensor product element data are of the + form (slow, fastest, faster, fast). These strides are not "C" or "F" order. + Hence, this specialized array context takes care of specifying the + particular strides required. + """ + + def transform_loopy_program(self, t_unit): + if len(t_unit.callables_table) == 1: + knl = t_unit.default_entrypoint + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) + + return super().transform_loopy_program(t_unit) +# }}} + +# {{{ Lazy tensor product array context +class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext): + def transform_dag(self, dag): + return super().transform_dag(dag) + + def transform_loopy_program(self, t_unit): + knl = t_unit.default_entrypoint + + # {{{ adjust strides according to tensor product structure + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + # }}} + + # {{{ prefetch + # }}} + + # {{{ tile + # }}} + + import loopy as lp + # FIXME: remove this (eventually) + knl = lp.set_options(knl, insert_gbarriers=True) + t_unit = t_unit.with_kernel(knl) + self.dev_code = lp.generate_code_v2(t_unit).device_code() + + return super().transform_loopy_program(t_unit) +# }}} + +# }}} + # vim: foldmethod=marker From c294df7fba0e1227aec8e156ea8264e1beaaeb69 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Wed, 25 Oct 2023 09:40:35 -0500 Subject: [PATCH 50/97] Add baseline MPI version of TP actx --- grudge/array_context.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index b65901595..b361e83d3 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -628,9 +628,11 @@ def clone(self): pass # }}} -# {{{ Tensor product array context +# {{{ Tensor product array contexts # {{{ Relevant tags + + class OutputIsTensorProductDOFArrayOrdered(Tag): """Signify that the strides will not be of order "C" or "F". See :class:`grudge.array_context.TensorProductArrayContext` for more details. @@ -639,6 +641,8 @@ class OutputIsTensorProductDOFArrayOrdered(Tag): # }}} # {{{ Eager TP array context + + class TensorProductArrayContext(_PyOpenCLArrayContextBase): """Specialized array context for use with tensor product elements. @@ -667,6 +671,12 @@ def transform_loopy_program(self, t_unit): t_unit = t_unit.with_kernel(knl) return super().transform_loopy_program(t_unit) + + +# {{{ Distributed eager tensor product array context +class TensorProductMPIPyOpenCLArrayContext(MPIPyOpenCLArrayContext, + TensorProductArrayContext): + pass # }}} # {{{ Lazy tensor product array context From a9dbb5a5df7755de9fd0efd9c0432731e35febb8 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 26 Oct 2023 14:51:12 -0500 Subject: [PATCH 51/97] Initial lazy array context for TP elements --- grudge/array_context.py | 83 +++++++++++++++++++++++++++-------------- 1 file changed, 56 insertions(+), 27 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index b361e83d3..320ae1d30 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -628,20 +628,21 @@ def clone(self): pass # }}} + # {{{ Tensor product array contexts # {{{ Relevant tags - class OutputIsTensorProductDOFArrayOrdered(Tag): """Signify that the strides will not be of order "C" or "F". See :class:`grudge.array_context.TensorProductArrayContext` for more details. """ pass + # }}} -# {{{ Eager TP array context +# {{{ Eager TP array contexts class TensorProductArrayContext(_PyOpenCLArrayContextBase): """Specialized array context for use with tensor product elements. @@ -653,33 +654,35 @@ class TensorProductArrayContext(_PyOpenCLArrayContextBase): """ def transform_loopy_program(self, t_unit): - if len(t_unit.callables_table) == 1: - knl = t_unit.default_entrypoint - if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): - new_args = [] - for arg in knl.args: - if arg.is_output: - arg = arg.copy(dim_tags=( - f"N{len(arg.shape)-1}," - + ",".join(f"N{i}" - for i in range(len(arg.shape)-1)) - )) + #if len(t_unit.callables_table) == 1: + knl = t_unit.default_entrypoint + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) - new_args.append(arg) + new_args.append(arg) - knl = knl.copy(args=new_args) - t_unit = t_unit.with_kernel(knl) + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) return super().transform_loopy_program(t_unit) -# {{{ Distributed eager tensor product array context class TensorProductMPIPyOpenCLArrayContext(MPIPyOpenCLArrayContext, TensorProductArrayContext): pass + # }}} -# {{{ Lazy tensor product array context + +# {{{ Lazy tensor product array contexts + class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext): def transform_dag(self, dag): return super().transform_dag(dag) @@ -703,21 +706,47 @@ def transform_loopy_program(self, t_unit): knl = knl.copy(args=new_args) # }}} - # {{{ prefetch - # }}} + return super().transform_loopy_program(t_unit) - # {{{ tile - # }}} +# }}} + + +# {{{ TP fusion actx + +from meshmode.array_context import FusionContractorArrayContext + + +class TensorProductFusionContractorArrayContext(FusionContractorArrayContext): + + def transform_loopy_program(self, t_unit): + if len(t_unit.callables_table) == 1: + knl = t_unit.default_entrypoint + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) - import loopy as lp - # FIXME: remove this (eventually) - knl = lp.set_options(knl, insert_gbarriers=True) - t_unit = t_unit.with_kernel(knl) - self.dev_code = lp.generate_code_v2(t_unit).device_code() + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) return super().transform_loopy_program(t_unit) + + +class TensorProductMPIFusionContractorArrayContext( + MPIPytatoArrayContextBase, TensorProductFusionContractorArrayContext): + pass + # }}} + # }}} + # vim: foldmethod=marker From 7beb2e8228d72c9ffcab6390892b1599c3b7f122 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 2 Nov 2023 23:37:09 -0500 Subject: [PATCH 52/97] Fix up some array context business. Make lazy TPMPIArrayContext a subclass of MPIPytatoArrayContext --- grudge/array_context.py | 34 ++++++++++++++++------------------ 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index 320ae1d30..c1964ce91 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -654,7 +654,6 @@ class TensorProductArrayContext(_PyOpenCLArrayContextBase): """ def transform_loopy_program(self, t_unit): - #if len(t_unit.callables_table) == 1: knl = t_unit.default_entrypoint if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): new_args = [] @@ -719,28 +718,27 @@ def transform_loopy_program(self, t_unit): class TensorProductFusionContractorArrayContext(FusionContractorArrayContext): def transform_loopy_program(self, t_unit): - if len(t_unit.callables_table) == 1: - knl = t_unit.default_entrypoint - if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): - new_args = [] - for arg in knl.args: - if arg.is_output: - arg = arg.copy(dim_tags=( - f"N{len(arg.shape)-1}," - + ",".join(f"N{i}" - for i in range(len(arg.shape)-1)) - )) - - new_args.append(arg) - - knl = knl.copy(args=new_args) - t_unit = t_unit.with_kernel(knl) + knl = t_unit.default_entrypoint + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) return super().transform_loopy_program(t_unit) class TensorProductMPIFusionContractorArrayContext( - MPIPytatoArrayContextBase, TensorProductFusionContractorArrayContext): + MPIPytatoArrayContext, TensorProductFusionContractorArrayContext): pass # }}} From fb42a02980966f5ddd3783775e6b1c889f32e61f Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 21 Dec 2023 09:43:46 -0600 Subject: [PATCH 53/97] Add fixed up einsums for TP --- grudge/op.py | 484 +++++++++++++++++++++++++-------------------------- 1 file changed, 235 insertions(+), 249 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index f913ebf38..aebfc465c 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -74,17 +74,20 @@ from functools import partial -from meshmode.dof_array import DOFArray +from meshmode.dof_array import DOFArray, warn from meshmode.transform_metadata import (FirstAxisIsElementsTag, DiscretizationDOFAxisTag, DiscretizationElementAxisTag, DiscretizationFaceAxisTag) from meshmode.discretization.poly_element import TensorProductElementGroupBase +from modepy.tools import ( + reshape_array_for_tensor_product_space as fold, + unreshape_array_for_tensor_product_space as unfold) + from grudge.discretization import DiscretizationCollection from grudge.dof_desc import as_dofdesc from grudge.array_context import ( - TensorProductArrayContext, OutputIsTensorProductDOFArrayOrdered) from pytools import keyed_memoize_in @@ -99,7 +102,7 @@ ) from grudge.interpolation import interp -from grudge.projection import project, volume_quadrature_project +from grudge.projection import project from grudge.reductions import ( norm, @@ -122,11 +125,8 @@ interior_trace_pair, interior_trace_pairs, local_interior_trace_pair, - connected_parts, - inter_volume_trace_pairs, - local_inter_volume_trace_pairs, + #connected_ranks, cross_rank_trace_pairs, - cross_rank_inter_volume_trace_pairs, bdry_trace_pair, bv_trace_pair ) @@ -134,7 +134,6 @@ __all__ = ( "project", - "volume_quadrature_project", "interp", "norm", @@ -155,11 +154,8 @@ "interior_trace_pair", "interior_trace_pairs", "local_interior_trace_pair", - "connected_parts", - "inter_volume_trace_pairs", - "local_inter_volume_trace_pairs", + "connected_ranks", "cross_rank_trace_pairs", - "cross_rank_inter_volume_trace_pairs", "bdry_trace_pair", "bv_trace_pair", @@ -190,20 +186,51 @@ def _single_axis_derivative_kernel( # - whether the chain rule terms ("inv_jac_mat") sit outside (strong) # or inside (weak) the matrix-vector product that carries out the # derivative, cf. "metric_in_matvec". + + + # {{{ tensor product single axis derivative + + # FIXME: actually implement single axis tensor product derivatives + def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, + xyz_axis, metric_in_matvec): + + + return compute_simplicial_derivative(actx, grp, grp, get_diff_mat, vec, + ijm, xyz_axis, metric_in_matvec) + + # }}} + + + # {{{ simplicial single axis derivative + + def compute_simplicial_derivative(actx, in_grp, out_grp, + get_diff_mat, vec_i, ijm_i, + xyz_axis, metric_in_matvec): + # r for rst axis + return actx.einsum( + "rej,rij,ej->ei" if metric_in_matvec else "rei,rij,ej->ei", + ijm_i[xyz_axis], + get_diff_mat( + actx, + out_element_group=out_grp, + in_element_group=in_grp), + vec_i, + arg_names=("inv_jac_t", "ref_stiffT_mat", "vec", ), + tagged=(FirstAxisIsElementsTag(),)) + + # }}} + + return DOFArray( actx, data=tuple( - # r for rst axis - actx.einsum("rej,rij,ej->ei" if metric_in_matvec else "rei,rij,ej->ei", - ijm_i[xyz_axis], - get_diff_mat( - actx, - out_element_group=out_grp, - in_element_group=in_grp), - vec_i, - arg_names=("inv_jac_t", "ref_stiffT_mat", "vec", ), - tagged=(FirstAxisIsElementsTag(),)) - + compute_tensor_product_derivative(actx, in_grp, out_grp, + get_diff_mat, vec_i, ijm_i, + xyz_axis, metric_in_matvec) + if isinstance(in_grp, TensorProductElementGroupBase) + else compute_simplicial_derivative(actx, in_grp, out_grp, + get_diff_mat, vec_i, ijm_i, + xyz_axis, metric_in_matvec) for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, inv_jac_mat))) @@ -214,142 +241,116 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. + + # {{{ tensor product gradient + def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, metric_in_matvec): - """Exploits tensor product structure to differentiate each coordinate - axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) + # TODO: add note about inverse mass simplification, point to + # op.inverse_mass (assuming this is where the explanation will live) """ - from modepy.tools import ( - reshape_array_for_tensor_product_space, - unreshape_array_for_tensor_product_space) + Exploits tensor product structure to reduce complexity. Applies a + differentiation operator containing 1D information to a tensor of DOF + data. For example, in the 2D strong form case, this computes partial + derivatives in a similar manner to - # reshape u to expose tensor product structure - vec = reshape_array_for_tensor_product_space(grp.space, vec) + .. math:: - # apply operators to function data - dim = grp.dim - diff_mat = get_diff_mat(actx, grp, grp) + \partial_x \mathbf{f}_{ij} = \sum_{\ell} \mathbf{J}^e_{ij} + \mathbf{D}_{i\ell} \mathbf{f}_{\ell j} - # weak form case: - # 3D weak_x: einsum("estu,ps,qt,ru->epqr", - # f, stiff_1D, mass_1D, mass_1D) - if metric_in_matvec: - stiff_1D, mass_1D = diff_mat + where $\mathbf{D}$ is a 1D differentiation operator, $\mathbf{f}$ is a + vector of function data, $\mathbf{J}^e$ is the element Jacobian matrix. + The weak form uses a 1D element mass operator and a 1D element stiffness + operator to perform the contraction - if dim == 3: - weak_x = actx.einsum( - "estu,ps,qt,ru->epqr", - vec, - stiff_1D, - mass_1D, - mass_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - - weak_y = actx.einsum( - "estu,ps,qt,ru->epqr", - vec, - mass_1D, - stiff_1D, - mass_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) + .. math:: - weak_z = actx.einsum( - "estu,ps,qt,ru->epqr", - vec, - mass_1D, - mass_1D, - stiff_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) + \partial_x \mathbf{f}_{ij} = \sum_{\ell,b} \mathbf{J}^e_{\ell b} + \mathbf{f}_{\ell b} \mathbf{S}^e_{i\ell} \mathbf{M}^e_{jb} + """ - grad = make_obj_array([ - weak_x, - weak_y, - weak_z - ]) - elif dim == 2: - weak_x = actx.einsum( - "est,ps,qt->epq", - vec, - stiff_1D, - mass_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) + if grp.dim > 3 and metric_in_matvec: + warn('Efficient tensor product weak ' + 'differentiation operators only ' + 'implemented for dimension 2 and 3. ' + 'Defaulting to inefficient version.') + return compute_simplicial_grad(actx, grp, grp, diff_mat, vec, ijm, + metric_in_matvec) - weak_y = actx.einsum( - "est,ps,qt->epq", - vec, - mass_1D, - stiff_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) + # reshape u to expose tensor product structure + vec = fold(grp.space, vec) + diff_mat = get_diff_mat(actx, grp, grp) - grad = make_obj_array([ - weak_x, - weak_y - ]) + # weak form case: + # 3D weak_x: einsum("estu,ps,qt,ru->epqr", + # f, stiff_1D, mass_1D, mass_1D) + # TODO:? make this more general, maybe offload to a function that + # generates argnames and einsum specs + if metric_in_matvec: + stiff_1D, mass_1D = diff_mat + grad = make_obj_array([ + actx.einsum( + f"e{'bd'[:i]}j{'bd'[i:grp.dim-1]}," + + "ij," + + ("ab,cd" if grp.dim == 3 else "ab") + + "->" + f"e{'ac'[:i]}i{'ac'[i:grp.dim-1]}", + vec, + stiff_1D, + *(mass_1D,)*(grp.dim-1), + arg_names=("vec", "stiff_1D", + *(("mass_1D_1", "mass_1D_2")[:grp.dim-1])), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + for i in range(grp.dim) + ]) - # strong form case: + # Carries out, e.g., 3D strong form contraction # x partial: einsum("il,eljk->eijk", D, f) else: grad = make_obj_array([ - actx.einsum( - f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", - diff_mat, - vec, - arg_names=("diff_mat", "vec"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - for i in range(dim) + actx.einsum( + "yz," + + f"e{'abcdfghijkl'[:i]}z{'mnopqstuvwx'[:grp.dim-i-1]}->" + + f"e{'abcdfghijkl'[:i]}y{'mnopqstuvwx'[:grp.dim-i-1]}", + diff_mat, + vec, + arg_names=("diff_mat", "vec"), + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered())) + for i in range(grp.dim) ]) - # unreshape grad to apply geometric factors - grad = make_obj_array([ - unreshape_array_for_tensor_product_space(grp.space, grad[i]) - for i in range(grp.dim) + # {{{ unreshape grad and apply geometric factors + + # TODO: Chain einsums together with geometric factors + grad = actx.np.stack([ + unfold(grp.space, grad[rst_axis]) + for rst_axis in range(grp.dim) ]) - # apply geometric factors in strong case - from arraycontext.metadata import NameHint - if metric_in_matvec: - grad = make_obj_array([ - actx.einsum( - "rei,ei->ei", - ijm[i], - grad[i], - arg_names=("inv_jac_t", "vec"), - tagged=FirstAxisIsElementsTag()) - for i in range(dim) - ]) - else: - grad = actx.np.stack([grad[i] for i in range(dim)]) - grad = actx.einsum( - "xrei,xei->xei", - ijm, - grad, - arg_names=("inv_jac_t", "vec"), - tagged=(FirstAxisIsElementsTag(), - NameHint("tp_gradient"),)) + grad = actx.einsum( + "xrej,rej->xej", + ijm, + grad, + arg_names=("inv_jac_mat", "grad"), + tagged=(FirstAxisIsElementsTag(),) + ) + + # }}} return grad - per_group_grads = [ + # }}} - compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i, - metric_in_matvec) - if isinstance(in_grp, TensorProductElementGroupBase) - # r for rst axis - # x for xyz axis - else actx.einsum( + # {{{ simplicial grad + + def compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec): + return actx.einsum( "xrej,rij,ej->xei" if metric_in_matvec else "xrei,rij,ej->xei", ijm_i, get_diff_mat( @@ -361,9 +362,20 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), tagged=(FirstAxisIsElementsTag(),)) + # }}} + + + per_group_grads = [ + compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i, + metric_in_matvec) + if isinstance(in_grp, TensorProductElementGroupBase) + else compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec) + for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, - inv_jac_mat)] + inv_jac_mat) + ] return make_obj_array([ DOFArray( @@ -376,150 +388,108 @@ def _divergence_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. + + # {{{ tensor product div + def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): - """Exploits tensor product structure to differentiate each coordinate - axis using a single differentiation matrix of shape (nnodes1d, nnodes1d) + """ + Exploits tensor product structure to reduce complexity. See + `_gradient_kernel.compute_tensor_product_grad` for more details. """ - from modepy.tools import ( - reshape_array_for_tensor_product_space, - unreshape_array_for_tensor_product_space) + if grp.dim > 3 and metric_in_matvec: + warn('Efficient tensor product weak ' + 'differentiation operators only ' + 'implemented for dimension 2 and 3. ' + 'Defaulting to inefficient version.') + return compute_simplicial_div(actx, grp, grp, diff_mat, vec, ijm, + metric_in_matvec) # reshape u to expose tensor product structure - vec = reshape_array_for_tensor_product_space(grp.space, vec) - - dim = grp.dim diff_mat = get_diff_mat(actx, grp, grp) + vec = make_obj_array([ + fold(grp.space, vec[xyz_axis]) + for xyz_axis in range(grp.dim) + ]) # weak form if metric_in_matvec: stiff_1D, mass_1D = diff_mat - - if dim == 3: - weak_x = actx.einsum( - "estu,ps,qt,ru->epqr", - vec[0], - stiff_1D, - mass_1D, - mass_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - - weak_y = actx.einsum( - "estu,ps,qt,ru->epqr", - vec[1], - mass_1D, - stiff_1D, - mass_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - - weak_z = actx.einsum( - "estu,ps,qt,ru->epqr", - vec[2], - mass_1D, - mass_1D, - stiff_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s", "mass_1D_t"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - - partials = make_obj_array([ - weak_x, weak_y, weak_z - ]) - - elif dim == 2: - weak_x = actx.einsum( - "est,ps,qt->epq", - vec[0], - stiff_1D, - mass_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - - weak_y = actx.einsum( - "est,ps,qt->epq", - vec[1], - mass_1D, + partials = make_obj_array([ + make_obj_array([ + actx.einsum( + f"e{'bd'[:i]}j{'bd'[i:grp.dim-1]}," + + "ij," + + ("ab,cd" if grp.dim == 3 else "ab") + + "->" + f"e{'ac'[:i]}i{'ac'[i:grp.dim-1]}", + vec[func_axis], stiff_1D, - arg_names=("vec", "stiff_1D_r", "mass_1D_s"), + *(mass_1D,)*(grp.dim-1), + arg_names=("vec", "stiff_1D", + *(("mass_1D_1", "mass_1D_2")[:grp.dim-1])), tagged=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered())) - - partials = make_obj_array([ - weak_x, weak_y + for i in range(grp.dim) ]) - - else: - raise Exception("Dimensions of 2 and 3 are supported by " - "tensor product elements. Found dim = {dim}") - - - partials = make_obj_array([ - unreshape_array_for_tensor_product_space(grp.space, partials[i]) - for i in range(dim) - ]) - - partials = actx.np.stack(partials) - - div = make_obj_array([ - actx.einsum("rei,ei->ei", - ijm[i], - partials[i], - arg_names=("inv_jac_t", "vec"), - tagged=(FirstAxisIsElementsTag(),)) - for i in range(dim) + for func_axis in range(grp.dim) ]) - ret = 0 - for i in range(dim): - ret += div[i] - return ret - # strong form else: partials = make_obj_array([ - actx.einsum( - f"ij,e{'kl'[:i]}j{'mn'[:dim-i-1]}->e{'kl'[:i]}i{'mn'[:dim-i-1]}", + make_obj_array([ + actx.einsum( + "yz," + + f"e{'abcdfghijkl'[:i]}z{'mnopqstuvwx'[:grp.dim-i-1]}->" + + f"e{'abcdfghijkl'[:i]}y{'mnopqstuvwx'[:grp.dim-i-1]}", diff_mat, - vec[i], + vec[func_axis], arg_names=("diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - for i in range(dim) + OutputIsTensorProductDOFArrayOrdered())) + for i in range(grp.dim) + ]) + for func_axis in range(grp.dim) ]) - # unreshape partials to apply geometric factors - # TODO: chain the einsum above with the einsum below - partials = make_obj_array([ - unreshape_array_for_tensor_product_space(grp.space, partials[i]) - for i in range(partials.shape[0]) - ]) + # {{{ unreshape, apply geometric factors, and sum over partials - # apply geometric factors - partials = actx.np.stack([partials[i] for i in range(dim)]) + # TODO: Chain einsums together with geometric factors + partials = actx.np.stack([ + unfold(grp.space, partials[xyz_axis][rst_axis]) + for xyz_axis in range(grp.dim) + for rst_axis in range(grp.dim) + ]) - div = actx.einsum( - "xrei,xei->ei", - ijm, - partials, - arg_names=("inv_jac_t", "vec"), - tagged=(FirstAxisIsElementsTag(),)) + try: + partials = partials.reshape( + grp.dim, grp.dim, partials.shape[1], partials.shape[2]) + except IndexError: + partials = partials.reshape( + grp.dim, grp.dim, partials.shape[1] + ) - return div + div = actx.einsum( + "xrej,xrej->ej", + ijm, + partials, + arg_names=("inv_jac_mat", "partials",), + tagged=(FirstAxisIsElementsTag(),) + ) + # }}} - per_group_divs = [ + return div - compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) - if isinstance(in_grp, TensorProductElementGroupBase) + # }}} - # r for rst axis - # x for xyz axis - else actx.einsum( + + # {{{ simplicial div + + def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec): + return actx.einsum( "xrej,rij,xej->ei" if metric_in_matvec else "xrei,rij,xej->ei", ijm_i, get_diff_mat( @@ -531,9 +501,23 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), tagged=(FirstAxisIsElementsTag(),)) + # }}} + + + per_group_divs = [ + + compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) + if isinstance(in_grp, TensorProductElementGroupBase) + + # r for rst axis + # x for xyz axis + else compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec) + for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, - inv_jac_mat)] + inv_jac_mat) + ] return DOFArray(actx, data=tuple(per_group_divs)) @@ -556,6 +540,7 @@ def get_ref_derivative_mats(grp): import modepy as mp import numpy.linalg as la + #FIXME: Can be gotten rid of by updating meshmode nodes1d = grp.unit_nodes_1d bases_1d = grp.bases_1d() @@ -754,6 +739,7 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): import modepy as mp import numpy.linalg as la + # FIXME: can be gotten rid of by updating meshmode operators basis_1d = out_grp.bases_1d() nodes_1d = out_grp.unit_nodes_1d From 06463efa5028499848bbe68a8123417c96b2a221 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Thu, 21 Dec 2023 09:49:59 -0600 Subject: [PATCH 54/97] Minor changes --- grudge/op.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index aebfc465c..d6cb05e13 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -102,7 +102,7 @@ ) from grudge.interpolation import interp -from grudge.projection import project +from grudge.projection import project, volume_quadrature_project from grudge.reductions import ( norm, @@ -125,8 +125,12 @@ interior_trace_pair, interior_trace_pairs, local_interior_trace_pair, + connected_parts, + inter_volume_trace_pairs, + local_inter_volume_trace_pairs, #connected_ranks, cross_rank_trace_pairs, + cross_rank_inter_volume_trace_pairs, bdry_trace_pair, bv_trace_pair ) @@ -134,6 +138,7 @@ __all__ = ( "project", + "volume_quadrature_project", "interp", "norm", @@ -154,7 +159,7 @@ "interior_trace_pair", "interior_trace_pairs", "local_interior_trace_pair", - "connected_ranks", + #"connected_ranks", "cross_rank_trace_pairs", "bdry_trace_pair", "bv_trace_pair", From 43543ece74f4604b3923ebe00147c4136c9f79df Mon Sep 17 00:00:00 2001 From: Michael Campbell Date: Thu, 21 Dec 2023 14:03:42 -0600 Subject: [PATCH 55/97] Restore some op changes --- grudge/op.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/grudge/op.py b/grudge/op.py index d6cb05e13..63e512435 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -160,7 +160,11 @@ "interior_trace_pairs", "local_interior_trace_pair", #"connected_ranks", + "connected_parts", + "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", From fbb936993bbb3583f010e624d63c95d6d3f8d7a4 Mon Sep 17 00:00:00 2001 From: Michael Campbell Date: Thu, 21 Dec 2023 14:06:14 -0600 Subject: [PATCH 56/97] Remove stale requirement --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index f56f10888..2107e5aeb 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/a-alveyblanc/meshmode.git@tensor-product-1d-nodes-and-1d-basis#egg=meshmode +git+https://github.com/inducer/meshmode.git#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 From 511358b4c6022d82e84926dbdaa8d43594a9b26d Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Thu, 1 Feb 2024 10:59:31 -0600 Subject: [PATCH 57/97] Update dt_geometric_factors to support tensor product elements. --- grudge/dt_utils.py | 34 ++++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 54a9343e8..0a9b223cf 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -229,20 +229,35 @@ def h_min_from_volume( def dt_geometric_factors( 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). + [Hesthaven_2008]_, section 6.4, For simplicial elemenents, this factor is + defined as the inradius (radius of an inscribed circle/sphere). For + non-simplicial elements, a mean length measure is returned. - Specifically, the inradius for each element is computed using the following - formula from [Shewchuk_2002]_, Table 1, for simplicial cells - (triangles/tetrahedra): + Specifically, the inradius for each simplicial element is computed using the + following formula from [Shewchuk_2002]_, Table 1 (triangles, tetrahedra): .. math:: - r_D = \frac{d V}{\sum_{i=1}^{N_{faces}} F_i}, + r_D = \frac{d~V}{\sum_{i=1}^{N_{faces}} F_i}, where :math:`d` is the topological dimension, :math:`V` is the cell volume, and :math:`F_i` are the areas of each face of the cell. + For non-simplicial elements, we use the following formula for a mean + cell size measure: + + .. math:: + + r_D = \frac{2~d~V}{\sum_{i=1}^{N_{faces}} F_i}, + + where :math:`d` is the topological dimension, :math:`V` is the cell volume, + and :math:`F_i` are the areas of each face of the cell. Other valid choices + here include the shortest, longest, average of the cell diagonals, or edges. + The value returned by this routine (i.e. the cell volume divided by the + average cell face area) is bounded by the extrema of the cell edge lengths, + is straightforward to calculate regardless of element shape, and jibes well + with the foregoing calculation for simplicial elements. + :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one. Defaults to the base volume discretization if not provided. :returns: a frozen :class:`~meshmode.dof_array.DOFArray` containing the @@ -256,11 +271,10 @@ def dt_geometric_factors( actx = dcoll._setup_actx volm_discr = dcoll.discr_from_dd(dd) + r_fac = dcoll.dim if any(not isinstance(grp, SimplexElementGroupBase) for grp in volm_discr.groups): - raise NotImplementedError( - "Geometric factors are only implemented for simplex element groups" - ) + r_fac = 2.0*r_fac if volm_discr.dim != volm_discr.ambient_dim: from warnings import warn @@ -342,7 +356,7 @@ def dt_geometric_factors( "e,ei->ei", 1/sae_i, actx.tag_axis(1, DiscretizationDOFAxisTag(), cv_i), - tagged=(FirstAxisIsElementsTag(),)) * dcoll.dim + tagged=(FirstAxisIsElementsTag(),)) * r_fac for cv_i, sae_i in zip(cell_vols, surface_areas))))) # }}} From 7967767abac5c12a9432d6632f255403b8a16e1c Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 5 Feb 2024 15:18:35 -0600 Subject: [PATCH 58/97] Add tpe to dt util testing --- test/mesh_data.py | 2 ++ test/test_dt_utils.py | 36 +++++++++++++++++++++++++++++------- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/test/mesh_data.py b/test/mesh_data.py index 0ccc369a1..08799d09e 100644 --- a/test/mesh_data.py +++ b/test/mesh_data.py @@ -86,6 +86,7 @@ def get_mesh(self, resolution, mesh_order): class BoxMeshBuilder(MeshBuilder): ambient_dim = 2 + group_cls = None mesh_order = 1 resolutions = [4, 8, 16] @@ -100,6 +101,7 @@ def get_mesh(self, resolution, mesh_order): return mgen.generate_regular_rect_mesh( a=self.a, b=self.b, nelements_per_axis=resolution, + group_cls=self.group_cls, order=mesh_order) diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py index 9322f7d28..24f7f991e 100644 --- a/test/test_dt_utils.py +++ b/test/test_dt_utils.py @@ -47,31 +47,52 @@ @pytest.mark.parametrize("name", ["interval", "box2d", "box3d"]) -def test_geometric_factors_regular_refinement(actx_factory, name): +@pytest.mark.parametrize("tpe", [False, True]) +def test_geometric_factors_regular_refinement(actx_factory, name, tpe): from grudge.dt_utils import dt_geometric_factors + import pyopencl as cl + from grudge.array_context import TensorProductArrayContext - actx = actx_factory() + if tpe: + ctx = cl.create_some_context() + queue = cl.CommandQueue(ctx) + actx = TensorProductArrayContext(queue) + else: + actx = actx_factory() # {{{ cases + from meshmode.mesh import TensorProductElementGroup + group_cls = TensorProductElementGroup if tpe else None + if name == "interval": from mesh_data import BoxMeshBuilder - builder = BoxMeshBuilder(ambient_dim=1) + builder = BoxMeshBuilder(ambient_dim=1, group_cls=group_cls) elif name == "box2d": from mesh_data import BoxMeshBuilder - builder = BoxMeshBuilder(ambient_dim=2) + builder = BoxMeshBuilder(ambient_dim=2, group_cls=group_cls) elif name == "box3d": from mesh_data import BoxMeshBuilder - builder = BoxMeshBuilder(ambient_dim=3) + builder = BoxMeshBuilder(ambient_dim=3, group_cls=group_cls) else: raise ValueError("unknown geometry name: %s" % name) # }}} + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as Lgl + + from grudge.dof_desc import DISCR_TAG_BASE + dtag_to_grp_fac = { + DISCR_TAG_BASE: Lgl(builder.order) + } if tpe else None + order = None if tpe else builder.order + min_factors = [] for resolution in builder.resolutions: mesh = builder.get_mesh(resolution, builder.mesh_order) - dcoll = DiscretizationCollection(actx, mesh, order=builder.order) + dcoll = DiscretizationCollection(actx, mesh, order=order, + discr_tag_to_group_factory=dtag_to_grp_fac) min_factors.append( actx.to_numpy( op.nodal_min(dcoll, "vol", actx.thaw(dt_geometric_factors(dcoll)))) @@ -85,7 +106,8 @@ def test_geometric_factors_regular_refinement(actx_factory, name): # Make sure it works with empty meshes mesh = builder.get_mesh(0, builder.mesh_order) - dcoll = DiscretizationCollection(actx, mesh, order=builder.order) + dcoll = DiscretizationCollection(actx, mesh, order=order, + discr_tag_to_group_factory=dtag_to_grp_fac) factors = actx.thaw(dt_geometric_factors(dcoll)) # noqa: F841 From a4dd3f048ba2581731f1519c4c88c67285f05e23 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 5 Feb 2024 16:02:56 -0600 Subject: [PATCH 59/97] Add TPE to dt estimate tests --- test/test_dt_utils.py | 37 +++++++++++++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py index 24f7f991e..884aa4f78 100644 --- a/test/test_dt_utils.py +++ b/test/test_dt_utils.py @@ -173,8 +173,23 @@ def rhs(x): @pytest.mark.parametrize("dim", [1, 2]) @pytest.mark.parametrize("degree", [2, 4]) -def test_wave_dt_estimate(actx_factory, dim, degree, visualize=False): - actx = actx_factory() +@pytest.mark.parametrize("tpe", [False, True]) +def test_wave_dt_estimate(actx_factory, dim, degree, tpe, visualize=False): + + import pyopencl as cl + from grudge.array_context import TensorProductArrayContext + + if tpe: + ctx = cl.create_some_context() + queue = cl.CommandQueue(ctx) + actx = TensorProductArrayContext(queue) + else: + actx = actx_factory() + + # {{{ cases + + from meshmode.mesh import TensorProductElementGroup + group_cls = TensorProductElementGroup if tpe else None import meshmode.mesh.generation as mgen @@ -182,10 +197,24 @@ def test_wave_dt_estimate(actx_factory, dim, degree, visualize=False): b = [1, 1, 1] mesh = mgen.generate_regular_rect_mesh( a=a[:dim], b=b[:dim], - nelements_per_axis=(3,)*dim) + nelements_per_axis=(3,)*dim, + group_cls=group_cls) + assert mesh.dim == dim - dcoll = DiscretizationCollection(actx, mesh, order=degree) + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as Lgl + + from grudge.dof_desc import DISCR_TAG_BASE + order = degree + dtag_to_grp_fac = None + if tpe: + order = None + dtag_to_grp_fac = { + DISCR_TAG_BASE: Lgl(degree) + } + dcoll = DiscretizationCollection(actx, mesh, order=order, + discr_tag_to_group_factory=dtag_to_grp_fac) from grudge.models.wave import WeakWaveOperator wave_op = WeakWaveOperator(dcoll, c=1) From 434109930be739711c7a2da37ffb5980eda3cb74 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sun, 11 Feb 2024 15:29:32 -0600 Subject: [PATCH 60/97] Update TP operators, TP operator tests, and add support for TP in base arraycontext classes --- grudge/array_context.py | 187 ++++++---------- grudge/op.py | 464 ++++++++++++++++++++++++---------------- test/test_op.py | 334 ++++++----------------------- 3 files changed, 419 insertions(+), 566 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index c1964ce91..0b7f4f4e2 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -109,8 +109,7 @@ class PyOpenCLArrayContext(_PyOpenCLArrayContextBase): """Inherits from :class:`meshmode.array_context.PyOpenCLArrayContext`. Extends it - to understand :mod:`grudge`-specific transform metadata. (Of which there isn't - any, for now.) + to understand :mod:`grudge`-specific transform metadata. """ def __init__(self, queue: "pyopencl.CommandQueue", allocator: Optional["pyopencl.tools.AllocatorBase"] = None, @@ -125,6 +124,30 @@ def __init__(self, queue: "pyopencl.CommandQueue", super().__init__(queue, allocator, wait_event_queue_length, force_device_scalars) + def transform_loopy_program(self, t_unit): + knl = t_unit.default_entrypoint + + # {{{ process tensor product specific metadata + + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) + + # }}} + + return super().transform_loopy_program(t_unit) + # }}} @@ -132,8 +155,7 @@ def __init__(self, queue: "pyopencl.CommandQueue", class PytatoPyOpenCLArrayContext(_PytatoPyOpenCLArrayContextBase): """Inherits from :class:`meshmode.array_context.PytatoPyOpenCLArrayContext`. - Extends it to understand :mod:`grudge`-specific transform metadata. (Of - which there isn't any, for now.) + Extends it to understand :mod:`grudge`-specific transform metadata. """ def __init__(self, queue, allocator=None, *, @@ -154,6 +176,29 @@ def __init__(self, queue, allocator=None, super().__init__(queue, allocator, compile_trace_callback=compile_trace_callback) + def transform_loopy_program(self, t_unit): + knl = t_unit.default_entrypoint + + # {{{ process tensor product specific metadata + + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + + # }}} + + return super().transform_loopy_program(t_unit) + # }}} @@ -605,146 +650,44 @@ def get_reasonable_array_context_class( # }}} - -# {{{ distributed + numpy -try: - from arraycontext import NumpyArrayContext - - class MPINumpyArrayContext(NumpyArrayContext, MPIBasedArrayContext): - """An array context for using distributed computation with :mod:`numpy` - eager evaluation. - .. autofunction:: __init__ - """ - - def __init__(self, mpi_communicator) -> None: - super().__init__() - self.mpi_communicator = mpi_communicator - - def clone(self): - return type(self)(self.mpi_communicator) - -except ImportError: - print("Failed to import numpy array context.") - pass -# }}} - - -# {{{ Tensor product array contexts - -# {{{ Relevant tags +# {{{ tensor product-specific machinery class OutputIsTensorProductDOFArrayOrdered(Tag): """Signify that the strides will not be of order "C" or "F". See :class:`grudge.array_context.TensorProductArrayContext` for more details. - """ - pass - -# }}} - - -# {{{ Eager TP array contexts - -class TensorProductArrayContext(_PyOpenCLArrayContextBase): - """Specialized array context for use with tensor product elements. The strides for the arrays containing tensor product element data are of the form (slow, fastest, faster, fast). These strides are not "C" or "F" order. Hence, this specialized array context takes care of specifying the particular strides required. """ - - def transform_loopy_program(self, t_unit): - knl = t_unit.default_entrypoint - if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): - new_args = [] - for arg in knl.args: - if arg.is_output: - arg = arg.copy(dim_tags=( - f"N{len(arg.shape)-1}," - + ",".join(f"N{i}" - for i in range(len(arg.shape)-1)) - )) - - new_args.append(arg) - - knl = knl.copy(args=new_args) - t_unit = t_unit.with_kernel(knl) - - return super().transform_loopy_program(t_unit) + pass -class TensorProductMPIPyOpenCLArrayContext(MPIPyOpenCLArrayContext, - TensorProductArrayContext): +class MassMatrix1d(Tag): + """Used in DAG transformation to realize algebraic simplification of 1D + inverse mass operator times mass operator. + """ pass -# }}} +class InverseMassMatrix1d(Tag): + """See MassMatrix1d. + """ +# }}} -# {{{ Lazy tensor product array contexts +# {{{ Eager TP array context +class TensorProductArrayContext(_PyOpenCLArrayContextBase): + """Specialized array context for use with tensor product elements. + """ +# }}} +# {{{ Lazy tensor product array context class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext): def transform_dag(self, dag): return super().transform_dag(dag) - - def transform_loopy_program(self, t_unit): - knl = t_unit.default_entrypoint - - # {{{ adjust strides according to tensor product structure - if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): - new_args = [] - for arg in knl.args: - if arg.is_output: - arg = arg.copy(dim_tags=( - f"N{len(arg.shape)-1}," - + ",".join(f"N{i}" - for i in range(len(arg.shape)-1)) - )) - - new_args.append(arg) - - knl = knl.copy(args=new_args) - # }}} - - return super().transform_loopy_program(t_unit) - # }}} - -# {{{ TP fusion actx - -from meshmode.array_context import FusionContractorArrayContext - - -class TensorProductFusionContractorArrayContext(FusionContractorArrayContext): - - def transform_loopy_program(self, t_unit): - knl = t_unit.default_entrypoint - if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): - new_args = [] - for arg in knl.args: - if arg.is_output: - arg = arg.copy(dim_tags=( - f"N{len(arg.shape)-1}," - + ",".join(f"N{i}" - for i in range(len(arg.shape)-1)) - )) - - new_args.append(arg) - - knl = knl.copy(args=new_args) - t_unit = t_unit.with_kernel(knl) - - return super().transform_loopy_program(t_unit) - - -class TensorProductMPIFusionContractorArrayContext( - MPIPytatoArrayContext, TensorProductFusionContractorArrayContext): - pass - # }}} - -# }}} - - # vim: foldmethod=marker diff --git a/grudge/op.py b/grudge/op.py index 63e512435..dd15db781 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -75,11 +75,13 @@ from functools import partial from meshmode.dof_array import DOFArray, warn +from meshmode.discretization.poly_element import ( + TensorProductElementGroupBase as TensorProductElementGroup, + SimplexElementGroupBase as SimplexElementGroup) from meshmode.transform_metadata import (FirstAxisIsElementsTag, DiscretizationDOFAxisTag, DiscretizationElementAxisTag, DiscretizationFaceAxisTag) -from meshmode.discretization.poly_element import TensorProductElementGroupBase from modepy.tools import ( reshape_array_for_tensor_product_space as fold, @@ -87,8 +89,7 @@ from grudge.discretization import DiscretizationCollection from grudge.dof_desc import as_dofdesc -from grudge.array_context import ( - OutputIsTensorProductDOFArrayOrdered) +from grudge.array_context import OutputIsTensorProductDOFArrayOrdered from pytools import keyed_memoize_in from pytools.obj_array import make_obj_array @@ -199,13 +200,62 @@ def _single_axis_derivative_kernel( # {{{ tensor product single axis derivative - # FIXME: actually implement single axis tensor product derivatives def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, xyz_axis, metric_in_matvec): + vec = fold(grp.space, vec) + + if metric_in_matvec: + stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) + + apply_mass_axes = set(range(grp.dim)) - {xyz_axis} - return compute_simplicial_derivative(actx, grp, grp, get_diff_mat, vec, - ijm, xyz_axis, metric_in_matvec) + for ax in apply_mass_axes: + vec_mass_applied = single_axis_operator_application( + actx, grp.dim, mass_1d, ax, vec, + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("mass_1d", "vec") + ) + + ref_weak_derivative = unfold( + grp.space, + single_axis_operator_application( + actx, grp.dim, stiff_1d, xyz_axis, vec_mass_applied, + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("stiff_1d", "vec_with_mass_applied")) + ) + + derivative = actx.einsum( + 'rej,ej->ej', + ijm[xyz_axis], + ref_weak_derivative, + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", "ref_weak_derivative") + ) + + else: + diff_mat = get_diff_mat(actx, grp, grp) + + ref_derivative = unfold( + grp.space, + single_axis_operator_application( + actx, grp.dim, diff_mat, xyz_axis, vec, + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("diff_mat", "vec")) + ) + + derivative = actx.einsum( + 'rej,ej->ej', + ijm[xyz_axis], + ref_derivative, + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", "ref_derivs") + ) + + return derivative # }}} @@ -213,17 +263,17 @@ def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, # {{{ simplicial single axis derivative def compute_simplicial_derivative(actx, in_grp, out_grp, - get_diff_mat, vec_i, ijm_i, + get_diff_mat, vec, ijm, xyz_axis, metric_in_matvec): # r for rst axis return actx.einsum( "rej,rij,ej->ei" if metric_in_matvec else "rei,rij,ej->ei", - ijm_i[xyz_axis], + ijm[xyz_axis], get_diff_mat( actx, out_element_group=out_grp, in_element_group=in_grp), - vec_i, + vec, arg_names=("inv_jac_t", "ref_stiffT_mat", "vec", ), tagged=(FirstAxisIsElementsTag(),)) @@ -233,10 +283,9 @@ def compute_simplicial_derivative(actx, in_grp, out_grp, return DOFArray( actx, data=tuple( - compute_tensor_product_derivative(actx, in_grp, out_grp, - get_diff_mat, vec_i, ijm_i, - xyz_axis, metric_in_matvec) - if isinstance(in_grp, TensorProductElementGroupBase) + compute_tensor_product_derivative(actx, in_grp, get_diff_mat, vec_i, + ijm_i, xyz_axis, metric_in_matvec) + if isinstance(in_grp, TensorProductElementGroup) else compute_simplicial_derivative(actx, in_grp, out_grp, get_diff_mat, vec_i, ijm_i, xyz_axis, metric_in_matvec) @@ -258,28 +307,8 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, # TODO: add note about inverse mass simplification, point to # op.inverse_mass (assuming this is where the explanation will live) """ - Exploits tensor product structure to reduce complexity. Applies a - differentiation operator containing 1D information to a tensor of DOF - data. For example, in the 2D strong form case, this computes partial - derivatives in a similar manner to - - .. math:: - - \partial_x \mathbf{f}_{ij} = \sum_{\ell} \mathbf{J}^e_{ij} - \mathbf{D}_{i\ell} \mathbf{f}_{\ell j} - - where $\mathbf{D}$ is a 1D differentiation operator, $\mathbf{f}$ is a - vector of function data, $\mathbf{J}^e$ is the element Jacobian matrix. - The weak form uses a 1D element mass operator and a 1D element stiffness - operator to perform the contraction - - .. math:: - - \partial_x \mathbf{f}_{ij} = \sum_{\ell,b} \mathbf{J}^e_{\ell b} - \mathbf{f}_{\ell b} \mathbf{S}^e_{i\ell} \mathbf{M}^e_{jb} """ - if grp.dim > 3 and metric_in_matvec: warn('Efficient tensor product weak ' 'differentiation operators only ' @@ -288,69 +317,69 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, return compute_simplicial_grad(actx, grp, grp, diff_mat, vec, ijm, metric_in_matvec) - # reshape u to expose tensor product structure + # reshape vector to expose tensor product structure vec = fold(grp.space, vec) - diff_mat = get_diff_mat(actx, grp, grp) - # weak form case: - # 3D weak_x: einsum("estu,ps,qt,ru->epqr", - # f, stiff_1D, mass_1D, mass_1D) - # TODO:? make this more general, maybe offload to a function that - # generates argnames and einsum specs if metric_in_matvec: - stiff_1D, mass_1D = diff_mat - grad = make_obj_array([ - actx.einsum( - f"e{'bd'[:i]}j{'bd'[i:grp.dim-1]}," + - "ij," + - ("ab,cd" if grp.dim == 3 else "ab") + - "->" - f"e{'ac'[:i]}i{'ac'[i:grp.dim-1]}", - vec, - stiff_1D, - *(mass_1D,)*(grp.dim-1), - arg_names=("vec", "stiff_1D", - *(("mass_1D_1", "mass_1D_2")[:grp.dim-1])), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - for i in range(grp.dim) - ]) - - # Carries out, e.g., 3D strong form contraction - # x partial: einsum("il,eljk->eijk", D, f) - else: - grad = make_obj_array([ - actx.einsum( - "yz," + - f"e{'abcdfghijkl'[:i]}z{'mnopqstuvwx'[:grp.dim-i-1]}->" + - f"e{'abcdfghijkl'[:i]}y{'mnopqstuvwx'[:grp.dim-i-1]}", - diff_mat, - vec, - arg_names=("diff_mat", "vec"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - for i in range(grp.dim) - ]) - - # {{{ unreshape grad and apply geometric factors - - # TODO: Chain einsums together with geometric factors - grad = actx.np.stack([ - unfold(grp.space, grad[rst_axis]) - for rst_axis in range(grp.dim) - ]) + stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) + + grad = [] + for xyz_axis in range(grp.dim): + grad.append(vec) + apply_mass_axes = set(range(grp.dim)) - {xyz_axis} + + # apply mass operators + for ax in apply_mass_axes: + grad[xyz_axis] = single_axis_operator_application( + actx, grp.dim, mass_1d, ax, grad[xyz_axis], + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("mass_1d", f"vec_{xyz_axis}") + ) - grad = actx.einsum( - "xrej,rej->xej", - ijm, - grad, - arg_names=("inv_jac_mat", "grad"), - tagged=(FirstAxisIsElementsTag(),) - ) + # apply stiffness operator and unfold + grad[xyz_axis] = unfold( + grp.space, + single_axis_operator_application( + actx, grp.dim, stiff_1d, xyz_axis, grad[xyz_axis], + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("stiff_1d", f"vec_{xyz_axis}")) + ) + + # apply metric terms + grad[xyz_axis] = actx.einsum( + 'rej,ej->ej', + ijm[xyz_axis], + grad[xyz_axis], + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", f"vec_{xyz_axis}") + ) + else: + diff_mat = get_diff_mat(actx, grp, grp) + + grad = [] + for xyz_axis in range(grp.dim): + grad.append(vec) + grad[xyz_axis] = unfold( + grp.space, + single_axis_operator_application( + actx, grp.dim, diff_mat, xyz_axis, grad[xyz_axis], + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("diff_mat", f"vec_{xyz_axis}") + ) + ) - # }}} + grad[xyz_axis] = actx.einsum( + "rej,ej->ej", + ijm[xyz_axis], + grad[xyz_axis], + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", f"vec_{xyz_axis}") + ) - return grad + return make_obj_array(grad) # }}} @@ -377,7 +406,7 @@ def compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, per_group_grads = [ compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i, metric_in_matvec) - if isinstance(in_grp, TensorProductElementGroupBase) + if isinstance(in_grp, TensorProductElementGroup) else compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, ijm_i, metric_in_matvec) @@ -414,83 +443,72 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): return compute_simplicial_div(actx, grp, grp, diff_mat, vec, ijm, metric_in_matvec) - # reshape u to expose tensor product structure - diff_mat = get_diff_mat(actx, grp, grp) vec = make_obj_array([ - fold(grp.space, vec[xyz_axis]) - for xyz_axis in range(grp.dim) + fold(grp.space, vec[func_axis]) + for func_axis in range(vec.shape[0]) ]) - # weak form if metric_in_matvec: - stiff_1D, mass_1D = diff_mat - partials = make_obj_array([ - make_obj_array([ - actx.einsum( - f"e{'bd'[:i]}j{'bd'[i:grp.dim-1]}," + - "ij," + - ("ab,cd" if grp.dim == 3 else "ab") + - "->" - f"e{'ac'[:i]}i{'ac'[i:grp.dim-1]}", - vec[func_axis], - stiff_1D, - *(mass_1D,)*(grp.dim-1), - arg_names=("vec", "stiff_1D", - *(("mass_1D_1", "mass_1D_2")[:grp.dim-1])), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - for i in range(grp.dim) - ]) - for func_axis in range(grp.dim) - ]) - - # strong form + stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) + + partials = [] + for func_axis in range(vec.shape[0]): + ref = [] + for xyz_axis in range(grp.dim): + ref.append(vec[func_axis]) + + apply_mass_axes = set(range(grp.dim)) - {xyz_axis} + for ax in apply_mass_axes: + ref[xyz_axis] = single_axis_operator_application( + actx, grp.dim, mass_1d, ax, ref[xyz_axis], + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("mass_1d", f"vec_{func_axis}_{xyz_axis}") + ) + + ref[xyz_axis] = single_axis_operator_application( + actx, grp.dim, stiff_1d, xyz_axis, ref[xyz_axis], + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("stiff_1d", f"vec_{func_axis}_{xyz_axis}") + ) + + partials.append(ref) + else: - partials = make_obj_array([ - make_obj_array([ - actx.einsum( - "yz," + - f"e{'abcdfghijkl'[:i]}z{'mnopqstuvwx'[:grp.dim-i-1]}->" + - f"e{'abcdfghijkl'[:i]}y{'mnopqstuvwx'[:grp.dim-i-1]}", - diff_mat, - vec[func_axis], - arg_names=("diff_mat", "vec"), - tagged=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered())) - for i in range(grp.dim) - ]) - for func_axis in range(grp.dim) - ]) - - # {{{ unreshape, apply geometric factors, and sum over partials - - # TODO: Chain einsums together with geometric factors + diff_mat = get_diff_mat(actx, grp, grp) + + partials = [] + for func_axis in range(vec.shape[0]): + ref = [] + for xyz_axis in range(grp.dim): + ref.append(vec[func_axis]) + + ref[xyz_axis] = single_axis_operator_application( + actx, grp.dim, diff_mat, xyz_axis, ref[xyz_axis], + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("diff_mat", f"vec_{func_axis}_{xyz_axis}") + ) + + partials.append(ref) + partials = actx.np.stack([ - unfold(grp.space, partials[xyz_axis][rst_axis]) + unfold(grp.space, partials[func_axis][xyz_axis]) + for func_axis in range(grp.dim) for xyz_axis in range(grp.dim) - for rst_axis in range(grp.dim) ]) - - try: - partials = partials.reshape( - grp.dim, grp.dim, partials.shape[1], partials.shape[2]) - except IndexError: - partials = partials.reshape( - grp.dim, grp.dim, partials.shape[1] - ) + partials = partials.reshape(grp.dim, grp.dim, *partials.shape[-2:]) div = actx.einsum( - "xrej,xrej->ej", + 'xrej,xrej->ej', ijm, partials, - arg_names=("inv_jac_mat", "partials",), + arg_names=("inv_jac_t", "partials"), tagged=(FirstAxisIsElementsTag(),) ) - # }}} - return div - # }}} @@ -516,7 +534,7 @@ def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, per_group_divs = [ compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) - if isinstance(in_grp, TensorProductElementGroupBase) + if isinstance(in_grp, TensorProductElementGroup) # r for rst axis # x for xyz axis @@ -545,7 +563,7 @@ def _reference_derivative_matrices(actx: ArrayContext, actx, _reference_derivative_matrices, lambda grp: grp.discretization_key()) def get_ref_derivative_mats(grp): - if isinstance(grp, TensorProductElementGroupBase): + if isinstance(grp, TensorProductElementGroup): import modepy as mp import numpy.linalg as la @@ -565,13 +583,18 @@ def get_ref_derivative_mats(grp): 1: DiscretizationDOFAxisTag()}, diff_mat))) - else: + elif isinstance(grp, SimplexElementGroup): from meshmode.discretization.poly_element import diff_matrices return actx.freeze( actx.tag_axis( 1, DiscretizationDOFAxisTag(), actx.from_numpy( np.asarray(diff_matrices(grp))))) + + else: + raise TypeError("grp must be either a TensorProductElementGroup or" + f" a SimplexElementGroup. Found {grp}") + return get_ref_derivative_mats(out_element_group) @@ -744,7 +767,7 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): # {{{ tensor product case - if isinstance(out_grp, TensorProductElementGroupBase): + if isinstance(out_grp, TensorProductElementGroup): import modepy as mp import numpy.linalg as la @@ -755,22 +778,24 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): vdm = mp.vandermonde(basis_1d.functions, nodes_1d) vdm_p = mp.vandermonde(basis_1d.gradients, nodes_1d)[0] - mass_1D = la.inv(vdm @ vdm.T) + mass_1d = la.inv(vdm @ vdm.T) diff_mat = la.solve(vdm.T, vdm_p.T).T - stiff_1D = actx.freeze( + stiff_1d = actx.freeze( actx.tag_axis(1, DiscretizationDOFAxisTag(), actx.from_numpy( np.asarray( - diff_mat.T @ mass_1D.T)))) + diff_mat.T @ mass_1d.T)))) - mass_1D = actx.freeze( - actx.tag_axis(1, DiscretizationDOFAxisTag(), - actx.from_numpy( - np.asarray( - mass_1D)))) + from grudge.array_context import MassMatrix1d + mass_1d = actx.freeze( + actx.tag_axis( + 1, (DiscretizationDOFAxisTag(),), + actx.from_numpy(np.asarray(mass_1d))) + ) + mass_1d = actx.tag(MassMatrix1d(), mass_1d) - return (stiff_1D, mass_1D) + return (stiff_1d, mass_1d) # }}} @@ -803,6 +828,7 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): ).copy() # contigify the array ) ) + return get_ref_stiffness_transpose_mat(out_element_group, in_element_group) @@ -1113,14 +1139,31 @@ def reference_inverse_mass_matrix(actx: ArrayContext, element_group): lambda grp: grp.discretization_key()) def get_ref_inv_mass_mat(grp): from modepy import inverse_mass_matrix - basis = grp.basis_obj() - return actx.freeze( - actx.tag_axis(0, DiscretizationDOFAxisTag(), - actx.from_numpy( - np.asarray( - inverse_mass_matrix(basis.functions, grp.unit_nodes), - order="C")))) + if isinstance(grp, TensorProductElementGroup): + basis_1d = grp.bases_1d() + nodes_1d = grp.unit_nodes_1d + inv_mass_1d = inverse_mass_matrix(basis_1d.functions, nodes_1d) + + from grudge.array_context import InverseMassMatrix1d + inv_mass_1d = actx.tag_axis(0, DiscretizationDOFAxisTag(), + actx.from_numpy(np.asarray(inv_mass_1d))) + inv_mass_1d = actx.freeze( + actx.tag(InverseMassMatrix1d(), inv_mass_1d)) + + return inv_mass_1d + elif isinstance(grp, SimplexElementGroup): + basis = grp.basis_obj() + + return actx.freeze( + actx.tag_axis(0, DiscretizationDOFAxisTag(), + actx.from_numpy( + np.asarray( + inverse_mass_matrix(basis.functions, grp.unit_nodes), + order="C")))) + else: + raise TypeError("grp must be either a TensorProductElementGroup or" + f" a SimplexElementGroup. Found {grp}") return get_ref_inv_mass_mat(element_group) @@ -1145,15 +1188,50 @@ def _apply_inverse_mass_operator( discr = dcoll.discr_from_dd(dd_in) inv_area_elements = 1./area_element(actx, dcoll, dd=dd_in, _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + + + def apply_to_tensor_product_elements(grp, jac_inv, vec, ref_inv_mass): + + vec = fold(grp.space, vec) + + for xyz_axis in range(grp.dim): + vec = single_axis_operator_application( + actx, grp.dim, ref_inv_mass, xyz_axis, vec, + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("ref_inv_mass_1d", "vec")) + + vec = unfold(grp.space, vec) + + return actx.einsum( + "ei,ei->ei", + jac_inv, + vec, + tagged=(FirstAxisIsElementsTag(),) + ) + + + def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): + + # Based on https://arxiv.org/pdf/1608.03836.pdf + # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv + return actx.einsum( + "ei,ij,ej->ei", + jac_inv, + ref_inv_mass, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + group_data = [ - # Based on https://arxiv.org/pdf/1608.03836.pdf - # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv - actx.einsum("ei,ij,ej->ei", - jac_inv, - reference_inverse_mass_matrix(actx, element_group=grp), - vec_i, - tagged=(FirstAxisIsElementsTag(),)) - for grp, jac_inv, vec_i in zip(discr.groups, inv_area_elements, vec)] + apply_to_tensor_product_elements( + grp, jac_inv, vec_i, + reference_inverse_mass_matrix(actx, element_group=grp)) + if isinstance(grp, TensorProductElementGroup) else + apply_to_simplicial_elements(jac_inv, vec_i, + reference_inverse_mass_matrix(actx, element_group=grp)) + for grp, jac_inv, vec_i in zip(discr.groups, inv_area_elements, vec) + ] return DOFArray(actx, data=tuple(group_data)) @@ -1400,4 +1478,30 @@ def face_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: # }}} +# {{{ general single axis operator application + +def single_axis_operator_application(actx, dim, operator, axis, data, + arg_names=None, tags=None): + """ + Used for applying 1D operators to a single axis of a tensor of DOF data. + """ + + if not isinstance(arg_names, tuple): + raise TypeError("arg_names must be a tuple.") + if not isinstance(tags, tuple): + raise TypeError("arg_names must be a tuple.") + + operator_spec = 'ij' + data_spec = f'e{"abcdefghklm"[:axis]}j{"nopqrstuvwxyz"[:dim-axis-1]}' + out_spec = f'e{"abcdefghklm"[:axis]}i{"nopqrstuvwxyz"[:dim-axis-1]}' + + spec = operator_spec + ',' + data_spec + '->' + out_spec + + return actx.einsum(spec, operator, data, + arg_names=arg_names, + tagged=tags) + +# }}} + + # vim: foldmethod=marker diff --git a/test/test_op.py b/test/test_op.py index dfb361b4f..8e8330622 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -21,6 +21,7 @@ """ +from meshmode.mesh.processing import affine_map import numpy as np import meshmode.mesh.generation as mgen @@ -30,9 +31,10 @@ from grudge import op, geometry as geo, DiscretizationCollection from grudge.dof_desc import DOFDesc +from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup + import pytest -from grudge.discretization import make_discretization_collection from grudge.array_context import PytestPyOpenCLArrayContextFactory from arraycontext import pytest_generate_tests_for_array_contexts pytest_generate_tests = pytest_generate_tests_for_array_contexts( @@ -45,6 +47,10 @@ # {{{ gradient +@pytest.mark.parametrize("group_cls", [ + SimplexElementGroup, + TensorProductElementGroup +]) @pytest.mark.parametrize("form", ["strong", "weak"]) @pytest.mark.parametrize("dim", [1, 2, 3]) @pytest.mark.parametrize("order", [2, 3]) @@ -54,7 +60,7 @@ (True, True) ]) def test_gradient(actx_factory, form, dim, order, vectorize, nested, - visualize=False): + group_cls, visualize=False): actx = actx_factory() from pytools.convergence import EOCRecorder @@ -62,144 +68,40 @@ def test_gradient(actx_factory, form, dim, order, vectorize, nested, for n in [4, 6, 8]: mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, b=(1,)*dim, - nelements_per_axis=(n,)*dim) + a=(-1,)*dim, b=(1,)*dim, + nelements_per_axis=(n,)*dim, + group_cls=group_cls) - dcoll = DiscretizationCollection(actx, mesh, order=order) + if group_cls is TensorProductElementGroup: + # no reason to test 1D tensor product elements + if dim == 1: + return - def f(x): - result = dcoll.zeros(actx) + 1 - for i in range(dim-1): - result = result * actx.np.sin(np.pi*x[i]) - result = result * actx.np.cos(np.pi/2*x[dim-1]) - return result + import grudge.dof_desc as dd + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL - def grad_f(x): - result = make_obj_array([dcoll.zeros(actx) + 1 for _ in range(dim)]) - for i in range(dim-1): - for j in range(i): - result[i] = result[i] * actx.np.sin(np.pi*x[j]) - result[i] = result[i] * np.pi*actx.np.cos(np.pi*x[i]) - for j in range(i+1, dim-1): - result[i] = result[i] * actx.np.sin(np.pi*x[j]) - result[i] = result[i] * actx.np.cos(np.pi/2*x[dim-1]) - for j in range(dim-1): - result[dim-1] = result[dim-1] * actx.np.sin(np.pi*x[j]) - result[dim-1] = result[dim-1] * (-np.pi/2*actx.np.sin(np.pi/2*x[dim-1])) - return result - - x = actx.thaw(dcoll.nodes()) - - if vectorize: - u = make_obj_array([(i+1)*f(x) for i in range(dim)]) - else: - u = f(x) - - def get_flux(u_tpair): - dd = u_tpair.dd - dd_allfaces = dd.with_dtag("all_faces") - normal = geo.normal(actx, dcoll, dd) - u_avg = u_tpair.avg - if vectorize: - if nested: - flux = make_obj_array([u_avg_i * normal for u_avg_i in u_avg]) - else: - flux = np.outer(u_avg, normal) - else: - flux = u_avg * normal - return op.project(dcoll, dd, dd_allfaces, flux) - - dd_allfaces = DOFDesc("all_faces") + dcoll = DiscretizationCollection( + actx, + mesh, + discr_tag_to_group_factory={ + dd.DISCR_TAG_BASE: LGL(order)}) - if form == "strong": - grad_u = ( - op.local_grad(dcoll, u, nested=nested) - # No flux terms because u doesn't have inter-el jumps - ) - elif form == "weak": - grad_u = op.inverse_mass(dcoll, - -op.weak_local_grad(dcoll, u, nested=nested) # pylint: disable=E1130 - + # noqa: W504 - op.face_mass(dcoll, - dd_allfaces, - # Note: no boundary flux terms here because u_ext == u_int == 0 - sum(get_flux(utpair) - for utpair in op.interior_trace_pairs(dcoll, u)) - ) - ) - else: - raise ValueError("Invalid form argument.") + elif group_cls is SimplexElementGroup: + dcoll = DiscretizationCollection(actx, mesh, order=order) - if vectorize: - expected_grad_u = make_obj_array( - [(i+1)*grad_f(x) for i in range(dim)]) - if not nested: - expected_grad_u = np.stack(expected_grad_u, axis=0) else: - expected_grad_u = grad_f(x) + raise AssertionError('Expecting TensorProductElementGroup or ' + f'SimplexElementGroup. Found {group_cls}') - if visualize: - from grudge.shortcuts import make_visualizer - vis = make_visualizer(dcoll, vis_order=order if dim == 3 else dim+3) - - filename = (f"test_gradient_{form}_{dim}_{order}" - f"{'_vec' if vectorize else ''}{'_nested' if nested else ''}.vtu") - vis.write_vtk_file(filename, [ - ("u", u), - ("grad_u", grad_u), - ("expected_grad_u", expected_grad_u), - ], overwrite=True) + alpha = 0.3 + rot_mat = np.array([ + [np.cos(alpha), np.sin(alpha), 0], + [-np.sin(alpha), np.cos(alpha), 0], + [0, 0, 1], + ])[:dim, :dim] - rel_linf_err = actx.to_numpy( - op.norm(dcoll, grad_u - expected_grad_u, np.inf) - / op.norm(dcoll, expected_grad_u, np.inf)) - eoc_rec.add_data_point(1./n, rel_linf_err) - - print("L^inf error:") - print(eoc_rec) - assert (eoc_rec.order_estimate() >= order - 0.5 - or eoc_rec.max_error() < 1e-11) - - -@pytest.mark.parametrize("form", ["strong", "weak"]) -@pytest.mark.parametrize("dim", [2, 3]) -@pytest.mark.parametrize("order", [2, 3]) -@pytest.mark.parametrize(("vectorize", "nested"), [ - (False, False), - (True, False), - (True, True) - ]) -def test_tensor_product_gradient(form, dim, order, vectorize, - nested, visualize=False): - """A "one-dimensional tensor product element" does not make sense, so the - one-dimensional case is excluded from this test. - """ - - import pyopencl as cl - from grudge.array_context import TensorProductArrayContext - - ctx = cl.create_some_context() - queue = cl.CommandQueue(ctx) - actx = TensorProductArrayContext(queue) - - from pytools.convergence import EOCRecorder - eoc_rec = EOCRecorder() - - from meshmode.mesh import TensorProductElementGroup - from meshmode.discretization.poly_element import \ - LegendreGaussLobattoTensorProductGroupFactory as LGL - for n in [4, 6, 8]: - mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, b=(1,)*dim, - nelements_per_axis=(n,)*dim, - group_cls=TensorProductElementGroup) - - import grudge.dof_desc as dd - dcoll = DiscretizationCollection( - actx, - mesh, - discr_tag_to_group_factory={ - dd.DISCR_TAG_BASE: LGL(order)}) + mesh = affine_map(mesh, A=rot_mat) def f(x): result = dcoll.zeros(actx) + 1 @@ -232,7 +134,7 @@ def grad_f(x): def get_flux(u_tpair): dd = u_tpair.dd dd_allfaces = dd.with_dtag("all_faces") - normal = actx.thaw(dcoll.normal(dd)) + normal = geo.normal(actx, dcoll, dd) u_avg = u_tpair.avg if vectorize: if nested: @@ -299,16 +201,20 @@ def get_flux(u_tpair): # {{{ divergence +@pytest.mark.parametrize("group_cls", [ + #SimplexElementGroup, + TensorProductElementGroup +]) @pytest.mark.parametrize("form", ["strong", "weak"]) -@pytest.mark.parametrize("dim", [1, 2, 3]) +@pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize(("vectorize", "nested"), [ (False, False), (True, False), (True, True) - ]) +]) def test_divergence(actx_factory, form, dim, order, vectorize, nested, - visualize=False): + group_cls, visualize=False): actx = actx_factory() from pytools.convergence import EOCRecorder @@ -316,141 +222,40 @@ def test_divergence(actx_factory, form, dim, order, vectorize, nested, for n in [4, 6, 8]: mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, b=(1,)*dim, - nelements_per_axis=(n,)*dim) - - dcoll = DiscretizationCollection(actx, mesh, order=order) - - def f(x): - result = make_obj_array([dcoll.zeros(actx) + (i+1) for i in range(dim)]) - for i in range(dim-1): - result = result * actx.np.sin(np.pi*x[i]) - result = result * actx.np.cos(np.pi/2*x[dim-1]) - return result - - def div_f(x): - result = dcoll.zeros(actx) - for i in range(dim-1): - deriv = dcoll.zeros(actx) + (i+1) - for j in range(i): - deriv = deriv * actx.np.sin(np.pi*x[j]) - deriv = deriv * np.pi*actx.np.cos(np.pi*x[i]) - for j in range(i+1, dim-1): - deriv = deriv * actx.np.sin(np.pi*x[j]) - deriv = deriv * actx.np.cos(np.pi/2*x[dim-1]) - result = result + deriv - deriv = dcoll.zeros(actx) + dim - for j in range(dim-1): - deriv = deriv * actx.np.sin(np.pi*x[j]) - deriv = deriv * (-np.pi/2*actx.np.sin(np.pi/2*x[dim-1])) - result = result + deriv - return result - - x = actx.thaw(dcoll.nodes()) - - if vectorize: - u = make_obj_array([(i+1)*f(x) for i in range(dim)]) - if not nested: - u = np.stack(u, axis=0) - else: - u = f(x) - - def get_flux(u_tpair): - dd = u_tpair.dd - dd_allfaces = dd.with_dtag("all_faces") - normal = geo.normal(actx, dcoll, dd) - flux = u_tpair.avg @ normal - return op.project(dcoll, dd, dd_allfaces, flux) - - dd_allfaces = DOFDesc("all_faces") - - if form == "strong": - div_u = ( - op.local_div(dcoll, u) - # No flux terms because u doesn't have inter-el jumps - ) - elif form == "weak": - div_u = op.inverse_mass(dcoll, - -op.weak_local_div(dcoll, u) - + # noqa: W504 - op.face_mass(dcoll, - dd_allfaces, - # Note: no boundary flux terms here because u_ext == u_int == 0 - sum(get_flux(utpair) - for utpair in op.interior_trace_pairs(dcoll, u)) - ) - ) - else: - raise ValueError("Invalid form argument.") + a=(-1,)*dim, b=(1,)*dim, + nelements_per_axis=(n,)*dim, + group_cls=group_cls) - if vectorize: - expected_div_u = make_obj_array([(i+1)*div_f(x) for i in range(dim)]) - else: - expected_div_u = div_f(x) + if group_cls is TensorProductElementGroup: + # no reason to test 1D tensor product elements + if dim == 1: + return - if visualize: - from grudge.shortcuts import make_visualizer - vis = make_visualizer(dcoll, vis_order=order if dim == 3 else dim+3) - - filename = (f"test_divergence_{form}_{dim}_{order}" - f"{'_vec' if vectorize else ''}{'_nested' if nested else ''}.vtu") - vis.write_vtk_file(filename, [ - ("u", u), - ("div_u", div_u), - ("expected_div_u", expected_div_u), - ], overwrite=True) + import grudge.dof_desc as dd + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL - rel_linf_err = actx.to_numpy( - op.norm(dcoll, div_u - expected_div_u, np.inf) - / op.norm(dcoll, expected_div_u, np.inf)) - eoc_rec.add_data_point(1./n, rel_linf_err) - - print("L^inf error:") - print(eoc_rec) - assert (eoc_rec.order_estimate() >= order - 0.5 - or eoc_rec.max_error() < 1e-11) - - -@pytest.mark.parametrize("form", ["strong", "weak"]) -@pytest.mark.parametrize("dim", [2, 3]) -@pytest.mark.parametrize("order", [2, 3]) -@pytest.mark.parametrize(("vectorize", "nested"), [ - (False, False), - (True, False), - (True, True) - ]) -def test_tensor_product_divergence(form, dim, order, vectorize, - nested, visualize=False): - """A "one-dimensional tensor product element" does not make sense, so the - one-dimensional case is excluded from this test. - """ - import pyopencl as cl - from grudge.array_context import TensorProductArrayContext - - ctx = cl.create_some_context() - queue = cl.CommandQueue(ctx) - actx = TensorProductArrayContext(queue) - - from pytools.convergence import EOCRecorder - eoc_rec = EOCRecorder() - - from meshmode.mesh import TensorProductElementGroup - from meshmode.discretization.poly_element import \ - LegendreGaussLobattoTensorProductGroupFactory as LGL - for n in [4, 6, 8]: - mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, - b=(1,)*dim, - nelements_per_axis=(n,)*dim, - group_cls=TensorProductElementGroup) - - import grudge.dof_desc as dd - dcoll = make_discretization_collection( + dcoll = DiscretizationCollection( actx, mesh, discr_tag_to_group_factory={ dd.DISCR_TAG_BASE: LGL(order)}) + elif group_cls is SimplexElementGroup: + dcoll = DiscretizationCollection(actx, mesh, order=order) + + else: + raise AssertionError('Expecting TensorProductElementGroup or ' + f'SimplexElementGroup. Found {group_cls}') + + alpha = 0.3 + rot_mat = np.array([ + [np.cos(alpha), np.sin(alpha), 0], + [-np.sin(alpha), np.cos(alpha), 0], + [0, 0, 1], + ])[:dim, :dim] + + mesh = affine_map(mesh, A=rot_mat) def f(x): result = make_obj_array([dcoll.zeros(actx) + (i+1) for i in range(dim)]) for i in range(dim-1): @@ -488,7 +293,7 @@ def div_f(x): def get_flux(u_tpair): dd = u_tpair.dd dd_allfaces = dd.with_dtag("all_faces") - normal = actx.thaw(dcoll.normal(dd)) + normal = geo.normal(actx, dcoll, dd) flux = u_tpair.avg @ normal return op.project(dcoll, dd, dd_allfaces, flux) @@ -539,6 +344,7 @@ def get_flux(u_tpair): print(eoc_rec) assert (eoc_rec.order_estimate() >= order - 0.5 or eoc_rec.max_error() < 1e-11) + # }}} From 999904f3342615d776cb6b6abf846e4a330c1eb7 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sun, 11 Feb 2024 15:49:18 -0600 Subject: [PATCH 61/97] Remove no longer used TP actx classes --- grudge/array_context.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index 0b7f4f4e2..d827274e7 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -650,6 +650,7 @@ def get_reasonable_array_context_class( # }}} + # {{{ tensor product-specific machinery class OutputIsTensorProductDOFArrayOrdered(Tag): @@ -676,18 +677,5 @@ class InverseMassMatrix1d(Tag): # }}} -# {{{ Eager TP array context -class TensorProductArrayContext(_PyOpenCLArrayContextBase): - """Specialized array context for use with tensor product elements. - """ -# }}} - -# {{{ Lazy tensor product array context -class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext): - def transform_dag(self, dag): - return super().transform_dag(dag) -# }}} - -# }}} # vim: foldmethod=marker From 79ae993c946eb1f7cda229637c0bffd85e9cccf4 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sun, 11 Feb 2024 16:15:31 -0600 Subject: [PATCH 62/97] Add in MPINumpyArrayContext --- grudge/array_context.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/grudge/array_context.py b/grudge/array_context.py index d827274e7..41164c54a 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -650,6 +650,30 @@ def get_reasonable_array_context_class( # }}} +# {{{ distributed + numpy + +try: + from arraycontext import NumpyArrayContext + + class MPINumpyArrayContext(NumpyArrayContext, MPIBasedArrayContext): + """An array context for using distributed computation with :mod:`numpy` + eager evaluation. + .. autofunction:: __init__ + """ + + def __init__(self, mpi_communicator) -> None: + super().__init__() + self.mpi_communicator = mpi_communicator + + def clone(self): + return type(self)(self.mpi_communicator) + +except ImportError: + print("Failed to import numpy array context.") + pass + +# }}} + # {{{ tensor product-specific machinery From d0ca536c2557b9cdf82f8955910d98001c7ecaa1 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sun, 11 Feb 2024 16:44:26 -0600 Subject: [PATCH 63/97] Clean up TP arraycontext imports --- .../tensor-product-examples/acoustic_pulse.py | 6 ++---- grudge/array_context.py | 3 +-- test/test_dt_utils.py | 19 ++----------------- 3 files changed, 5 insertions(+), 23 deletions(-) diff --git a/examples/tensor-product-examples/acoustic_pulse.py b/examples/tensor-product-examples/acoustic_pulse.py index 13c2194cf..577d0da4e 100644 --- a/examples/tensor-product-examples/acoustic_pulse.py +++ b/examples/tensor-product-examples/acoustic_pulse.py @@ -216,14 +216,12 @@ def main(ctx_factory, order=3, final_time=1, resolution=16, queue = cl.CommandQueue(cl_ctx) if lazy: - from grudge.array_context import PytatoTensorProductArrayContext - actx = PytatoTensorProductArrayContext( + actx = PytatoPyOpenCLArrayContext( queue, allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)), ) else: - from grudge.array_context import TensorProductArrayContext - actx = TensorProductArrayContext( + actx = PyOpenCLArrayContext( queue, allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)), force_device_scalars=True, diff --git a/grudge/array_context.py b/grudge/array_context.py index 41164c54a..f387d56a7 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -678,8 +678,7 @@ def clone(self): # {{{ tensor product-specific machinery class OutputIsTensorProductDOFArrayOrdered(Tag): - """Signify that the strides will not be of order "C" or "F". See - :class:`grudge.array_context.TensorProductArrayContext` for more details. + """Signify that the strides will not be of order "C" or "F". The strides for the arrays containing tensor product element data are of the form (slow, fastest, faster, fast). These strides are not "C" or "F" order. diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py index 884aa4f78..b0e6cf4e1 100644 --- a/test/test_dt_utils.py +++ b/test/test_dt_utils.py @@ -50,15 +50,8 @@ @pytest.mark.parametrize("tpe", [False, True]) def test_geometric_factors_regular_refinement(actx_factory, name, tpe): from grudge.dt_utils import dt_geometric_factors - import pyopencl as cl - from grudge.array_context import TensorProductArrayContext - if tpe: - ctx = cl.create_some_context() - queue = cl.CommandQueue(ctx) - actx = TensorProductArrayContext(queue) - else: - actx = actx_factory() + actx = actx_factory() # {{{ cases @@ -176,15 +169,7 @@ def rhs(x): @pytest.mark.parametrize("tpe", [False, True]) def test_wave_dt_estimate(actx_factory, dim, degree, tpe, visualize=False): - import pyopencl as cl - from grudge.array_context import TensorProductArrayContext - - if tpe: - ctx = cl.create_some_context() - queue = cl.CommandQueue(ctx) - actx = TensorProductArrayContext(queue) - else: - actx = actx_factory() + actx = actx_factory() # {{{ cases From cb447382649f6d32d7bf6e3ee2c6f7e1a0efb2d7 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Sun, 11 Feb 2024 17:00:04 -0600 Subject: [PATCH 64/97] Add back simplicial test in test_op --- test/test_op.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_op.py b/test/test_op.py index 8e8330622..48d93a5a1 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -202,17 +202,17 @@ def get_flux(u_tpair): # {{{ divergence @pytest.mark.parametrize("group_cls", [ - #SimplexElementGroup, + SimplexElementGroup, TensorProductElementGroup ]) @pytest.mark.parametrize("form", ["strong", "weak"]) -@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("dim", [1, 2, 3]) @pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize(("vectorize", "nested"), [ (False, False), (True, False), (True, True) -]) + ]) def test_divergence(actx_factory, form, dim, order, vectorize, nested, group_cls, visualize=False): actx = actx_factory() From d9fea00ac6c1cc993b9d31ec213fa1ac6c410ba7 Mon Sep 17 00:00:00 2001 From: Addison Alvey-Blanco Date: Mon, 12 Feb 2024 13:26:14 -0600 Subject: [PATCH 65/97] Change how geometric factors are applied --- grudge/op.py | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index dd15db781..93b55b11f 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -347,14 +347,6 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, arg_names=("stiff_1d", f"vec_{xyz_axis}")) ) - # apply metric terms - grad[xyz_axis] = actx.einsum( - 'rej,ej->ej', - ijm[xyz_axis], - grad[xyz_axis], - tagged=(FirstAxisIsElementsTag(),), - arg_names=("inv_jac_t", f"vec_{xyz_axis}") - ) else: diff_mat = get_diff_mat(actx, grp, grp) @@ -371,15 +363,14 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, ) ) - grad[xyz_axis] = actx.einsum( - "rej,ej->ej", - ijm[xyz_axis], - grad[xyz_axis], - tagged=(FirstAxisIsElementsTag(),), - arg_names=("inv_jac_t", f"vec_{xyz_axis}") - ) - - return make_obj_array(grad) + grad = actx.np.stack(grad) + return actx.einsum( + "xrej,rej->xej", + ijm, + grad, + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", f"grad") + ) # }}} From 74fea5201b1091dba5e5886ee751fd6673a1d368 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 12 Feb 2024 15:50:02 -0600 Subject: [PATCH 66/97] Skip 1d dt test for tpe, debugging print. --- test/test_dt_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py index 884aa4f78..d9ea7aef8 100644 --- a/test/test_dt_utils.py +++ b/test/test_dt_utils.py @@ -180,6 +180,8 @@ def test_wave_dt_estimate(actx_factory, dim, degree, tpe, visualize=False): from grudge.array_context import TensorProductArrayContext if tpe: + if dim == 1: + pytest.skip() ctx = cl.create_some_context() queue = cl.CommandQueue(ctx) actx = TensorProductArrayContext(queue) @@ -239,6 +241,7 @@ def test_wave_dt_estimate(actx_factory, dim, degree, tpe, visualize=False): RK4MethodBuilder.output_coeffs)) dt_est = actx.to_numpy(wave_op.estimate_rk4_timestep(actx, dcoll)) + print(f"{dt_est=}") if visualize: re, im = np.mgrid[-4:1:30j, -5:5:30j] From bb82f724899a746a2afa48dc07be2be37ab9f9d3 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 19 Feb 2024 18:50:59 -0600 Subject: [PATCH 67/97] Tag axes for normal to fix fallbacks for 1d cases. --- grudge/geometry/metrics.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py index 260f219b3..59f55daa0 100644 --- a/grudge/geometry/metrics.py +++ b/grudge/geometry/metrics.py @@ -72,6 +72,9 @@ DD_VOLUME_ALL, DOFDesc, DISCR_TAG_BASE ) +from meshmode.transform_metadata import (DiscretizationDOFAxisTag, + DiscretizationElementAxisTag) + from meshmode.transform_metadata import (DiscretizationAmbientDimAxisTag, DiscretizationTopologicalDimAxisTag) @@ -660,11 +663,11 @@ def area_element( @memoize_in(dcoll, (area_element, dd, _use_geoderiv_connection)) def _area_elements(): - result = actx.np.sqrt( - pseudoscalar( - actx, dcoll, dd=dd, - _use_geoderiv_connection=_use_geoderiv_connection).norm_squared()) - + res = pseudoscalar( + actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection + ).norm_squared() + result = actx.np.sqrt(tag_axes(actx, {0: DiscretizationElementAxisTag(), + 1: DiscretizationDOFAxisTag()}, res)) return actx.freeze( actx.tag(NameHint(f"area_el_{dd.as_identifier()}"), result)) From cdd2632377726994566e85e81be68645c6ff32f8 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 26 Feb 2024 14:32:19 -0600 Subject: [PATCH 68/97] make array context fallbacks optional (disabled by default) --- grudge/array_context.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index f387d56a7..7d6e9106f 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -445,9 +445,11 @@ def to_output_template(keys, _): class MPIPytatoArrayContextBase(MPIBasedArrayContext): def __init__( - self, mpi_communicator, queue, *, mpi_base_tag, allocator=None, - compile_trace_callback: Optional[Callable[[Any, str, Any], None]] - = None) -> None: + self, mpi_communicator, queue, *, + mpi_base_tag, allocator=None, + compile_trace_callback: Optional[Callable[[Any, str, Any], None]] = None, + use_axis_tag_inference_fallback: bool = False, + use_einsum_inference_fallback: bool = False) -> None: """ :arg compile_trace_callback: A function of three arguments *(what, stage, ir)*, where *what* identifies the object @@ -462,7 +464,9 @@ def __init__( "to reduce device allocations)") super().__init__(queue, allocator, - compile_trace_callback=compile_trace_callback) + compile_trace_callback=compile_trace_callback, + use_axis_tag_inference_fallback=use_axis_tag_inference_fallback, + use_einsum_inference_fallback=use_einsum_inference_fallback) self.mpi_communicator = mpi_communicator self.mpi_base_tag = mpi_base_tag @@ -477,7 +481,9 @@ def clone(self): # pylint: disable=no-member return type(self)(self.mpi_communicator, self.queue, mpi_base_tag=self.mpi_base_tag, - allocator=self.allocator) + allocator=self.allocator, + use_axis_tag_inference_fallback=self.use_axis_tag_inference_fallback, + use_einsum_inference_fallback=self.use_einsum_inference_fallback) # }}} From 5599defcdcc01850b1204b4fc87ed4d84d586ac3 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 28 Feb 2024 13:04:37 -0600 Subject: [PATCH 69/97] Use vol/max(face_area) for geometric_factor of TPE --- grudge/dt_utils.py | 105 ++++++++++++++++++++++++++++----------------- 1 file changed, 65 insertions(+), 40 deletions(-) diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 0a9b223cf..1eba29639 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -272,8 +272,11 @@ def dt_geometric_factors( volm_discr = dcoll.discr_from_dd(dd) r_fac = dcoll.dim - if any(not isinstance(grp, SimplexElementGroupBase) - for grp in volm_discr.groups): + # assumes !simplex = tpe + tpe = any(not isinstance(grp, SimplexElementGroupBase) + fro grp in volm_discr.groups) + + if tpe: # old way: use average face area r_fac = 2.0*r_fac if volm_discr.dim != volm_discr.ambient_dim: @@ -301,52 +304,74 @@ def dt_geometric_factors( op.elementwise_integral( dcoll, dd_face, face_discr.zeros(actx) + 1.0 ) - ) + } + + if tpe: + # Reshape from (nelements*nfaces, 1) to (nfaces, nelements, 1) + # to get per-element data + node_data_per_group = [] + for igrp, group in enumerate(volm_discr.mesh.groups): + nelements = group.nelements + nfaces = group.nfaces + el_face_areas = face_areas[igrp].reshape(nfaces, nelements, + face_areas[igrp].shape[1]) + if actx.supports_nonscalar_broadcasting: + el_data = actx.np.max(el_face_areas, axis=0)[:, 0:1] + node_data = actx.np.broadcast_to( + el_data, dcoll.zeros(actx, dd=dd)[igrp].shape) + else: + el_data_np = np.max(actx.to_numpy(el_face_areas), axis=0)[:, 0:1] + node_data_np = np.ascontiguousarray( + np.broadcast_to(el_data_np, dcoll.zeros(actx, dd=dd)[igrp].shape)) + node_data = actx.from_numpy(node_data_np) + + node_data_per_group.append(node_data) + surface_areas = rfac * DOFArray(actx, node_data_per_group) - if actx.supports_nonscalar_broadcasting: - # Compute total surface area of an element by summing over the - # individual face areas - surface_areas = DOFArray( - actx, - data=tuple( - actx.einsum( - "fej->e", - tag_axes(actx, { - 0: DiscretizationFaceAxisTag(), - 1: DiscretizationElementAxisTag(), - 2: DiscretizationDOFAxisTag() + else: + if actx.supports_nonscalar_broadcasting: + # Compute total surface area of an element by summing over the + # individual face areas + surface_areas = DOFArray( + actx, + data=tuple( + actx.einsum( + "fej->e", + tag_axes(actx, { + 0: DiscretizationFaceAxisTag(), + 1: DiscretizationElementAxisTag(), + 2: DiscretizationDOFAxisTag() }, + face_ae_i.reshape( + vgrp.mesh_el_group.nfaces, + vgrp.nelements, + face_ae_i.shape[-1])), + tagged=(FirstAxisIsElementsTag(),)) + for vgrp, face_ae_i in zip(volm_discr.groups, face_areas))) + else: + surface_areas = DOFArray( + actx, + data=tuple( + # NOTE: Whenever the array context can't perform nonscalar + # broadcasting, elementwise reductions + # (like `elementwise_integral`) repeat the *same* scalar value of + # the reduction at each degree of freedom. To get a single + # value for the face area (per face), + # we simply average over the nodes, which gives the desired result. + actx.einsum( + "fej->e", face_ae_i.reshape( vgrp.mesh_el_group.nfaces, vgrp.nelements, - face_ae_i.shape[-1])), + face_ae_i.shape[-1] + ) / afgrp.nunit_dofs, tagged=(FirstAxisIsElementsTag(),)) - for vgrp, face_ae_i in zip(volm_discr.groups, face_areas))) - else: - surface_areas = DOFArray( - actx, - data=tuple( - # NOTE: Whenever the array context can't perform nonscalar - # broadcasting, elementwise reductions - # (like `elementwise_integral`) repeat the *same* scalar value of - # the reduction at each degree of freedom. To get a single - # value for the face area (per face), - # we simply average over the nodes, which gives the desired result. - actx.einsum( - "fej->e", - face_ae_i.reshape( - vgrp.mesh_el_group.nfaces, - vgrp.nelements, - face_ae_i.shape[-1] - ) / afgrp.nunit_dofs, - tagged=(FirstAxisIsElementsTag(),)) - - for vgrp, afgrp, face_ae_i in zip(volm_discr.groups, - face_discr.groups, - face_areas) + for vgrp, afgrp, face_ae_i in zip(volm_discr.groups, + face_discr.groups, + face_areas) + ) ) - ) return actx.freeze( actx.tag(NameHint(f"dt_geometric_{dd.as_identifier()}"), From de152130f94b3e9095233865db2813e0b1522b12 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 28 Feb 2024 15:49:50 -0600 Subject: [PATCH 70/97] Makes lazy work, eager still gives old result --- grudge/dt_utils.py | 74 ++++++++++++++++++++++++++++------------------ 1 file changed, 45 insertions(+), 29 deletions(-) diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 1eba29639..4ea9a76af 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -274,7 +274,7 @@ def dt_geometric_factors( r_fac = dcoll.dim # assumes !simplex = tpe tpe = any(not isinstance(grp, SimplexElementGroupBase) - fro grp in volm_discr.groups) + for grp in volm_discr.groups) if tpe: # old way: use average face area r_fac = 2.0*r_fac @@ -304,35 +304,52 @@ def dt_geometric_factors( op.elementwise_integral( dcoll, dd_face, face_discr.zeros(actx) + 1.0 ) - } - + ) + print(f"{face_areas=}") if tpe: - # Reshape from (nelements*nfaces, 1) to (nfaces, nelements, 1) - # to get per-element data - node_data_per_group = [] - for igrp, group in enumerate(volm_discr.mesh.groups): - nelements = group.nelements - nfaces = group.nfaces - el_face_areas = face_areas[igrp].reshape(nfaces, nelements, - face_areas[igrp].shape[1]) - if actx.supports_nonscalar_broadcasting: - el_data = actx.np.max(el_face_areas, axis=0)[:, 0:1] - node_data = actx.np.broadcast_to( - el_data, dcoll.zeros(actx, dd=dd)[igrp].shape) - else: - el_data_np = np.max(actx.to_numpy(el_face_areas), axis=0)[:, 0:1] - node_data_np = np.ascontiguousarray( - np.broadcast_to(el_data_np, dcoll.zeros(actx, dd=dd)[igrp].shape)) - node_data = actx.from_numpy(node_data_np) - - node_data_per_group.append(node_data) - surface_areas = rfac * DOFArray(actx, node_data_per_group) + if actx.supports_nonscalar_broadcasting: + surface_areas = r_fac * DOFArray( + actx, + data=tuple( + actx.np.max( + tag_axes(actx, { + 0: DiscretizationFaceAxisTag(), + 1: DiscretizationElementAxisTag(), + }, + face_ae_i.reshape( + vgrp.mesh_el_group.nfaces, + vgrp.nelements)), + axis=0) + for vgrp, face_ae_i in zip(volm_discr.groups, face_areas))) + else: + surface_areas = DOFArray( + actx, + data=tuple( + # NOTE: Whenever the array context can't perform nonscalar + # broadcasting, elementwise reductions + # (like `elementwise_integral`) repeat the *same* scalar value of + # the reduction at each degree of freedom. To get a single + # value for the face area (per face), + # we simply average over the nodes, which gives the desired result. + actx.einsum( + "fej->e", + face_ae_i.reshape( + vgrp.mesh_el_group.nfaces, + vgrp.nelements, + face_ae_i.shape[-1] + ) / afgrp.nunit_dofs, + tagged=(FirstAxisIsElementsTag(),)) + for vgrp, afgrp, face_ae_i in zip(volm_discr.groups, + face_discr.groups, + face_areas) + ) + ) else: if actx.supports_nonscalar_broadcasting: # Compute total surface area of an element by summing over the # individual face areas - surface_areas = DOFArray( + surface_areas = r_fac * DOFArray( actx, data=tuple( actx.einsum( @@ -342,10 +359,10 @@ def dt_geometric_factors( 1: DiscretizationElementAxisTag(), 2: DiscretizationDOFAxisTag() }, - face_ae_i.reshape( - vgrp.mesh_el_group.nfaces, - vgrp.nelements, - face_ae_i.shape[-1])), + face_ae_i.reshape( + vgrp.mesh_el_group.nfaces, + vgrp.nelements, + face_ae_i.shape[-1])), tagged=(FirstAxisIsElementsTag(),)) for vgrp, face_ae_i in zip(volm_discr.groups, face_areas))) else: @@ -383,7 +400,6 @@ def dt_geometric_factors( actx.tag_axis(1, DiscretizationDOFAxisTag(), cv_i), tagged=(FirstAxisIsElementsTag(),)) * r_fac for cv_i, sae_i in zip(cell_vols, surface_areas))))) - # }}} From 41aba76e313d92dee68b790b15b378528c7fd494 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Thu, 29 Feb 2024 10:02:11 -0600 Subject: [PATCH 71/97] Fix up max face area snippet for non-lazy actx. --- grudge/dt_utils.py | 35 ++++++++++++----------------------- 1 file changed, 12 insertions(+), 23 deletions(-) diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 4ea9a76af..11583f270 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -322,29 +322,18 @@ def dt_geometric_factors( axis=0) for vgrp, face_ae_i in zip(volm_discr.groups, face_areas))) else: - surface_areas = DOFArray( - actx, - data=tuple( - # NOTE: Whenever the array context can't perform nonscalar - # broadcasting, elementwise reductions - # (like `elementwise_integral`) repeat the *same* scalar value of - # the reduction at each degree of freedom. To get a single - # value for the face area (per face), - # we simply average over the nodes, which gives the desired result. - actx.einsum( - "fej->e", - face_ae_i.reshape( - vgrp.mesh_el_group.nfaces, - vgrp.nelements, - face_ae_i.shape[-1] - ) / afgrp.nunit_dofs, - tagged=(FirstAxisIsElementsTag(),)) - - for vgrp, afgrp, face_ae_i in zip(volm_discr.groups, - face_discr.groups, - face_areas) - ) - ) + el_data_per_group = [] + for igrp, group in enumerate(volm_discr.mesh.groups): + nelements = group.nelements + nfaces = group.nfaces + el_face_data = face_areas[igrp].reshape(nfaces, nelements, + face_areas[igrp].shape[1]) + el_data_np = np.ascontiguousarray( + np.max(actx.to_numpy(el_face_data), axis=0)[:, 0:1]) + el_data = actx.from_numpy(el_data_np) + el_data = el_data.reshape(nelements) + el_data_per_group.append(el_data) + surface_areas = DOFArray(actx, tuple(el_data_per_group)) else: if actx.supports_nonscalar_broadcasting: # Compute total surface area of an element by summing over the From 7ef92bcaaf8ee925fa4505f76a90eb2fb880dfa3 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Thu, 29 Feb 2024 10:12:24 -0600 Subject: [PATCH 72/97] Bring back missing scale factor for face areas. --- grudge/dt_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 11583f270..260d59716 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -305,7 +305,7 @@ def dt_geometric_factors( dcoll, dd_face, face_discr.zeros(actx) + 1.0 ) ) - print(f"{face_areas=}") + if tpe: if actx.supports_nonscalar_broadcasting: surface_areas = r_fac * DOFArray( @@ -333,7 +333,7 @@ def dt_geometric_factors( el_data = actx.from_numpy(el_data_np) el_data = el_data.reshape(nelements) el_data_per_group.append(el_data) - surface_areas = DOFArray(actx, tuple(el_data_per_group)) + surface_areas = r_fac * DOFArray(actx, tuple(el_data_per_group)) else: if actx.supports_nonscalar_broadcasting: # Compute total surface area of an element by summing over the From e315767f33a4622350eaa4b0ece07c90cc9b3a7b Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Sat, 2 Mar 2024 12:20:37 -0600 Subject: [PATCH 73/97] Correct misfac --- grudge/dt_utils.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 260d59716..246b938c0 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -271,13 +271,11 @@ def dt_geometric_factors( actx = dcoll._setup_actx volm_discr = dcoll.discr_from_dd(dd) - r_fac = dcoll.dim # assumes !simplex = tpe tpe = any(not isinstance(grp, SimplexElementGroupBase) for grp in volm_discr.groups) - if tpe: # old way: use average face area - r_fac = 2.0*r_fac + r_fac = 1.0 if tpe else dcoll.dim if volm_discr.dim != volm_discr.ambient_dim: from warnings import warn @@ -308,7 +306,7 @@ def dt_geometric_factors( if tpe: if actx.supports_nonscalar_broadcasting: - surface_areas = r_fac * DOFArray( + surface_areas = DOFArray( actx, data=tuple( actx.np.max( @@ -333,12 +331,12 @@ def dt_geometric_factors( el_data = actx.from_numpy(el_data_np) el_data = el_data.reshape(nelements) el_data_per_group.append(el_data) - surface_areas = r_fac * DOFArray(actx, tuple(el_data_per_group)) + surface_areas = DOFArray(actx, tuple(el_data_per_group)) else: if actx.supports_nonscalar_broadcasting: # Compute total surface area of an element by summing over the # individual face areas - surface_areas = r_fac * DOFArray( + surface_areas = DOFArray( actx, data=tuple( actx.einsum( From a5a58ea0892de65567900b0d14da36346d7c57a9 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Sat, 2 Mar 2024 14:15:12 -0600 Subject: [PATCH 74/97] Tweak more --- grudge/dt_utils.py | 2 +- test/test_dt_utils.py | 64 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py index 246b938c0..9754dc6e0 100644 --- a/grudge/dt_utils.py +++ b/grudge/dt_utils.py @@ -275,7 +275,7 @@ def dt_geometric_factors( tpe = any(not isinstance(grp, SimplexElementGroupBase) for grp in volm_discr.groups) - r_fac = 1.0 if tpe else dcoll.dim + r_fac = 0.5 if tpe else dcoll.dim if volm_discr.dim != volm_discr.ambient_dim: from warnings import warn diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py index be7352028..7fee076b7 100644 --- a/test/test_dt_utils.py +++ b/test/test_dt_utils.py @@ -258,6 +258,70 @@ def test_wave_dt_estimate(actx_factory, dim, degree, tpe, visualize=False): assert not stable_dt_factors or max(stable_dt_factors) < 1.5, stable_dt_factors +@pytest.mark.parametrize("dim", [2]) +@pytest.mark.parametrize("degree", [1, 2]) +@pytest.mark.parametrize("tpe", [True]) +def test_charlen(actx_factory, dim, degree, tpe, visualize=False): + + from grudge.dt_utils import ( + dt_geometric_factors, + dt_non_geometric_factors, + h_min_from_volume, + h_max_from_volume + ) + actx = actx_factory() + + if tpe: + if dim == 1: + pytest.skip() + + # {{{ cases + + from meshmode.mesh import TensorProductElementGroup + group_cls = TensorProductElementGroup if tpe else None + + import meshmode.mesh.generation as mgen + + a = [0, 0, 0] + b = [1, 1, 1] + nels1d = [2, 3, 4] + + for nel1d in nels1d: + print(f"{dim=},{nel1d=},{degree=}") + mesh = mgen.generate_regular_rect_mesh( + a=a[:dim], b=b[:dim], + nelements_per_axis=(nel1d,)*dim, + group_cls=group_cls) + print(f"{mesh=}") + assert mesh.dim == dim + + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as Lgl + + from grudge.dof_desc import DISCR_TAG_BASE + order = degree + dtag_to_grp_fac = None + if tpe: + order = None + dtag_to_grp_fac = { + DISCR_TAG_BASE: Lgl(degree) + } + + dcoll = DiscretizationCollection(actx, mesh, order=order, + discr_tag_to_group_factory=dtag_to_grp_fac) + + h_min = actx.to_numpy(h_min_from_volume(dcoll)) + h_max = actx.to_numpy(h_max_from_volume(dcoll)) + gfac = actx.to_numpy(dt_geometric_factors(dcoll)) + ngfac = dt_non_geometric_factors(dcoll) + + print(f"{h_min=}") + print(f"{h_max=}") + print(f"{gfac=}") + print(f"{ngfac=}") + + assert False + # You can test individual routines by typing # $ python test_grudge.py 'test_routine()' From e7af8e844fe1a403ceefe56e743fe42e13d8fb4d Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 1 May 2024 13:57:51 -0500 Subject: [PATCH 75/97] Update dt tpe --- test/test_dt_utils.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py index 7fee076b7..66f72e726 100644 --- a/test/test_dt_utils.py +++ b/test/test_dt_utils.py @@ -33,7 +33,7 @@ [PytestPyOpenCLArrayContextFactory, PytestPytatoPyOpenCLArrayContextFactory]) -from grudge import DiscretizationCollection +from grudge import make_discretization_collection import grudge.op as op @@ -84,8 +84,8 @@ def test_geometric_factors_regular_refinement(actx_factory, name, tpe): min_factors = [] for resolution in builder.resolutions: mesh = builder.get_mesh(resolution, builder.mesh_order) - dcoll = DiscretizationCollection(actx, mesh, order=order, - discr_tag_to_group_factory=dtag_to_grp_fac) + dcoll = make_discretization_collection(actx, mesh, order=order, + discr_tag_to_group_factory=dtag_to_grp_fac) min_factors.append( actx.to_numpy( op.nodal_min(dcoll, "vol", actx.thaw(dt_geometric_factors(dcoll)))) @@ -99,8 +99,8 @@ def test_geometric_factors_regular_refinement(actx_factory, name, tpe): # Make sure it works with empty meshes mesh = builder.get_mesh(0, builder.mesh_order) - dcoll = DiscretizationCollection(actx, mesh, order=order, - discr_tag_to_group_factory=dtag_to_grp_fac) + dcoll = make_discretization_collection(actx, mesh, order=order, + discr_tag_to_group_factory=dtag_to_grp_fac) factors = actx.thaw(dt_geometric_factors(dcoll)) # noqa: F841 @@ -130,7 +130,7 @@ def test_non_geometric_factors(actx_factory, name): degrees = list(range(1, 8)) for degree in degrees: mesh = builder.get_mesh(1, degree) - dcoll = DiscretizationCollection(actx, mesh, order=degree) + dcoll = make_discretization_collection(actx, mesh, order=degree) factors.append(min(dt_non_geometric_factors(dcoll))) # Crude estimate, factors should behave like 1/N**2 @@ -149,7 +149,7 @@ def test_build_jacobian(actx_factory): mesh = mgen.generate_regular_rect_mesh(a=[0], b=[1], nelements_per_axis=(3,)) assert mesh.dim == 1 - dcoll = DiscretizationCollection(actx, mesh, order=1) + dcoll = make_discretization_collection(actx, mesh, order=1) def rhs(x): return 3*x**2 + 2*x + 5 @@ -202,7 +202,7 @@ def test_wave_dt_estimate(actx_factory, dim, degree, tpe, visualize=False): dtag_to_grp_fac = { DISCR_TAG_BASE: Lgl(degree) } - dcoll = DiscretizationCollection(actx, mesh, order=order, + dcoll = make_discretization_collection(actx, mesh, order=order, discr_tag_to_group_factory=dtag_to_grp_fac) from grudge.models.wave import WeakWaveOperator @@ -307,7 +307,7 @@ def test_charlen(actx_factory, dim, degree, tpe, visualize=False): DISCR_TAG_BASE: Lgl(degree) } - dcoll = DiscretizationCollection(actx, mesh, order=order, + dcoll = make_discretization_collection(actx, mesh, order=order, discr_tag_to_group_factory=dtag_to_grp_fac) h_min = actx.to_numpy(h_min_from_volume(dcoll)) From 4c72f9cd4f5c9b1538d88f8c8ccd06aae8ce5182 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 1 May 2024 18:42:13 -0500 Subject: [PATCH 76/97] Update TPE tests - esp for wave rk4 estimate --- grudge/models/wave.py | 8 +++++++- test/test_dt_utils.py | 9 +++++---- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/grudge/models/wave.py b/grudge/models/wave.py index 4fc726870..dbe45abe7 100644 --- a/grudge/models/wave.py +++ b/grudge/models/wave.py @@ -189,7 +189,13 @@ def max_characteristic_velocity(self, actx, t=None, fields=None): def estimate_rk4_timestep(self, actx, dcoll, **kwargs): # FIXME: Sketchy, empirically determined fudge factor - return 0.38 * super().estimate_rk4_timestep(actx, dcoll, **kwargs) + from meshmode.discretization.poly_element import SimplexElementGroupBase + from grudge.dof_desc import DD_VOLUME_ALL + volm_discr = dcoll.discr_from_dd(DD_VOLUME_ALL) + tpe = any(not isinstance(grp, SimplexElementGroupBase) + for grp in volm_discr.groups) + fudge_fac = 0.38 if not tpe else 0.23 + return fudge_fac * super().estimate_rk4_timestep(actx, dcoll, **kwargs) # }}} diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py index 66f72e726..0c3c3d3cb 100644 --- a/test/test_dt_utils.py +++ b/test/test_dt_utils.py @@ -98,10 +98,11 @@ def test_geometric_factors_regular_refinement(actx_factory, name, tpe): assert np.all(np.isclose(ratios, 2)) # Make sure it works with empty meshes - mesh = builder.get_mesh(0, builder.mesh_order) - dcoll = make_discretization_collection(actx, mesh, order=order, - discr_tag_to_group_factory=dtag_to_grp_fac) - factors = actx.thaw(dt_geometric_factors(dcoll)) # noqa: F841 + if not tpe: + mesh = builder.get_mesh(0, builder.mesh_order) + dcoll = make_discretization_collection(actx, mesh, order=order, + discr_tag_to_group_factory=dtag_to_grp_fac) + factors = actx.thaw(dt_geometric_factors(dcoll)) # noqa: F841 @pytest.mark.parametrize("name", ["interval", "box2d", "box3d"]) From 811b3c09d16e6445d38f27bfc39bf5f8cf605bd0 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Thu, 13 Jun 2024 15:29:08 -0500 Subject: [PATCH 77/97] Add TPE option, reorder term calc in euler to expose issue with face normals in lazy --- examples/euler/acoustic_pulse.py | 28 +++++++++++++++++----------- grudge/models/euler.py | 12 ++++++------ 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/examples/euler/acoustic_pulse.py b/examples/euler/acoustic_pulse.py index 779062910..d5bbc0375 100644 --- a/examples/euler/acoustic_pulse.py +++ b/examples/euler/acoustic_pulse.py @@ -37,6 +37,7 @@ EulerOperator, InviscidWallBC ) +from meshmode.mesh import TensorProductElementGroup from grudge.shortcuts import rk4_step from meshmode.mesh import BTAG_ALL @@ -112,7 +113,8 @@ def run_acoustic_pulse(actx, final_time=1, resolution=16, overintegration=False, - visualize=False): + visualize=False, + tpe=False): # eos-related parameters gamma = 1.4 @@ -124,16 +126,18 @@ def run_acoustic_pulse(actx, dim = 2 box_ll = -0.5 box_ur = 0.5 + group_cls = TensorProductElementGroup if tpe else None mesh = generate_regular_rect_mesh( a=(box_ll,)*dim, b=(box_ur,)*dim, - nelements_per_axis=(resolution,)*dim) + nelements_per_axis=(resolution,)*dim, + group_cls=group_cls) from grudge import DiscretizationCollection from grudge.dof_desc import DISCR_TAG_BASE, DISCR_TAG_QUAD - from meshmode.discretization.poly_element import \ - (default_simplex_group_factory, - QuadratureSimplexGroupFactory) + from meshmode.discretization.poly_element import ( + InterpolatoryEdgeClusteredGroupFactory, + QuadratureGroupFactory) exp_name = f"fld-acoustic-pulse-N{order}-K{resolution}" if overintegration: @@ -145,9 +149,8 @@ def run_acoustic_pulse(actx, dcoll = DiscretizationCollection( actx, mesh, discr_tag_to_group_factory={ - DISCR_TAG_BASE: default_simplex_group_factory( - base_dim=mesh.dim, order=order), - DISCR_TAG_QUAD: QuadratureSimplexGroupFactory(2*order) + DISCR_TAG_BASE: InterpolatoryEdgeClusteredGroupFactory(order), + DISCR_TAG_QUAD: QuadratureGroupFactory(2*order) } ) @@ -212,7 +215,8 @@ def rhs(t, q): def main(ctx_factory, order=3, final_time=1, resolution=16, - overintegration=False, visualize=False, lazy=False): + overintegration=False, visualize=False, lazy=False, + tpe=False): cl_ctx = ctx_factory() queue = cl.CommandQueue(cl_ctx) @@ -234,7 +238,7 @@ def main(ctx_factory, order=3, final_time=1, resolution=16, resolution=resolution, overintegration=overintegration, final_time=final_time, - visualize=visualize + visualize=visualize, tpe=tpe ) @@ -251,6 +255,8 @@ def main(ctx_factory, order=3, final_time=1, resolution=16, help="write out vtk output") parser.add_argument("--lazy", action="store_true", help="switch to a lazy computation mode") + parser.add_argument("--tpe", action="store_true", + help="use tensor product elements") args = parser.parse_args() logging.basicConfig(level=logging.INFO) @@ -260,4 +266,4 @@ def main(ctx_factory, order=3, final_time=1, resolution=16, resolution=args.resolution, overintegration=args.oi, visualize=args.visualize, - lazy=args.lazy) + lazy=args.lazy, tpe=args.tpe) diff --git a/grudge/models/euler.py b/grudge/models/euler.py index 1b6eb569c..25d26e780 100644 --- a/grudge/models/euler.py +++ b/grudge/models/euler.py @@ -322,12 +322,6 @@ def operator(self, t, q): def interp_to_quad(u): return op.project(dcoll, "vol", dq, u) - # Compute volume fluxes - volume_fluxes = op.weak_local_div( - dcoll, dq, - interp_to_quad(euler_volume_flux(dcoll, q, gamma=gamma)) - ) - # Compute interior interface fluxes interface_fluxes = ( sum( @@ -357,6 +351,12 @@ def interp_to_quad(u): ) interface_fluxes = interface_fluxes + bc_fluxes + # Compute volume fluxes + volume_fluxes = op.weak_local_div( + dcoll, dq, + interp_to_quad(euler_volume_flux(dcoll, q, gamma=gamma)) + ) + return op.inverse_mass( dcoll, volume_fluxes - op.face_mass(dcoll, df, interface_fluxes) From 5b563de6fc84d7fed2d543944bdd89410aeb9d64 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 17 Jun 2024 17:39:59 -0500 Subject: [PATCH 78/97] Use geoderiv_connection only for Simplices --- grudge/geometry/metrics.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py index 260f219b3..3e8177cca 100644 --- a/grudge/geometry/metrics.py +++ b/grudge/geometry/metrics.py @@ -86,6 +86,11 @@ register_multivector_as_array_container() +def _has_geoderiv_connection(grp): + from modepy.shapes import Simplex + return grp.is_affine and issubclass(grp._modepy_shape_cls, Simplex) + + def _geometry_to_quad_if_requested( dcoll, inner_dd, dd, vec, _use_geoderiv_connection): @@ -105,7 +110,7 @@ def to_quad(vec): return DOFArray( vec.array_context, tuple( - geoderiv_vec_i if megrp.is_affine else all_quad_vec_i + geoderiv_vec_i if _has_geoderiv_connection(megrp) else all_quad_vec_i for megrp, geoderiv_vec_i, all_quad_vec_i in zip( dcoll.discr_from_dd(inner_dd).mesh.groups, dcoll._base_to_geoderiv_connection(inner_dd)(vec), From 6df78a7c57a2e4bf0e3aaf9fb89d19762649a94c Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 17 Jun 2024 17:40:54 -0500 Subject: [PATCH 79/97] Allow generation of TPE mesh type --- test/mesh_data.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/test/mesh_data.py b/test/mesh_data.py index c71950bb9..6658f495e 100644 --- a/test/mesh_data.py +++ b/test/mesh_data.py @@ -4,6 +4,7 @@ from meshmode.mesh.io import read_gmsh import numpy as np import meshmode.mesh.generation as mgen +from meshmode.mesh import TensorProductElementGroup class MeshBuilder(ABC): @@ -111,6 +112,7 @@ def get_mesh(self, resolution, mesh_order=4): class _BoxMeshBuilderBase(MeshBuilder): resolutions = [4, 8, 16] mesh_order = 1 + group_cls = None a = (-0.5, -0.5, -0.5) b = (+0.5, +0.5, +0.5) @@ -122,20 +124,32 @@ def get_mesh(self, resolution, mesh_order=4): return mgen.generate_regular_rect_mesh( a=self.a, b=self.b, nelements_per_axis=resolution, - order=mesh_order) + order=mesh_order, group_cls=self.group_cls) class BoxMeshBuilder1D(_BoxMeshBuilderBase): ambient_dim = 1 + def __init__(self, tpe=False): + if tpe: + self.group_cls = TensorProductElementGroup + class BoxMeshBuilder2D(_BoxMeshBuilderBase): ambient_dim = 2 + def __init__(self, tpe=False): + if tpe: + self.group_cls = TensorProductElementGroup + class BoxMeshBuilder3D(_BoxMeshBuilderBase): ambient_dim = 2 + def __init__(self, tpe=False): + if tpe: + self.group_cls = TensorProductElementGroup + class WarpedRectMeshBuilder(MeshBuilder): resolutions = [4, 6, 8] From 00b6ca45efc2b24daa6caf4190addc75945ed34b Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 17 Jun 2024 17:41:26 -0500 Subject: [PATCH 80/97] Extend some tests to hit TPEs --- test/test_grudge.py | 48 +++++++++++++++++++++++++++++++-------------- 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/test/test_grudge.py b/test/test_grudge.py index ff42eb63f..c1432f37d 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -61,10 +61,10 @@ # {{{ mass operator trig integration +@pytest.mark.parametrize("tpe", [True, False]) @pytest.mark.parametrize("ambient_dim", [1, 2, 3]) -@pytest.mark.parametrize("discr_tag", [dof_desc.DISCR_TAG_BASE, - dof_desc.DISCR_TAG_QUAD]) -def test_mass_mat_trig(actx_factory, ambient_dim, discr_tag): +@pytest.mark.parametrize("use_overint", [False, True]) +def test_mass_mat_trig(actx_factory, tpe, ambient_dim, use_overint): """Check the integral of some trig functions on an interval using the mass matrix. """ @@ -75,22 +75,26 @@ def test_mass_mat_trig(actx_factory, ambient_dim, discr_tag): a = -4.0 * np.pi b = +9.0 * np.pi + true_integral = 13*np.pi/2 * (b - a)**(ambient_dim - 1) + discr_tag = dof_desc.DISCR_TAG_QUAD if use_overint else dof_desc.DISCR_TAG_BASE - from meshmode.discretization.poly_element import QuadratureSimplexGroupFactory dd_quad = dof_desc.DOFDesc(dof_desc.DTAG_VOLUME_ALL, discr_tag) + discr_order = order if discr_tag is dof_desc.DISCR_TAG_BASE: discr_tag_to_group_factory = {} else: + discr_order = None discr_tag_to_group_factory = { - discr_tag: QuadratureSimplexGroupFactory(order=2*order) + dof_desc.DISCR_TAG_BASE: InterpolatoryEdgeClusteredGroupFactory(order), + dof_desc.DISCR_TAG_QUAD: QuadratureGroupFactory(order=2*order) } mesh = mgen.generate_regular_rect_mesh( a=(a,)*ambient_dim, b=(b,)*ambient_dim, nelements_per_axis=(nel_1d,)*ambient_dim, order=1) dcoll = DiscretizationCollection( - actx, mesh, order=order, + actx, mesh, order=discr_order, discr_tag_to_group_factory=discr_tag_to_group_factory ) @@ -160,10 +164,11 @@ def _spheroid_surface_area(radius, aspect_ratio): return 2.0 * np.pi * radius**2 * (1 + (c/a)**2 / e * np.arctanh(e)) +@pytest.mark.parametrize("tpe", [True, False]) @pytest.mark.parametrize("name", [ "2-1-ellipse", "spheroid", "box2d", "box3d" ]) -def test_mass_surface_area(actx_factory, name): +def test_mass_surface_area(actx_factory, tpe, name): actx = actx_factory() # {{{ cases @@ -171,16 +176,20 @@ def test_mass_surface_area(actx_factory, name): order = 4 if name == "2-1-ellipse": + if tpe: + pytest.skip() builder = mesh_data.EllipseMeshBuilder(radius=3.1, aspect_ratio=2.0) surface_area = _ellipse_surface_area(builder.radius, builder.aspect_ratio) elif name == "spheroid": + if tpe: + pytest.skip() builder = mesh_data.SpheroidMeshBuilder() surface_area = _spheroid_surface_area(builder.radius, builder.aspect_ratio) elif name == "box2d": - builder = mesh_data.BoxMeshBuilder2D() + builder = mesh_data.BoxMeshBuilder2D(tpe) surface_area = 1.0 elif name == "box3d": - builder = mesh_data.BoxMeshBuilder3D() + builder = mesh_data.BoxMeshBuilder3D(tpe) surface_area = 1.0 else: raise ValueError("unknown geometry name: %s" % name) @@ -976,11 +985,11 @@ def rhs(t, w): # {{{ models: variable coefficient advection oversampling @pytest.mark.parametrize("order", [2, 3, 4]) -def test_improvement_quadrature(actx_factory, order): +@pytest.mark.parametrize("tpe", [False, True]) +def test_improvement_quadrature(actx_factory, order, tpe): """Test whether quadrature improves things and converges""" from grudge.models.advection import VariableCoefficientAdvectionOperator from pytools.convergence import EOCRecorder - from meshmode.discretization.poly_element import QuadratureSimplexGroupFactory from meshmode.mesh import BTAG_ALL actx = actx_factory() @@ -1002,23 +1011,30 @@ def conv_test(descr, use_quad): else: qtag = None + group_cls = TensorProductElementGroup if tpe else None + qfac = 2 if tpe else 4 ns = [20, 25] + discr_order = order for n in ns: mesh = mgen.generate_regular_rect_mesh( a=(-0.5,)*dims, b=(0.5,)*dims, nelements_per_axis=(n,)*dims, - order=order) + order=order, group_cls=group_cls) if use_quad: discr_tag_to_group_factory = { - qtag: QuadratureSimplexGroupFactory(order=4*order) + dof_desc.DISCR_TAG_BASE: + InterpolatoryEdgeClusteredGroupFactory(order), + dof_desc.DISCR_TAG_QUAD: + QuadratureGroupFactory(order=qfac*order) } + discr_order = None else: discr_tag_to_group_factory = {} dcoll = DiscretizationCollection( - actx, mesh, order=order, + actx, mesh, order=discr_order, discr_tag_to_group_factory=discr_tag_to_group_factory ) @@ -1050,9 +1066,11 @@ def zero_inflow(dtag, t=0): eoc, errs = conv_test("no quadrature", False) q_eoc, q_errs = conv_test("with quadrature", True) - assert q_eoc > eoc assert (q_errs < errs).all() assert q_eoc > order - 0.1 + # Fails for all tensor-product element types + assert q_eoc > eoc + # }}} From 95d68340cd3997edefc1b0ae5912409e91394334 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 19 Jun 2024 10:22:09 -0500 Subject: [PATCH 81/97] Disable tpe div for now --- grudge/op.py | 60 ++++++++++++++++++++++------------------------------ 1 file changed, 25 insertions(+), 35 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 93b55b11f..878468092 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -197,7 +197,6 @@ def _single_axis_derivative_kernel( # or inside (weak) the matrix-vector product that carries out the # derivative, cf. "metric_in_matvec". - # {{{ tensor product single axis derivative def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, @@ -228,7 +227,7 @@ def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, ) derivative = actx.einsum( - 'rej,ej->ej', + "rej,ej->ej", ijm[xyz_axis], ref_weak_derivative, tagged=(FirstAxisIsElementsTag(),), @@ -248,7 +247,7 @@ def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, ) derivative = actx.einsum( - 'rej,ej->ej', + "rej,ej->ej", ijm[xyz_axis], ref_derivative, tagged=(FirstAxisIsElementsTag(),), @@ -259,7 +258,6 @@ def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, # }}} - # {{{ simplicial single axis derivative def compute_simplicial_derivative(actx, in_grp, out_grp, @@ -279,7 +277,6 @@ def compute_simplicial_derivative(actx, in_grp, out_grp, # }}} - return DOFArray( actx, data=tuple( @@ -299,7 +296,6 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. - # {{{ tensor product gradient def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, @@ -310,10 +306,10 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, """ if grp.dim > 3 and metric_in_matvec: - warn('Efficient tensor product weak ' - 'differentiation operators only ' - 'implemented for dimension 2 and 3. ' - 'Defaulting to inefficient version.') + warn("Efficient tensor product weak " + "differentiation operators only " + "implemented for dimension 2 and 3. " + "Defaulting to inefficient version.") return compute_simplicial_grad(actx, grp, grp, diff_mat, vec, ijm, metric_in_matvec) @@ -334,8 +330,7 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, actx, grp.dim, mass_1d, ax, grad[xyz_axis], tags=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("mass_1d", f"vec_{xyz_axis}") - ) + arg_names=("mass_1d", f"vec_{xyz_axis}")) # apply stiffness operator and unfold grad[xyz_axis] = unfold( @@ -369,12 +364,11 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, ijm, grad, tagged=(FirstAxisIsElementsTag(),), - arg_names=("inv_jac_t", f"grad") + arg_names=("inv_jac_t", "grad") ) # }}} - # {{{ simplicial grad def compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, @@ -393,7 +387,6 @@ def compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, # }}} - per_group_grads = [ compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i, metric_in_matvec) @@ -417,7 +410,6 @@ def _divergence_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. - # {{{ tensor product div def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): @@ -427,10 +419,10 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): """ if grp.dim > 3 and metric_in_matvec: - warn('Efficient tensor product weak ' - 'differentiation operators only ' - 'implemented for dimension 2 and 3. ' - 'Defaulting to inefficient version.') + warn("Efficient tensor product weak " + "differentiation operators only " + "implemented for dimension 2 and 3. " + "Defaulting to inefficient version.") return compute_simplicial_div(actx, grp, grp, diff_mat, vec, ijm, metric_in_matvec) @@ -492,7 +484,7 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): partials = partials.reshape(grp.dim, grp.dim, *partials.shape[-2:]) div = actx.einsum( - 'xrej,xrej->ej', + "xrej,xrej->ej", ijm, partials, arg_names=("inv_jac_t", "partials"), @@ -502,7 +494,6 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): return div # }}} - # {{{ simplicial div def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, @@ -521,16 +512,17 @@ def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, # }}} - per_group_divs = [ - compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) - if isinstance(in_grp, TensorProductElementGroup) + # Disable for now: TPE div not working properly + # compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) + # if isinstance(in_grp, TensorProductElementGroup) # r for rst axis # x for xyz axis - else compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, - ijm_i, metric_in_matvec) + # else + compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec) for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, @@ -569,9 +561,10 @@ def get_ref_derivative_mats(grp): from arraycontext.metadata import NameHint return actx.freeze( - actx.tag(NameHint("tp_diff_mat_1d"), - tag_axes(actx, { - 1: DiscretizationDOFAxisTag()}, + actx.tag( + NameHint("tp_diff_mat_1d"), + tag_axes(actx, { + 1: DiscretizationDOFAxisTag()}, diff_mat))) elif isinstance(grp, SimplexElementGroup): @@ -1180,7 +1173,6 @@ def _apply_inverse_mass_operator( inv_area_elements = 1./area_element(actx, dcoll, dd=dd_in, _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) - def apply_to_tensor_product_elements(grp, jac_inv, vec, ref_inv_mass): vec = fold(grp.space, vec) @@ -1201,7 +1193,6 @@ def apply_to_tensor_product_elements(grp, jac_inv, vec, ref_inv_mass): tagged=(FirstAxisIsElementsTag(),) ) - def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): # Based on https://arxiv.org/pdf/1608.03836.pdf @@ -1213,7 +1204,6 @@ def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): vec, tagged=(FirstAxisIsElementsTag(),)) - group_data = [ apply_to_tensor_product_elements( grp, jac_inv, vec_i, @@ -1482,11 +1472,11 @@ def single_axis_operator_application(actx, dim, operator, axis, data, if not isinstance(tags, tuple): raise TypeError("arg_names must be a tuple.") - operator_spec = 'ij' + operator_spec = "ij" data_spec = f'e{"abcdefghklm"[:axis]}j{"nopqrstuvwxyz"[:dim-axis-1]}' out_spec = f'e{"abcdefghklm"[:axis]}i{"nopqrstuvwxyz"[:dim-axis-1]}' - spec = operator_spec + ',' + data_spec + '->' + out_spec + spec = operator_spec + "," + data_spec + "->" + out_spec return actx.einsum(spec, operator, data, arg_names=arg_names, From 85c0eece0b8734102faa96e34e072b27a2c95e71 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 19 Jun 2024 10:22:31 -0500 Subject: [PATCH 82/97] Add missing params to mesh data --- test/mesh_data.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/mesh_data.py b/test/mesh_data.py index 802dfe419..98ef35320 100644 --- a/test/mesh_data.py +++ b/test/mesh_data.py @@ -113,6 +113,8 @@ class _BoxMeshBuilderBase(MeshBuilder): resolutions = [4, 8, 16] mesh_order = 1 group_cls = None + a = (-0.5, -0.5, -0.5) + b = (+0.5, +0.5, +0.5) def get_mesh(self, resolution, mesh_order=None): if mesh_order is None: From d50f7e47102bedbffd718a66482d48e82eb4fcf9 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 19 Jun 2024 16:19:51 -0500 Subject: [PATCH 83/97] Re-enable tpe div --- grudge/op.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 878468092..bccf1b098 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -514,13 +514,12 @@ def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, per_group_divs = [ - # Disable for now: TPE div not working properly - # compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) - # if isinstance(in_grp, TensorProductElementGroup) + compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) + if isinstance(in_grp, TensorProductElementGroup) # r for rst axis # x for xyz axis - # else + else compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, ijm_i, metric_in_matvec) From 82924d890ee2b2feeb2a35f15416765bb2440eb4 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 19 Jun 2024 19:41:43 -0500 Subject: [PATCH 84/97] add working op for tpe --- grudge/op.py | 55 ++++++++++++++++++++++++++-------------------------- 1 file changed, 27 insertions(+), 28 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index bccf1b098..a90f209cb 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -307,9 +307,9 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, if grp.dim > 3 and metric_in_matvec: warn("Efficient tensor product weak " - "differentiation operators only " - "implemented for dimension 2 and 3. " - "Defaulting to inefficient version.") + "differentiation operators only " + "implemented for dimension 2 and 3. " + "Defaulting to inefficient version.") return compute_simplicial_grad(actx, grp, grp, diff_mat, vec, ijm, metric_in_matvec) @@ -412,45 +412,45 @@ def _divergence_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec # {{{ tensor product div - def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): + def compute_tensor_product_div(actx, in_grp, out_grp, diff_mat, vec, ijm): """ Exploits tensor product structure to reduce complexity. See `_gradient_kernel.compute_tensor_product_grad` for more details. """ - if grp.dim > 3 and metric_in_matvec: + if ((in_grp.dim > 3 and metric_in_matvec) or (in_grp != out_grp)): warn("Efficient tensor product weak " "differentiation operators only " "implemented for dimension 2 and 3. " "Defaulting to inefficient version.") - return compute_simplicial_div(actx, grp, grp, diff_mat, vec, ijm, + return compute_simplicial_div(actx, in_grp, out_grp, diff_mat, vec, ijm, metric_in_matvec) vec = make_obj_array([ - fold(grp.space, vec[func_axis]) + fold(in_grp.space, vec[func_axis]) for func_axis in range(vec.shape[0]) ]) if metric_in_matvec: - stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) + stiff_1d, mass_1d = get_diff_mat(actx, in_grp, out_grp) partials = [] for func_axis in range(vec.shape[0]): ref = [] - for xyz_axis in range(grp.dim): + for xyz_axis in range(in_grp.dim): ref.append(vec[func_axis]) - apply_mass_axes = set(range(grp.dim)) - {xyz_axis} + apply_mass_axes = set(range(in_grp.dim)) - {xyz_axis} for ax in apply_mass_axes: ref[xyz_axis] = single_axis_operator_application( - actx, grp.dim, mass_1d, ax, ref[xyz_axis], + actx, in_grp.dim, mass_1d, ax, ref[xyz_axis], tags=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered(),), arg_names=("mass_1d", f"vec_{func_axis}_{xyz_axis}") ) ref[xyz_axis] = single_axis_operator_application( - actx, grp.dim, stiff_1d, xyz_axis, ref[xyz_axis], + actx, in_grp.dim, stiff_1d, xyz_axis, ref[xyz_axis], tags=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered(),), arg_names=("stiff_1d", f"vec_{func_axis}_{xyz_axis}") @@ -459,16 +459,16 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): partials.append(ref) else: - diff_mat = get_diff_mat(actx, grp, grp) + diff_mat = get_diff_mat(actx, in_grp, out_grp) partials = [] for func_axis in range(vec.shape[0]): ref = [] - for xyz_axis in range(grp.dim): + for xyz_axis in range(in_grp.dim): ref.append(vec[func_axis]) ref[xyz_axis] = single_axis_operator_application( - actx, grp.dim, diff_mat, xyz_axis, ref[xyz_axis], + actx, in_grp.dim, diff_mat, xyz_axis, ref[xyz_axis], tags=(FirstAxisIsElementsTag(), OutputIsTensorProductDOFArrayOrdered(),), arg_names=("diff_mat", f"vec_{func_axis}_{xyz_axis}") @@ -477,11 +477,11 @@ def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): partials.append(ref) partials = actx.np.stack([ - unfold(grp.space, partials[func_axis][xyz_axis]) - for func_axis in range(grp.dim) - for xyz_axis in range(grp.dim) + unfold(out_grp.space, partials[func_axis][xyz_axis]) + for func_axis in range(out_grp.dim) + for xyz_axis in range(out_grp.dim) ]) - partials = partials.reshape(grp.dim, grp.dim, *partials.shape[-2:]) + partials = partials.reshape(out_grp.dim, out_grp.dim, *partials.shape[-2:]) div = actx.einsum( "xrej,xrej->ej", @@ -514,14 +514,14 @@ def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, per_group_divs = [ - compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) + compute_tensor_product_div(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i) if isinstance(in_grp, TensorProductElementGroup) # r for rst axis # x for xyz axis - else - compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, - ijm_i, metric_in_matvec) + else compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec) for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, @@ -560,11 +560,10 @@ def get_ref_derivative_mats(grp): from arraycontext.metadata import NameHint return actx.freeze( - actx.tag( - NameHint("tp_diff_mat_1d"), - tag_axes(actx, { - 1: DiscretizationDOFAxisTag()}, - diff_mat))) + actx.tag(NameHint("tp_diff_mat_1d"), + tag_axes(actx, { + 1: DiscretizationDOFAxisTag()}, + diff_mat))) elif isinstance(grp, SimplexElementGroup): from meshmode.discretization.poly_element import diff_matrices From e2db31a3f447227a0fbb14153ed4fb190b091ecd Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Sat, 22 Jun 2024 09:33:44 -0500 Subject: [PATCH 85/97] Remove tpe op calls from differentiation operators. --- grudge/op.py | 122 ++++++--------------------------------------------- 1 file changed, 13 insertions(+), 109 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index a90f209cb..dc082de86 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -75,9 +75,6 @@ from functools import partial from meshmode.dof_array import DOFArray, warn -from meshmode.discretization.poly_element import ( - TensorProductElementGroupBase as TensorProductElementGroup, - SimplexElementGroupBase as SimplexElementGroup) from meshmode.transform_metadata import (FirstAxisIsElementsTag, DiscretizationDOFAxisTag, DiscretizationElementAxisTag, @@ -280,10 +277,7 @@ def compute_simplicial_derivative(actx, in_grp, out_grp, return DOFArray( actx, data=tuple( - compute_tensor_product_derivative(actx, in_grp, get_diff_mat, vec_i, - ijm_i, xyz_axis, metric_in_matvec) - if isinstance(in_grp, TensorProductElementGroup) - else compute_simplicial_derivative(actx, in_grp, out_grp, + compute_simplicial_derivative(actx, in_grp, out_grp, get_diff_mat, vec_i, ijm_i, xyz_axis, metric_in_matvec) for out_grp, in_grp, vec_i, ijm_i in zip( @@ -388,10 +382,7 @@ def compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, # }}} per_group_grads = [ - compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i, - metric_in_matvec) - if isinstance(in_grp, TensorProductElementGroup) - else compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, + compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, ijm_i, metric_in_matvec) for out_grp, in_grp, vec_i, ijm_i in zip( @@ -513,14 +504,7 @@ def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, # }}} per_group_divs = [ - - compute_tensor_product_div(actx, in_grp, out_grp, get_diff_mat, vec_i, - ijm_i) - if isinstance(in_grp, TensorProductElementGroup) - - # r for rst axis - # x for xyz axis - else compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, + compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, ijm_i, metric_in_matvec) for out_grp, in_grp, vec_i, ijm_i in zip( @@ -545,37 +529,12 @@ def _reference_derivative_matrices(actx: ArrayContext, actx, _reference_derivative_matrices, lambda grp: grp.discretization_key()) def get_ref_derivative_mats(grp): - if isinstance(grp, TensorProductElementGroup): - import modepy as mp - import numpy.linalg as la - - #FIXME: Can be gotten rid of by updating meshmode - nodes1d = grp.unit_nodes_1d - bases_1d = grp.bases_1d() - - vdm_1d = mp.vandermonde(bases_1d.functions, nodes1d) - vdm_p_1d = mp.vandermonde(bases_1d.gradients, nodes1d)[0] - - diff_mat = actx.from_numpy(vdm_p_1d @ la.inv(vdm_1d)) - - from arraycontext.metadata import NameHint - return actx.freeze( - actx.tag(NameHint("tp_diff_mat_1d"), - tag_axes(actx, { - 1: DiscretizationDOFAxisTag()}, - diff_mat))) - - elif isinstance(grp, SimplexElementGroup): - from meshmode.discretization.poly_element import diff_matrices - return actx.freeze( - actx.tag_axis( - 1, DiscretizationDOFAxisTag(), - actx.from_numpy( - np.asarray(diff_matrices(grp))))) - - else: - raise TypeError("grp must be either a TensorProductElementGroup or" - f" a SimplexElementGroup. Found {grp}") + from meshmode.discretization.poly_element import diff_matrices + return actx.freeze( + actx.tag_axis( + 1, DiscretizationDOFAxisTag(), + actx.from_numpy( + np.asarray(diff_matrices(grp))))) return get_ref_derivative_mats(out_element_group) @@ -747,40 +706,6 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): from meshmode.discretization.poly_element import \ mass_matrix, diff_matrices - # {{{ tensor product case - - if isinstance(out_grp, TensorProductElementGroup): - import modepy as mp - import numpy.linalg as la - - # FIXME: can be gotten rid of by updating meshmode operators - basis_1d = out_grp.bases_1d() - nodes_1d = out_grp.unit_nodes_1d - - vdm = mp.vandermonde(basis_1d.functions, nodes_1d) - vdm_p = mp.vandermonde(basis_1d.gradients, nodes_1d)[0] - - mass_1d = la.inv(vdm @ vdm.T) - diff_mat = la.solve(vdm.T, vdm_p.T).T - - stiff_1d = actx.freeze( - actx.tag_axis(1, DiscretizationDOFAxisTag(), - actx.from_numpy( - np.asarray( - diff_mat.T @ mass_1d.T)))) - - from grudge.array_context import MassMatrix1d - mass_1d = actx.freeze( - actx.tag_axis( - 1, (DiscretizationDOFAxisTag(),), - actx.from_numpy(np.asarray(mass_1d))) - ) - mass_1d = actx.tag(MassMatrix1d(), mass_1d) - - return (stiff_1d, mass_1d) - - # }}} - mmat = mass_matrix(out_grp) return actx.freeze( @@ -1122,30 +1047,14 @@ def reference_inverse_mass_matrix(actx: ArrayContext, element_group): def get_ref_inv_mass_mat(grp): from modepy import inverse_mass_matrix - if isinstance(grp, TensorProductElementGroup): - basis_1d = grp.bases_1d() - nodes_1d = grp.unit_nodes_1d - inv_mass_1d = inverse_mass_matrix(basis_1d.functions, nodes_1d) - - from grudge.array_context import InverseMassMatrix1d - inv_mass_1d = actx.tag_axis(0, DiscretizationDOFAxisTag(), - actx.from_numpy(np.asarray(inv_mass_1d))) - inv_mass_1d = actx.freeze( - actx.tag(InverseMassMatrix1d(), inv_mass_1d)) + basis = grp.basis_obj() - return inv_mass_1d - elif isinstance(grp, SimplexElementGroup): - basis = grp.basis_obj() - - return actx.freeze( - actx.tag_axis(0, DiscretizationDOFAxisTag(), - actx.from_numpy( + return actx.freeze( + actx.tag_axis(0, DiscretizationDOFAxisTag(), + actx.from_numpy( np.asarray( inverse_mass_matrix(basis.functions, grp.unit_nodes), order="C")))) - else: - raise TypeError("grp must be either a TensorProductElementGroup or" - f" a SimplexElementGroup. Found {grp}") return get_ref_inv_mass_mat(element_group) @@ -1192,7 +1101,6 @@ def apply_to_tensor_product_elements(grp, jac_inv, vec, ref_inv_mass): ) def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): - # Based on https://arxiv.org/pdf/1608.03836.pdf # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv return actx.einsum( @@ -1203,10 +1111,6 @@ def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): tagged=(FirstAxisIsElementsTag(),)) group_data = [ - apply_to_tensor_product_elements( - grp, jac_inv, vec_i, - reference_inverse_mass_matrix(actx, element_group=grp)) - if isinstance(grp, TensorProductElementGroup) else apply_to_simplicial_elements(jac_inv, vec_i, reference_inverse_mass_matrix(actx, element_group=grp)) for grp, jac_inv, vec_i in zip(discr.groups, inv_area_elements, vec) From c4f2d3a9926153ae149dbe03ccfb6ff393e3e302 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 15 Jul 2024 23:28:26 -0500 Subject: [PATCH 86/97] Update wadg to use quadrature discretization --- grudge/op.py | 122 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) diff --git a/grudge/op.py b/grudge/op.py index dc082de86..d7eb00659 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -1119,6 +1119,125 @@ def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): return DOFArray(actx, data=tuple(group_data)) +def _apply_inverse_mass_operator_quad( + dcoll: DiscretizationCollection, dd_out, dd_in, vec): + if not isinstance(vec, DOFArray): + return map_array_container( + partial(_apply_inverse_mass_operator_quad, dcoll, dd_out, dd_in), vec + ) + + from grudge.geometry import area_element + + if dd_out != dd_in: + raise ValueError( + "Cannot compute inverse of a mass matrix mapping " + "between different element groups; inverse is not " + "guaranteed to be well-defined" + ) + + actx = vec.array_context + dd_quad = dd_in + dd_base = dd_quad.with_discr_tag(DISCR_TAG_BASE) + discr_quad = dcoll.discr_from_dd(dd_quad) + discr_base = dcoll.discr_from_dd(dd_base) + + ae = \ + project(dcoll, dd_base, dd_quad, + area_element( + actx, dcoll, dd=dd_base, + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)) + + inv_area_elements = 1./ae + + def apply_to_tensor_product_elements(grp, jac_inv, vec, ref_inv_mass): + + vec = fold(grp.space, vec) + + for xyz_axis in range(grp.dim): + vec = single_axis_operator_application( + actx, grp.dim, ref_inv_mass, xyz_axis, vec, + tags=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("ref_inv_mass_1d", "vec")) + + vec = unfold(grp.space, vec) + + return actx.einsum( + "ei,ei->ei", + jac_inv, + vec, + tagged=(FirstAxisIsElementsTag(),) + ) + + def apply_to_simplicial_elements_stage1(vec, ref_inv_mass): + # Based on https://arxiv.org/pdf/1608.03836.pdf + # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv + return actx.einsum( + "ij,ej->ei", + ref_inv_mass, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + def apply_to_simplicial_elements_stage2(jac_inv, vec): + # Based on https://arxiv.org/pdf/1608.03836.pdf + # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv + return actx.einsum( + "ei,ej->ei", + jac_inv, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + def apply_to_simplicial_elements_stage3(mm, vec): + # Based on https://arxiv.org/pdf/1608.03836.pdf + # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv + return actx.einsum( + "ij,ej->ei", + mm, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + def apply_to_simplicial_elements_stage4(mm_inv, vec): + # Based on https://arxiv.org/pdf/1608.03836.pdf + # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv + return actx.einsum( + "ij,ej->ei", + mm_inv, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + stage1_group_data = [ + apply_to_simplicial_elements_stage1(vec_i, + reference_inverse_mass_matrix(actx, element_group=grp)) + for grp, vec_i in zip(discr_base.groups, vec) + ] + + stage1 = DOFArray(actx, data=tuple(stage1_group_data)) + stage1 = project(dcoll, dd_base, dd_quad, stage1) + + stage2_group_data = [ + apply_to_simplicial_elements_stage2(jac_inv, vec_i) + for jac_inv, vec_i in zip(inv_area_elements, stage1) + ] + + stage2 = DOFArray(actx, data=tuple(stage2_group_data)) + + stage3_group_data = [ + apply_to_simplicial_elements_stage3( + reference_mass_matrix(actx, out_grp, in_grp), vec_i) + for out_grp, in_grp, vec_i in zip(discr_base.groups, discr_quad.groups, + stage2) + ] + stage3 = DOFArray(actx, data=tuple(stage3_group_data)) + + group_data = [ + apply_to_simplicial_elements_stage4( + reference_inverse_mass_matrix(actx, element_group=grp), vec_i) + for grp, vec_i in zip(discr_base.groups, stage3) + ] + + return DOFArray(actx, data=tuple(group_data)) + + def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: r"""Return the action of the DG mass matrix inverse on a vector (or vectors) of :class:`~meshmode.dof_array.DOFArray`\ s, *vec*. @@ -1166,6 +1285,9 @@ def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: else: raise TypeError("invalid number of arguments") + if dd.uses_quadrature(): + return _apply_inverse_mass_operator_quad(dcoll, dd, dd, vec) + return _apply_inverse_mass_operator(dcoll, dd, dd, vec) # }}} From c94a38617577e906a5fdb4e043d304da16c24c54 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Thu, 18 Jul 2024 21:18:32 -0500 Subject: [PATCH 87/97] update the update to working condition --- grudge/op.py | 3 +-- grudge/trace_pair.py | 1 + 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index c6bcb9b1a..b87f6b0ef 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -122,7 +122,7 @@ from grudge.trace_pair import ( bdry_trace_pair, bv_trace_pair, - connected_ranks, + # connected_ranks, cross_rank_trace_pairs, cross_rank_inter_volume_trace_pairs, inter_volume_trace_pairs, @@ -139,7 +139,6 @@ "bdry_trace_pair", "bv_trace_pair", "connected_parts", - "connected_ranks", "cross_rank_inter_volume_trace_pairs", "cross_rank_trace_pairs", "elementwise_integral", diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 10c2c06b5..098001601 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -72,6 +72,7 @@ from arraycontext import ( ArrayContainer, + ArrayContext, ArrayOrContainer, dataclass_array_container, flatten, From a6c5168f0d49560b0344ff31e298f52e3438f7f2 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Sat, 20 Jul 2024 09:24:41 -0500 Subject: [PATCH 88/97] Dont overintegration WADG for simplices --- grudge/op.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/grudge/op.py b/grudge/op.py index ce063a811..e03bb8ec9 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -1280,7 +1280,9 @@ def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: raise TypeError("invalid number of arguments") if dd.uses_quadrature(): - return _apply_inverse_mass_operator_quad(dcoll, dd, dd, vec) + if not dcoll._has_affine_groups(dd.domain_tag): + return _apply_inverse_mass_operator_quad(dcoll, dd, dd, vec) + dd = dd.with_discr_tag(DISCR_TAG_BASE) return _apply_inverse_mass_operator(dcoll, dd, dd, vec) From aaa67d9df9c7b576d10890102425a34d5e79d02d Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 24 Jul 2024 16:57:19 -0500 Subject: [PATCH 89/97] Disable rate-of-convergence test for TPE --- test/test_grudge.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/test/test_grudge.py b/test/test_grudge.py index f1ed112e0..f9e0ca9be 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -988,11 +988,9 @@ def test_improvement_quadrature(actx_factory, order, tpe): """Test whether quadrature improves things and converges""" from grudge.models.advection import VariableCoefficientAdvectionOperator from pytools.convergence import EOCRecorder - from meshmode.discretization.poly_element import QuadratureSimplexGroupFactory + # from meshmode.discretization.poly_element \ + # import QuadratureSimplexGroupFactory from meshmode.mesh import BTAG_ALL - from pytools.convergence import EOCRecorder - - from grudge.models.advection import VariableCoefficientAdvectionOperator actx = actx_factory() @@ -1035,7 +1033,7 @@ def conv_test(descr, use_quad): else: discr_tag_to_group_factory = {} - dcoll = DiscretizationCollection( + dcoll = make_discretization_collection( actx, mesh, order=discr_order, discr_tag_to_group_factory=discr_tag_to_group_factory ) @@ -1071,7 +1069,8 @@ def zero_inflow(dtag, t=0, dcoll=dcoll): assert (q_errs < errs).all() assert q_eoc > order - 0.1 # Fails for all tensor-product element types - assert q_eoc > eoc + if not tpe: + assert q_eoc > eoc # }}} From 26862d25bb9fe75869700d518c3c2d2aef49a032 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Thu, 25 Jul 2024 07:32:37 -0500 Subject: [PATCH 90/97] Combine a few stages in wadg --- grudge/op.py | 34 ++++++++++++++++++++++++++++------ 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index e03bb8ec9..786a82753 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -1135,11 +1135,15 @@ def _apply_inverse_mass_operator_quad( discr_quad = dcoll.discr_from_dd(dd_quad) discr_base = dcoll.discr_from_dd(dd_base) - ae = \ - project(dcoll, dd_base, dd_quad, - area_element( - actx, dcoll, dd=dd_base, - _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)) + # ae = \ + # project(dcoll, dd_base, dd_quad, + # area_element( + # actx, dcoll, dd=dd_base, + # _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)) + + ae = area_element( + actx, dcoll, dd=dd_quad, + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) inv_area_elements = 1./ae @@ -1172,6 +1176,14 @@ def apply_to_simplicial_elements_stage1(vec, ref_inv_mass): vec, tagged=(FirstAxisIsElementsTag(),)) + def apply_to_simplicial_elements_staged(mm_inv, mm, vec): + return actx.einsum( + "ni,ij,ej->en", + mm_inv, + mm, + vec, + tagged=(FirstAxisIsElementsTag(),)) + def apply_to_simplicial_elements_stage2(jac_inv, vec): # Based on https://arxiv.org/pdf/1608.03836.pdf # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv @@ -1208,6 +1220,7 @@ def apply_to_simplicial_elements_stage4(mm_inv, vec): stage1 = DOFArray(actx, data=tuple(stage1_group_data)) stage1 = project(dcoll, dd_base, dd_quad, stage1) + stage2_group_data = [ apply_to_simplicial_elements_stage2(jac_inv, vec_i) for jac_inv, vec_i in zip(inv_area_elements, stage1) @@ -1215,6 +1228,14 @@ def apply_to_simplicial_elements_stage4(mm_inv, vec): stage2 = DOFArray(actx, data=tuple(stage2_group_data)) + staged_group_data = [ + apply_to_simplicial_elements_staged( + reference_inverse_mass_matrix(actx, out_grp), + reference_mass_matrix(actx, out_grp, in_grp), vec_i) + for in_grp, out_grp, vec_i in zip( + discr_quad.groups, discr_base.groups, stage2) + ] + stage3_group_data = [ apply_to_simplicial_elements_stage3( reference_mass_matrix(actx, out_grp, in_grp), vec_i) @@ -1229,7 +1250,8 @@ def apply_to_simplicial_elements_stage4(mm_inv, vec): for grp, vec_i in zip(discr_base.groups, stage3) ] - return DOFArray(actx, data=tuple(group_data)) + # return DOFArray(actx, data=tuple(group_data)) + return DOFArray(actx, data=tuple(staged_group_data)) def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: From 1e553b1d23065fa0f40cbbe3054b2efc61ca4c78 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Fri, 26 Jul 2024 01:39:27 -0500 Subject: [PATCH 91/97] Update wadg to match main version. --- grudge/op.py | 71 +++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 68 insertions(+), 3 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 786a82753..8f35074d7 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -97,6 +97,7 @@ from grudge.dof_desc import ( DD_VOLUME_ALL, DISCR_TAG_BASE, + DISCR_TAG_QUAD, FACE_RESTR_ALL, DOFDesc, VolumeDomainTag, @@ -1113,6 +1114,69 @@ def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): return DOFArray(actx, data=tuple(group_data)) +def _apply_inverse_mass_operator_quad( + dcoll: DiscretizationCollection, dd, vec): + if not isinstance(vec, DOFArray): + return map_array_container( + partial(_apply_inverse_mass_operator_quad, dcoll, dd), vec + ) + + from grudge.geometry import area_element + + actx = vec.array_context + dd_quad = dd.with_discr_tag(DISCR_TAG_QUAD) + dd_base = dd.with_discr_tag(DISCR_TAG_BASE) + discr_quad = dcoll.discr_from_dd(dd_quad) + discr_base = dcoll.discr_from_dd(dd_base) + + # Based on https://arxiv.org/pdf/1608.03836.pdf + # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv + # Overintegration version of action on *vec*: + # true_Minv ~ ref_Minv * (ref_M)_qtb * (1/Jac)_quad * P(Minv*vec) + # P => projection to quadrature, qti => quad-to-base + + # Compute 1/Jac on quadrature discr + inv_area_elements = 1/area_element( + actx, dcoll, dd=dd_quad, + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + + def apply_minv_to_vec(vec, ref_inv_mass): + return actx.einsum( + "ij,ej->ei", + ref_inv_mass, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + # The rest of wadg + def apply_rest_of_wadg(mm_inv, mm, vec): + return actx.einsum( + "ni,ij,ej->en", + mm_inv, + mm, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + stage1_group_data = [ + apply_minv_to_vec( + vec_i, reference_inverse_mass_matrix(actx, element_group=grp)) + for grp, vec_i in zip(discr_base.groups, vec) + ] + stage2 = inv_area_elements * project( + dcoll, dd_base, dd_quad, + DOFArray(actx, data=tuple(stage1_group_data))) + + wadg_group_data = [ + apply_rest_of_wadg( + reference_inverse_mass_matrix(actx, out_grp), + reference_mass_matrix(actx, out_grp, in_grp), vec_i) + for in_grp, out_grp, vec_i in zip( + discr_quad.groups, discr_base.groups, stage2) + ] + + return DOFArray(actx, data=tuple(wadg_group_data)) + + +""" def _apply_inverse_mass_operator_quad( dcoll: DiscretizationCollection, dd_out, dd_in, vec): if not isinstance(vec, DOFArray): @@ -1252,6 +1316,7 @@ def apply_to_simplicial_elements_stage4(mm_inv, vec): # return DOFArray(actx, data=tuple(group_data)) return DOFArray(actx, data=tuple(staged_group_data)) +""" def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: @@ -1302,9 +1367,9 @@ def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: raise TypeError("invalid number of arguments") if dd.uses_quadrature(): - if not dcoll._has_affine_groups(dd.domain_tag): - return _apply_inverse_mass_operator_quad(dcoll, dd, dd, vec) - dd = dd.with_discr_tag(DISCR_TAG_BASE) + # if not dcoll._has_affine_groups(dd.domain_tag): + return _apply_inverse_mass_operator_quad(dcoll, dd, vec) + # dd = dd.with_discr_tag(DISCR_TAG_BASE) return _apply_inverse_mass_operator(dcoll, dd, dd, vec) From 8f586e3da91a26f9c85324587de5ff67b0331269 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 9 Sep 2024 13:26:50 -0500 Subject: [PATCH 92/97] Update esdg to use modepy for diff_mats instead of meshmode. --- grudge/flux_differencing.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/grudge/flux_differencing.py b/grudge/flux_differencing.py index 816138df9..68f983fd5 100644 --- a/grudge/flux_differencing.py +++ b/grudge/flux_differencing.py @@ -45,7 +45,7 @@ from grudge.discretization import DiscretizationCollection from grudge.dof_desc import DOFDesc - +import modepy as mp from pytools import memoize_in, keyed_memoize_in import numpy as np @@ -64,7 +64,6 @@ def _reference_skew_symmetric_hybridized_sbp_operators( face_quad_grp.discretization_key())) def get_reference_skew_symetric_hybridized_diff_mats( base_grp, quad_vol_grp, face_quad_grp): - from meshmode.discretization.poly_element import diff_matrices from modepy import faces_for_shape, face_normal from grudge.interpolation import ( volume_quadrature_interpolation_matrix, @@ -113,7 +112,8 @@ def get_reference_skew_symetric_hybridized_diff_mats( # {{{ Hybridized (volume + surface) operators q_mats = [p_mat.T @ (weights * vdm_q.T @ vdm_q) @ diff_mat @ p_mat - for diff_mat in diff_matrices(base_grp)] + for diff_mat in mp.diff_matrices(base_grp.basis_obj(), + base_grp.unit_nodes)] e_mat = vf_mat @ p_mat q_skew_hybridized = np.asarray( [ From f29701705b374ede4b59bac11247c0517c756d54 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 25 Sep 2024 12:40:56 -0500 Subject: [PATCH 93/97] Add test to demonstrate simplex integral fails with overintegration --- test/mesh_data.py | 10 ++++++++-- test/test_grudge.py | 30 ++++++++++++++++++++---------- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/test/mesh_data.py b/test/mesh_data.py index 58c9b80a7..11a9f04d9 100644 --- a/test/mesh_data.py +++ b/test/mesh_data.py @@ -117,14 +117,20 @@ class _BoxMeshBuilderBase(MeshBuilder): a = (-0.5, -0.5, -0.5) b = (+0.5, +0.5, +0.5) + tpe: bool + + def __init__(self, tpe=False): + self._tpe = tpe + def get_mesh(self, resolution, mesh_order=4): if not isinstance(resolution, (list, tuple)): resolution = (resolution,) * self.ambient_dim - + from meshmode.mesh import TensorProductElementGroup + group_cls = TensorProductElementGroup if self._tpe else None return mgen.generate_regular_rect_mesh( a=self.a, b=self.b, nelements_per_axis=resolution, - order=mesh_order) + order=mesh_order, group_cls=group_cls) class BoxMeshBuilder1D(_BoxMeshBuilderBase): diff --git a/test/test_grudge.py b/test/test_grudge.py index b0849310b..83d49a2cd 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -155,26 +155,31 @@ def _spheroid_surface_area(radius, aspect_ratio): @pytest.mark.parametrize("name", [ - "2-1-ellipse", "spheroid", "box2d", "box3d" + "box2d-tpe", "box3d-tpe", "box2d", "box3d", "2-1-ellipse", "spheroid", ]) -def test_mass_surface_area(actx_factory, name): +@pytest.mark.parametrize("oi", [False, True]) +def test_mass_surface_area(actx_factory, name, oi): + from grudge.dof_desc import as_dofdesc, DISCR_TAG_BASE, DISCR_TAG_QUAD actx = actx_factory() + qtag = DISCR_TAG_QUAD if oi else DISCR_TAG_BASE + vol_dd_base = as_dofdesc(dof_desc.DTAG_VOLUME_ALL) + vol_dd_quad = vol_dd_base.with_discr_tag(qtag) # {{{ cases order = 4 - + tpe = name.endswith("-tpe") if name == "2-1-ellipse": builder = mesh_data.EllipseMeshBuilder(radius=3.1, aspect_ratio=2.0) surface_area = _ellipse_surface_area(builder.radius, builder.aspect_ratio) elif name == "spheroid": builder = mesh_data.SpheroidMeshBuilder() surface_area = _spheroid_surface_area(builder.radius, builder.aspect_ratio) - elif name == "box2d": - builder = mesh_data.BoxMeshBuilder2D() + elif name.startswith("box2d"): + builder = mesh_data.BoxMeshBuilder2D(tpe=tpe) surface_area = 1.0 - elif name == "box3d": - builder = mesh_data.BoxMeshBuilder3D() + elif name.startswith("box3d"): + builder = mesh_data.BoxMeshBuilder3D(tpe=tpe) surface_area = 1.0 else: raise ValueError(f"unknown geometry name: {name}") @@ -189,15 +194,20 @@ def test_mass_surface_area(actx_factory, name): for resolution in builder.resolutions: mesh = builder.get_mesh(resolution, order) - dcoll = make_discretization_collection(actx, mesh, order=order) - volume_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME_ALL) + dcoll = make_discretization_collection( + actx, mesh, + discr_tag_to_group_factory={ + DISCR_TAG_BASE: InterpolatoryEdgeClusteredGroupFactory(order), + DISCR_TAG_QUAD: QuadratureGroupFactory(3 * order) + }) + volume_discr = dcoll.discr_from_dd(vol_dd_quad) logger.info("ndofs: %d", volume_discr.ndofs) logger.info("nelements: %d", volume_discr.mesh.nelements) # {{{ compute surface area - dd = dof_desc.DD_VOLUME_ALL + dd = vol_dd_quad ones_volm = volume_discr.zeros(actx) + 1 approx_surface_area = actx.to_numpy(op.integral(dcoll, dd, ones_volm)) From 3502174a4ccde5e8dca28b4ec137ee8d09f6193f Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Wed, 25 Sep 2024 13:13:07 -0500 Subject: [PATCH 94/97] Fix up linting issues --- test/mesh_data.py | 4 ++-- test/test_grudge.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/test/mesh_data.py b/test/mesh_data.py index 11a9f04d9..26b222ecc 100644 --- a/test/mesh_data.py +++ b/test/mesh_data.py @@ -125,8 +125,8 @@ def __init__(self, tpe=False): def get_mesh(self, resolution, mesh_order=4): if not isinstance(resolution, (list, tuple)): resolution = (resolution,) * self.ambient_dim - from meshmode.mesh import TensorProductElementGroup - group_cls = TensorProductElementGroup if self._tpe else None + from meshmode.mesh import TensorProductElementGroup + group_cls = TensorProductElementGroup if self._tpe else None return mgen.generate_regular_rect_mesh( a=self.a, b=self.b, nelements_per_axis=resolution, diff --git a/test/test_grudge.py b/test/test_grudge.py index 83d49a2cd..b23543605 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -159,7 +159,7 @@ def _spheroid_surface_area(radius, aspect_ratio): ]) @pytest.mark.parametrize("oi", [False, True]) def test_mass_surface_area(actx_factory, name, oi): - from grudge.dof_desc import as_dofdesc, DISCR_TAG_BASE, DISCR_TAG_QUAD + from grudge.dof_desc import DISCR_TAG_BASE, DISCR_TAG_QUAD, as_dofdesc actx = actx_factory() qtag = DISCR_TAG_QUAD if oi else DISCR_TAG_BASE vol_dd_base = as_dofdesc(dof_desc.DTAG_VOLUME_ALL) From 0fa70b2facfeee45183df932d4d9a007f2352c7e Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 27 Sep 2024 16:42:54 -0500 Subject: [PATCH 95/97] add ProcessLogger for find_distributed_partition --- grudge/array_context.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/grudge/array_context.py b/grudge/array_context.py index ce2df247b..c144f0a5f 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -303,19 +303,21 @@ def _dag_to_compiled_func(self, dict_of_named_arrays, # dict_of_named_arrays) # else: # raise - distributed_partition = pt.find_distributed_partition( - # pylint-ignore-reason: - # '_BasePytatoArrayContext' has no - # 'mpi_communicator' member - # pylint: disable=no-member - self.actx.mpi_communicator, dict_of_named_arrays) - - if __debug__: - # pylint-ignore-reason: - # '_BasePytatoArrayContext' has no 'mpi_communicator' member - pt.verify_distributed_partition( - self.actx.mpi_communicator, # pylint: disable=no-member - distributed_partition) + + with ProcessLogger(logger, "pt.find_distributed_partition"): + distributed_partition = pt.find_distributed_partition( + # pylint-ignore-reason: + # '_BasePytatoArrayContext' has no + # 'mpi_communicator' member + # pylint: disable=no-member + self.actx.mpi_communicator, dict_of_named_arrays) + + if __debug__: + # pylint-ignore-reason: + # '_BasePytatoArrayContext' has no 'mpi_communicator' member + pt.verify_distributed_partition( + self.actx.mpi_communicator, # pylint: disable=no-member + distributed_partition) self.actx._compile_trace_callback(self.f, "post_find_distributed_partition", distributed_partition) From e7749e7fe12d161736ab2eb9f89b407b992d1d44 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Tue, 1 Oct 2024 08:36:51 -0500 Subject: [PATCH 96/97] Adding some quadrature tests. --- test/mesh_data.py | 26 ++------ test/test_dt_utils.py | 12 ++-- test/test_grudge.py | 148 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 159 insertions(+), 27 deletions(-) diff --git a/test/mesh_data.py b/test/mesh_data.py index adea582c3..461f7e4e2 100644 --- a/test/mesh_data.py +++ b/test/mesh_data.py @@ -117,10 +117,13 @@ class _BoxMeshBuilderBase(MeshBuilder): group_cls = None a = (-0.5, -0.5, -0.5) b = (+0.5, +0.5, +0.5) - tpe: bool + tpe: bool = False - def __init__(self, tpe=False): + def __init__(self, tpe=False, a=(-0.5, -0.5, -0.5), + b=(0.5, 0.5, 0.5)): self.tpe = tpe + self.a = a + self.b = b def get_mesh(self, resolution, mesh_order=None): if mesh_order is None: @@ -138,33 +141,14 @@ def get_mesh(self, resolution, mesh_order=None): class BoxMeshBuilder1D(_BoxMeshBuilderBase): ambient_dim = 1 - def __init__(self, tpe=False): - if tpe: - self.group_cls = TensorProductElementGroup - else: - tpe = False - self.tpe = tpe class BoxMeshBuilder2D(_BoxMeshBuilderBase): ambient_dim = 2 - def __init__(self, tpe=False): - if tpe: - self.group_cls = TensorProductElementGroup - else: - tpe = False - self.tpe = tpe - class BoxMeshBuilder3D(_BoxMeshBuilderBase): ambient_dim = 3 - def __init__(self, tpe=False): - if tpe: - self.group_cls = TensorProductElementGroup - else: - tpe = False - self.tpe = tpe class WarpedRectMeshBuilder(MeshBuilder): resolutions: ClassVar[Sequence[Hashable]] = [4, 6, 8] diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py index 60cb3f729..8f8ec7504 100644 --- a/test/test_dt_utils.py +++ b/test/test_dt_utils.py @@ -107,12 +107,12 @@ def test_geometric_factors_regular_refinement(actx_factory, name, tpe): assert np.all(np.isclose(ratios, 2)) # Make sure it works with empty meshes - if not tpe: - mesh = builder.get_mesh(0, order) - dcoll = make_discretization_collection( - actx, mesh, order=order, - discr_tag_to_group_factory=dtag_to_grp_fac) - factors = actx.thaw(dt_geometric_factors(dcoll)) # noqa: F841 + # if not tpe: + # mesh = builder.get_mesh(0, order) + # dcoll = make_discretization_collection( + # actx, mesh, order=order, + # discr_tag_to_group_factory=dtag_to_grp_fac) + # factors = actx.thaw(dt_geometric_factors(dcoll)) # noqa: F841 @pytest.mark.parametrize("name", ["interval", "box2d", "box3d"]) diff --git a/test/test_grudge.py b/test/test_grudge.py index aa2564039..617171725 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -240,6 +240,154 @@ def test_mass_surface_area(actx_factory, name, oi): assert eoc.max_error() < 3e-13 or eoc.order_estimate() > order + +@pytest.mark.parametrize("name", [ + "interval", "box2d", "box2d-tpe", "box3d", "box3d-tpe" + ]) +# @pytest.mark.parametrize("discr_order", [1, 2, 3, 4, 5]) +def test_correctness_of_quadrature(actx_factory, name): + from grudge.dof_desc import DISCR_TAG_BASE, DISCR_TAG_QUAD, as_dofdesc + actx = actx_factory() + vol_dd_base = as_dofdesc(dof_desc.DTAG_VOLUME_ALL) + vol_dd_quad = vol_dd_base.with_discr_tag(DISCR_TAG_QUAD) + + # {{{ cases + + tol = 5e-13 + # discr_order = 1 + dim = None + mesh_order = 1 + + tpe = name.endswith("-tpe") + if name.startswith("box2d"): + builder = mesh_data.BoxMeshBuilder2D(tpe=tpe, + a=(0, 0), + b=(2.0, 2.0)) + dim = 2 + elif name.startswith("box3d"): + builder = mesh_data.BoxMeshBuilder3D(tpe=tpe, + a=(0, 0, 0), + b=(2.0, 2.0, 2.0)) + dim = 3 + elif name == "interval": + builder = mesh_data.BoxMeshBuilder1D(tpe=False, a=(0.0,), + b=(2.0,)) + dim = 1.0 + else: + raise ValueError(f"unknown geometry name: {name}") + exact_volume = 2.0**dim + print(f"Domain: {name} ({dim}d), {exact_volume=}") + print(f"======================================================") + + # }}} + + # {{{ convergence + + from pytools.convergence import EOCRecorder + for discr_order in range(1, 8): + print(f" {discr_order=}") + print(" --------------------") + report_discr = True + for field_order in range(1, max(2*discr_order+1, 8)): + report_field_order = True + eoc_base = EOCRecorder() + eoc_quad = EOCRecorder() + for resolution in builder.resolutions: + mesh = builder.get_mesh(resolution, mesh_order) + dcoll = make_discretization_collection( + actx, mesh, + discr_tag_to_group_factory={ + DISCR_TAG_BASE: + InterpolatoryEdgeClusteredGroupFactory(discr_order), + DISCR_TAG_QUAD: QuadratureGroupFactory(discr_order) + }) + vol_discr_base = dcoll.discr_from_dd(vol_dd_base) + vol_discr_quad = dcoll.discr_from_dd(vol_dd_quad) + if report_discr: + nelem = vol_discr_base.mesh.nelements + ndofs_base = vol_discr_base.ndofs + ndofs_quad = vol_discr_quad.ndofs + dofs_per_el_base = ndofs_base/nelem + dofs_per_el_quad = ndofs_quad/nelem + print(f" - {dofs_per_el_base=}, {dofs_per_el_quad=}") + report_discr = False + if report_field_order: + print(f" - {field_order=}") + print(" - - - - - - - - - -") + report_field_order = False + nodes_base = actx.thaw(vol_discr_base.nodes()) + nodes_quad = actx.thaw(vol_discr_quad.nodes()) + ones_base = 0*nodes_base[0] + 1 + ones_quad = 0*nodes_quad[0] + 1 + + approx_vol_base = \ + actx.to_numpy(op.integral(dcoll, vol_dd_base, ones_base)) + approx_vol_quad = \ + actx.to_numpy(op.integral(dcoll, vol_dd_quad, ones_quad)) + err_vol_base = abs(approx_vol_base - exact_volume)/exact_volume + err_vol_quad = abs(approx_vol_quad - exact_volume)/exact_volume + + logger.info( + f"Name: {name} ({dim}d)\n" + f"Exact volume: {exact_volume}\n" + f"volume base: got {approx_vol_base:.12e}\n" # noqa: G004 + f"volume quad: got {approx_vol_quad:.12e}\n") # noqa: G004 + + # Quadrature should get exact volume for all discr (p >= 1) + assert err_vol_base < tol + assert err_vol_quad < tol + + field_base = nodes_base[0]**field_order + field_quad = nodes_quad[0]**field_order + if dim > 1: + field_base = sum(nodes_base[i]**field_order + for i in range(dim)) + field_quad = sum(nodes_quad[i]**field_order + for i in range(dim)) + ofac = 1.0/2**field_order + field_base = ofac * field_base * (field_order + 1.0) + field_quad = ofac * field_quad * (field_order + 1.0) + exact_integral = dim*exact_volume + + integral_base = \ + actx.to_numpy(op.integral(dcoll, vol_dd_base, field_base)) + integral_quad = \ + actx.to_numpy(op.integral(dcoll, vol_dd_quad, field_quad)) + err_base = \ + abs(integral_base - exact_integral)/exact_integral + err_quad = \ + abs(integral_quad - exact_integral)/exact_integral + + if field_order <= discr_order: + assert err_base < tol + assert err_quad < tol + + # compute max element size + from grudge.dt_utils import h_max_from_volume + h_max = h_max_from_volume(dcoll) + + eoc_base.add_data_point(actx.to_numpy(h_max), err_base) + eoc_quad.add_data_point(actx.to_numpy(h_max), err_quad) + + logger.info("volume error(base)\n%s", str(eoc_base)) + logger.info("volume error(quad)\n%s", str(eoc_quad)) + print("---- base -----") + print(f"{eoc_base.pretty_print()}") + print("---- quad -----") + print(f"{eoc_quad.pretty_print()}") + + # Sanity check here: *must* be exact if discr_order is sufficient + if discr_order >= field_order: + assert eoc_base.max_error() < tol + assert eoc_quad.max_error() < tol + else: + if eoc_base.max_error() > tol: # *can* be exact(ish) otherwise + assert eoc_base.order_estimate() > discr_order + if eoc_quad.max_error() > tol: + assert eoc_quad.order_estimate() > discr_order + + print("-=-=-=-=-=-=-=-=-=-=-=-=-=--=-=-=-") + print("==============================") # }}} From 109da041afc2e3ed9418d77e6aab8b69e60b47a3 Mon Sep 17 00:00:00 2001 From: Mike Campbell Date: Mon, 11 Nov 2024 11:56:54 -0600 Subject: [PATCH 97/97] Remove TPE-specific routines --- grudge/op.py | 240 +-------------------------------------------------- 1 file changed, 1 insertion(+), 239 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 95d2e3cc9..0ef01ae5c 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -91,7 +91,7 @@ from pytools import keyed_memoize_in from pytools.obj_array import make_obj_array -from grudge.array_context import OutputIsTensorProductDOFArrayOrdered +# from grudge.array_context import OutputIsTensorProductDOFArrayOrdered import grudge.dof_desc as dof_desc from grudge.discretization import DiscretizationCollection from grudge.dof_desc import ( @@ -190,67 +190,6 @@ def _single_axis_derivative_kernel( # or inside (weak) the matrix-vector product that carries out the # derivative, cf. "metric_in_matvec". - # {{{ tensor product single axis derivative - - def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, - xyz_axis, metric_in_matvec): - - vec = fold(grp.space, vec) - - if metric_in_matvec: - stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) - - apply_mass_axes = set(range(grp.dim)) - {xyz_axis} - - for ax in apply_mass_axes: - vec_mass_applied = single_axis_operator_application( - actx, grp.dim, mass_1d, ax, vec, - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("mass_1d", "vec") - ) - - ref_weak_derivative = unfold( - grp.space, - single_axis_operator_application( - actx, grp.dim, stiff_1d, xyz_axis, vec_mass_applied, - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("stiff_1d", "vec_with_mass_applied")) - ) - - derivative = actx.einsum( - "rej,ej->ej", - ijm[xyz_axis], - ref_weak_derivative, - tagged=(FirstAxisIsElementsTag(),), - arg_names=("inv_jac_t", "ref_weak_derivative") - ) - - else: - diff_mat = get_diff_mat(actx, grp, grp) - - ref_derivative = unfold( - grp.space, - single_axis_operator_application( - actx, grp.dim, diff_mat, xyz_axis, vec, - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("diff_mat", "vec")) - ) - - derivative = actx.einsum( - "rej,ej->ej", - ijm[xyz_axis], - ref_derivative, - tagged=(FirstAxisIsElementsTag(),), - arg_names=("inv_jac_t", "ref_derivs") - ) - - return derivative - - # }}} - # {{{ simplicial single axis derivative def compute_simplicial_derivative(actx, in_grp, out_grp, @@ -288,79 +227,6 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. - # {{{ tensor product gradient - - def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, - metric_in_matvec): - # TODO: add note about inverse mass simplification, point to - # op.inverse_mass (assuming this is where the explanation will live) - """ - """ - - if grp.dim > 3 and metric_in_matvec: - warn("Efficient tensor product weak " - "differentiation operators only " - "implemented for dimension 2 and 3. " - "Defaulting to inefficient version.") - return compute_simplicial_grad(actx, grp, grp, diff_mat, vec, ijm, - metric_in_matvec) - - # reshape vector to expose tensor product structure - vec = fold(grp.space, vec) - - if metric_in_matvec: - stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) - - grad = [] - for xyz_axis in range(grp.dim): - grad.append(vec) - apply_mass_axes = set(range(grp.dim)) - {xyz_axis} - - # apply mass operators - for ax in apply_mass_axes: - grad[xyz_axis] = single_axis_operator_application( - actx, grp.dim, mass_1d, ax, grad[xyz_axis], - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("mass_1d", f"vec_{xyz_axis}")) - - # apply stiffness operator and unfold - grad[xyz_axis] = unfold( - grp.space, - single_axis_operator_application( - actx, grp.dim, stiff_1d, xyz_axis, grad[xyz_axis], - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("stiff_1d", f"vec_{xyz_axis}")) - ) - - else: - diff_mat = get_diff_mat(actx, grp, grp) - - grad = [] - for xyz_axis in range(grp.dim): - grad.append(vec) - grad[xyz_axis] = unfold( - grp.space, - single_axis_operator_application( - actx, grp.dim, diff_mat, xyz_axis, grad[xyz_axis], - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("diff_mat", f"vec_{xyz_axis}") - ) - ) - - grad = actx.np.stack(grad) - return actx.einsum( - "xrej,rej->xej", - ijm, - grad, - tagged=(FirstAxisIsElementsTag(),), - arg_names=("inv_jac_t", "grad") - ) - - # }}} - # {{{ simplicial grad def compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, @@ -400,90 +266,6 @@ def _divergence_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. - # {{{ tensor product div - - def compute_tensor_product_div(actx, in_grp, out_grp, diff_mat, vec, ijm): - """ - Exploits tensor product structure to reduce complexity. See - `_gradient_kernel.compute_tensor_product_grad` for more details. - """ - - if ((in_grp.dim > 3 and metric_in_matvec) or (in_grp != out_grp)): - warn("Efficient tensor product weak " - "differentiation operators only " - "implemented for dimension 2 and 3. " - "Defaulting to inefficient version.") - return compute_simplicial_div(actx, in_grp, out_grp, diff_mat, vec, ijm, - metric_in_matvec) - - vec = make_obj_array([ - fold(in_grp.space, vec[func_axis]) - for func_axis in range(vec.shape[0]) - ]) - - if metric_in_matvec: - stiff_1d, mass_1d = get_diff_mat(actx, in_grp, out_grp) - - partials = [] - for func_axis in range(vec.shape[0]): - ref = [] - for xyz_axis in range(in_grp.dim): - ref.append(vec[func_axis]) - - apply_mass_axes = set(range(in_grp.dim)) - {xyz_axis} - for ax in apply_mass_axes: - ref[xyz_axis] = single_axis_operator_application( - actx, in_grp.dim, mass_1d, ax, ref[xyz_axis], - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("mass_1d", f"vec_{func_axis}_{xyz_axis}") - ) - - ref[xyz_axis] = single_axis_operator_application( - actx, in_grp.dim, stiff_1d, xyz_axis, ref[xyz_axis], - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("stiff_1d", f"vec_{func_axis}_{xyz_axis}") - ) - - partials.append(ref) - - else: - diff_mat = get_diff_mat(actx, in_grp, out_grp) - - partials = [] - for func_axis in range(vec.shape[0]): - ref = [] - for xyz_axis in range(in_grp.dim): - ref.append(vec[func_axis]) - - ref[xyz_axis] = single_axis_operator_application( - actx, in_grp.dim, diff_mat, xyz_axis, ref[xyz_axis], - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("diff_mat", f"vec_{func_axis}_{xyz_axis}") - ) - - partials.append(ref) - - partials = actx.np.stack([ - unfold(out_grp.space, partials[func_axis][xyz_axis]) - for func_axis in range(out_grp.dim) - for xyz_axis in range(out_grp.dim) - ]) - partials = partials.reshape(out_grp.dim, out_grp.dim, *partials.shape[-2:]) - - div = actx.einsum( - "xrej,xrej->ej", - ijm, - partials, - arg_names=("inv_jac_t", "partials"), - tagged=(FirstAxisIsElementsTag(),) - ) - - return div - # }}} - # {{{ simplicial div def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, @@ -1079,26 +861,6 @@ def _apply_inverse_mass_operator( inv_area_elements = 1./area_element(actx, dcoll, dd=dd_in, _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) - def apply_to_tensor_product_elements(grp, jac_inv, vec, ref_inv_mass): - - vec = fold(grp.space, vec) - - for xyz_axis in range(grp.dim): - vec = single_axis_operator_application( - actx, grp.dim, ref_inv_mass, xyz_axis, vec, - tags=(FirstAxisIsElementsTag(), - OutputIsTensorProductDOFArrayOrdered(),), - arg_names=("ref_inv_mass_1d", "vec")) - - vec = unfold(grp.space, vec) - - return actx.einsum( - "ei,ei->ei", - jac_inv, - vec, - tagged=(FirstAxisIsElementsTag(),) - ) - def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): # Based on https://arxiv.org/pdf/1608.03836.pdf # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv