From 914936abb527d497660430ff4fae7ea0fb4bb9ee Mon Sep 17 00:00:00 2001 From: "Joshua C. Macdonald" Date: Wed, 14 Oct 2026 04:16:33 -0400 Subject: [PATCH 1/2] Honor floating reference precision in transport matrix builders --- .github/workflows/ci.yaml | 2 +- CHANGELOG.md | 5 ++ docs/guides/backends.md | 6 ++ src/op_engine/matrix_ops.py | 12 ++- tests/op_engine/test_operator_precision.py | 94 ++++++++++++++++++++++ 5 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 tests/op_engine/test_operator_precision.py diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index f2dbc9c..51a54fb 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -34,7 +34,7 @@ jobs: uv sync --dev uv pip install 'torch>=2.6' --index-url https://download.pytorch.org/whl/cpu - name: 'Run PyTorch contracts' - run: 'uv run --no-sync pytest tests/op_engine/test_array_namespace.py -k torch' + run: 'uv run --no-sync pytest tests/op_engine/test_array_namespace.py tests/op_engine/test_operator_precision.py -k torch' build-check: name: 'Build Check (${{ matrix.python-version }})' runs-on: 'ubuntu-latest' diff --git a/CHANGELOG.md b/CHANGELOG.md index 91e6d3e..9b09291 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ versioning while the project remains pre-1.0. ### Fixed +- Namespace-aware advection and diffusion builders honor a floating reference + dtype, avoiding float32 promotion on NumPy and JAX with x64 enabled. Integer + references retain floating coefficient construction; stencil and native + coefficient derivatives are checked across CPU namespaces (#212). + - `ModelCore.set_initial_state` resolves its configured NumPy-compatible dtype in the initial state's array namespace, so native PyTorch tensors work with `CoreSolver.run`. CPU Torch CI covers fixed and adaptive explicit runs, diff --git a/docs/guides/backends.md b/docs/guides/backends.md index 69edcaf..fd84126 100644 --- a/docs/guides/backends.md +++ b/docs/guides/backends.md @@ -20,6 +20,12 @@ using NumPy-compatible values such as `ModelCoreOptions(dtype=np.float32)`; integration and complete method qualification remain provider/backend capabilities rather than consequences of namespace selection. +The namespace-aware `build_advection_matrix` and `build_diffusion_matrix` +builders use a floating `reference` array's dtype directly, including float32 +when JAX x64 is enabled. Integer references select a floating result so physical +coefficients are not truncated. Geometry remains static; coefficients can remain +dynamic for JAX JIT/AD and CPU Torch AD. + ## Portable fixed-step methods With JAX state, fixed-step Euler, Heun, RK4, Dormand--Prince, dense IMEX diff --git a/src/op_engine/matrix_ops.py b/src/op_engine/matrix_ops.py index 81142bf..9df1f0d 100644 --- a/src/op_engine/matrix_ops.py +++ b/src/op_engine/matrix_ops.py @@ -291,7 +291,11 @@ def build_advection_matrix( namespace_source = np.asarray(namespace_source) xp = _namespace_of(namespace_source) source_dtype = cast("Any", namespace_source).dtype - dtype = xp.result_type(xp.asarray(0.0).dtype, source_dtype) + dtype = ( + source_dtype + if xp.isdtype(source_dtype, "real floating") + else xp.result_type(xp.asarray(0.0).dtype, source_dtype) + ) velocity_array = cast("Array", xp.asarray(velocity, dtype=dtype)) if velocity_array.shape != (): msg = f"advection velocity must be scalar; got shape {velocity_array.shape}." @@ -474,7 +478,11 @@ def build_diffusion_matrix( # noqa: PLR0913 namespace_source = np.asarray(namespace_source) xp = _namespace_of(namespace_source) source_dtype = cast("Any", namespace_source).dtype - dtype = xp.result_type(xp.asarray(0.0).dtype, source_dtype) + dtype = ( + source_dtype + if xp.isdtype(source_dtype, "real floating") + else xp.result_type(xp.asarray(0.0).dtype, source_dtype) + ) coefficient_array = cast("Array", xp.asarray(coefficient, dtype=dtype)) if coefficient_array.shape != (): msg = ( diff --git a/tests/op_engine/test_operator_precision.py b/tests/op_engine/test_operator_precision.py new file mode 100644 index 0000000..6a98aca --- /dev/null +++ b/tests/op_engine/test_operator_precision.py @@ -0,0 +1,94 @@ +"""Reference-controlled precision for namespace-aware finite-volume builders.""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import Any + +import numpy as np +import pytest + +from op_engine import array_namespace, build_advection_matrix, build_diffusion_matrix + +BACKENDS = ("numpy", "jax", "torch") +STENCILS = { + "advection": np.asarray([[-1, 0, 0], [1, -1, 0], [0, 1, 0]]), + "diffusion": np.asarray([[-1, 1, 0], [1, -2, 1], [0, 1, -1]]), +} +BUILDERS = {"advection": build_advection_matrix, "diffusion": build_diffusion_matrix} + + +def _host(value: Any) -> np.ndarray: # noqa: ANN401 + """Return eager values for independent assertions. + + Returns: + Host NumPy values. + """ + return np.asarray(value.detach().numpy() if hasattr(value, "detach") else value) + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +@pytest.mark.parametrize("kind", ["advection", "diffusion"]) +def test_floating_reference_selects_operator_precision( + backend: str, dtype: str, kind: str +) -> None: + """Known reflecting stencils retain the reference's precision and namespace.""" + module = ( + np + if backend == "numpy" + else pytest.importorskip("jax.numpy" if backend == "jax" else "torch") + ) + precision = ( + pytest.importorskip("jax").enable_x64() if backend == "jax" else nullcontext() + ) + with precision: + reference = module.asarray(np.ones((2, 3, 3), dtype=dtype)) + result = BUILDERS[kind](3, 1.0, 0.4, bc="reflecting", reference=reference) + assert result.dtype == reference.dtype + assert result.shape == (3, 3) + assert array_namespace(result) is array_namespace(reference) + np.testing.assert_allclose(_host(result), 0.4 * STENCILS[kind], rtol=2e-7) + + +@pytest.mark.parametrize("backend", ["jax", "torch"]) +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +@pytest.mark.parametrize("kind", ["advection", "diffusion"]) +def test_reference_precision_preserves_dynamic_coefficient_derivatives( + backend: str, dtype: str, kind: str +) -> None: + """Coefficient derivatives match the known stencil's quadratic objective.""" + module = pytest.importorskip("jax.numpy" if backend == "jax" else "torch") + precision = ( + pytest.importorskip("jax").enable_x64() if backend == "jax" else nullcontext() + ) + with precision: + reference = module.asarray(np.ones(3, dtype=dtype)) + coefficient = module.asarray(np.asarray(0.4, dtype=dtype)) + + def objective(value: Any) -> Any: # noqa: ANN401 + operator = BUILDERS[kind]( + 3, 1.0, value, bc="reflecting", reference=reference + ) + return module.sum(operator * operator) + + if backend == "jax": + jax = pytest.importorskip("jax") + result, gradient = jax.jit(jax.value_and_grad(objective))(coefficient) + else: + result = objective(coefficient) + gradient = module.func.jacrev(objective)(coefficient) + assert result.dtype == reference.dtype + assert gradient.dtype == reference.dtype + expected = 0.8 * sum(float(value) ** 2 for value in STENCILS[kind].ravel()) + np.testing.assert_allclose(_host(gradient), expected, rtol=2e-7) + + +@pytest.mark.parametrize("kind", ["advection", "diffusion"]) +def test_integer_reference_still_constructs_floating_operator(kind: str) -> None: + """Integer geometry/state does not truncate coefficients or stencil values.""" + result = BUILDERS[kind]( + 3, 1.0, 0.4, bc="reflecting", reference=np.ones(3, dtype=np.int64) + ) + assert np.issubdtype(result.dtype, np.floating) + np.testing.assert_allclose(result, 0.4 * STENCILS[kind]) From d2830b47daa71ca8c7c12a85ee5dc522612a4352 Mon Sep 17 00:00:00 2001 From: "Joshua C. Macdonald" Date: Wed, 14 Oct 2026 04:18:44 -0400 Subject: [PATCH 2/2] Use the existing scoped x64 fallback in precision regressions --- tests/op_engine/test_operator_precision.py | 31 +++++++++++++++------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/tests/op_engine/test_operator_precision.py b/tests/op_engine/test_operator_precision.py index 6a98aca..b71d735 100644 --- a/tests/op_engine/test_operator_precision.py +++ b/tests/op_engine/test_operator_precision.py @@ -2,14 +2,31 @@ from __future__ import annotations -from contextlib import nullcontext -from typing import Any +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any import numpy as np import pytest from op_engine import array_namespace, build_advection_matrix, build_diffusion_matrix +if TYPE_CHECKING: + from collections.abc import Iterator + + +@contextmanager +def _precision(backend: str) -> Iterator[None]: + if backend == "jax": + jax = pytest.importorskip("jax") + enable = getattr(jax, "enable_x64", None) + if enable is None: + enable = jax.experimental.enable_x64 + with enable(): + yield + else: + yield + + BACKENDS = ("numpy", "jax", "torch") STENCILS = { "advection": np.asarray([[-1, 0, 0], [1, -1, 0], [0, 1, 0]]), @@ -39,10 +56,7 @@ def test_floating_reference_selects_operator_precision( if backend == "numpy" else pytest.importorskip("jax.numpy" if backend == "jax" else "torch") ) - precision = ( - pytest.importorskip("jax").enable_x64() if backend == "jax" else nullcontext() - ) - with precision: + with _precision(backend): reference = module.asarray(np.ones((2, 3, 3), dtype=dtype)) result = BUILDERS[kind](3, 1.0, 0.4, bc="reflecting", reference=reference) assert result.dtype == reference.dtype @@ -59,10 +73,7 @@ def test_reference_precision_preserves_dynamic_coefficient_derivatives( ) -> None: """Coefficient derivatives match the known stencil's quadratic objective.""" module = pytest.importorskip("jax.numpy" if backend == "jax" else "torch") - precision = ( - pytest.importorskip("jax").enable_x64() if backend == "jax" else nullcontext() - ) - with precision: + with _precision(backend): reference = module.asarray(np.ones(3, dtype=dtype)) coefficient = module.asarray(np.asarray(0.4, dtype=dtype))