Preserve floating reference dtype in transport matrix builders

.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'

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,

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

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 = (

tests/op_engine/test_operator_precision.py

@@ -0,0 +1,105 @@
+"""Reference-controlled precision for namespace-aware finite-volume builders."""
+
+from __future__ import annotations
+
+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]]),
+    "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")
+    )
+    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
+        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")
+    with _precision(backend):
+        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])