Purge .join() from mirgecom

mirgecom/euler.py

@@ -57,12 +57,8 @@
     inviscid_flux,
     inviscid_facial_flux
 )
-from grudge.eager import (
-    interior_trace_pair,
-    cross_rank_trace_pairs
-)
-from grudge.trace_pair import TracePair
-from mirgecom.fluid import make_conserved
+
+from grudge.trace_pair import interior_trace_pairs
 from mirgecom.operators import div_operator
 
 
@@ -100,28 +96,21 @@ def euler_operator(discr, eos, boundaries, cv, time=0.0):
         Agglomerated object array of DOF arrays representing the RHS of the Euler
         flow equations.
     """
+    # Compute volume contributions
     inviscid_flux_vol = inviscid_flux(discr, eos, cv)
+    interior_cv = interior_trace_pairs(discr, cv)
+    # Compute interface contributions
     inviscid_flux_bnd = (
-        inviscid_facial_flux(discr, eos=eos, cv_tpair=interior_trace_pair(discr, cv))
-        + sum(inviscid_facial_flux(
-            discr, eos=eos, cv_tpair=TracePair(
-                part_tpair.dd, interior=make_conserved(discr.dim, q=part_tpair.int),
-                exterior=make_conserved(discr.dim, q=part_tpair.ext)))
-              for part_tpair in cross_rank_trace_pairs(discr, cv.join()))
-        + sum(boundaries[btag].inviscid_divergence_flux(discr, btag=btag, cv=cv,
-                                                        eos=eos, time=time)
+        # Interior faces
+        + sum(inviscid_facial_flux(discr, eos=eos, cv_tpair=part_tpair)
+              for part_tpair in interior_cv)
+        # Domain boundaries
+        + sum(boundaries[btag].inviscid_boundary_flux(discr, btag=btag, cv=cv,
+                                                      eos=eos, time=time)
               for btag in boundaries)
     )
-    q = -div_operator(discr, inviscid_flux_vol.join(), inviscid_flux_bnd.join())
-    return make_conserved(discr.dim, q=q)
-
 
-def inviscid_operator(discr, eos, boundaries, q, t=0.0):
-    """Interface :function:`euler_operator` with backwards-compatible API."""
-    from warnings import warn
-    warn("Do not call inviscid_operator; it is now called euler_operator. This"
-         "function will disappear August 1, 2021", DeprecationWarning, stacklevel=2)
-    return euler_operator(discr, eos, boundaries, make_conserved(discr.dim, q=q), t)
+    return -div_operator(discr, inviscid_flux_vol, inviscid_flux_bnd)
 
 
 # By default, run unitless

mirgecom/filter.py

@@ -48,9 +48,13 @@
 import numpy as np
 import grudge.dof_desc as dof_desc
 
+from arraycontext import map_array_container
+
+from functools import partial
+
 from meshmode.dof_array import DOFArray
+
 from pytools import keyed_memoize_in
-from pytools.obj_array import obj_array_vectorized_n_args
 
 
 def exponential_mode_response_function(mode, alpha, cutoff, nfilt, filter_order):
@@ -163,7 +167,6 @@ def apply_spectral_filter(actx, modal_field, discr, cutoff,
     )
 
 
-@obj_array_vectorized_n_args
 def filter_modally(dcoll, dd, cutoff, mode_resp_func, field):
     """Stand-alone procedural interface to spectral filtering.
 
@@ -189,21 +192,24 @@ def filter_modally(dcoll, dd, cutoff, mode_resp_func, field):
         Mode below which *field* will not be filtered
     mode_resp_func:
         Modal response function returns a filter coefficient for input mode id
-    field: :class:`numpy.ndarray`
-        DOFArray or object array of DOFArrays
+    field: :class:`mirgecom.fluid.ConservedVars`
+        An array container containing the relevant field(s) to filter.
 
     Returns
     -------
