Merge remote-tracking branch 'origin/main' into multi-volume

grudge/op.py

@@ -98,7 +98,11 @@
     tracepair_with_discr_tag,
     interior_trace_pair,
     interior_trace_pairs,
+    local_interior_trace_pair,
+    inter_volume_trace_pairs,
+    local_inter_volume_trace_pairs,
     cross_rank_trace_pairs,
+    cross_rank_inter_volume_trace_pairs,
     bdry_trace_pair,
     bv_trace_pair
 )
@@ -126,8 +130,10 @@
     "interior_trace_pair",
     "interior_trace_pairs",
     "local_interior_trace_pair",
-    "connected_ranks",
+    "inter_volume_trace_pairs",
+    "local_inter_volume_trace_pairs",
     "cross_rank_trace_pairs",
+    "cross_rank_inter_volume_trace_pairs",
     "bdry_trace_pair",
     "bv_trace_pair",
 
@@ -207,90 +213,6 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec,
 # }}}
 
 
-# {{{ common derivative "helpers"
-
-def _div_helper(dcoll, diff_func, *args):
-    if len(args) == 1:
-        vecs, = args
-        dd = DD_VOLUME_ALL
-    elif len(args) == 2:
-        dd, vecs = args
-    else:
-        raise TypeError("invalid number of arguments")
-
-    if not isinstance(vecs, np.ndarray):
-        # vecs is not an object array -> treat as array container
-        return map_array_container(
-            partial(_div_helper, dcoll, diff_func, dd), vecs)
-
-    assert vecs.dtype == object
-
-    if vecs.size:
-        sample_vec = vecs[(0,)*vecs.ndim]
-
-        if isinstance(sample_vec, np.ndarray):
-            assert sample_vec.dtype == object
-            # vecs is an object array containing further object arrays
-            # -> treat as array container
-            return map_array_container(
-                partial(_div_helper, dcoll, diff_func, dd), vecs)
-
-    if vecs.shape[-1] != dcoll.ambient_dim:
-        raise ValueError("last/innermost dimension of *vecs* argument doesn't match "
-                "ambient dimension")
-
-    div_result_shape = vecs.shape[:-1]
-
-    if len(div_result_shape) == 0:
-        return sum(diff_func(dd, i, vec_i) for i, vec_i in enumerate(vecs))
-    else:
-        result = np.zeros(div_result_shape, dtype=object)
-        for idx in np.ndindex(div_result_shape):
-            result[idx] = sum(
-                    diff_func(dd, i, vec_i) for i, vec_i in enumerate(vecs[idx]))
-        return result
-
-
-def _grad_helper(dcoll, scalar_grad, *args, nested):
-    if len(args) == 1:
-        vec, = args
-        dd_in = dof_desc.DD_VOLUME_ALL
-    elif len(args) == 2:
-        dd_in, vec = args
-    else:
-        raise TypeError("invalid number of arguments")
-
-    if isinstance(vec, np.ndarray):
-        # Occasionally, data structures coming from *mirgecom* will
-        # contain empty object arrays as placeholders for fields.
-        # For example, species mass fractions is an empty object array when
-        # running in a single-species configuration.
-        # This hack here ensures that these empty arrays, at the very least,
-        # have their shape updated when applying the gradient operator
-        if vec.size == 0:
-            return vec.reshape(vec.shape + (dcoll.ambient_dim,))
-
-        # For containers with ndarray data (such as momentum/velocity),
-        # the gradient is matrix-valued, so we compute the gradient for
-        # each component. If requested (via 'not nested'), return a matrix of
-        # derivatives by stacking the results.
-        grad = obj_array_vectorize(
-            lambda el: _grad_helper(
-                dcoll, scalar_grad, dd_in, el, nested=nested), vec)
-        if nested:
-            return grad
-        else:
-            return np.stack(grad, axis=0)
-
-    if not isinstance(vec, DOFArray):
-        return map_array_container(
-            partial(_grad_helper, dcoll, scalar_grad, dd_in, nested=nested), vec)
-
-    return scalar_grad(dcoll, dd_in, vec)
-
-# }}}
-
-
 # {{{ Derivative operators
 
 def _reference_derivative_matrices(actx: ArrayContext,
@@ -347,7 +269,7 @@ def local_grad(
         :class:`~meshmode.dof_array.DOFArray`\ s or
         :class:`~arraycontext.container.ArrayContainer`\ of object arrays.
     """
-    dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE)
+    dd_in = DD_VOLUME_ALL
     from grudge.tools import rec_map_subarrays
     return rec_map_subarrays(
         partial(_strong_scalar_grad, dcoll, dd_in),
@@ -502,7 +424,7 @@ def weak_local_grad(
     """
     if len(args) == 1:
         vecs, = args
-        dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE)
+        dd_in = DD_VOLUME_ALL
     elif len(args) == 2:
         dd_in, vecs = args
     else:
@@ -606,7 +528,7 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT:
     """
     if len(args) == 1:
         vecs, = args
-        dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE)
+        dd_in = DD_VOLUME_ALL
     elif len(args) == 2:
         dd_in, vecs = args
     else:

grudge/trace_pair.py

@@ -7,8 +7,13 @@
 
 .. autoclass:: TracePair
 
-Boundary traces
----------------
+.. currentmodule:: grudge.op
+
+.. autoclass:: project_tracepair
+.. autoclass:: tracepair_with_discr_tag
+
+Boundary trace functions
+------------------------
 
 .. autofunction:: bdry_trace_pair
 .. autofunction:: bv_trace_pair