make compatible with non-coupling-enabled versions of grudge v3 (#1048)

Co-authored-by: Matthew Smith <[email protected]>

doc/support/arrays.rst

@@ -43,24 +43,25 @@ To demonstrate the effect of this, first we need some setup:
 .. doctest::
 
    >>> import pyopencl as cl
-   >>> from arraycontext import PyOpenCLArrayContext, thaw
+   >>> from arraycontext import thaw
+   >>> from meshmode.array_context import PyOpenCLArrayContext
    >>> ctx = cl.create_some_context()
    >>> queue = cl.CommandQueue(ctx)
    >>> actx = PyOpenCLArrayContext(queue)
    >>> from meshmode.mesh.generation import generate_regular_rect_mesh
    >>> mesh = generate_regular_rect_mesh(a=(0, 0), b=(1, 1), nelements_per_axis=(10, 10))
-   >>> from grudge import DiscretizationCollection
-   >>> dcoll = DiscretizationCollection(actx, mesh, order=5)
+   >>> from grudge import make_discretization_collection
+   >>> dcoll = make_discretization_collection(actx, mesh, order=5)
 
 Most quantities that are maintained by the discretization will be frozen. For example,
 if one wanted to grab the nodes of the mesh or normals of a named surface in the mesh:
 
 .. doctest::
 
-   >>> from grudge.dof_desc import DOFDesc
-   >>> nodes = thaw(dcoll.nodes(), actx)
-   >>> dd = DOFDesc("all_faces")
-   >>> nhat = thaw(dcoll.normal(dd), actx)
+   >>> from grudge.dof_desc import as_dofdesc
+   >>> nodes = actx.thaw(dcoll.nodes())
+   >>> dd = as_dofdesc("all_faces")
+   >>> nhat = actx.thaw(dcoll.normal(dd))
 
 What can go wrong?  Attempts to operate on frozen data will yield errors similar to
 the following:
@@ -76,7 +77,7 @@ Fortunately, recovering from this is straightforward:
 
 .. doctest::
 
-   >>> nodes = thaw(dcoll.nodes(), actx)
+   >>> nodes = actx.thaw(dcoll.nodes())
    >>> result = nodes * 5
 
 Array Containers

examples/advdiff-tpe.py

@@ -198,7 +198,15 @@ def vol_max(x):
         return actx.to_numpy(nodal_max(dcoll, "vol", x))[()]
 
     from grudge.dt_utils import characteristic_lengthscales
-    dx = characteristic_lengthscales(actx, dcoll)
+
+    try:
+        dx = characteristic_lengthscales(actx, dcoll)
+    except NotImplementedError:
+        from warnings import warn
+        warn("This example requires https://github.com/inducer/grudge/pull/338 . "
+             "Exiting.")
+        return
+
     dx_min, dx_max = vol_min(dx), vol_max(dx)
 
     print(f"DX: ({dx_min}, {dx_max})")

examples/poiseuille-tpe.py

@@ -426,8 +426,14 @@ def my_rhs(t, state):
     from mirgecom.simutil import force_evaluation
     current_state = force_evaluation(actx, current_state)
 
-    current_dt = get_sim_timestep(dcoll, current_state, current_t, current_dt,
+    try:
+        current_dt = get_sim_timestep(dcoll, current_state, current_t, current_dt,
                                   current_cfl, t_final, constant_cfl)
+    except NotImplementedError:
+        from warnings import warn
+        warn("This example requires https://github.com/inducer/grudge/pull/338 . "
+             "Exiting.")
+        return
 
     current_step, current_t, current_cv = \
         advance_state(rhs=my_rhs, timestepper=timestepper,

examples/run_examples.sh

@@ -16,14 +16,19 @@ function endgroup {
 
 # }}}
 
+python -c "from grudge.array_context import MPINumpyArrayConext" && numpy_actx_available=numpy || numpy_actx_available=
+
+echo "Numpy array context available: $numpy_actx_available"
+
+
 run_path=$(pwd)
 examples_dir=${1:-"./"}
 shift
 
 if [[ ! -d "${examples_dir}" ]]; then
     echo "Usage: run_examples.sh <examples directory> [list of examples]"
     printf "\nThis script runs examples on 2 MPI ranks (where appropriate) and\n"
-    printf "compares the results it gets from running with eager, lazy, and numpy\n"
+    printf "compares the results it gets from running with eager, lazy, and numpy (if available)\n"
     printf "array contexts. Users may optionally provide a list of which examples\n"
     printf "to run.\n\nArguments:\n"
     printf "\n<examples directory>: defaults to the current working directory.\n"
@@ -93,14 +98,6 @@ do
     TOL_LAZY=1e-9
     TOL_NUMPY=1e-9
 
-    # FIXME: Have to ignore "rhs", and "gradients" to make
-    # this example pass.
-    # if [[ $example == "thermally-coupled.py" ]]; then
-    #    echo "Setting tolerance=1 for $example"
-    #    TOL_LAZY=1
-    #    TOL_NUMPY=1
-    # fi
-
     date
     printf "***\n***\n"
 
@@ -110,7 +107,7 @@ do
     test_results=()
 
     # Run example with Eager, Lazy, and Numpy arraycontexts
-    for actx in eager lazy numpy; do
+    for actx in eager lazy $numpy_actx_available; do
         test_name="${example_name}_$actx"
         startgroup "**** Running $test_name"
         set -x

examples/thermally-coupled-tpe.py

@@ -90,6 +90,13 @@ def main(actx_class, use_esdg=False, use_overintegration=False,
     if casename is None:
         casename = "mirgecom"
 
+    try:
+        from grudge.discretization import PartID  # noqa: F401
+    except ImportError:
+        from warnings import warn
+        warn("This example requires a coupling-enabled branch of grudge; exiting.")
+        return
+
     from mpi4py import MPI
     comm = MPI.COMM_WORLD
     rank = comm.Get_rank()

examples/thermally-coupled.py

@@ -90,6 +90,13 @@ def main(actx_class, use_esdg=False, use_overintegration=False,
     if casename is None:
         casename = "mirgecom"
 
+    try:
+        from grudge.discretization import PartID  # noqa: F401
+    except ImportError:
+        from warnings import warn
+        warn("This example requires a coupling-enabled branch of grudge; exiting.")
+        return
+
     from mpi4py import MPI
     comm = MPI.COMM_WORLD
     rank = comm.Get_rank()

mirgecom/array_context.py

@@ -59,10 +59,12 @@ def get_reasonable_array_context_class(*, lazy: bool, distributed: bool,
         warn("The NumpyArrayContext is still under development")
 
         if distributed:
-            from grudge.array_context import MPINumpyArrayContext
+            from grudge.array_context import MPINumpyArrayContext  \
+                # pylint: disable=no-name-in-module
             return MPINumpyArrayContext
         else:
-            from grudge.array_context import NumpyArrayContext
+            from grudge.array_context import NumpyArrayContext  \
+                # pylint: disable=no-name-in-module
             return NumpyArrayContext
 
     if profiling:
@@ -108,6 +110,13 @@ def actx_class_is_numpy(actx_class: Type[ArrayContext]) -> bool:
         return False
 
 
+def actx_class_has_fallback_args(actx_class: Type[ArrayContext]) -> bool:
+    """Return True if *actx_class* has fallback arguments."""
+    import inspect
+    spec = inspect.getfullargspec(actx_class.__init__)
+    return "use_axis_tag_inference_fallback" in spec.args
+
+
 def _check_cache_dirs_node() -> None:
     """Check whether multiple ranks share cache directories on the same node."""
     from mpi4py import MPI
@@ -271,15 +280,17 @@ def initialize_actx(
         use_einsum_inference_fallback: bool = False) -> ArrayContext:
     """Initialize a new :class:`~arraycontext.ArrayContext` based on *actx_class*."""
     from grudge.array_context import (MPIPyOpenCLArrayContext,
-                                      MPIPytatoArrayContext,
-                                      MPINumpyArrayContext)
+                                      MPIPytatoArrayContext
+                                      )
 
     actx_kwargs: Dict[str, Any] = {}
 
     if comm:
         actx_kwargs["mpi_communicator"] = comm
 
     if actx_class_is_numpy(actx_class):
+        from grudge.array_context import MPINumpyArrayContext  \
+            # pylint: disable=no-name-in-module
         if comm:
             assert issubclass(actx_class, MPINumpyArrayContext)
         else:
@@ -299,10 +310,12 @@ def initialize_actx(
 
         if actx_class_is_lazy(actx_class):
             assert issubclass(actx_class, PytatoPyOpenCLArrayContext)
-            actx_kwargs["use_axis_tag_inference_fallback"] = \
-                use_axis_tag_inference_fallback
-            actx_kwargs["use_einsum_inference_fallback"] = \
-                use_einsum_inference_fallback
+
+            if actx_class_has_fallback_args(actx_class):
+                actx_kwargs["use_axis_tag_inference_fallback"] = \
+                    use_axis_tag_inference_fallback
+                actx_kwargs["use_einsum_inference_fallback"] = \
+                    use_einsum_inference_fallback
             if comm:
                 assert issubclass(actx_class, MPIPytatoArrayContext)
                 actx_kwargs["mpi_base_tag"] = 12000

mirgecom/euler.py

@@ -93,8 +93,6 @@
     interior_trace_pairs,
     tracepair_with_discr_tag
 )
-from grudge.projection import volume_quadrature_project
-from grudge.flux_differencing import volume_flux_differencing
 
 import grudge.op as op
 
@@ -174,6 +172,9 @@ def entropy_stable_euler_operator(
     gamma_quad = gas_model.eos.gamma(state_quad.cv, state_quad.temperature)
 
     # Compute the projected (nodal) entropy variables
+    from grudge.projection import volume_quadrature_project  \
+        # pylint: disable=no-name-in-module
+
     entropy_vars = volume_quadrature_project(
         dcoll, dd_vol_quad,
         # Map to entropy variables
@@ -207,6 +208,9 @@ def _reshape(shape, ary):
         _reshape((-1, 1), modified_conserved_fluid_state))
 
     # Compute volume derivatives using flux differencing
+    from grudge.flux_differencing import volume_flux_differencing  \
+        # pylint: disable=no-name-in-module,import-error
+
     inviscid_vol_term = \
         -volume_flux_differencing(dcoll, dd_vol_quad, dd_allfaces_quad,
                                   flux_matrices)

mirgecom/gas_model.py

@@ -933,7 +933,8 @@ def replace_fluid_state(
 def make_entropy_projected_fluid_state(
         discr, dd_vol, dd_faces, state, entropy_vars, gamma, gas_model):
     """Projects the entropy vars to target manifold, computes the CV from that."""
-    from grudge.interpolation import volume_and_surface_quadrature_interpolation
+    from grudge.interpolation import volume_and_surface_quadrature_interpolation  \
+        # pylint: disable=no-name-in-module
 
     # Interpolate to the volume and surface (concatenated) quadrature
     # discretizations: v = [v_vol, v_surf]

mirgecom/multiphysics/thermally_coupled_fluid_wall.py

@@ -74,7 +74,6 @@
 from meshmode.dof_array import DOFArray
 from grudge.trace_pair import (
     TracePair,
-    inter_volume_trace_pairs
 )
 from grudge.dof_desc import (
     DISCR_TAG_BASE,
@@ -787,6 +786,10 @@ def _get_interface_trace_pairs_no_grad(
         (fluid_dd, wall_dd): (
             _make_thermal_data(fluid_kappa, fluid_temperature),
             _make_thermal_data(wall_kappa, wall_temperature))}
+
+    from grudge.trace_pair import inter_volume_trace_pairs  \
+        # pylint: disable=no-name-in-module
+
     return inter_volume_trace_pairs(
         dcoll, pairwise_thermal_data,
         comm_tag=(_ThermalDataNoGradInterVolTag, comm_tag))
@@ -811,6 +814,9 @@ def _get_interface_trace_pairs(
                 wall_temperature,
                 wall_grad_temperature))}
 
+    from grudge.trace_pair import inter_volume_trace_pairs  \
+        # pylint: disable=no-name-in-module
+
     return inter_volume_trace_pairs(
         dcoll, pairwise_thermal_data,
         comm_tag=(_ThermalDataInterVolTag, comm_tag))

mirgecom/simutil.py

@@ -84,7 +84,7 @@
     DiscretizationDOFAxisTag
 )
 from arraycontext import flatten, map_array_container
-from grudge.discretization import DiscretizationCollection, PartID
+from grudge.discretization import DiscretizationCollection
 from grudge.dof_desc import DD_VOLUME_ALL
 from meshmode.dof_array import DOFArray
 
@@ -981,6 +981,9 @@ def get_rank_to_mesh_data():
                     for rank in range(num_ranks)]
 
             else:
+                from grudge.discretization import \
+                    PartID  # pylint: disable=no-name-in-module
+
                 tag_to_volume = {
                     tag: vol
                     for vol, tags in volume_to_tags.items()

test/test_av.py

@@ -29,11 +29,11 @@
 import numpy as np
 import pyopencl as cl
 import pytest
-from meshmode.array_context import (  # noqa
-    PyOpenCLArrayContext,
-    pytest_generate_tests_for_pyopencl_array_context
-    as pytest_generate_tests
-)
+from meshmode.array_context import PyOpenCLArrayContext
+
+from meshmode.array_context import PytestPyOpenCLArrayContextFactory
+from arraycontext import pytest_generate_tests_for_array_contexts
+
 from meshmode.mesh import BTAG_ALL
 from meshmode.discretization.connection import FACE_RESTR_ALL
 import grudge.op as op
@@ -51,14 +51,14 @@
 from mirgecom.eos import IdealSingleGas
 
 from mirgecom.discretization import create_discretization_collection
-from pyopencl.tools import (  # noqa
-    pytest_generate_tests_for_pyopencl as pytest_generate_tests,
-)
 from mirgecom.simutil import get_box_mesh
 from pytools.obj_array import make_obj_array
 
 logger = logging.getLogger(__name__)
 
+pytest_generate_tests = pytest_generate_tests_for_array_contexts(
+    [PytestPyOpenCLArrayContextFactory])
+
 
 # NOTE:  Testing of this av_laplacian_operator is currently
 # pretty limited.  This fact is somewhat indicative of the

test/test_bc.py

@@ -52,9 +52,10 @@
 import grudge.op as op
 from mirgecom.simutil import get_box_mesh
 
-from meshmode.array_context import (  # noqa
-    pytest_generate_tests_for_pyopencl_array_context
-    as pytest_generate_tests)
+from meshmode.array_context import PytestPyOpenCLArrayContextFactory
+from arraycontext import pytest_generate_tests_for_array_contexts
+pytest_generate_tests = pytest_generate_tests_for_array_contexts(
+    [PytestPyOpenCLArrayContextFactory])
 
 logger = logging.getLogger(__name__)
 

test/test_chemistry.py

@@ -34,9 +34,8 @@
 
 from meshmode.array_context import PyOpenCLArrayContext
 from meshmode.mesh.generation import generate_regular_rect_mesh
-from meshmode.array_context import (  # noqa
-    pytest_generate_tests_for_pyopencl_array_context
-    as pytest_generate_tests)
+from meshmode.array_context import PytestPyOpenCLArrayContextFactory
+from arraycontext import pytest_generate_tests_for_array_contexts
 
 from mirgecom.fluid import make_conserved
 from mirgecom.eos import PyrometheusMixture
@@ -47,6 +46,9 @@
     get_pyrometheus_wrapper_class
 )
 
+pytest_generate_tests = pytest_generate_tests_for_array_contexts(
+    [PytestPyOpenCLArrayContextFactory])
+
 
 @pytest.mark.parametrize(("mechname", "fuel", "rate_tol"),
                          [("uiuc_7sp", "C2H4", 1e-11),

test/test_diffusion.py

@@ -33,9 +33,10 @@
 from grudge.dof_desc import BoundaryDomainTag, DISCR_TAG_BASE, DISCR_TAG_QUAD
 from grudge.shortcuts import make_visualizer
 from meshmode.dof_array import DOFArray
-from meshmode.array_context import (  # noqa
-    pytest_generate_tests_for_pyopencl_array_context
-    as pytest_generate_tests)
+
+from meshmode.array_context import PytestPyOpenCLArrayContextFactory
+from arraycontext import pytest_generate_tests_for_array_contexts
+
 from mirgecom.symbolic import (
     diff as sym_diff,
     grad as sym_grad,
@@ -55,6 +56,9 @@
 
 logger = logging.getLogger(__name__)
 
+pytest_generate_tests = pytest_generate_tests_for_array_contexts(
+    [PytestPyOpenCLArrayContextFactory])
+
 
 def test_diffusion_boundary_conditions(actx_factory):
     """Checks the boundary conditions for diffusion operator."""