def flat_norm(ary, ord=None) -> float:
r"""Return an element-wise :math:`\ell^{\text{ord}}` norm of *ary*.
:arg ary: may be a :class:`DOFArray` or a
:class:`~arraycontext.ArrayContainer` containing them.
"""
from numbers import Number
if isinstance(ary, Number):
return abs(ary)
if ord is None:
ord = 2
from arraycontext import is_array_container
import numpy.linalg as la
if isinstance(ary, DOFArray):
actx = ary.array_context
return la.norm([
actx.np.linalg.norm(actx.np.ravel(ary, order="A"), ord=ord)
for _, subary in serialize_container(ary)
],
ord=ord)
elif is_array_container(ary):
return la.norm([
flat_norm(subary, ord=ord)
for _, subary in serialize_container(ary)
],
ord=ord)
raise TypeError(
f"unsupported array type passed to flat_norm: '{type(ary).__name__}'")
def _unflatten_like(_ary, _prototype):
if isinstance(_prototype, DOFArray):
group_shapes = [subary.shape for subary in _prototype]
group_sizes = [subary.size for subary in _prototype]
group_starts = np.cumsum([0] + group_sizes)
return _unflatten_dof_array(actx,
_ary,
group_shapes,
group_starts,
strict=True)
elif is_array_container(_prototype):
assert type(_ary) is type(_prototype)
return deserialize_container(_prototype, [
(_same_key(key1, key2), _unflatten_like(subary, subprototype))
for (key1, subary), (key2, subprototype) in zip(
serialize_container(_ary), serialize_container(_prototype))
])
else:
if strict:
raise ValueError(
"cannot unflatten array "
f"with prototype '{type(_prototype).__name__}'; "
"use 'strict=False' to leave the array unchanged")
assert type(_ary) is type(_prototype)
return _ary
def nodal_sum_loc(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
r"""Return the rank-local nodal sum of a vector of degrees of freedom *vec*.
:arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value
convertible to one.
:arg vec: a :class:`~meshmode.dof_array.DOFArray` or an
:class:`~arraycontext.container.ArrayContainer` of them.
:returns: a scalar denoting the rank-local nodal sum.
"""
if not isinstance(vec, DOFArray):
return sum(
nodal_sum_loc(dcoll, dd, comp)
for _, comp in serialize_container(vec))
actx = vec.array_context
return sum([actx.np.sum(grp_ary) for grp_ary in vec])
def nodal_max_loc(dcoll: DiscretizationCollection, dd, vec) -> "DeviceScalar":
r"""Return the rank-local nodal maximum of a vector of degrees
of freedom *vec*.
:arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value
convertible to one.
:arg vec: a :class:`~meshmode.dof_array.DOFArray` or an
:class:`~arraycontext.container.ArrayContainer`.
:returns: a scalar denoting the rank-local nodal maximum.
"""
if not isinstance(vec, DOFArray):
return max(
nodal_max_loc(dcoll, dd, comp)
for _, comp in serialize_container(vec))
actx = vec.array_context
return reduce(
lambda acc, grp_ary: actx.np.maximum(acc, actx.np.max(grp_ary)), vec,
actx.from_numpy(np.array(-np.inf)))