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 rand():
from meshmode.dof_array import DOFArray
return DOFArray(
actx,
tuple(
actx.from_numpy(np.random.rand(grp.nelements, grp.nunit_dofs))
for grp in discr.discr_from_dd("vol").groups))
def test_flatten_unflatten(actx_factory):
actx = actx_factory()
ambient_dim = 2
from meshmode.mesh.generation import generate_regular_rect_mesh
mesh = generate_regular_rect_mesh(a=(-0.5, ) * ambient_dim,
b=(+0.5, ) * ambient_dim,
n=(3, ) * ambient_dim,
order=1)
discr = Discretization(actx, mesh, PolynomialWarpAndBlendGroupFactory(3))
a = np.random.randn(discr.ndofs)
from meshmode.dof_array import flatten, unflatten
a_round_trip = actx.to_numpy(
flatten(unflatten(actx, discr, actx.from_numpy(a))))
assert np.array_equal(a, a_round_trip)
from meshmode.dof_array import flatten_to_numpy, unflatten_from_numpy
a_round_trip = flatten_to_numpy(actx, unflatten_from_numpy(actx, discr, a))
assert np.array_equal(a, a_round_trip)
x = thaw(discr.nodes(), actx)
avg_mass = DOFArray(
actx,
tuple([(np.pi + actx.zeros((grp.nelements, 1), a.dtype))
for grp in discr.groups]))
c = MyContainer(name="flatten",
mass=avg_mass,
momentum=make_obj_array([x, x, x]),
enthalpy=x)
from meshmode.dof_array import unflatten_like
c_round_trip = unflatten_like(actx, flatten(c), c)
assert flat_norm(c - c_round_trip) < 1.0e-8
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_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 _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 _interleave_dof_arrays(self, ary1, ary2):
if not isinstance(ary1, DOFArray) or not isinstance(ary2, DOFArray):
raise TypeError("non-array passed to connection")
@memoize_in(self.array_context,
(CenterGranularityConnection, "interleave"))
def prg():
from meshmode.array_context import make_loopy_program
return make_loopy_program(
"""{[iel, idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}""",
"""
dst[iel, 2*idof] = src1[iel, idof]
dst[iel, 2*idof + 1] = src2[iel, idof]
""", [
lp.GlobalArg("src1", shape="(nelements, nunit_dofs)"),
lp.GlobalArg("src2", shape="(nelements, nunit_dofs)"),
lp.GlobalArg("dst", shape="(nelements, 2*nunit_dofs)"),
"...",
],
name="interleave")
results = []
for grp, src1, src2 in zip(self.discr.groups, ary1, ary2):
if src1.dtype != src2.dtype:
raise ValueError("dtype mismatch in inputs")
result = self.array_context.empty(
(grp.nelements, 2 * grp.nunit_dofs), dtype=src1.dtype)
self.array_context.call_loopy(prg(),
src1=src1,
src2=src2,
dst=result,
nelements=grp.nelements,
nunit_dofs=grp.nunit_dofs)
results.append(result)
return DOFArray(self.array_context, tuple(results))
def nodes(self):
r"""
:returns: object array of shape ``(ambient_dim,)`` containing
:class:`~meshmode.dof_array.DOFArray`\ s of node coordinates.
"""
actx = self._setup_actx
@memoize_in(actx, (Discretization, "nodes_prg"))
def prg():
return make_loopy_program("""{[iel,idof,j]:
0<=iel<nelements and
0<=idof<ndiscr_nodes and
0<=j<nmesh_nodes}""",
"""
result[iel, idof] = \
sum(j, resampling_mat[idof, j] * nodes[iel, j])
""",
name="nodes")
return make_obj_array([
_DOFArray.from_list(None, [
actx.freeze(
actx.call_loopy(
prg(),
resampling_mat=actx.from_numpy(
grp.from_mesh_interp_matrix()),
nodes=actx.from_numpy(
grp.mesh_el_group.nodes[iaxis]))["result"])
for grp in self.groups
]) for iaxis in range(self.ambient_dim)
])
def evaluate_sphere_eigf(actx, discr, m: int, n: int) -> DOFArray:
assert discr.ambient_dim == 3
# {{{ get spherical coordinates
from arraycontext import thaw
x, y, z = thaw(discr.nodes(), actx)
theta = actx.np.arctan2(actx.np.sqrt(x**2 + y**2), z)
phi = actx.np.arctan2(y, x)
# }}}
# {{{ evaluate Y^m_n
from scipy.special import sph_harm # pylint: disable=no-name-in-module
y_mn = []
for gtheta, gphi in zip(theta, phi):
result = sph_harm(m, n, actx.to_numpy(gphi), actx.to_numpy(gtheta))
y_mn.append(actx.from_numpy(result.real.copy()))
# }}}
return DOFArray(actx, tuple(y_mn))
def num_reference_derivative(self, ref_axes, vec):
actx = vec.array_context
@memoize_in(actx, (Discretization, "reference_derivative_prg"))
def prg():
return make_loopy_program(
"""{[iel,idof,j]:
0<=iel<nelements and
0<=idof,j<nunit_dofs}""",
"result[iel,idof] = sum(j, diff_mat[idof, j] * vec[iel, j])",
name="diff")
def get_mat(grp):
mat = None
for ref_axis in ref_axes:
next_mat = grp.diff_matrices()[ref_axis]
if mat is None:
mat = next_mat
else:
mat = np.dot(next_mat, mat)
return mat
return _DOFArray.from_list(actx, [
actx.call_loopy(prg(),
diff_mat=actx.from_numpy(get_mat(grp)),
vec=vec[grp.index])["result"]
for grp in self.groups
])
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 test_elementwise_reductions(actx_factory):
actx = actx_factory()
from mesh_data import BoxMeshBuilder
builder = BoxMeshBuilder(ambient_dim=1)
nelements = 4
mesh = builder.get_mesh(nelements, builder.mesh_order)
dcoll = DiscretizationCollection(actx, mesh, order=builder.order)
x = thaw(dcoll.nodes(), actx)
def f(x):
return actx.np.sin(x[0])
field = f(x)
mins = []
maxs = []
sums = []
for gidx, grp_f in enumerate(field):
min_res = np.empty(grp_f.shape)
max_res = np.empty(grp_f.shape)
sum_res = np.empty(grp_f.shape)
for eidx in range(dcoll._volume_discr.groups[gidx].nelements):
element_data = actx.to_numpy(grp_f[eidx])
min_res[eidx, :] = np.min(element_data)
max_res[eidx, :] = np.max(element_data)
sum_res[eidx, :] = np.sum(element_data)
mins.append(actx.from_numpy(min_res))
maxs.append(actx.from_numpy(max_res))
sums.append(actx.from_numpy(sum_res))
from meshmode.dof_array import DOFArray, flat_norm
ref_mins = DOFArray(actx, data=tuple(mins))
ref_maxs = DOFArray(actx, data=tuple(maxs))
ref_sums = DOFArray(actx, data=tuple(sums))
elem_mins = op.elementwise_min(dcoll, field)
elem_maxs = op.elementwise_max(dcoll, field)
elem_sums = op.elementwise_sum(dcoll, field)
assert flat_norm(elem_mins - ref_mins, ord=np.inf) < 1.e-15
assert flat_norm(elem_maxs - ref_maxs, ord=np.inf) < 1.e-15
assert flat_norm(elem_sums - ref_sums, ord=np.inf) < 1.e-15
def test_modal_coefficients_by_projection(actx_factory, quad_group_factory):
group_cls = SimplexElementGroup
modal_group_factory = ModalSimplexGroupFactory
actx = actx_factory()
order = 10
m_order = 5
# Make a regular rectangle mesh
mesh = mgen.generate_regular_rect_mesh(a=(0, 0),
b=(5, 3),
npoints_per_axis=(10, 6),
order=order,
group_cls=group_cls)
# Make discretizations
nodal_disc = Discretization(actx, mesh, quad_group_factory(order))
modal_disc = Discretization(actx, mesh, modal_group_factory(m_order))
# Make connections one using quadrature projection
nodal_to_modal_conn_quad = NodalToModalDiscretizationConnection(
nodal_disc, modal_disc, allow_approximate_quad=True)
def f(x):
return 2 * actx.np.sin(5 * x)
x_nodal = thaw(nodal_disc.nodes()[0], actx)
nodal_f = f(x_nodal)
# Compute modal coefficients we expect to get
import modepy as mp
grp, = nodal_disc.groups
shape = mp.Simplex(grp.dim)
space = mp.space_for_shape(shape, order=m_order)
basis = mp.orthonormal_basis_for_space(space, shape)
quad = grp.quadrature_rule()
nodal_f_data = actx.to_numpy(nodal_f[0])
vdm = mp.vandermonde(basis.functions, quad.nodes)
w_diag = np.diag(quad.weights)
modal_data = []
for _, nodal_data in enumerate(nodal_f_data):
# Compute modal data in each element: V.T * W * nodal_data
elem_modal_f = np.dot(vdm.T, np.dot(w_diag, nodal_data))
modal_data.append(elem_modal_f)
modal_data = actx.from_numpy(np.asarray(modal_data))
modal_f_expected = DOFArray(actx, data=(modal_data, ))
# Map nodal coefficients using the quadrature-based projection
modal_f_computed = nodal_to_modal_conn_quad(nodal_f)
err = flat_norm(modal_f_expected - modal_f_computed)
assert err <= 1e-13
def _compute_characteristic_lengthscales():
return freeze(
DOFArray(
actx,
data=tuple(
# Scale each group array of geometric factors by the
# corresponding group non-geometric factor
cng * geo_facts for cng, geo_facts in zip(
dt_non_geometric_factors(dcoll),
thaw(dt_geometric_factors(dcoll), actx)))))
def _get_ref_data(field):
mins = []
maxs = []
sums = []
for grp_f in field:
min_res = np.empty(grp_f.shape)
max_res = np.empty(grp_f.shape)
sum_res = np.empty(grp_f.shape)
for eidx in range(dcoll.mesh.nelements):
element_data = actx.to_numpy(grp_f[eidx])
min_res[eidx, :] = np.min(element_data)
max_res[eidx, :] = np.max(element_data)
sum_res[eidx, :] = np.sum(element_data)
mins.append(actx.from_numpy(min_res))
maxs.append(actx.from_numpy(max_res))
sums.append(actx.from_numpy(sum_res))
min_field = DOFArray(actx, data=tuple(mins))
max_field = DOFArray(actx, data=tuple(maxs))
sums_field = DOFArray(actx, data=tuple(sums))
return min_field, max_field, sums_field
def _new_array(self, actx, creation_func, dtype=None):
if dtype is None:
dtype = self.real_dtype
elif dtype == "c":
dtype = self.complex_dtype
else:
dtype = np.dtype(dtype)
return _DOFArray.from_list(actx, [
creation_func(shape=(grp.nelements, grp.nunit_dofs), dtype=dtype)
for grp in self.groups
])
def _apply_elementwise_reduction(op_name: str, dcoll: DiscretizationCollection,
*args) -> DOFArray:
r"""Returns a vector of DOFs with all entries on each element set
to the reduction operation *op_name* over all degrees of freedom.
:arg \*args: Arguments for the reduction operator, such as *dd* and *vec*.
:returns: a :class:`~meshmode.dof_array.DOFArray` or object arrary of
:class:`~meshmode.dof_array.DOFArray`s.
"""
if len(args) == 1:
vec, = args
dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE)
elif len(args) == 2:
dd, vec = args
else:
raise TypeError("invalid number of arguments")
dd = dof_desc.as_dofdesc(dd)
if isinstance(vec, np.ndarray):
return obj_array_vectorize(
lambda vi: _apply_elementwise_reduction(op_name, dcoll, dd, vi),
vec)
actx = vec.array_context
@memoize_in(actx,
(_apply_elementwise_reduction, "elementwise_%s_prg" % op_name))
def elementwise_prg():
# FIXME: This computes the reduction value redundantly for each
# output DOF.
t_unit = make_loopy_program([
"{[iel]: 0 <= iel < nelements}",
"{[idof, jdof]: 0 <= idof, jdof < ndofs}"
],
"""
result[iel, idof] = %s(jdof, operand[iel, jdof])
""" % op_name,
name="grudge_elementwise_%s_knl" % op_name)
import loopy as lp
from meshmode.transform_metadata import (ConcurrentElementInameTag,
ConcurrentDOFInameTag)
return lp.tag_inames(t_unit, {
"iel": ConcurrentElementInameTag(),
"idof": ConcurrentDOFInameTag()
})
return DOFArray(
actx,
data=tuple(
actx.call_loopy(elementwise_prg(), operand=vec_i)["result"]
for vec_i in vec))
def _map_elementwise_reduction(self, reduction_name, expr):
@memoize_in(self.places, "elementwise_node_" + reduction_name)
def node_knl():
from meshmode.array_context import make_loopy_program
return make_loopy_program("""{[iel, idof, jdof]:
0<=iel<nelements and
0<=idof, jdof<ndofs}""",
"""
result[iel, idof] = %s(jdof, operand[iel, jdof])
""" % reduction_name,
name="nodewise_reduce")
@memoize_in(self.places, "elementwise_" + reduction_name)
def element_knl():
from meshmode.array_context import make_loopy_program
return make_loopy_program("""{[iel, jdof]:
0<=iel<nelements and
0<=jdof<ndofs}
""",
"""
result[iel, 0] = %s(jdof, operand[iel, jdof])
""" % reduction_name,
name="elementwise_reduce")
discr = self.places.get_discretization(expr.dofdesc.geometry,
expr.dofdesc.discr_stage)
operand = self.rec(expr.operand)
assert operand.shape == (len(discr.groups), )
def _reduce(knl, result):
for grp in discr.groups:
self.array_context.call_loopy(knl,
operand=operand[grp.index],
result=result[grp.index])
return result
dtype = operand.entry_dtype
granularity = expr.dofdesc.granularity
if granularity is sym.GRANULARITY_NODE:
return _reduce(node_knl(),
discr.empty(self.array_context, dtype=dtype))
elif granularity is sym.GRANULARITY_ELEMENT:
result = DOFArray(
self.array_context,
tuple([
self.array_context.empty((grp.nelements, 1), dtype=dtype)
for grp in discr.groups
]))
return _reduce(element_knl(), result)
else:
raise ValueError(f"unsupported granularity: {granularity}")
def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec):
if isinstance(vec, np.ndarray):
return obj_array_vectorize(
lambda vi: _apply_face_mass_operator(dcoll, dd, vi), 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
@memoize_in(actx, (_apply_face_mass_operator, "face_mass_knl"))
def prg():
t_unit = make_loopy_program([
"{[iel]: 0 <= iel < nelements}", "{[f]: 0 <= f < nfaces}",
"{[idof]: 0 <= idof < nvol_nodes}",
"{[jdof]: 0 <= jdof < nface_nodes}"
],
"""
result[iel, idof] = sum(f, sum(jdof, mat[idof, f, jdof]
* jac_surf[f, iel, jdof]
* vec[f, iel, jdof]))
""",
name="face_mass")
import loopy as lp
from meshmode.transform_metadata import (ConcurrentElementInameTag,
ConcurrentDOFInameTag)
return lp.tag_inames(t_unit, {
"iel": ConcurrentElementInameTag(),
"idof": ConcurrentDOFInameTag()
})
assert len(face_discr.groups) == len(volm_discr.groups)
surf_area_elements = area_element(actx, dcoll, dd=dd)
return DOFArray(
actx,
data=tuple(
actx.call_loopy(
prg(),
mat=reference_face_mass_matrix(actx,
face_element_group=afgrp,
vol_element_group=vgrp,
dtype=dtype),
jac_surf=surf_ae_i.reshape(vgrp.mesh_el_group.nfaces,
vgrp.nelements, afgrp.nunit_dofs),
vec=vec_i.reshape(vgrp.mesh_el_group.nfaces, vgrp.nelements,
afgrp.nunit_dofs))["result"] for vgrp, afgrp,
vec_i, surf_ae_i in zip(volm_discr.groups, face_discr.groups, vec,
surf_area_elements)))
def map_insn_loopy_kernel(self, insn, profile_data=None):
kdescr = insn.kernel_descriptor
discr = self.discrwb.discr_from_dd(kdescr.governing_dd)
dof_array_kwargs = {}
other_kwargs = {}
for name, expr in kdescr.input_mappings.items():
v = self.rec(expr)
if isinstance(v, DOFArray):
dof_array_kwargs[name] = v
else:
other_kwargs[name] = v
for name in kdescr.scalar_args():
v = other_kwargs[name]
if isinstance(v, (int, float)):
other_kwargs[name] = discr.real_dtype.type(v)
elif isinstance(v, complex):
other_kwargs[name] = discr.complex_dtype.type(v)
elif isinstance(v, np.number):
pass
else:
raise ValueError(
"unrecognized scalar type for variable '%s': %s" %
(name, type(v)))
result = {}
for grp in discr.groups:
kwargs = other_kwargs.copy()
kwargs["nelements"] = grp.nelements
kwargs["nunit_dofs"] = grp.nunit_dofs
for name, ary in dof_array_kwargs.items():
kwargs[name] = ary[grp.index]
knl_result = self.array_context.call_loopy(kdescr.loopy_kernel,
**kwargs)
for name, val in knl_result.items():
result.setdefault(name, []).append(val)
result = {
name: DOFArray(self.array_context, tuple(val))
for name, val in result.items()
}
return list(result.items()), []
def _interleave_dof_arrays(self, ary1, ary2):
if not isinstance(ary1, DOFArray) or not isinstance(ary2, DOFArray):
raise TypeError("non-array passed to connection")
if ary1.array_context is not ary2.array_context:
raise ValueError("array context of the two arguments must match")
actx = ary1.array_context
@memoize_in(actx, (CenterGranularityConnection, "interleave"))
def prg():
from arraycontext import make_loopy_program
t_unit = make_loopy_program(
"{[iel, idof]: 0 <= iel < nelements and 0 <= idof < nunit_dofs}",
"""
result[iel, 2*idof] = ary1[iel, idof]
result[iel, 2*idof + 1] = ary2[iel, idof]
""", [
lp.GlobalArg("ary1", shape="(nelements, nunit_dofs)"),
lp.GlobalArg("ary2", shape="(nelements, nunit_dofs)"),
lp.GlobalArg("result", shape="(nelements, 2*nunit_dofs)"),
...
],
name="interleave")
from meshmode.transform_metadata import (ConcurrentElementInameTag,
ConcurrentDOFInameTag)
return lp.tag_inames(
t_unit, {
"iel": ConcurrentElementInameTag(),
"idof": ConcurrentDOFInameTag()
})
results = []
for grp, subary1, subary2 in zip(self.discr.groups, ary1, ary2):
if subary1.dtype != subary2.dtype:
raise ValueError(
"dtype mismatch in inputs: "
f"'{subary1.dtype.name}' and '{subary2.dtype.name}'")
result = actx.call_loopy(prg(),
ary1=subary1,
ary2=subary2,
nelements=grp.nelements,
nunit_dofs=grp.nunit_dofs)["result"]
results.append(result)
return DOFArray(actx, tuple(results))
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 _visualize_refinement(actx: PyOpenCLArrayContext,
discr,
niter,
stage_nr,
stage_name,
flags,
visualize=False):
if not visualize:
return
if stage_nr not in (1, 2):
raise ValueError("unexpected stage number")
flags = flags.get()
logger.info("for stage %s: splitting %d/%d stage-%d elements", stage_name,
np.sum(flags), discr.mesh.nelements, stage_nr)
from meshmode.discretization.visualization import make_visualizer
vis = make_visualizer(actx, discr, 3)
assert len(flags) == discr.mesh.nelements
flags = flags.astype(bool)
nodes_flags_template = discr.zeros(actx)
nodes_flags = []
for grp in discr.groups:
meg = grp.mesh_el_group
nodes_flags_grp = actx.to_numpy(nodes_flags_template[grp.index])
nodes_flags_grp[flags[meg.element_nr_base:meg.nelements +
meg.element_nr_base]] = 1
nodes_flags.append(actx.from_numpy(nodes_flags_grp))
nodes_flags = DOFArray(actx, tuple(nodes_flags))
vis_data = [
("refine_flags", nodes_flags),
]
if 0:
from pytential import sym, bind
bdry_normals = bind(discr, sym.normal(
discr.ambient_dim))(actx).as_vector(dtype=object)
vis_data.append(("bdry_normals", bdry_normals), )
vis.write_vtk_file(f"refinement-{stage_name}-{niter:03d}.vtu", vis_data)
def inverse_mass(self, vec):
if not isinstance(vec, DOFArray):
return map_array_container(self.inverse_mass, vec)
actx = vec.array_context
dtype = vec.entry_dtype
discr = self.volume_discr
return DOFArray(
actx,
data=tuple(
actx.einsum("ij,ej->ei",
self.get_inverse_mass_matrix(grp, dtype),
vec_i,
arg_names=("mass_inv_mat", "vec"),
tagged=(FirstAxisIsElementsTag(), ))
for grp, vec_i in zip(discr.groups, vec))) / thaw(
self.vol_jacobian(), actx)
def map_insn_diff_batch_assign(self, insn, profile_data=None):
field = self.rec(insn.field)
repr_op = insn.operators[0]
assert repr_op.dd_in.domain_tag == repr_op.dd_out.domain_tag
in_discr = self.discrwb.discr_from_dd(repr_op.dd_in)
out_discr = self.discrwb.discr_from_dd(repr_op.dd_out)
prg = self._elwise_linear_loopy_prg()
result = []
for name, op in zip(insn.names, insn.operators):
group_results = []
for in_grp, out_grp in zip(in_discr.groups, out_discr.groups):
assert in_grp.nelements == out_grp.nelements
if in_grp.nelements == 0:
continue
# Cache operator
cache_key = "diff_batch", in_grp, out_grp, tuple(insn.operators),\
field.entry_dtype
try:
matrices_dev = self.bound_op.operator_data_cache[cache_key]
except KeyError:
matrices_dev = [
self.array_context.from_numpy(mat)
for mat in repr_op.matrices(out_grp, in_grp)
]
self.bound_op.operator_data_cache[cache_key] = matrices_dev
group_results.append(
self.array_context.call_loopy(
prg,
mat=matrices_dev[op.rst_axis],
vec=field[in_grp.index])["result"])
result.append(
(name, DOFArray(self.array_context, tuple(group_results))))
return result, []
def quad_weights(self):
""":returns: A :class:`~meshmode.dof_array.DOFArray` with quadrature weights.
"""
actx = self._setup_actx
@memoize_in(actx, (Discretization, "quad_weights_prg"))
def prg():
return make_loopy_program(
"{[iel,idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}",
"result[iel,idof] = weights[idof]",
name="quad_weights")
return _DOFArray.from_list(None, [
actx.freeze(
actx.call_loopy(
prg(),
weights=actx.from_numpy(grp.weights),
nelements=grp.nelements,
)["result"]) for grp in self.groups
])
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 test_node_reduction(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
actx = PyOpenCLArrayContext(queue)
# {{{ build discretization
target_order = 4
nelements = 32
mesh = make_curve_mesh(starfish,
np.linspace(0.0, 1.0, nelements + 1),
target_order)
discr = Discretization(actx, mesh,
InterpolatoryQuadratureSimplexGroupFactory(target_order))
# }}}
# {{{ test
# create a shuffled [1, nelements + 1] array
ary = []
el_nr_base = 0
for grp in discr.groups:
x = 1 + np.arange(el_nr_base, grp.nelements)
np.random.shuffle(x)
ary.append(actx.freeze(actx.from_numpy(x.reshape(-1, 1))))
el_nr_base += grp.nelements
from meshmode.dof_array import DOFArray
ary = DOFArray(actx, tuple(ary))
for func, expected in [
(sym.NodeSum, nelements * (nelements + 1) // 2),
(sym.NodeMax, nelements),
(sym.NodeMin, 1),
]:
r = bind(discr, func(sym.var("x")))(actx, x=ary)
assert abs(r - expected) < 1.0e-15, r
def face_mass(self, vec):
if not isinstance(vec, DOFArray):
return map_array_container(self.face_mass, vec)
actx = vec.array_context
dtype = vec.entry_dtype
@memoize_in(self, "face_mass_knl")
def knl():
return make_loopy_program(
"""{[iel,idof,f,j]:
0<=iel<nelements and
0<=f<nfaces and
0<=idof<nvol_nodes and
0<=j<nface_nodes}""", "result[iel,idof] = "
"sum(f, sum(j, mat[idof, f, j] * vec[f, iel, j]))",
name="face_mass")
all_faces_conn = self.get_connection("vol", "all_faces")
all_faces_discr = all_faces_conn.to_discr
vol_discr = all_faces_conn.from_discr
fj = thaw(self.face_jacobian("all_faces"), vec.array_context)
vec = vec * fj
assert len(all_faces_discr.groups) == len(vol_discr.groups)
return DOFArray(
actx,
data=tuple(
actx.call_loopy(
knl(),
mat=self.get_local_face_mass_matrix(afgrp, volgrp, dtype),
vec=vec_i.reshape(volgrp.mesh_el_group.nfaces, volgrp.
nelements, afgrp.nunit_dofs))["result"]
for afgrp, volgrp, vec_i in zip(all_faces_discr.groups,
vol_discr.groups, vec)))