-    result: numpy.ndarray
-        Filtered version of *field*.
+    result: :class:`mirgecom.fluid.ConservedVars`
+        An array container containing the filtered field(s).
     """
+    if not isinstance(field, DOFArray):
+        return map_array_container(
+            partial(filter_modally, dcoll, dd, cutoff, mode_resp_func), field
+        )
+
+    actx = field.array_context
     dd = dof_desc.as_dofdesc(dd)
     dd_modal = dof_desc.DD_VOLUME_MODAL
     discr = dcoll.discr_from_dd(dd)
 
-    assert isinstance(field, DOFArray)
-    actx = field.array_context
-
     modal_map = dcoll.connection_from_dds(dd, dd_modal)
     nodal_map = dcoll.connection_from_dds(dd_modal, dd)
     field = modal_map(field)

mirgecom/simutil.py

@@ -12,6 +12,8 @@
 --------------------
 
 .. autofunction:: compare_fluid_solutions
+.. autofunction:: componentwise_norms
+.. autofunction:: max_component_norm
 .. autofunction:: check_naninf_local
 .. autofunction:: check_range_local
 
@@ -48,6 +50,13 @@
 import numpy as np
 import grudge.op as op
 
+from arraycontext import map_array_container, flatten
+
+from functools import partial
+
+from meshmode.dof_array import DOFArray
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -290,7 +299,37 @@ def compare_fluid_solutions(discr, red_state, blue_state):
     """
     actx = red_state.array_context
     resid = red_state - blue_state
-    return [actx.to_numpy(discr.norm(v, np.inf)) for v in resid.join()]
+    resid_errs = actx.to_numpy(
+        flatten(componentwise_norms(discr, resid, order=np.inf), actx))
+
+    return resid_errs.tolist()
+
+
+def componentwise_norms(discr, fields, order=np.inf):
+    """Return the *order*-norm for each component of *fields*.
+
+    .. note::
+        This is a collective routine and must be called by all MPI ranks.
+    """
+    if not isinstance(fields, DOFArray):
+        return map_array_container(
+            partial(componentwise_norms, discr, order=order), fields)
+    if len(fields) > 0:
+        return discr.norm(fields, order)
+    else:
+        # FIXME: This work-around for #575 can go away after #569
+        return 0
+
+
+def max_component_norm(discr, fields, order=np.inf):
+    """Return the max *order*-norm over the components of *fields*.
+
+    .. note::
+        This is a collective routine and must be called by all MPI ranks.
+    """
+    actx = fields.array_context
+    return max(actx.to_numpy(flatten(
+        componentwise_norms(discr, fields, order), actx)))
 
 
 def generate_and_distribute_mesh(comm, generate_mesh):

test/test_euler.py

@@ -53,6 +53,7 @@
     pytest_generate_tests_for_pyopencl_array_context
     as pytest_generate_tests)
 
+from mirgecom.simutil import max_component_norm
 
 from grudge.shortcuts import make_visualizer
 from mirgecom.inviscid import get_inviscid_timestep
@@ -234,7 +235,8 @@ def test_vortex_rhs(actx_factory, order):
             discr, eos=IdealSingleGas(), boundaries=boundaries,
             cv=vortex_soln, time=0.0)
 
-        err_max = actx.to_numpy(discr.norm(inviscid_rhs.join(), np.inf))
+        err_max = max_component_norm(discr, inviscid_rhs, np.inf)
+
         eoc_rec.add_data_point(1.0 / nel_1d, err_max)
 
     logger.info(
@@ -293,8 +295,7 @@ def test_lump_rhs(actx_factory, dim, order):
         )
         expected_rhs = lump.exact_rhs(discr, cv=lump_soln, time=0)
 
-        err_max = actx.to_numpy(
-            discr.norm((inviscid_rhs-expected_rhs).join(), np.inf))
+        err_max = max_component_norm(discr, inviscid_rhs-expected_rhs, np.inf)
         if err_max > maxxerr:
             maxxerr = err_max
 
@@ -371,7 +372,7 @@ def test_multilump_rhs(actx_factory, dim, order, v0):
         print(f"expected_rhs = {expected_rhs}")
 
         err_max = actx.to_numpy(
-            discr.norm((inviscid_rhs-expected_rhs).join(), np.inf))
+            discr.norm((inviscid_rhs-expected_rhs), np.inf))
         if err_max > maxxerr:
             maxxerr = err_max
 
