def __call__(self, arys):
r"""
:arg arys: either a single :class:`~meshmode.dof_array.DOFArray`
or a list/tuple with exactly 2 entries that are both
:class:`~meshmode.dof_array.DOFArray`\ s.
Additionally, this function vectorizes over object arrays of
:class:`~meshmode.dof_array.DOFArrays`\ s.
:return: an interleaved array or list of :class:`pyopencl.array.Array`s.
If *vecs* was a pair of arrays :math:`(x, y)`, they are
interleaved as :math:`[x_1, y_1, x_2, y_2, \ddots, x_n, y_n]`.
A single array is simply interleaved with itself.
"""
if isinstance(arys, DOFArray):
arys = (arys, arys)
if isinstance(arys, (list, tuple)):
assert len(arys) == 2
else:
raise ValueError("cannot interleave arrays")
if isinstance(arys[0], DOFArray):
return self._interleave_dof_arrays(*arys)
else:
from pytools.obj_array import obj_array_vectorize_n_args
return obj_array_vectorize_n_args(self._interleave_dof_arrays,
*arys)
def structured_vdot(x, y):
# vdot() implementation that is aware of scalars and host or
# PyOpenCL arrays. It also recurses down nested object arrays.
if (isinstance(x, Number)
or (isinstance(x, np.ndarray) and x.dtype.char != "O")):
return np.vdot(x, y)
elif isinstance(x, cl.array.Array):
return cl.array.vdot(x, y).get()
else:
assert isinstance(x, np.ndarray) and x.dtype.char == "O"
return sum(obj_array_vectorize_n_args(structured_vdot, x, y))
def _comparison(self, operator_func, other):
from numbers import Number
if isinstance(other, DOFArray):
return obj_array_vectorize_n_args(operator_func, self, other)
elif isinstance(other, Number):
return obj_array_vectorize(
lambda self_entry: operator_func(self_entry, other), self)
else:
# fall back to "best effort" (i.e. likley failure)
return operator_func(self, other)
def flatten_if_needed(actx: PyOpenCLArrayContext, ary: np.ndarray):
from pytools.obj_array import obj_array_vectorize_n_args
from meshmode.dof_array import DOFArray, thaw, flatten
if (isinstance(ary, np.ndarray) and ary.dtype.char == "O"
and not isinstance(ary, DOFArray)):
return obj_array_vectorize_n_args(flatten_if_needed, actx, ary)
if not isinstance(ary, DOFArray):
return ary
if ary.array_context is None:
ary = thaw(actx, ary)
return flatten(ary)
def obj_or_dof_array_vectorize_n_args(f, *args):
r"""Apply the function *f* elementwise to all entries of any
object arrays or :class:`DOFArray`\ s in *args*. All such arrays are expected
to have the same shape (but this is not checked).
Equivalent to an appropriately-looped execution of::
result[idx] = f(obj_array_arg1[idx], arg2, obj_array_arg3[idx])
Return an array of the same shape as the arguments consisting of the
return values of *f*. If the elements of arrays found in *args* are
further object arrays, recurse. If a :class:`DOFArray` is found, apply
*f* to its entries. If non-object-arrays are found, apply *f* to those.
"""
dofarray_arg_indices = [
i for i, arg in enumerate(args)
if isinstance(arg, DOFArray)]
if not dofarray_arg_indices:
if any(isinstance(arg, np.ndarray) and arg.dtype.char == "O"
for i, arg in enumerate(args)):
from pytools.obj_array import obj_array_vectorize_n_args
return obj_array_vectorize_n_args(
partial(obj_or_dof_array_vectorize_n_args, f), *args)
else:
return f(*args)
leading_da_index = dofarray_arg_indices[0]
template_ary = args[leading_da_index]
result = []
new_args = list(args)
for igrp in range(len(template_ary)):
for arg_i in dofarray_arg_indices:
new_args[arg_i] = args[arg_i][igrp]
result.append(f(*new_args))
return DOFArray(template_ary.array_context, tuple(result))
def __call__(self, *exprs):
from pytools.obj_array import obj_array_vectorize_n_args
return obj_array_vectorize_n_args(
super(FunctionSymbol, self).__call__, *exprs)
def diffusion_operator(discr,
quad_tag,
alpha,
boundaries,
u,
return_grad_u=False):
r"""
Compute the diffusion operator.
The diffusion operator is defined as
$\nabla\cdot(\alpha\nabla u)$, where $\alpha$ is the diffusivity and
$u$ is a scalar field.
Uses unstabilized central numerical fluxes.
Parameters
----------
discr: grudge.eager.EagerDGDiscretization
the discretization to use
quad_tag:
quadrature tag indicating which discretization in *discr* to use for
overintegration
alpha: numbers.Number or meshmode.dof_array.DOFArray
the diffusivity value(s)
boundaries:
dictionary (or list of dictionaries) mapping boundary tags to
:class:`DiffusionBoundary` instances
u: meshmode.dof_array.DOFArray or numpy.ndarray
the DOF array (or object array of DOF arrays) to which the operator should be
applied
return_grad_u: bool
an optional flag indicating whether $\nabla u$ should also be returned
Returns
-------
diff_u: meshmode.dof_array.DOFArray or numpy.ndarray
the diffusion operator applied to *u*
grad_u: numpy.ndarray
the gradient of *u*; only returned if *return_grad_u* is True
"""
if isinstance(u, np.ndarray):
if not isinstance(boundaries, list):
raise TypeError(
"boundaries must be a list if u is an object array")
if len(boundaries) != len(u):
raise TypeError("boundaries must be the same length as u")
return obj_array_vectorize_n_args(
lambda boundaries_i, u_i: diffusion_operator(discr,
quad_tag,
alpha,
boundaries_i,
u_i,
return_grad_u=
return_grad_u),
make_obj_array(boundaries), u)
for btag, bdry in boundaries.items():
if not isinstance(bdry, DiffusionBoundary):
raise TypeError(f"Unrecognized boundary type for tag {btag}. "
"Must be an instance of DiffusionBoundary.")
dd_quad = DOFDesc("vol", quad_tag)
dd_allfaces_quad = DOFDesc("all_faces", quad_tag)
grad_u = discr.inverse_mass(
discr.weak_grad(-u) - # noqa: W504
discr.face_mass(
dd_allfaces_quad,
gradient_flux(discr, quad_tag, interior_trace_pair(discr, u)) +
sum(
bdry.get_gradient_flux(discr, quad_tag, as_dofdesc(btag),
alpha, u)
for btag, bdry in boundaries.items()) + sum(
gradient_flux(discr, quad_tag, u_tpair)
for u_tpair in cross_rank_trace_pairs(discr, u))))
alpha_quad = discr.project("vol", dd_quad, alpha)
grad_u_quad = discr.project("vol", dd_quad, grad_u)
diff_u = discr.inverse_mass(
discr.weak_div(dd_quad, -alpha_quad * grad_u_quad) - # noqa: W504
discr.face_mass(
dd_allfaces_quad,
diffusion_flux(discr, quad_tag, interior_trace_pair(discr, alpha),
interior_trace_pair(discr, grad_u)) +
sum(
bdry.get_diffusion_flux(discr, quad_tag, as_dofdesc(btag),
alpha, grad_u)
for btag, bdry in boundaries.items()) + sum(
diffusion_flux(discr, quad_tag, alpha_tpair, grad_u_tpair)
for alpha_tpair, grad_u_tpair in zip(
cross_rank_trace_pairs(discr, alpha),
cross_rank_trace_pairs(discr, grad_u)))))
if return_grad_u:
return diff_u, grad_u
else:
return diff_u
def diffusion_operator(discr, alpha, boundaries, u):
r"""
Compute the diffusion operator.
The diffusion operator is defined as
$\nabla\cdot(\alpha\nabla u)$, where $\alpha$ is the diffusivity and
$u$ is a scalar field.
Parameters
----------
discr: grudge.eager.EagerDGDiscretization
the discretization to use
alpha: float
the (constant) diffusivity
boundaries:
dictionary (or list of dictionaries) mapping boundary tags to
:class:`DiffusionBoundary` instances
u: meshmode.dof_array.DOFArray or numpy.ndarray
the DOF array (or object array of DOF arrays) to which the operator should be
applied
Returns
-------
meshmode.dof_array.DOFArray or numpy.ndarray
the diffusion operator applied to *u*
"""
if isinstance(u, np.ndarray):
if not isinstance(boundaries, list):
raise TypeError(
"boundaries must be a list if u is an object array")
if len(boundaries) != len(u):
raise TypeError("boundaries must be the same length as u")
return obj_array_vectorize_n_args(
lambda boundaries_i, u_i: diffusion_operator(
discr, alpha, boundaries_i, u_i), make_obj_array(boundaries),
u)
for btag, bdry in boundaries.items():
if not isinstance(bdry, DiffusionBoundary):
raise TypeError(f"Unrecognized boundary type for tag {btag}. "
"Must be an instance of DiffusionBoundary.")
q = discr.inverse_mass(
-math.sqrt(alpha) * discr.weak_grad(u) + # noqa: W504
discr.face_mass(
_q_flux(discr, alpha=alpha, u_tpair=interior_trace_pair(discr, u))
+ sum(
bdry.get_q_flux(discr, alpha=alpha, dd=btag, u=u)
for btag, bdry in boundaries.items()) + sum(
_q_flux(discr, alpha=alpha, u_tpair=tpair)
for tpair in cross_rank_trace_pairs(discr, u))))
return (discr.inverse_mass(
-math.sqrt(alpha) * discr.weak_div(q) + # noqa: W504
discr.face_mass(
_u_flux(discr, alpha=alpha, q_tpair=interior_trace_pair(discr, q))
+ sum(
bdry.get_u_flux(discr, alpha=alpha, dd=btag, q=q)
for btag, bdry in boundaries.items()) + sum(
_u_flux(discr, alpha=alpha, q_tpair=tpair)
for tpair in cross_rank_trace_pairs(discr, q)))))
