diff --git a/grudge/reductions.py b/grudge/reductions.py
index 5337349e1c1701b6389292ffdcde67f47b9a72eb..d9b0a6bad37a84bd33f4482dea8d6ba939e5c40f 100644
--- a/grudge/reductions.py
+++ b/grudge/reductions.py
@@ -60,10 +60,7 @@ THE SOFTWARE.
from numbers import Number
from functools import reduce
-from arraycontext import (
- ArrayContext,
- make_loopy_program
-)
+from arraycontext import make_loopy_program
from grudge.discretization import DiscretizationCollection
@@ -260,36 +257,15 @@ def integral(dcoll: DiscretizationCollection, dd, vec) -> float:
# {{{ Elementwise reductions
-def _map_elementwise_reduction(actx: ArrayContext, op_name):
- @memoize_in(actx, (_map_elementwise_reduction,
- "elementwise_%s_prg" % op_name))
- def prg():
- return 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
- )
- return prg()
-
-
-def elementwise_sum(dcoll: DiscretizationCollection, *args) -> DOFArray:
+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 sum of DOFs on that element.
+ to the reduction operation *op_name* over all degrees of freedom.
- May be called with ``(dcoll, vec)`` or ``(dcoll, dd, vec)``.
-
- :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
- Defaults to the base volume discretization if not provided.
- :arg vec: a :class:`~meshmode.dof_array.DOFArray`
- :returns: a :class:`~meshmode.dof_array.DOFArray` whose entries
- denote the element-wise sum of *vec*.
+ :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)
@@ -302,64 +278,63 @@ def elementwise_sum(dcoll: DiscretizationCollection, *args) -> DOFArray:
if isinstance(vec, np.ndarray):
return obj_array_vectorize(
- lambda vi: elementwise_sum(dcoll, dd, vi), vec
+ 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():
+ return 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
+ )
+
return DOFArray(
actx,
data=tuple(
- actx.call_loopy(
- _map_elementwise_reduction(actx, "sum"),
- operand=vec_i
- )["result"]
+ actx.call_loopy(elementwise_prg(), operand=vec_i)["result"]
for vec_i in vec
)
)
-def elementwise_max(dcoll: DiscretizationCollection, *args) -> DOFArray:
+def elementwise_sum(dcoll: DiscretizationCollection, *args) -> DOFArray:
r"""Returns a vector of DOFs with all entries on each element set
- to the maximum over all DOFs on that element.
+ to the sum of DOFs on that element.
May be called with ``(dcoll, vec)`` or ``(dcoll, dd, vec)``.
- :arg dcoll: a :class:`grudge.discretization.DiscretizationCollection`.
:arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
Defaults to the base volume discretization if not provided.
:arg vec: a :class:`~meshmode.dof_array.DOFArray`
:returns: a :class:`~meshmode.dof_array.DOFArray` whose entries
- denote the element-wise max of *vec*.
+ denote the element-wise sum of *vec*.
"""
+ return _apply_elementwise_reduction("sum", dcoll, *args)
- 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: elementwise_max(dcoll, dd, vi), vec
- )
+def elementwise_max(dcoll: DiscretizationCollection, *args) -> DOFArray:
+ r"""Returns a vector of DOFs with all entries on each element set
+ to the maximum over all DOFs on that element.
- actx = vec.array_context
+ May be called with ``(dcoll, vec)`` or ``(dcoll, dd, vec)``.
- return DOFArray(
- actx,
- tuple(
- actx.call_loopy(
- _map_elementwise_reduction(actx, "max"),
- operand=vec_i
- )["result"]
- for vec_i in vec
- )
- )
+ :arg dcoll: a :class:`grudge.discretization.DiscretizationCollection`.
+ :arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
+ Defaults to the base volume discretization if not provided.
+ :arg vec: a :class:`~meshmode.dof_array.DOFArray`
+ :returns: a :class:`~meshmode.dof_array.DOFArray` whose entries
+ denote the element-wise max of *vec*.
+ """
+ return _apply_elementwise_reduction("max", dcoll, *args)
def elementwise_min(dcoll: DiscretizationCollection, *args) -> DOFArray:
@@ -375,34 +350,7 @@ def elementwise_min(dcoll: DiscretizationCollection, *args) -> DOFArray:
:returns: a :class:`~meshmode.dof_array.DOFArray` whose entries
denote the element-wise min of *vec*.
"""
-
- 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: elementwise_min(dcoll, dd, vi), vec
- )
-
- actx = vec.array_context
-
- return DOFArray(
- actx,
- tuple(
- actx.call_loopy(
- _map_elementwise_reduction(actx, "min"),
- operand=vec_i
- )["result"]
- for vec_i in vec
- )
- )
+ return _apply_elementwise_reduction("min", dcoll, *args)
def elementwise_integral(dcoll: DiscretizationCollection, dd, vec) -> DOFArray: