diff --git a/CHANGELOG.md b/CHANGELOG.md index c9b30d0..ab5ce07 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,12 @@ versioning while the project remains pre-1.0. The capability matrix separates method, loop, sparse, and hardware evidence and documents matching `just qualify` commands (#206). +- Live adaptive schedules carry lightweight method/layout/kernel provenance and + descriptive discovery controls, namespace, precision and caller regime labels. + Replay rejects known structural mismatches while retaining portable dynamic + parameter/state reuse and compatible manual meshes. A structured analytic + example demonstrates accuracy-driven refresh outside differentiation (#205). + ### Fixed - Namespace-aware advection and diffusion builders honor a floating reference diff --git a/docs/guides/backends.md b/docs/guides/backends.md index 783748a..04f5600 100644 --- a/docs/guides/backends.md +++ b/docs/guides/backends.md @@ -65,8 +65,11 @@ The resulting derivative is conditional on the recorded mesh. This is also the branchwise meaning of an eager live-adaptive gradient: neither mode differentiates the discrete accept/reject decision. Refresh the schedule as optimization parameters move, and always refresh it after a material model or tolerance -change. A schedule is validated against the output grid, while matching the -method and other configuration is the caller's responsibility. +change. A recorded schedule checks producing method, state layout, relevant kernel +configuration, and optional structural model identity before replay. Backend, +dtype and fitted values can change; manually constructed meshes remain +unverified unless provenance is required. See the +[replay contract and refresh example](replay-provenance.md). Core schedule replay retains the eager Python loop by default. Select `RunConfig(replay_loop="auto")` to use a namespace's registered scan driver, diff --git a/docs/guides/replay-provenance.md b/docs/guides/replay-provenance.md new file mode 100644 index 0000000..9c089d2 --- /dev/null +++ b/docs/guides/replay-provenance.md @@ -0,0 +1,65 @@ +# Accepted meshes: compatibility and accuracy + +A frozen accepted mesh is static controller output. Replay invokes its numerical +steps and bypasses adaptive error estimates and acceptance decisions. Its AD +results are derivatives **conditional on that mesh**. A compatible mesh may +still become inaccurate when model parameters or initial conditions move. + +The core contract distinguishes these questions before testing accuracy: + +| Metadata | Replay treatment | +| --- | --- | +| Finite/increasing output grid and positive interval steps summing to each interval | Always validated by `AdaptiveStepSchedule` | +| Producing resolved method, state shape/axis order, relevant operator axis, TR-BDF2 gamma, nonlinear stage solver/configuration | Known mismatches reject before stepping | +| Caller-supplied structural model identity | Compared when recording/replay use an identity; it must describe equations/geometry/layout, not fitted parameter values | +| Array namespace and dtype | Descriptive; portable meshes may change NumPy/JAX/Torch or precision | +| Discovery tolerances, controller limits, and reference parameter regime | Descriptive; changing fitted values or requested accuracy requires an external accuracy decision, not an automatic incompatibility | +| Callback bodies and arbitrary custom solver internals | No automatic semantic identity guarantee; use the structural identity and independently qualify changed implementations | + +`AdaptiveScheduleProvenance` is attached to schedules created by live adaptive +runs. `RunConfig.replay_model_id` and `replay_reference_regime` are optional +static labels supplied by the caller. Record concrete nominal values or a +stable parameter-set identifier in the reference label at an eager boundary; +replay does not inspect or convert dynamic parameter arrays. + +Manually constructed meshes and legacy schedules with `provenance=None` remain +accepted by default after mesh/output-grid validation. Their method/model +compatibility is unverified and remains the caller's responsibility. Use +`RunConfig(replay_require_provenance=True)` to reject unverified schedules. +Adding a provenance record is not an accuracy certificate, and copying labels +from another model does not make that model compatible. + +Controller settings are recorded to explain discovery. Replay ignores error +control even if the supplied configuration requests tighter tolerances. A +provider may impose stronger configuration/context checks on its own artifact; +these core rules do not remove those provider checks. + +A practical reuse policy is: + +1. Discover a live mesh at a documented nominal parameter regime. +2. Reuse it while changing dynamic values, retaining structural model identity. +3. Compare a replay trajectory to a fresh accurate solve or to a refined mesh + at representative new parameter values. Check the scientific objective and + needed derivatives as well as state error. +4. If the independent discrepancy exceeds the intended accuracy, rediscover + the mesh at the new regime, recheck agreement, and freeze that mesh for the + next differentiation phase. + +The structured example accompanying this guide executes that policy and shows +an accuracy failure despite compatible provenance. Discrete refresh decisions +remain outside the differentiated numerical solve. + +Run the [structured refresh example](https://github.com/ACCIDDA/op_engine/blob/main/examples/replay_mesh_refresh.py): + +```shell +uv run python examples/replay_mesh_refresh.py +``` + +It transfers mass between two generic pools with `A'=-k*A^2`, `B'=+k*A^2`, +a shared rate and group-local multipliers. The independent solution is +`A(t)=A0/(1+k*A0*t)`, with conserved `A+B`. A mesh discovered at rate 0.05 +remains accurate nearby; rate 3.0 is structurally compatible but exceeds the +example's `1e-4` state-error policy on the old mesh. Fresh discovery restores +that accuracy. JSON reports both discrepancies and accepted-step counts. +This policy checks this fixture's state error; different scientific objectives +and derivative accuracy need their own independent checks. diff --git a/examples/replay_mesh_refresh.py b/examples/replay_mesh_refresh.py new file mode 100644 index 0000000..eaf4708 --- /dev/null +++ b/examples/replay_mesh_refresh.py @@ -0,0 +1,122 @@ +"""Check and refresh a compatible mesh after a nonlinear parameter change. + +Run ``uv run python examples/replay_mesh_refresh.py``. A shared transfer rate +and two group-local multipliers drive a generic closed two-pool model. The +analytic solution independently checks replay accuracy before any mesh reuse +is trusted. Refresh is an eager decision outside a differentiated solve. +""" + +from __future__ import annotations + +import json +from dataclasses import replace +from typing import cast + +import numpy as np + +from op_engine import Array, CoreSolver, ModelCore, array_namespace +from op_engine.core_solver import AdaptiveConfig, AdaptiveStepSchedule, RunConfig +from op_engine.model_core import ModelCoreOptions + +INITIAL = np.asarray([[0.8, 0.6], [0.2, 0.4]]) +LOCAL = np.asarray([1.0, 1.3]) +CONFIG = RunConfig( + method="heun", + adaptive=True, + adaptive_cfg=AdaptiveConfig(rtol=1e-4, atol=1e-7, dt_init=0.5), + replay_model_id="quadratic-transfer:pool,group:v1", + replay_require_provenance=True, +) + + +def trajectory( + rate: float, schedule: AdaptiveStepSchedule | None = None +) -> tuple[np.ndarray, AdaptiveStepSchedule]: + """Discover or replay this structured model at a shared transfer rate. + + Returns: + Final state and its accepted mesh. + + Raises: + RuntimeError: If live discovery did not return a mesh. + """ + core = ModelCore( + 2, + 2, + np.asarray([0.0, 1.0]), + options=ModelCoreOptions( + axis_names=("pool", "group"), + ), + ) + core.set_initial_state(INITIAL) + solver = CoreSolver(core) + + def rhs(_time: object, state: Array) -> Array: + xp = array_namespace(state) + flow = rate * xp.asarray(LOCAL, dtype=state.dtype) * state[0] ** 2 + return cast("Array", xp.stack((-flow, flow))) + + config = replace( + CONFIG, replay_reference_regime=f"shared_rate={rate}; local=[1,1.3]; A0=[.8,.6]" + ) + if schedule is None: + solver.run(rhs, config=config) + else: + solver.replay_adaptive_schedule(rhs, schedule, config=config) + recorded = solver.last_adaptive_schedule + if recorded is None: + msg = "Adaptive solve did not return an accepted mesh" + raise RuntimeError(msg) + return np.asarray(core.get_current_state()), recorded + + +def exact(rate: float) -> np.ndarray: + """Evaluate the independent closed-form two-pool solution at time one. + + Returns: + Final state from A'=-k*A^2 and the conserved total A+B. + """ + first = INITIAL[0] / (1 + rate * LOCAL * INITIAL[0]) + return np.stack((first, INITIAL.sum(axis=0) - first)) + + +def assess_and_refresh() -> dict[str, float | int | bool | str]: + """Apply an independent state-error policy at a changed parameter regime. + + Returns: + Accuracy and mesh-count evidence for the reusable and refreshed cases. + + Raises: + RuntimeError: If the selected mesh fails the independent accuracy check. + """ + _, nominal = trajectory(0.05) + near, _ = trajectory(0.055, nominal) + changed_rate = 3.0 + changed, _ = trajectory(changed_rate, nominal) + reference = exact(changed_rate) + maximum_error = float(np.max(np.abs(changed - reference))) + accuracy_limit = 1e-4 + selected = nominal + result = changed + if maximum_error > accuracy_limit: + result, selected = trajectory(changed_rate) + refreshed_error = float(np.max(np.abs(result - reference))) + if refreshed_error > accuracy_limit: + msg = "The selected mesh still fails the independent accuracy check" + raise RuntimeError(msg) + return { + "model_id": CONFIG.replay_model_id or "", + "nominal_rate": 0.05, + "changed_rate": changed_rate, + "accuracy_limit": accuracy_limit, + "nearby_error": float(np.max(np.abs(near - exact(0.055)))), + "old_mesh_error": maximum_error, + "selected_mesh_error": refreshed_error, + "nominal_steps": sum(map(len, nominal.step_sizes)), + "selected_steps": sum(map(len, selected.step_sizes)), + "refreshed": selected is not nominal, + } + + +if __name__ == "__main__": + print(json.dumps(assess_and_refresh(), indent=2)) # noqa: T201 diff --git a/justfile b/justfile index a36cb20..dd0059b 100644 --- a/justfile +++ b/justfile @@ -102,6 +102,7 @@ build-test-core: uv pip install --python "${CLEANROOM}/venv/bin/python" -r "${CLEANROOM}/dev-requirements.txt" cp pyproject.toml "${CLEANROOM}/pyproject.toml" cp -R tests "${CLEANROOM}/tests" + cp -R examples "${CLEANROOM}/examples" cd "${CLEANROOM}" "${CLEANROOM}/venv/bin/pytest" --import-mode=importlib tests --quiet --exitfirst diff --git a/mkdocs.yml b/mkdocs.yml index d2c8945..d9bc7c8 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -57,6 +57,8 @@ nav: - Biogeochemical Network: guides/biogeochemical-network.md - Trait-Structured Chemostat: guides/chemostat-pde.md - Backend Qualification: guides/backend-qualification.md + + - Replay Provenance: guides/replay-provenance.md - Backend Strategies: guides/backends.md - API Reference: - Adaptive Tau: api-reference/adaptive_tau.md diff --git a/src/op_engine/__init__.py b/src/op_engine/__init__.py index 32e39bc..d75411e 100644 --- a/src/op_engine/__init__.py +++ b/src/op_engine/__init__.py @@ -9,6 +9,7 @@ from ._typing import Scalar as Scalar from .adaptive_tau import AdaptiveTauLeapingConfig, AdaptiveTauLeapingSolver from .core_solver import ( + AdaptiveScheduleProvenance, CoreSolver, NonlinearIntegrationConvergenceError, NonlinearIntegrationDiagnostics, @@ -88,6 +89,7 @@ ) __all__ = [ + "AdaptiveScheduleProvenance", "AdaptiveTauLeapingConfig", "AdaptiveTauLeapingSolver", "Array", diff --git a/src/op_engine/core_solver.py b/src/op_engine/core_solver.py index 6e47a2d..354138a 100644 --- a/src/op_engine/core_solver.py +++ b/src/op_engine/core_solver.py @@ -378,6 +378,27 @@ def __post_init__(self) -> None: raise ValueError(msg) +class AdaptiveScheduleProvenance(NamedTuple): + """Static compatibility fields and descriptive adaptive-discovery metadata. + + Backend, dtype, discovery controls, and reference regime are descriptive. + Replay checks method, state layout, relevant operator/kernel settings, and + the caller's structural model identity. Callback semantics are not inferred. + """ + + method: str + state_shape: tuple[int, ...] + axis_names: tuple[str, ...] + operator_axis: int | None + gamma: float | None + nonlinear_signature: tuple[str, ...] + discovery_settings: tuple[tuple[str, str], ...] + array_namespace: str + dtype: str + model_id: str | None + reference_regime: str | None + + @dataclass(slots=True, frozen=True) class AdaptiveStepSchedule: """Accepted step sizes for replaying one adaptive solve. @@ -389,17 +410,26 @@ class AdaptiveStepSchedule: Attributes: output_times: Output grid used to create the schedule. step_sizes: Accepted internal step sizes for each output interval. + provenance: Optional producing method/layout/configuration metadata. + None retains manually constructed/legacy mesh compatibility. """ output_times: tuple[float, ...] step_sizes: tuple[tuple[float, ...], ...] + provenance: AdaptiveScheduleProvenance | None = None def __post_init__(self) -> None: """Normalize and validate the recorded mesh. Raises: + TypeError: If provenance has the wrong record type. ValueError: If times or step sizes do not define a valid mesh. """ + if self.provenance is not None and not isinstance( + self.provenance, AdaptiveScheduleProvenance + ): + msg = "schedule provenance must be an AdaptiveScheduleProvenance or None" + raise TypeError(msg) output_times = tuple(float(value) for value in self.output_times) step_sizes = tuple( tuple(float(step_size) for step_size in interval) @@ -740,6 +770,12 @@ class RunConfig: ``scan`` requires one. SDIRK2 diagnostics retain eager replay. replay_checkpoint: Rematerialize the scan body during reverse mode. Used only when replay selects an adapter. + replay_model_id: Optional static structural model label, checked during + replay. Keep fitted parameter values out of this identity. + replay_reference_regime: Optional descriptive label for the nominal + parameter/initial-state regime used to discover a mesh. + replay_require_provenance: Reject manually constructed/legacy meshes + without provenance when True. False preserves compatibility. """ method: str = "heun" @@ -754,6 +790,9 @@ class RunConfig: fixed_max_step: float | None = None replay_loop: Literal["auto", "scan", "unroll"] = "unroll" replay_checkpoint: bool = False + replay_model_id: str | None = None + replay_reference_regime: str | None = None + replay_require_provenance: bool = False def __post_init__(self) -> None: """Normalize the method and validate context-free configuration. @@ -799,7 +838,7 @@ def _validate_replay_options(self) -> None: """Validate namespace-loop options independently of the chosen method. Raises: - TypeError: If the checkpoint flag is not boolean. + TypeError: If a replay flag or static label has the wrong type. ValueError: If the loop strategy is unknown. """ if self.replay_loop not in {"auto", "scan", "unroll"}: @@ -808,6 +847,16 @@ def _validate_replay_options(self) -> None: if not isinstance(self.replay_checkpoint, bool): msg = "replay_checkpoint must be boolean" raise TypeError(msg) + if not isinstance(self.replay_require_provenance, bool): + msg = "replay_require_provenance must be boolean" + raise TypeError(msg) + for name in ("replay_model_id", "replay_reference_regime"): + label = getattr(self, name) + if label is not None: + if not isinstance(label, str) or not label.strip(): + msg = f"{name} must be a non-empty static string or None" + raise TypeError(msg) + object.__setattr__(self, name, label.strip()) # noqa: PLC2801 @dataclass(slots=True, frozen=True) @@ -1138,6 +1187,104 @@ def last_nonlinear_diagnostics( """Return nonlinear diagnostics from the latest SDIRK run or replay.""" return self._last_nonlinear_diagnostics + def _schedule_provenance( + self, config: RunConfig, plan: RunPlan + ) -> AdaptiveScheduleProvenance: + """Capture static metadata without inspecting fitted array values. + + Returns: + Producing compatibility fields and descriptive discovery settings. + """ + state = self.core.get_current_state() + adaptive = config.adaptive_cfg + atol = adaptive.atol + atol_description = ( + str(atol) + if isinstance(atol, (float, int, np.floating, np.integer)) + else f"array shape={atol.shape}, dtype={atol.dtype}" + ) + settings = ( + ("rtol", str(adaptive.rtol)), + ("atol", atol_description), + ("dt_init", str(adaptive.dt_init)), + ("max_reject", str(adaptive.max_reject)), + ("max_steps", str(adaptive.max_steps)), + *( + (name, str(getattr(config.dt_controller, name))) + for name in ( + "dt_min", + "dt_max", + "safety", + "fac_min", + "fac_max", + ) + ), + ) + nonlinear: tuple[str, ...] = () + if plan.method == "sdirk2" and plan.nonlinear is not None: + solver = plan.nonlinear.solver + nonlinear = (f"{type(solver).__module__}.{type(solver).__qualname__}",) + if isinstance(solver, DenseNewtonSolver): + nonlinear += tuple( + str(getattr(solver.config, name)) + for name in ( + "max_iterations", + "rtol", + "atol", + "damping", + ) + ) + operator_axis = ( + None + if plan.method in _EXPLICIT_METHODS or plan.method == "sdirk2" + else self.core.axis_index(self._op_axis) + ) + return AdaptiveScheduleProvenance( + plan.method, + self.state_shape, + self.core.axis_names, + operator_axis, + plan.gamma, + nonlinear, + settings, + _namespace_of(state).__name__, + str(state.dtype), + config.replay_model_id, + config.replay_reference_regime, + ) + + def _validate_schedule_provenance( + self, schedule: AdaptiveStepSchedule, config: RunConfig, plan: RunPlan + ) -> None: + """Reject known structural incompatibility before any replay step. + + Raises: + ValueError: If required provenance is absent or a checked field differs. + """ + recorded = schedule.provenance + if recorded is None: + if config.replay_require_provenance: + msg = "Adaptive schedule has no provenance; discover a verified mesh" + raise ValueError(msg) + return + active = self._schedule_provenance(config, plan) + for name in ( + "method", + "state_shape", + "axis_names", + "operator_axis", + "gamma", + "nonlinear_signature", + "model_id", + ): + if getattr(recorded, name) != getattr(active, name): + msg = ( + f"Adaptive schedule provenance {name} mismatch: " + f"recorded {getattr(recorded, name)!r}, " + f"configured {getattr(active, name)!r}; discover a fresh mesh" + ) + raise ValueError(msg) + # ------------------------------------------------------------------ # Axis / operator helpers # ------------------------------------------------------------------ @@ -3853,6 +4000,7 @@ def _run_explicit( self._last_adaptive_schedule = AdaptiveStepSchedule( output_times=tuple(float(value) for value in time_grid), step_sizes=tuple(schedule_steps), + provenance=self._schedule_provenance(config, plan), ) def _advance_array_implicit_adaptive_to_time( # noqa: PLR0913 @@ -3993,6 +4141,7 @@ def _run_array_implicit( self._last_adaptive_schedule = AdaptiveStepSchedule( output_times=tuple(float(value) for value in time_grid), step_sizes=tuple(schedule_steps), + provenance=self._schedule_provenance(config, plan), ) def _advance_nonadaptive_to_time( # noqa: PLR0913 @@ -4366,6 +4515,7 @@ def _run_sdirk( self._last_adaptive_schedule = AdaptiveStepSchedule( output_times=tuple(float(value) for value in time_grid), step_sizes=tuple(schedule_steps), + provenance=self._schedule_provenance(config, plan), ) return self._update_nonlinear_diagnostics( reference, @@ -4637,7 +4787,8 @@ def replay_adaptive_schedule( Raises: TypeError: If schedule has the wrong type or the state changes array ecosystems during replay. - ValueError: If config is not adaptive or the output grid differs. + ValueError: If config is not adaptive, the output grid differs, + required provenance is absent, or structural provenance differs. """ if not isinstance(schedule, AdaptiveStepSchedule): msg = "schedule must be an AdaptiveStepSchedule" @@ -4650,6 +4801,7 @@ def replay_adaptive_schedule( self._last_nonlinear_diagnostics = None self._validate_schedule_time_grid(schedule) plan = self._resolve_run_plan(config) + self._validate_schedule_provenance(schedule, config, plan) adapter = None if config.replay_loop != "unroll": @@ -4791,5 +4943,6 @@ def run( self._last_adaptive_schedule = AdaptiveStepSchedule( output_times=tuple(float(value) for value in time_grid), step_sizes=tuple(schedule_steps), + provenance=self._schedule_provenance(cfg, plan), ) return None diff --git a/tests/op_engine/test_replay_provenance.py b/tests/op_engine/test_replay_provenance.py new file mode 100644 index 0000000..23815ea --- /dev/null +++ b/tests/op_engine/test_replay_provenance.py @@ -0,0 +1,316 @@ +"""Separate mesh validity, structural compatibility, and accuracy-driven refresh.""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from runpy import run_path +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +import pytest + +from op_engine import ( + AdaptiveScheduleProvenance, + Array, + CoreSolver, + DenseNewtonSolver, + ModelCore, + NewtonConfig, + NonlinearMethodConfig, + array_namespace, +) +from op_engine.core_solver import ( + AdaptiveConfig, + AdaptiveStepSchedule, + OperatorSpecs, + RunConfig, +) +from op_engine.model_core import ModelCoreOptions + +if TYPE_CHECKING: + from op_engine.core_solver import RHSFunction + +INITIAL = np.asarray([[0.8, 0.6], [0.2, 0.4]]) +LOCAL = np.asarray([1.0, 1.3]) + + +def _transfer(rate: Any) -> RHSFunction: # noqa: ANN401 + def rhs(_time: object, state: Array) -> Array: + xp = array_namespace(state) + flow = rate * xp.asarray(LOCAL, dtype=state.dtype) * state[0] ** 2 + return cast("Array", xp.stack((-flow, flow))) + + return rhs + + +def _core( + initial: Array | None = None, *, axes: tuple[str, str] = ("pool", "group") +) -> ModelCore: + core = ModelCore( + 2, 2, np.asarray([0.0, 1.0]), options=ModelCoreOptions(axis_names=axes) + ) + core.set_initial_state(INITIAL if initial is None else initial) + return core + + +def _config() -> RunConfig: + return RunConfig( + method="heun", + adaptive=True, + adaptive_cfg=AdaptiveConfig(rtol=1e-4, atol=1e-7, dt_init=0.5), + replay_model_id="quadratic-transfer:pool,group:v1", + replay_reference_regime="rate=.05; local=[1,1.3]; A0=[.8,.6]", + ) + + +def _record() -> tuple[AdaptiveStepSchedule, RunConfig]: + config = _config() + solver = CoreSolver(_core()) + solver.run(_transfer(0.05), config=config) + assert solver.last_adaptive_schedule is not None + return solver.last_adaptive_schedule, config + + +def _reference(rate: float, initial: np.ndarray = INITIAL) -> np.ndarray: + # A'=-k*A^2, B'=+k*A^2 gives A(t)=A0/(1+k*A0*t), A+B constant. + a = initial[0] / (1 + rate * LOCAL * initial[0]) + return np.stack((a, initial.sum(axis=0) - a)) + + +def test_live_mesh_captures_descriptive_and_structural_metadata() -> None: + """A live mesh identifies its method, layout, and nominal regime.""" + schedule, config = _record() + provenance = schedule.provenance + assert isinstance(provenance, AdaptiveScheduleProvenance) + assert provenance.method == "heun" + assert provenance.state_shape == INITIAL.shape + assert provenance.axis_names == ("pool", "group") + assert provenance.operator_axis is None + assert provenance.model_id == config.replay_model_id + assert provenance.reference_regime == config.replay_reference_regime + assert dict(provenance.discovery_settings)["rtol"] == "0.0001" + assert "numpy" in provenance.array_namespace + assert provenance.dtype == "float64" + + +@pytest.mark.parametrize( + ("field", "change"), + [ + ("method", {"method": "rk4"}), + ("model_id", {"replay_model_id": "different-equations:v2"}), + ], +) +def test_known_method_and_model_mismatch_reject_before_advancing( + field: str, change: dict[str, Any] +) -> None: + """Known structural changes fail while retaining the untouched initial state.""" + schedule, config = _record() + solver = CoreSolver(_core()) + with pytest.raises(ValueError, match=f"provenance {field} mismatch"): + solver.replay_adaptive_schedule( + _transfer(0.05), schedule, config=replace(config, **change) + ) + assert solver.core.current_step == 0 + np.testing.assert_array_equal(solver.core.get_current_state(), INITIAL) + + +@pytest.mark.parametrize("layout", ["shape", "axis_names"]) +def test_state_layout_mismatch_rejects(layout: str) -> None: + """Cell shape and named axis order are structural compatibility constraints.""" + schedule, config = _record() + core = ( + _core(axes=("state", "group")) + if layout == "axis_names" + else ModelCore(2, 3, np.asarray([0.0, 1.0])) + ) + field = "axis_names" if layout == "axis_names" else "state_shape" + with pytest.raises( + ValueError, + match=f"provenance {field} mismatch", + ): + CoreSolver(core).replay_adaptive_schedule( + _transfer(0.05), schedule, config=config + ) + assert core.current_step == 0 + + +def test_legacy_mesh_is_explicitly_unverified_and_optionally_rejected() -> None: + """Legacy constructors retain mesh-only behavior unless provenance is required.""" + manual = AdaptiveStepSchedule((0.0, 1.0), ((0.5, 0.5),)) + assert manual.provenance is None + solver = CoreSolver(_core()) + solver.replay_adaptive_schedule(_transfer(0.05), manual, config=_config()) + assert solver.core.current_step == 1 + strict = CoreSolver(_core()) + with pytest.raises(ValueError, match="has no provenance"): + strict.replay_adaptive_schedule( + _transfer(0.05), + manual, + config=replace(_config(), replay_require_provenance=True), + ) + assert strict.core.current_step == 0 + with pytest.raises(ValueError, match="sum to each output interval"): + AdaptiveStepSchedule((0.0, 1.0), ((0.4, 0.4),)) + + +def test_compatible_parameter_reuse_can_need_accuracy_refresh() -> None: + """Independent analytic error checks distinguish reuse from an accuracy refresh.""" + schedule, config = _record() + near = CoreSolver(_core()) + near.replay_adaptive_schedule( + _transfer(0.055), + schedule, + config=replace(config, replay_reference_regime="nearby=.055"), + ) + np.testing.assert_allclose( + near.core.get_current_state(), _reference(0.055), rtol=1e-5, atol=1e-6 + ) + far = CoreSolver(_core()) + far.replay_adaptive_schedule(_transfer(3.0), schedule, config=config) + discrepancy = np.max(np.abs(far.core.get_current_state() - _reference(3.0))) + assert discrepancy > 1e-3 + fresh = CoreSolver(_core()) + refreshed_config = replace(config, replay_reference_regime="rate=3.0") + fresh.run(_transfer(3.0), config=refreshed_config) + refreshed = fresh.last_adaptive_schedule + assert refreshed is not None + assert len(refreshed.step_sizes[0]) > len(schedule.step_sizes[0]) + assert refreshed.provenance is not None + assert refreshed.provenance.reference_regime == "rate=3.0" + np.testing.assert_allclose( + fresh.core.get_current_state(), _reference(3.0), rtol=1e-4, atol=1e-5 + ) + + +def test_portable_jax_reuse_preserves_dynamic_state_and_tolerance_arrays() -> None: + """Static compatibility checks preserve traced arrays across a backend change.""" + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + enable = getattr(jax, "enable_x64", None) + if enable is None: + enable = jax.experimental.enable_x64 + schedule, config = _record() + with enable(): + + def objective(gain: Any, atol: Any) -> Any: # noqa: ANN401 + solver = CoreSolver(_core(jnp.asarray(INITIAL) * gain)) + replay_config = replace( + config, + replay_loop="scan", + adaptive_cfg=AdaptiveConfig(rtol=1e-5, atol=atol), + replay_reference_regime="dynamic initial scale", + ) + solver.replay_adaptive_schedule( + _transfer(0.05), schedule, config=replay_config + ) + result = solver.core.get_current_state() + return jnp.sum(result**2) + + gain = jnp.asarray(1.0) + value, derivative = jax.jit(jax.value_and_grad(objective))( + gain, jnp.full((2, 2), 1e-8) + ) + step = 1e-4 + expected = ( + np.sum(_reference(0.05, INITIAL * (1 + step)) ** 2) + - np.sum(_reference(0.05, INITIAL * (1 - step)) ** 2) + ) / (2 * step) + assert np.isfinite(float(value)) + np.testing.assert_allclose(derivative, expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("method", ["ros2", "imex-trbdf2", "sdirk2"]) +def test_implicit_discovery_records_relevant_kernel_identity(method: str) -> None: + """Resolved implicit kernels reject axis, gamma, or nonlinear iteration changes.""" + core = _core() + xp = array_namespace(core.get_current_state()) + identity = xp.eye(2) + + def implicit(dt: float, scale: float, _context: object) -> tuple[Array, Array]: + return cast("Array", identity * (1 + 0.2 * dt * scale)), cast("Array", identity) + + def trapezoidal(dt: float, scale: float, _context: object) -> tuple[Array, Array]: + return cast("Array", identity * (1 + 0.1 * dt * scale)), cast( + "Array", identity * (1 - 0.1 * dt * scale) + ) + + config = replace( + _config(), + method=method, + jacobian=lambda _t, _y: cast("Array", -0.2 * identity), + nonlinear=NonlinearMethodConfig( + rhs_jacobian=lambda _t, _y: cast("Array", -0.2 * xp.eye(4)) + ) + if method == "sdirk2" + else None, + operators=OperatorSpecs(tr=trapezoidal, bdf2=implicit) + if method == "imex-trbdf2" + else OperatorSpecs(), + ) + rhs = ( + (lambda _t, y: cast("Array", y - y)) + if method == "imex-trbdf2" + else (lambda _t, y: cast("Array", -0.2 * y)) + ) + solver = CoreSolver(core, operator_axis="pool") + solver.run(rhs, config=config) + mesh = solver.last_adaptive_schedule + assert mesh is not None + assert mesh.provenance is not None + assert mesh.provenance.method == method + if method == "imex-trbdf2": + assert mesh.provenance.gamma is not None + changed = replace(config, gamma=0.7) + field = "gamma" + elif method == "sdirk2": + assert config.nonlinear is not None + changed = replace( + config, + nonlinear=replace( + config.nonlinear, + solver=DenseNewtonSolver(NewtonConfig(max_iterations=4)), + ), + ) + field = "nonlinear_signature" + else: + changed = config + field = "operator_axis" + target = CoreSolver(_core(), operator_axis="group" if method == "ros2" else "pool") + with pytest.raises(ValueError, match=f"provenance {field} mismatch"): + target.replay_adaptive_schedule(rhs, mesh, config=changed) + + +def test_array_implicit_discovery_records_jax_metadata() -> None: + """The native implicit recording path attaches the same compatibility contract.""" + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + enable = getattr(jax, "enable_x64", None) + if enable is None: + enable = jax.experimental.enable_x64 + with enable(): + solver = CoreSolver(_core(jnp.asarray(INITIAL)), operator_axis="pool") + config = replace( + _config(), method="ros2", jacobian=lambda _t, _y: -0.2 * jnp.eye(2) + ) + solver.run(lambda _t, y: -0.2 * y, config=config) + mesh = solver.last_adaptive_schedule + assert mesh is not None + assert mesh.provenance is not None + assert mesh.provenance.array_namespace == "jax.numpy" + assert mesh.provenance.operator_axis == 0 + np.testing.assert_allclose( + solver.core.get_current_state(), INITIAL * np.exp(-0.2), rtol=1e-4 + ) + + +def test_published_structured_refresh_example() -> None: + """The runnable example makes and verifies an independent accuracy decision.""" + script = Path(__file__).resolve().parents[2] / "examples/replay_mesh_refresh.py" + namespace = run_path(str(script)) + report = namespace["assess_and_refresh"]() + assert report["refreshed"] is True + assert float(report["nearby_error"]) < float(report["accuracy_limit"]) + assert float(report["old_mesh_error"]) > float(report["accuracy_limit"]) + assert float(report["selected_mesh_error"]) < float(report["accuracy_limit"]) + assert int(report["selected_steps"]) > int(report["nominal_steps"])