@@ -437,7 +438,7 @@ def _euler_flow_stepper(actx, parameters):
     def write_soln(state, write_status=True):
         dv = eos.dependent_vars(cv=state)
         expected_result = initializer(nodes, t=t)
-        result_resid = (state - expected_result).join()
+        result_resid = state - expected_result
         maxerr = [np.max(np.abs(result_resid[i].get())) for i in range(dim + 2)]
         mindv = [np.min(dvfld.get()) for dvfld in dv]
         maxdv = [np.max(dvfld.get()) for dvfld in dv]
@@ -492,9 +493,7 @@ def rhs(t, q):
                 write_soln(state=cv)
 
         cv = rk4_step(cv, t, dt, rhs)
-        cv = make_conserved(
-            dim, q=filter_modally(discr, "vol", cutoff, frfunc, cv.join())
-        )
+        cv = filter_modally(discr, "vol", cutoff, frfunc, cv)
 
         t += dt
         istep += 1
@@ -506,7 +505,7 @@ def rhs(t, q):
         maxerr = max(write_soln(cv, False))
     else:
         expected_result = initializer(nodes, time=t)
-        maxerr = actx.to_numpy(discr.norm((cv - expected_result).join(), np.inf))
+        maxerr = max_component_norm(discr, cv-expected_result, np.inf)
 
     logger.info(f"Max Error: {maxerr}")
     if maxerr > exittol:

test/test_filter.py

@@ -186,13 +186,14 @@ def test_filter_function(actx_factory, dim, order, do_viz=False):
     # the filter unharmed.
     from mirgecom.initializers import Uniform
     initr = Uniform(dim=dim)
-    uniform_soln = initr(t=0, x_vec=nodes).join()
+    uniform_soln = initr(t=0, x_vec=nodes)
 
     from mirgecom.filter import filter_modally
     filtered_soln = filter_modally(discr, "vol", cutoff,
                                    frfunc, uniform_soln)
     soln_resid = uniform_soln - filtered_soln
-    max_errors = [actx.to_numpy(discr.norm(v, np.inf)) for v in soln_resid]
+    from mirgecom.simutil import componentwise_norms
+    max_errors = componentwise_norms(discr, soln_resid, np.inf)
 
     tol = 1e-14
 

test/test_fluid.py

@@ -77,8 +77,7 @@ def test_velocity_gradient_sanity(actx_factory, dim, mass_exp, vel_fac):
 
     cv = make_conserved(dim, mass=mass, energy=energy, momentum=mom)
     from grudge.op import local_grad
-    grad_cv = make_conserved(dim,
-                             q=local_grad(discr, cv.join()))
+    grad_cv = local_grad(discr, cv)
 
     grad_v = velocity_gradient(discr, cv, grad_cv)
 
@@ -123,8 +122,7 @@ def test_velocity_gradient_eoc(actx_factory, dim):
 
         cv = make_conserved(dim, mass=mass, energy=energy, momentum=mom)
         from grudge.op import local_grad
-        grad_cv = make_conserved(dim,
-                                 q=local_grad(discr, cv.join()))
+        grad_cv = local_grad(discr, cv)
         grad_v = velocity_gradient(discr, cv, grad_cv)
 
         def exact_grad_row(xdata, gdim, dim):
@@ -178,8 +176,7 @@ def test_velocity_gradient_structure(actx_factory):
 
     cv = make_conserved(dim, mass=mass, energy=energy, momentum=mom)
     from grudge.op import local_grad
-    grad_cv = make_conserved(dim,
-                             q=local_grad(discr, cv.join()))
+    grad_cv = local_grad(discr, cv)
     grad_v = velocity_gradient(discr, cv, grad_cv)
 
     tol = 1e-11
@@ -233,8 +230,8 @@ def test_species_mass_gradient(actx_factory, dim):
     cv = make_conserved(dim, mass=mass, energy=energy, momentum=mom,
                         species_mass=species_mass)
     from grudge.op import local_grad
