diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 6329c02..f2dbc9c 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -15,6 +15,26 @@ permissions: contents: 'read' jobs: + torch-cpu: + name: 'PyTorch CPU (3.12)' + runs-on: 'ubuntu-latest' + steps: + - name: 'Checkout repository' + uses: 'actions/checkout@v6' + with: + persist-credentials: false + - name: 'Setup uv' + uses: 'astral-sh/setup-uv@v7.2.0' + with: + enable-cache: true + cache-dependency-glob: 'pyproject.toml' + python-version: '3.12' + - name: 'Install test dependencies and CPU PyTorch' + run: | + 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' build-check: name: 'Build Check (${{ matrix.python-version }})' runs-on: 'ubuntu-latest' diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c3189f..fc06162 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,13 @@ versioning while the project remains pre-1.0. ## [Unreleased] +### Fixed + +- `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, + float32/float64, and optional history (#208). + ## [0.6.0] - 2026-10-09 ### Added diff --git a/docs/guides/backends.md b/docs/guides/backends.md index af04f4f..69edcaf 100644 --- a/docs/guides/backends.md +++ b/docs/guides/backends.md @@ -12,10 +12,13 @@ Namespace portability and automatic differentiation are related but distinct. The Array API does not define `grad`, `jit`, or traced control flow. A namespace such as JAX can nevertheless differentiate the ordinary `op_engine` methods because their fixed-step numerical operations remain in that namespace. -The optional `torch` dependency group qualifies representative eager -fixed-step kernels with PyTorch autograd. Compiler integration and complete -method qualification remain provider/backend capabilities rather than -consequences of namespace selection. +CPU PyTorch tests qualify `ModelCore` initialization and complete fixed and +adaptive explicit runs with float32/float64 and optional history, plus a +representative fixed-step kernel with PyTorch autograd. Configure the dtype +using NumPy-compatible values such as `ModelCoreOptions(dtype=np.float32)`; +`ModelCore` resolves that dtype in the initial state's namespace. Compiler +integration and complete method qualification remain provider/backend +capabilities rather than consequences of namespace selection. ## Portable fixed-step methods diff --git a/src/op_engine/model_core.py b/src/op_engine/model_core.py index e4513aa..35ee19c 100644 --- a/src/op_engine/model_core.py +++ b/src/op_engine/model_core.py @@ -87,7 +87,8 @@ class ModelCoreOptions: This is metadata only and enables non-uniform grids and FV/FDM operator builders. store_history: Whether to store the full time history. - dtype: Floating-point dtype for internal arrays. + dtype: NumPy-compatible dtype for internal arrays, resolved to the + initial state's namespace when that state is supplied. """ other_axes: tuple[int, ...] = () @@ -425,7 +426,7 @@ def set_initial_state(self, initial_state: Array) -> None: xp = _namespace_of(initial_state) initial_state_arr = cast( "Array", - xp.asarray(initial_state, dtype=self.dtype), + xp.asarray(initial_state, dtype=getattr(xp, self.dtype.name)), ) if initial_state_arr.shape != self.state_shape: raise ValueError( diff --git a/tests/op_engine/test_array_namespace.py b/tests/op_engine/test_array_namespace.py index 66d105c..5bddd04 100644 --- a/tests/op_engine/test_array_namespace.py +++ b/tests/op_engine/test_array_namespace.py @@ -152,6 +152,59 @@ def test_model_core_options_no_longer_store_backend_module() -> None: assert "xp" not in {field.name for field in fields(ModelCoreOptions)} +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("store_history", [False, True]) +@pytest.mark.parametrize( + ("method", "adaptive"), + [("rk4", False), ("dopri5", False), ("dopri5", True), ("heun", True)], +) +def test_core_solver_run_preserves_torch_state_and_matches_numpy( + dtype: type[np.float32 | np.float64], + method: str, + *, + adaptive: bool, + store_history: bool, +) -> None: + """The complete run resolves NumPy dtype options into native Torch dtypes.""" + torch = pytest.importorskip("torch") + times = np.asarray([0.0, 0.2, 0.7, 1.0]) + initial = np.arange(1.0, 7.0).reshape(2, 3) + options = ModelCoreOptions(dtype=dtype, store_history=store_history) + numpy_core = ModelCore(2, 3, times, options=options) + torch_core = ModelCore(2, 3, times, options=options) + numpy_core.set_initial_state(initial) + torch_core.set_initial_state(torch.tensor(initial, dtype=torch.float64)) + + def rhs(_time: Scalar, state: Array) -> Array: + return cast("Array", array_namespace(state).multiply(state, -0.3)) + + config = RunConfig( + method=method, + adaptive=adaptive, + adaptive_cfg=AdaptiveConfig(rtol=1e-5, atol=1e-7, dt_init=0.1), + ) + CoreSolver(numpy_core).run(rhs, config=config) + CoreSolver(torch_core).run(rhs, config=config) + + result = torch_core.get_current_state() + assert isinstance(result, torch.Tensor) + assert result.dtype == getattr(torch, np.dtype(dtype).name) + assert result.device.type == "cpu" + tolerance = 10 * np.finfo(dtype).eps + np.testing.assert_allclose( + result.numpy(), numpy_core.get_current_state(), rtol=tolerance + ) + np.testing.assert_allclose(result.numpy(), initial * np.exp(-0.3), rtol=3e-6) + if store_history: + assert isinstance(torch_core.state_array, torch.Tensor) + assert torch_core.state_array.dtype == result.dtype + np.testing.assert_allclose( + torch_core.state_array.numpy(), numpy_core.state_array, rtol=tolerance + ) + else: + assert torch_core.state_array is None + + def test_model_core_preserves_jax_namespace_for_updates_and_history() -> None: """Immutable JAX state/history updates never fall back to NumPy.""" jnp = pytest.importorskip("jax.numpy")