diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py
index e245ad2592972d90c2a89255adddfeb95d26821d..a9fd8cf43ae3c5b3239b51658fa6093dfa5aa5a3 100644
--- a/grudge/dt_utils.py
+++ b/grudge/dt_utils.py
@@ -45,8 +45,11 @@ THE SOFTWARE.
import numpy as np
-from arraycontext import ArrayContext, thaw, freeze, Scalar
-from meshmode.transform_metadata import FirstAxisIsElementsTag
+from arraycontext import ArrayContext, thaw, freeze, Scalar, tag_axes
+from meshmode.transform_metadata import (FirstAxisIsElementsTag,
+ DiscretizationDOFAxisTag,
+ DiscretizationFaceAxisTag,
+ DiscretizationElementAxisTag)
from grudge.dof_desc import DD_VOLUME, DOFDesc, as_dofdesc
from grudge.discretization import DiscretizationCollection
@@ -287,15 +290,18 @@ def dt_geometric_factors(
data=tuple(
actx.einsum(
"fej->e",
- face_ae_i.reshape(
- vgrp.mesh_el_group.nfaces,
- vgrp.nelements,
- face_ae_i.shape[-1]),
+ 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)
- )
- )
+ for vgrp, face_ae_i in zip(volm_discr.groups, face_areas)))
else:
surface_areas = DOFArray(
actx,
@@ -316,20 +322,17 @@ def dt_geometric_factors(
tagged=(FirstAxisIsElementsTag(),))
for vgrp, afgrp, face_ae_i in zip(volm_discr.groups,
- face_discr.groups,
- face_areas)
+ face_discr.groups,
+ face_areas)
)
)
return freeze(DOFArray(
actx,
data=tuple(
- actx.einsum("e,ei->ei",
- 1/sae_i,
- cv_i,
- tagged=(FirstAxisIsElementsTag(),)) * dcoll.dim
-
- for cv_i, sae_i in zip(cell_vols, surface_areas)
+ actx.einsum("e,ei->ei", 1/sae_i, cv_i,
+ tagged=(FirstAxisIsElementsTag(),)) * dcoll.dim
+ for cv_i, sae_i, vgrp in zip(cell_vols, surface_areas, volm_discr.groups)
)
))
diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py
index 492d8b7f783476ff9b4dca6359dfb5324158a596..d065a4360e1830149130859868a25becbd29b7b7 100644
--- a/grudge/geometry/metrics.py
+++ b/grudge/geometry/metrics.py
@@ -60,7 +60,7 @@ THE SOFTWARE.
import numpy as np
-from arraycontext import thaw, freeze, ArrayContext
+from arraycontext import thaw, freeze, ArrayContext, tag_axes
from meshmode.dof_array import DOFArray
from grudge import DiscretizationCollection
@@ -70,6 +70,10 @@ from grudge.dof_desc import (
DD_VOLUME, DOFDesc, DISCR_TAG_BASE
)
+from meshmode.transform_metadata import (DiscretizationAmbientDimAxisTag,
+ DiscretizationTopologicalDimAxisTag)
+
+
from pymbolic.geometric_algebra import MultiVector
from pytools.obj_array import make_obj_array
@@ -513,15 +517,19 @@ def inverse_surface_metric_derivative_mat(
multiplier = 1
mat = actx.np.stack([
- actx.np.stack([
- multiplier
- * inverse_surface_metric_derivative(actx, dcoll,
- rst_axis, xyz_axis, dd=dd,
- _use_geoderiv_connection=_use_geoderiv_connection)
- for rst_axis in range(dcoll.dim)])
- for xyz_axis in range(dcoll.ambient_dim)])
-
- return freeze(mat, actx)
+ actx.np.stack([
+ multiplier
+ * inverse_surface_metric_derivative(
+ actx, dcoll,
+ rst_axis, xyz_axis, dd=dd,
+ _use_geoderiv_connection=_use_geoderiv_connection)
+ for rst_axis in range(dcoll.dim)])
+ for xyz_axis in range(dcoll.ambient_dim)])
+
+ return freeze(tag_axes(actx, {
+ 0: DiscretizationAmbientDimAxisTag(),
+ 1: DiscretizationTopologicalDimAxisTag()},
+ mat))
return thaw(_inv_surf_metric_deriv(), actx)
diff --git a/grudge/op.py b/grudge/op.py
index d544d7ca605da4be76ad7f538638c7ad2b941592..0f4d6305244d78b74063d015bec0a8d988e2d558 100644
--- a/grudge/op.py
+++ b/grudge/op.py
@@ -57,13 +57,16 @@ THE SOFTWARE.
"""
-from arraycontext import ArrayContext, map_array_container
+from arraycontext import ArrayContext, map_array_container, tag_axes
from arraycontext.container import ArrayOrContainerT
from functools import partial
from meshmode.dof_array import DOFArray
-from meshmode.transform_metadata import FirstAxisIsElementsTag
+from meshmode.transform_metadata import (FirstAxisIsElementsTag,
+ DiscretizationDOFAxisTag,
+ DiscretizationElementAxisTag,
+ DiscretizationFaceAxisTag)
from grudge.discretization import DiscretizationCollection
@@ -168,8 +171,7 @@ def _single_axis_derivative_kernel(
get_diff_mat(
actx,
out_element_group=out_grp,
- in_element_group=in_grp
- ),
+ in_element_group=in_grp),
vec_i,
arg_names=("inv_jac_t", "ref_stiffT_mat", "vec", ),
tagged=(FirstAxisIsElementsTag(),))
@@ -222,12 +224,10 @@ def _reference_derivative_matrices(actx: ArrayContext,
def get_ref_derivative_mats(grp):
from meshmode.discretization.poly_element import diff_matrices
return actx.freeze(
- actx.from_numpy(
- np.asarray(
- [dfmat for dfmat in diff_matrices(grp)]
- )
- )
- )
+ actx.tag_axis(
+ 1, DiscretizationDOFAxisTag(),
+ actx.from_numpy(
+ np.asarray([dfmat for dfmat in diff_matrices(grp)]))))
return get_ref_derivative_mats(out_element_group)
@@ -346,12 +346,10 @@ def _reference_stiffness_transpose_matrix(
mmat = mass_matrix(out_grp)
return actx.freeze(
- actx.from_numpy(
- np.asarray(
- [dmat.T @ mmat.T for dmat in diff_matrices(out_grp)]
- )
- )
- )
+ actx.tag_axis(1, DiscretizationDOFAxisTag(),
+ actx.from_numpy(
+ np.asarray(
+ [dmat.T @ mmat.T for dmat in diff_matrices(out_grp)]))))
from modepy import vandermonde
basis = out_grp.basis_obj()
@@ -568,13 +566,11 @@ def reference_mass_matrix(actx: ArrayContext, out_element_group, in_element_grou
weights = in_grp.quadrature_rule().weights
return actx.freeze(
- actx.from_numpy(
- np.asarray(
- np.einsum("j,ik,jk->ij", weights, vand_inv_t, o_vand),
- order="C"
- )
- )
- )
+ actx.tag_axis(0, DiscretizationDOFAxisTag(),
+ actx.from_numpy(
+ np.asarray(
+ np.einsum("j,ik,jk->ij", weights, vand_inv_t, o_vand),
+ order="C"))))
return get_ref_mass_mat(out_element_group, in_element_group)
@@ -598,15 +594,15 @@ def _apply_mass_operator(
actx,
data=tuple(
actx.einsum("ij,ej,ej->ei",
- reference_mass_matrix(
- actx,
- out_element_group=out_grp,
- in_element_group=in_grp
- ),
- ae_i,
- vec_i,
- arg_names=("mass_mat", "jac", "vec"),
- tagged=(FirstAxisIsElementsTag(),))
+ reference_mass_matrix(
+ actx,
+ out_element_group=out_grp,
+ in_element_group=in_grp
+ ),
+ ae_i,
+ vec_i,
+ arg_names=("mass_mat", "jac", "vec"),
+ tagged=(FirstAxisIsElementsTag(),))
for in_grp, out_grp, ae_i, vec_i in zip(
in_discr.groups, out_discr.groups, area_elements, vec)
@@ -664,13 +660,11 @@ def reference_inverse_mass_matrix(actx: ArrayContext, element_group):
basis = grp.basis_obj()
return actx.freeze(
- actx.from_numpy(
- np.asarray(
- inverse_mass_matrix(basis.functions, grp.unit_nodes),
- order="C"
- )
- )
- )
+ actx.tag_axis(0, DiscretizationDOFAxisTag(),
+ actx.from_numpy(
+ np.asarray(
+ inverse_mass_matrix(basis.functions, grp.unit_nodes),
+ order="C"))))
return get_ref_inv_mass_mat(element_group)
@@ -695,21 +689,15 @@ 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)
- group_data = []
- for grp, jac_inv, vec_i in zip(discr.groups, inv_area_elements, vec):
-
- ref_mass_inverse = reference_inverse_mass_matrix(actx,
- element_group=grp)
-
- group_data.append(
+ 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,
- ref_mass_inverse,
+ 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)]
return DOFArray(actx, data=tuple(group_data))
@@ -840,7 +828,12 @@ def reference_face_mass_matrix(
vol_grp.unit_nodes,
)
- return actx.freeze(actx.from_numpy(matrix))
+ return actx.freeze(
+ tag_axes(actx, {
+ 0: DiscretizationDOFAxisTag(),
+ 2: DiscretizationDOFAxisTag()
+ },
+ actx.from_numpy(matrix)))
return get_ref_face_mass_mat(face_element_group, vol_element_group)
@@ -867,27 +860,27 @@ def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec):
data=tuple(
actx.einsum("ifj,fej,fej->ei",
reference_face_mass_matrix(
- actx,
- face_element_group=afgrp,
- vol_element_group=vgrp,
- dtype=dtype),
- surf_ae_i.reshape(
+ actx,
+ face_element_group=afgrp,
+ vol_element_group=vgrp,
+ dtype=dtype),
+ actx.tag_axis(1, DiscretizationElementAxisTag(),
+ surf_ae_i.reshape(
vgrp.mesh_el_group.nfaces,
vgrp.nelements,
- surf_ae_i.shape[-1]),
- vec_i.reshape(
+ surf_ae_i.shape[-1])),
+ actx.tag_axis(0, DiscretizationFaceAxisTag(),
+ vec_i.reshape(
vgrp.mesh_el_group.nfaces,
vgrp.nelements,
- afgrp.nunit_dofs),
+ afgrp.nunit_dofs)),
arg_names=("ref_face_mass_mat", "jac_surf", "vec"),
tagged=(FirstAxisIsElementsTag(),))
for vgrp, afgrp, vec_i, surf_ae_i in zip(volm_discr.groups,
face_discr.groups,
vec,
- surf_area_elements)
- )
- )
+ surf_area_elements)))
def face_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT:
diff --git a/grudge/reductions.py b/grudge/reductions.py
index 1efeff4e9aa5dbe5dbc9d6caae1087e49cc174b9..936467222b3652bcb33aa315c7efee0d74c29457 100644
--- a/grudge/reductions.py
+++ b/grudge/reductions.py
@@ -72,6 +72,7 @@ from grudge.discretization import DiscretizationCollection
from pytools import memoize_in
from meshmode.dof_array import DOFArray
+from meshmode.transform_metadata import DiscretizationDOFAxisTag
import numpy as np
import grudge.dof_desc as dof_desc
@@ -364,13 +365,12 @@ def _apply_elementwise_reduction(
"iel": ConcurrentElementInameTag(),
"idof": ConcurrentDOFInameTag()})
- return DOFArray(
- actx,
- data=tuple(
- actx.call_loopy(elementwise_prg(), operand=vec_i)["result"]
- for vec_i in vec
- )
- )
+ return actx.tag_axis(1, DiscretizationDOFAxisTag(),
+ DOFArray(
+ actx,
+ data=tuple(
+ actx.call_loopy(elementwise_prg(), operand=vec_i)["result"]
+ for vec_i in vec)))
def elementwise_sum(