-    grad_cv = make_conserved(dim,
-                             q=local_grad(discr, cv.join()))
+    grad_cv = local_grad(discr, cv)
+
     from mirgecom.fluid import species_mass_fraction_gradient
     grad_y = species_mass_fraction_gradient(discr, cv, grad_cv)
 

test/test_lazy.py

@@ -22,7 +22,7 @@
 
 import numpy as np
 from functools import partial
-from pytools.obj_array import make_obj_array, obj_array_vectorize
+from pytools.obj_array import make_obj_array, obj_array_vectorize  # noqa
 import pyopencl as cl
 import pyopencl.tools as cl_tools
 import pyopencl.array as cla  # noqa
@@ -75,19 +75,16 @@ def get_discr(order):
 
 # Mimics math.isclose for state arrays
 def _isclose(discr, x, y, rel_tol=1e-9, abs_tol=0, return_operands=False):
-    def componentwise_norm(a):
-        from mirgecom.fluid import ConservedVars
-        if isinstance(a, ConservedVars):
-            return componentwise_norm(a.join())
-        from arraycontext import get_container_context_recursively
-        actx = get_container_context_recursively(a)
-        return obj_array_vectorize(lambda b: actx.to_numpy(discr.norm(b, np.inf)), a)
-
-    lhs = componentwise_norm(x - y)
+
+    from mirgecom.simutil import componentwise_norms
+    from arraycontext import flatten
+    actx = x.array_context
+    lhs = actx.to_numpy(flatten(componentwise_norms(discr, x - y, np.inf), actx))
+
     rhs = np.maximum(
         rel_tol * np.maximum(
-            componentwise_norm(x),
-            componentwise_norm(y)),
+            actx.to_numpy(flatten(componentwise_norms(discr, x, np.inf), actx)),
+            actx.to_numpy(flatten(componentwise_norms(discr, y, np.inf), actx))),
         abs_tol)
 
     is_close = np.all(lhs <= rhs)

test/test_operators.py

@@ -217,13 +217,12 @@ def sym_eval(expr, x_vec):
         test_data = test_func(nodes)
         exact_grad = grad_test_func(nodes)
 
-        def inf_norm(x):
-            return actx.to_numpy(discr.norm(x, np.inf))
+        from mirgecom.simutil import componentwise_norms
+        from arraycontext import flatten
+
+        err_scale = max(flatten(componentwise_norms(discr, exact_grad, np.inf),
+                                actx))
 
-        if isinstance(test_data, ConservedVars):
-            err_scale = inf_norm(exact_grad.join())
-        else:
-            err_scale = inf_norm(exact_grad)
         if err_scale <= 1e-16:
             err_scale = 1
 
@@ -236,18 +235,12 @@ def inf_norm(x):
                                          test_data_int_tpair, boundaries)
 
         from mirgecom.operators import grad_operator
-        if isinstance(test_data, ConservedVars):
-            test_grad = make_conserved(
-                dim=dim, q=grad_operator(discr, test_data.join(),
-                                         test_data_flux_bnd.join())
-                )
-            grad_err = inf_norm((test_grad - exact_grad).join())/err_scale
-        else:
-            test_grad = grad_operator(discr, test_data, test_data_flux_bnd)
-            grad_err = inf_norm(test_grad - exact_grad)/err_scale
-
-        print(f"{test_grad=}")
-        eoc.add_data_point(actx.to_numpy(h_max), grad_err)
+        test_grad = grad_operator(discr, test_data, test_data_flux_bnd)
+        grad_err = \
+            max(flatten(componentwise_norms(discr, test_grad - exact_grad, np.inf),
+                        actx)) / err_scale
+
+        eoc.add_data_point(actx.to_numpy(h_max), actx.to_numpy(grad_err))
 
     assert (
         eoc.order_estimate() >= order - 0.5

test/test_restart.py

@@ -76,4 +76,5 @@ def test_restart_cv(actx_factory, nspecies):
     restart_data = read_restart_data(actx, rst_filename)
 
     resid = test_state - restart_data["state"]
-    assert actx.to_numpy(discr.norm(resid.join(), np.inf)) == 0
+    from mirgecom.simutil import max_component_norm
+    assert max_component_norm(discr, resid, np.inf) == 0