def _single_axis_derivative_kernel(actx, out_discr, in_discr, get_diff_mat,
inv_jac_mat, xyz_axis, vec, *,
metric_in_matvec):
# This gets used from both the strong and the weak derivative. These differ
# in three ways:
# - which differentiation matrix gets used,
# - whether inv_jac_mat is pre-multiplied by a factor that includes the
# area element, and
# - 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".
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(), ))
for out_grp, in_grp, vec_i, ijm_i in zip(
out_discr.groups, in_discr.groups, vec, inv_jac_mat)))
def _apply_mass_operator(dcoll: DiscretizationCollection, dd_out, dd_in, vec):
if not isinstance(vec, DOFArray):
return map_array_container(
partial(_apply_mass_operator, dcoll, dd_out, dd_in), vec)
from grudge.geometry import area_element
in_discr = dcoll.discr_from_dd(dd_in)
out_discr = dcoll.discr_from_dd(dd_out)
actx = vec.array_context
area_elements = area_element(
actx,
dcoll,
dd=dd_in,
_use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
return DOFArray(
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(), ))
for in_grp, out_grp, ae_i, vec_i in zip(
in_discr.groups, out_discr.groups, area_elements, vec)))
def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec):
if not isinstance(vec, DOFArray):
return map_array_container(
partial(_apply_face_mass_operator, dcoll, dd), vec)
from grudge.geometry import area_element
volm_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)
face_discr = dcoll.discr_from_dd(dd)
dtype = vec.entry_dtype
actx = vec.array_context
assert len(face_discr.groups) == len(volm_discr.groups)
surf_area_elements = area_element(
actx,
dcoll,
dd=dd,
_use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
return DOFArray(
actx,
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(vgrp.mesh_el_group.nfaces,
vgrp.nelements, -1),
vec_i.reshape(vgrp.mesh_el_group.nfaces,
vgrp.nelements, 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)))
def _apply_inverse_mass_operator(dcoll: DiscretizationCollection, dd_out,
dd_in, vec):
if isinstance(vec, np.ndarray):
return obj_array_vectorize(
lambda vi: _apply_inverse_mass_operator(dcoll, dd_out, dd_in, vi),
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
discr = dcoll.discr_from_dd(dd_in)
inv_area_elements = 1. / area_element(actx, dcoll, dd=dd_in)
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(
# 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,
vec_i,
tagged=(FirstAxisIsElementsTag(), )))
return DOFArray(actx, data=tuple(group_data))
def _apply_mass_operator(dcoll: DiscretizationCollection, dd_out, dd_in, vec):
if isinstance(vec, np.ndarray):
return obj_array_vectorize(
lambda vi: _apply_mass_operator(dcoll, dd_out, dd_in, vi), vec)
from grudge.geometry import area_element
in_discr = dcoll.discr_from_dd(dd_in)
out_discr = dcoll.discr_from_dd(dd_out)
actx = vec.array_context
area_elements = area_element(actx, dcoll, dd=dd_in)
return DOFArray(
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(), ))
for in_grp, out_grp, ae_i, vec_i in zip(
in_discr.groups, out_discr.groups, area_elements, vec)))
def _apply_stiffness_transpose_operator(dcoll: DiscretizationCollection,
dd_out, dd_in, vec, xyz_axis):
from grudge.geometry import \
inverse_surface_metric_derivative, area_element
in_discr = dcoll.discr_from_dd(dd_in)
out_discr = dcoll.discr_from_dd(dd_out)
actx = vec.array_context
area_elements = area_element(actx, dcoll, dd=dd_in)
inverse_jac_t = actx.np.stack([
inverse_surface_metric_derivative(actx,
dcoll,
rst_axis,
xyz_axis,
dd=dd_in)
for rst_axis in range(dcoll.dim)
])
return DOFArray(
actx,
data=tuple(
actx.einsum(
"dij,ej,ej,dej->ei",
reference_stiffness_transpose_matrix(
actx, out_element_group=out_grp, in_element_group=in_grp),
ae_i,
vec_i,
inv_jac_t_i,
arg_names=("ref_stiffT_mat", "jac", "vec", "inv_jac_t"),
tagged=(FirstAxisIsElementsTag(), ))
for out_grp, in_grp, vec_i, ae_i, inv_jac_t_i in zip(
out_discr.groups, in_discr.groups, vec, area_elements,
inverse_jac_t)))
def apply_spectral_filter(actx, modal_field, discr, cutoff,
mode_response_function):
r"""Apply the spectral filter, defined by the *mode_response_function*.
This routine returns filtered data in the modal basis, which has
been applied using a user-provided *mode_response_function*
to dampen modes beyond the user-provided *cutoff*.
Parameters
----------
actx: :class:`arraycontext.ArrayContext`
A :class:`arraycontext.ArrayContext` associated with
an array of degrees of freedom
modal_field: numpy.ndarray
DOFArray or object array of DOFArrays denoting the modal data
discr: :class:`meshmode.discretization.Discretization`
A :class:`meshmode.discretization.Discretization` describing
the volume discretization the *modal_field* comes from.
cutoff: int
Mode cutoff beyond which the filter will be applied, and below which
the filter will preserve.
mode_response_function:
A function that returns a filter weight for for each mode id.
Returns
-------
modal_field: :class:`meshmode.dof_array.DOFArray`
DOFArray or object array of DOFArrays
"""
from meshmode.transform_metadata import FirstAxisIsElementsTag
return DOFArray(
actx,
tuple(actx.einsum("j,ej->ej",
make_spectral_filter(
actx,
group=grp,
cutoff=cutoff,
mode_response_function=mode_response_function
),
vec_i,
arg_names=("filter", "vec"),
tagged=(FirstAxisIsElementsTag(),))
for grp, vec_i in zip(discr.groups, modal_field))
)
def _compute_local_gradient(dcoll: DiscretizationCollection, vec, xyz_axis):
from grudge.geometry import inverse_surface_metric_derivative
discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)
actx = vec.array_context
inverse_jac_t = actx.np.stack([
inverse_surface_metric_derivative(actx, dcoll, rst_axis, xyz_axis)
for rst_axis in range(dcoll.dim)
])
return DOFArray(actx,
data=tuple(
actx.einsum("dei,dij,ej->ei",
inv_jac_t_i,
reference_derivative_matrices(actx, grp),
vec_i,
arg_names=("inv_jac_t", "ref_diff_mat",
"vec"),
tagged=(FirstAxisIsElementsTag(), ))
for grp, vec_i, inv_jac_t_i in zip(
discr.groups, vec, inverse_jac_t)))
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.
per_group_grads = [
# 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(), )) for out_grp, in_grp, vec_i,
ijm_i in zip(out_discr.groups, in_discr.groups, vec, inv_jac_mat)
]
return make_obj_array([
DOFArray(actx,
data=tuple([pgg_i[xyz_axis] for pgg_i in per_group_grads]))
for xyz_axis in range(out_discr.ambient_dim)
])
def dt_geometric_factors(dcoll: DiscretizationCollection, dd=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).
Specifically, the inradius for each element is computed using the following
formula from [Shewchuk_2002]_, Table 1, for simplicial cells
(triangles/tetrahedra):
.. math::
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.
: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
geometric factors for each cell and at each nodal location.
"""
from meshmode.discretization.poly_element import SimplexElementGroupBase
if dd is None:
dd = DD_VOLUME
actx = dcoll._setup_actx
volm_discr = dcoll.discr_from_dd(dd)
if any(not isinstance(grp, SimplexElementGroupBase)
for grp in volm_discr.groups):
raise NotImplementedError(
"Geometric factors are only implemented for simplex element groups"
)
if volm_discr.dim != volm_discr.ambient_dim:
from warnings import warn
warn(
"The geometric factor for the characteristic length scale in "
"time step estimation is not necessarily valid for non-volume-"
"filling discretizations. Continuing anyway.",
stacklevel=3)
cell_vols = abs(
op.elementwise_integral(dcoll, dd,
volm_discr.zeros(actx) + 1.0))
if dcoll.dim == 1:
return freeze(cell_vols)
dd_face = DOFDesc("all_faces", dd.discretization_tag)
face_discr = dcoll.discr_from_dd(dd_face)
# To get a single value for the total surface area of a cell, we
# take the sum over the averaged face areas on each face.
# NOTE: The face areas are the *same* at each face nodal location.
# This assumes there are the *same* number of face nodes on each face.
surface_areas = abs(
op.elementwise_integral(dcoll, dd_face,
face_discr.zeros(actx) + 1.0))
surface_areas = DOFArray(
actx,
data=tuple(
actx.einsum("fej->e",
face_ae_i.reshape(vgrp.mesh_el_group.nfaces,
vgrp.nelements, afgrp.nunit_dofs),
tagged=(FirstAxisIsElementsTag(), )) / afgrp.nunit_dofs
for vgrp, afgrp, face_ae_i in zip(
volm_discr.groups, face_discr.groups, surface_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))))