diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py
index 24c5f8a28fe9320993a82b0d876b2a1bf542abde..5535725464383999347bb40953ae6a5dd507308d 100644
--- a/grudge/geometry/metrics.py
+++ b/grudge/geometry/metrics.py
@@ -379,7 +379,8 @@ def inverse_metric_derivative(
def inverse_surface_metric_derivative(
actx: ArrayContext, dcoll: DiscretizationCollection,
- rst_axis, xyz_axis, dd=None):
+ rst_axis, xyz_axis, dd=None,
+ *, _use_geoderiv_connection=False):
r"""Computes the inverse surface metric derivative of the physical
coordinate enumerated by *xyz_axis* with respect to the
reference axis *rst_axis*. These geometric terms are used in the
@@ -391,6 +392,12 @@ def inverse_surface_metric_derivative(
:arg xyz_axis: an integer denoting the physical coordinate axis.
:arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
Defaults to the base volume discretization.
+ :arg _use_geoderiv_connection: If *True*, process returned
+ :class:`~meshmode.dof_array.DOFArray`\ s through
+ :meth:`~grudge.DiscretizationCollection._base_to_geoderiv_connection`.
+ This should be set based on whether the code using the result of this
+ function is able to make use of these arrays. (This is an internal
+ argument and not intended for use outside :mod:`grudge`.)
:returns: a :class:`~meshmode.dof_array.DOFArray` containing the
inverse metric derivative at each nodal coordinate.
"""
@@ -402,20 +409,25 @@ def inverse_surface_metric_derivative(
dd = dof_desc.as_dofdesc(dd)
if ambient_dim == dim:
- return inverse_metric_derivative(
+ result = inverse_metric_derivative(
actx, dcoll, rst_axis, xyz_axis, dd=dd
)
else:
inv_form1 = inverse_first_fundamental_form(actx, dcoll, dd=dd)
- return sum(
+ result = sum(
inv_form1[rst_axis, d]*forward_metric_nth_derivative(
actx, dcoll, xyz_axis, d, dd=dd
) for d in range(dim))
+ if _use_geoderiv_connection:
+ result = dcoll._base_to_geoderiv_connection(dd)(result)
+
+ return result
+
def inverse_surface_metric_derivative_mat(
actx: ArrayContext, dcoll: DiscretizationCollection, dd=None,
- *, times_area_element=False):
+ *, times_area_element=False, _use_geoderiv_connection=False):
r"""Computes the matrix of inverse surface metric derivatives, indexed by
``(xyz_axis, rst_axis)``. It returns all values of
:func:`inverse_surface_metric_derivative_mat` in cached matrix form.
@@ -427,24 +439,33 @@ def inverse_surface_metric_derivative_mat(
:arg times_area_element: If *True*, each entry of the matrix is premultiplied
with the value of :func:`area_element`, reflecting the typical use
of the matrix in integrals evaluating weak derivatives.
+ :arg _use_geoderiv_connection: If *True*, process returned
+ :class:`~meshmode.dof_array.DOFArray`\ s through
+ :meth:`~grudge.DiscretizationCollection._base_to_geoderiv_connection`.
+ This should be set based on whether the code using the result of this
+ function is able to make use of these arrays. (This is an internal
+ argument and not intended for use outside :mod:`grudge`.)
:returns: a :class:`~meshmode.dof_array.DOFArray` containing the
inverse metric derivatives in per-group arrays of shape
``(xyz_dimension, rst_dimension, nelements, ndof)``.
"""
@memoize_in(dcoll, (inverse_surface_metric_derivative_mat, dd,
- times_area_element))
+ times_area_element, _use_geoderiv_connection))
def _inv_surf_metric_deriv():
if times_area_element:
- multiplier = area_element(actx, dcoll, dd=dd)
+ multiplier = area_element(actx, dcoll, dd=dd,
+ _use_geoderiv_connection=_use_geoderiv_connection)
else:
multiplier = 1
mat = actx.np.stack([
- actx.np.stack(
- [multiplier * inverse_surface_metric_derivative(actx, dcoll,
- rst_axis, xyz_axis, dd=dd)
- for rst_axis in range(dcoll.dim)])
+ 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)
@@ -524,7 +545,8 @@ def pseudoscalar(actx: ArrayContext, dcoll: DiscretizationCollection,
def area_element(
- actx: ArrayContext, dcoll: DiscretizationCollection, dd=None
+ actx: ArrayContext, dcoll: DiscretizationCollection, dd=None,
+ *, _use_geoderiv_connection=False
) -> DOFArray:
r"""Computes the scale factor used to transform integrals from reference
to global space.
@@ -533,16 +555,27 @@ def area_element(
:arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
Defaults to the base volume discretization.
+ :arg _use_geoderiv_connection: If *True*, process returned
+ :class:`~meshmode.dof_array.DOFArray`\ s through
+ :meth:`~grudge.DiscretizationCollection._base_to_geoderiv_connection`.
+ This should be set based on whether the code using the result of this
+ function is able to make use of these arrays. (This is an internal
+ argument and not intended for use outside :mod:`grudge`.)
:returns: a :class:`~meshmode.dof_array.DOFArray` containing the transformed
volumes for each element.
"""
if dd is None:
dd = DD_VOLUME
- @memoize_in(dcoll, (area_element, dd, actx.supports_nonscalar_broadcasting))
+ @memoize_in(dcoll, (area_element, dd, _use_geoderiv_connection))
def _area_elements():
- return freeze(actx.np.sqrt(
- pseudoscalar(actx, dcoll, dd=dd).norm_squared()), actx)
+ result = actx.np.sqrt(
+ pseudoscalar(actx, dcoll, dd=dd).norm_squared())
+
+ if _use_geoderiv_connection:
+ result = dcoll._base_to_geoderiv_connection(dd)(result)
+
+ return freeze(result, actx)
return thaw(_area_elements(), actx)
@@ -573,8 +606,9 @@ def rel_mv_normal(
def mv_normal(
actx: ArrayContext, dcoll: DiscretizationCollection, dd,
+ *, _use_geoderiv_connection=False
) -> MultiVector:
- """Exterior unit normal as a :class:`~pymbolic.geometric_algebra.MultiVector`.
+ r"""Exterior unit normal as a :class:`~pymbolic.geometric_algebra.MultiVector`.
This supports both volume discretizations
(where ambient == topological dimension) and surface discretizations
(where ambient == topological dimension + 1). In the latter case, extra
@@ -584,12 +618,22 @@ def mv_normal(
This function caches its results.
:arg dd: a :class:`~grudge.dof_desc.DOFDesc` as the surface discretization.
+ :arg _use_geoderiv_connection: If *True*, process returned
+ :class:`~meshmode.dof_array.DOFArray`\ s through
+ :meth:`~grudge.DiscretizationCollection._base_to_geoderiv_connection`.
+ This should be set based on whether the code using the result of this
+ function is able to make use of these arrays. The default value agrees
+ with ``actx.supports_nonscalar_broadcasting``. (This is an internal
+ argument and not intended for use outside :mod:`grudge`.)
:returns: a :class:`~pymbolic.geometric_algebra.MultiVector`
containing the unit normals.
"""
dd = dof_desc.as_dofdesc(dd)
- @memoize_in(dcoll, (mv_normal, dd))
+ if _use_geoderiv_connection is None:
+ _use_geoderiv_connection = actx.supports_nonscalar_broadcasting
+
+ @memoize_in(dcoll, (mv_normal, dd, _use_geoderiv_connection))
def _normal():
dim = dcoll.discr_from_dd(dd).dim
ambient_dim = dcoll.ambient_dim
@@ -627,13 +671,17 @@ def mv_normal(
result = mv / actx.np.sqrt(mv.norm_squared())
+ if _use_geoderiv_connection:
+ result = dcoll._base_to_geoderiv_connection(dd)(result)
+
return freeze(result, actx)
n = _normal()
return thaw(n, actx)
-def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd):
+def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd,
+ *, _use_geoderiv_connection=None):
"""Get the unit normal to the specified surface discretization, *dd*.
This supports both volume discretizations
(where ambient == topological dimension) and surface discretizations
@@ -644,10 +692,17 @@ def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd):
This function may be treated as if it caches its results.
:arg dd: a :class:`~grudge.dof_desc.DOFDesc` as the surface discretization.
+ :arg _use_geoderiv_connection: See :func:`mv_normal` for a full description.
+ (This is an internal argument and not intended for use outside
+ :mod:`grudge`.)
+
:returns: an object array of :class:`~meshmode.dof_array.DOFArray`
containing the unit normals at each nodal location.
"""
- return mv_normal(actx, dcoll, dd).as_vector(dtype=object)
+ return mv_normal(
+ actx, dcoll, dd,
+ _use_geoderiv_connection=_use_geoderiv_connection
+ ).as_vector(dtype=object)
# }}}
diff --git a/grudge/op.py b/grudge/op.py
index ed2a29b3fa3611388000dbe1a37fd4cccfc5015d..473e099d26f730e5d020f3746049b86fd1f17589 100644
--- a/grudge/op.py
+++ b/grudge/op.py
@@ -207,7 +207,8 @@ def local_grad(
discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)
actx = vec.array_context
- inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll)
+ inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll,
+ _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
return _gradient_kernel(actx, discr, discr,
_reference_derivative_matrices, inverse_jac_mat, vec,
metric_in_matvec=False)
@@ -230,7 +231,8 @@ def local_d_dx(dcoll: DiscretizationCollection, xyz_axis, vec):
actx = vec.array_context
from grudge.geometry import inverse_surface_metric_derivative_mat
- inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll)
+ inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll,
+ _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
return _single_axis_derivative_kernel(
actx, discr, discr,
@@ -374,7 +376,8 @@ def weak_local_grad(dcoll: DiscretizationCollection, *args, nested=False):
actx = vec.array_context
inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, dd=dd_in,
- times_area_element=True)
+ times_area_element=True,
+ _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
return _gradient_kernel(actx, out_discr, in_discr,
_reference_stiffness_transpose_matrix, inverse_jac_mat, vec,
@@ -422,7 +425,8 @@ def weak_local_d_dx(dcoll: DiscretizationCollection, *args):
actx = vec.array_context
inverse_jac_mat = inverse_surface_metric_derivative_mat(actx, dcoll, dd=dd_in,
- times_area_element=True)
+ times_area_element=True,
+ _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
return _single_axis_derivative_kernel(
actx, out_discr, in_discr, _reference_stiffness_transpose_matrix,
@@ -526,7 +530,8 @@ def _apply_mass_operator(
out_discr = dcoll.discr_from_dd(dd_out)
actx = vec.array_context
- area_elements = area_element(actx, dcoll, dd=dd_in)
+ area_elements = area_element(actx, dcoll, dd=dd_in,
+ _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
return DOFArray(
actx,
data=tuple(
@@ -629,7 +634,8 @@ def _apply_inverse_mass_operator(
actx = vec.array_context
discr = dcoll.discr_from_dd(dd_in)
- inv_area_elements = 1./area_element(actx, dcoll, 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):
@@ -796,7 +802,8 @@ def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec):
actx = vec.array_context
assert len(face_discr.groups) == len(volm_discr.groups)
- surf_area_elements = area_element(actx, dcoll, dd=dd)
+ surf_area_elements = area_element(actx, dcoll, dd=dd,
+ _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
return DOFArray(
actx,