Skip to content

Commit

Permalink
Merge branch 'optional-coupling-v3' into build-production
Browse files Browse the repository at this point in the history
  • Loading branch information
MTCam committed Jul 26, 2024
2 parents 4f18db1 + 0238cbe commit 864ff04
Show file tree
Hide file tree
Showing 35 changed files with 252 additions and 128 deletions.
17 changes: 9 additions & 8 deletions doc/support/arrays.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down
10 changes: 9 additions & 1 deletion examples/advdiff-tpe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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})")
Expand Down
8 changes: 7 additions & 1 deletion examples/poiseuille-tpe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
17 changes: 7 additions & 10 deletions examples/run_examples.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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"

Expand All @@ -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
Expand Down
7 changes: 7 additions & 0 deletions examples/thermally-coupled-tpe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
7 changes: 7 additions & 0 deletions examples/thermally-coupled.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
29 changes: 21 additions & 8 deletions mirgecom/array_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
8 changes: 6 additions & 2 deletions mirgecom/euler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
3 changes: 2 additions & 1 deletion mirgecom/gas_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
8 changes: 7 additions & 1 deletion mirgecom/multiphysics/thermally_coupled_fluid_wall.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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))
Expand All @@ -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))
Expand Down
5 changes: 4 additions & 1 deletion mirgecom/simutil.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down
16 changes: 8 additions & 8 deletions test/test_av.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
7 changes: 4 additions & 3 deletions test/test_bc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down
8 changes: 5 additions & 3 deletions test/test_chemistry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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),
Expand Down
10 changes: 7 additions & 3 deletions test/test_diffusion.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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."""
Expand Down
Loading

0 comments on commit 864ff04

Please sign in to comment.