Refactor elementwise reductions

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: