diff --git a/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py b/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py index f7817b4..1216cd3 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py @@ -17,7 +17,12 @@ ) from op_engine.model_core import ModelCore, ModelCoreOptions -from .config import OpEngineEngineConfig, _coerce_operator_specs, _has_operator_specs +from .config import ( + OpEngineEngineConfig, + SolverMethod, + _coerce_operator_specs, + _has_operator_specs, +) if TYPE_CHECKING: from collections.abc import Callable @@ -106,9 +111,11 @@ class OpEngineFlepimop2Engine(ModuleModel, EngineABC): config: OpEngineEngineConfig = Field(default_factory=OpEngineEngineConfig) def validate_system(self, system: SystemABC) -> list[ValidationIssue] | None: - """Validate system compatibility against the engine state-change mode.""" + """Validate system compatibility with engine config.""" + issues: list[ValidationIssue] = [] + if system.state_change != self.state_change: - return [ + issues.append( ValidationIssue( msg=( f"Engine state change type, '{self.state_change}', is not " @@ -116,9 +123,43 @@ def validate_system(self, system: SystemABC) -> list[ValidationIssue] | None: f"'{system.state_change}'." ), kind="incompatible_system", + ), + ) + + method = self.config.method + is_imex = method.is_imex + + if is_imex and not _has_operator_specs( + _coerce_operator_specs(self.config.operators), + ): + sys_ops = system.option("operators", None) + if not _has_operator_specs(_coerce_operator_specs(sys_ops)): + issues.append( + ValidationIssue( + msg=( + f"IMEX method '{method}' requires operator matrices, " + "but neither the engine config nor " + "system.option('operators') provides them." + ), + kind="missing_operators", + ), ) - ] - return None + + if method.is_implicit: + jac = system.option("jacobian", None) + if jac is None: + issues.append( + ValidationIssue( + msg=( + f"Implicit/Rosenbrock method '{method}' requires a " + "Jacobian callable, but system.option('jacobian') " + "is not provided." + ), + kind="missing_jacobian", + ), + ) + + return issues or None def run( self, @@ -137,7 +178,8 @@ def run( n_state = int(y0.size) run_cfg = self.config.to_run_config() - is_imex = run_cfg.method.startswith("imex-") + method = self.config.method + is_imex = method.is_imex operators = run_cfg.operators if is_imex and not _has_operator_specs(operators): @@ -153,6 +195,11 @@ def run( ) raise ValueError(msg) + if method.is_implicit: + jacobian = system.option("jacobian", None) + if callable(jacobian): + run_cfg = replace(run_cfg, jacobian=jacobian) + operator_axis = self.config.operator_axis if operator_axis == "state": system_axis = system.option("operator_axis", None) @@ -180,4 +227,4 @@ def run( return np.asarray(np.column_stack((times, states)), dtype=np.float64) -__all__ = ["OpEngineEngineConfig", "OpEngineFlepimop2Engine"] +__all__ = ["OpEngineEngineConfig", "OpEngineFlepimop2Engine", "SolverMethod"] diff --git a/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py b/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py index b7233ec..08254a8 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py @@ -2,7 +2,8 @@ from __future__ import annotations -from typing import Any, Literal +from enum import StrEnum +from typing import Any from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -32,14 +33,46 @@ def _coerce_operator_specs(specs: object) -> OperatorSpecs | None: return None +class SolverMethod(StrEnum): + """Solver method identifiers for op_engine integration.""" + + EULER = "euler" + HEUN = "heun" + IMEX_EULER = "imex-euler" + IMEX_HEUN_TR = "imex-heun-tr" + IMEX_TRBDF2 = "imex-trbdf2" + IMPLICIT_EULER = "implicit-euler" + TRAPEZOIDAL = "trapezoidal" + BDF2 = "bdf2" + ROS2 = "ros2" + + @property + def is_imex(self) -> bool: + """Whether this method is an IMEX method.""" + return self.value.startswith("imex-") + + @property + def is_implicit(self) -> bool: + """Whether this method requires a Jacobian.""" + return self in _IMPLICIT_METHODS + + +_IMPLICIT_METHODS: frozenset[SolverMethod] = frozenset( + { + SolverMethod.IMPLICIT_EULER, + SolverMethod.TRAPEZOIDAL, + SolverMethod.BDF2, + SolverMethod.ROS2, + }, +) + + class OpEngineEngineConfig(BaseModel): """Configuration schema for op_engine when used as a flepimop2 engine.""" model_config = ConfigDict(extra="allow") - method: Literal["euler", "heun", "imex-euler", "imex-heun-tr", "imex-trbdf2"] = ( - "heun" - ) + method: SolverMethod = SolverMethod.HEUN adaptive: bool = False strict: bool = True rtol: float = Field(default=1e-6, ge=0.0) @@ -91,4 +124,9 @@ def to_run_config(self) -> RunConfig: ) -__all__ = ["OpEngineEngineConfig", "_coerce_operator_specs", "_has_operator_specs"] +__all__ = [ + "OpEngineEngineConfig", + "SolverMethod", + "_coerce_operator_specs", + "_has_operator_specs", +] diff --git a/flepimop2-op_engine/src/flepimop2/engine/op_engine/py.typed b/flepimop2-op_engine/src/flepimop2/engine/op_engine/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/flepimop2-op_engine/tests/test_config.py b/flepimop2-op_engine/tests/test_config.py index 20e06e8..910acff 100644 --- a/flepimop2-op_engine/tests/test_config.py +++ b/flepimop2-op_engine/tests/test_config.py @@ -8,7 +8,7 @@ from op_engine.core_solver import OperatorSpecs, RunConfig # noqa: E402 from pydantic import ValidationError # noqa: E402 -from flepimop2.engine.op_engine import OpEngineEngineConfig # noqa: E402 +from flepimop2.engine.op_engine import OpEngineEngineConfig, SolverMethod # noqa: E402 def _has_any_operator_specs(specs: OperatorSpecs) -> bool: @@ -47,7 +47,7 @@ def test_engine_config_defaults_to_run_config() -> None: def test_engine_config_round_trips_selected_fields() -> None: """Engine config round-trips selected fields correctly.""" cfg = OpEngineEngineConfig( - method="euler", + method=SolverMethod.EULER, adaptive=True, strict=False, rtol=1e-4, @@ -77,7 +77,7 @@ def test_engine_config_round_trips_selected_fields() -> None: def test_engine_config_allows_unknown_fields() -> None: """Engine config should allow unknown fields without error.""" cfg = OpEngineEngineConfig( # type: ignore[call-arg] - method="heun", + method=SolverMethod.HEUN, adaptive=False, some_unknown_key=123, nested_unknown={"a": 1}, @@ -97,7 +97,7 @@ def test_engine_config_gamma_bounds_validation() -> None: """Engine config validates gamma bounds for imex-trbdf2 method.""" # IMEX requires operators at parse-time. cfg = OpEngineEngineConfig( - method="imex-trbdf2", + method=SolverMethod.IMEX_TRBDF2, gamma=0.6, operators={"default": "sentinel"}, ) @@ -108,28 +108,28 @@ def test_engine_config_gamma_bounds_validation() -> None: # invalid: gamma must be in (0, 1) with pytest.raises(ValidationError): OpEngineEngineConfig( - method="imex-trbdf2", + method=SolverMethod.IMEX_TRBDF2, gamma=0.0, operators={"default": "sentinel"}, ) with pytest.raises(ValidationError): OpEngineEngineConfig( - method="imex-trbdf2", + method=SolverMethod.IMEX_TRBDF2, gamma=1.0, operators={"default": "sentinel"}, ) with pytest.raises(ValidationError): OpEngineEngineConfig( - method="imex-trbdf2", + method=SolverMethod.IMEX_TRBDF2, gamma=-0.1, operators={"default": "sentinel"}, ) with pytest.raises(ValidationError): OpEngineEngineConfig( - method="imex-trbdf2", + method=SolverMethod.IMEX_TRBDF2, gamma=1.1, operators={"default": "sentinel"}, ) @@ -137,7 +137,7 @@ def test_engine_config_gamma_bounds_validation() -> None: def test_engine_config_imex_allows_deferred_operators() -> None: """IMEX methods may omit operators to defer to system options at runtime.""" - cfg = OpEngineEngineConfig(method="imex-euler") + cfg = OpEngineEngineConfig(method=SolverMethod.IMEX_EULER) run = cfg.to_run_config() assert run.method == "imex-euler" assert isinstance(run.operators, OperatorSpecs) @@ -147,13 +147,13 @@ def test_engine_config_imex_allows_deferred_operators() -> None: def test_engine_config_imex_rejects_explicitly_empty_operator_block() -> None: """Providing an empty operator block should raise validation errors.""" with pytest.raises(ValidationError): - OpEngineEngineConfig(method="imex-heun-tr", operators={}) + OpEngineEngineConfig(method=SolverMethod.IMEX_HEUN_TR, operators={}) def test_engine_config_imex_with_operators_still_valid() -> None: """Providing IMEX operators explicitly should still validate.""" cfg = OpEngineEngineConfig( - method="imex-euler", + method=SolverMethod.IMEX_EULER, operators={"default": "sentinel"}, ) run = cfg.to_run_config() diff --git a/flepimop2-op_engine/tests/test_engine.py b/flepimop2-op_engine/tests/test_engine.py index ddb3a48..a9cb543 100644 --- a/flepimop2-op_engine/tests/test_engine.py +++ b/flepimop2-op_engine/tests/test_engine.py @@ -8,8 +8,13 @@ import numpy as np import pytest from flepimop2.system.abc import SystemABC +from flepimop2.typing import StateChangeEnum -from flepimop2.engine.op_engine import OpEngineFlepimop2Engine +from flepimop2.engine.op_engine import ( + OpEngineEngineConfig, + OpEngineFlepimop2Engine, + SolverMethod, +) if TYPE_CHECKING: from flepimop2.typing import IdentifierString, SystemProtocol @@ -37,7 +42,7 @@ class _GoodSystem(SystemABC): """SystemABC implementation exposing a valid stepper via bind().""" module = "flepimop2.system.test_good" - state_change = "flow" + state_change = StateChangeEnum.FLOW def __init__(self) -> None: super().__init__() @@ -58,7 +63,7 @@ class _DeltaSystem(_GoodSystem): """SystemABC implementation with incompatible state_change.""" module = "flepimop2.system.test_delta" - state_change = "delta" + state_change = StateChangeEnum.DELTA # ----------------------------------------------------------------------------- @@ -68,7 +73,7 @@ class _DeltaSystem(_GoodSystem): def test_public_engine_wrapper_defines_module() -> None: """Public engine wrapper satisfies flepimop2's concrete module contract.""" - engine = OpEngineFlepimop2Engine(state_change="flow") + engine = OpEngineFlepimop2Engine(state_change=StateChangeEnum.FLOW) assert isinstance(engine, OpEngineFlepimop2Engine) assert engine.module == "flepimop2.engine.op_engine" @@ -81,7 +86,7 @@ def test_public_engine_wrapper_defines_module() -> None: def test_engine_run_basic_shape_and_dtype() -> None: """Engine returns correctly shaped float64 output array.""" - engine = OpEngineFlepimop2Engine(state_change="flow") + engine = OpEngineFlepimop2Engine(state_change=StateChangeEnum.FLOW) system = _GoodSystem() times = np.array([0.0, 0.5, 1.0], dtype=np.float64) @@ -102,7 +107,7 @@ def test_engine_run_identity_rhs_behavior() -> None: This test validates wiring correctness, not numerical accuracy. """ - engine = OpEngineFlepimop2Engine(state_change="flow") + engine = OpEngineFlepimop2Engine(state_change=StateChangeEnum.FLOW) system = _GoodSystem() times = np.array([0.0, 0.1, 0.2], dtype=np.float64) @@ -124,7 +129,7 @@ def test_engine_run_identity_rhs_behavior() -> None: def test_engine_rejects_non_increasing_times() -> None: """Engine rejects non-strictly-increasing time grids.""" - engine = OpEngineFlepimop2Engine(state_change="flow") + engine = OpEngineFlepimop2Engine(state_change=StateChangeEnum.FLOW) system = _GoodSystem() times = np.array([0.0, 0.0, 1.0], dtype=np.float64) @@ -138,7 +143,7 @@ def test_engine_rejects_non_increasing_times() -> None: def test_engine_rejects_non_1d_initial_state() -> None: """Engine rejects non-1D initial state arrays.""" - engine = OpEngineFlepimop2Engine(state_change="flow") + engine = OpEngineFlepimop2Engine(state_change=StateChangeEnum.FLOW) system = _GoodSystem() times = np.array([0.0, 1.0], dtype=np.float64) @@ -152,7 +157,7 @@ def test_engine_rejects_non_1d_initial_state() -> None: def test_validate_system_checks_state_change() -> None: """Engine validates state_change compatibility via validate_system.""" - engine = OpEngineFlepimop2Engine(state_change="flow") + engine = OpEngineFlepimop2Engine(state_change=StateChangeEnum.FLOW) good = _GoodSystem() assert engine.validate_system(good) is None @@ -162,6 +167,123 @@ def test_validate_system_checks_state_change() -> None: assert issues[0].kind == "incompatible_system" +# ----------------------------------------------------------------------------- +# validate_system: IMEX + missing operators +# ----------------------------------------------------------------------------- + + +def test_validate_imex_missing_operators() -> None: + """IMEX method without operators in config or system → missing_operators.""" + engine = OpEngineFlepimop2Engine( + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=SolverMethod.IMEX_EULER), + ) + system = _GoodSystem() + system.options = {} + + issues = engine.validate_system(system) + assert issues is not None + kinds = [i.kind for i in issues] + assert "missing_operators" in kinds + + +def test_validate_imex_system_provides_operators() -> None: + """IMEX method + system.option('operators') provided → no operator warning.""" + engine = OpEngineFlepimop2Engine( + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=SolverMethod.IMEX_EULER), + ) + system = _GoodSystem() + # _GoodSystem already has operators in options + + issues = engine.validate_system(system) + assert issues is None + + +def test_validate_imex_config_provides_operators() -> None: + """IMEX method + operators in engine config → no operator warning.""" + engine = OpEngineFlepimop2Engine( + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig( + method=SolverMethod.IMEX_EULER, + operators={ + "default": [np.eye(1).tolist(), np.eye(1).tolist()], + }, + ), + ) + system = _GoodSystem() + system.options = {} + + issues = engine.validate_system(system) + assert issues is None + + +# ----------------------------------------------------------------------------- +# validate_system: implicit/Rosenbrock + missing jacobian +# ----------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "method", + [ + SolverMethod.IMPLICIT_EULER, + SolverMethod.TRAPEZOIDAL, + SolverMethod.BDF2, + SolverMethod.ROS2, + ], +) +def test_validate_implicit_missing_jacobian(method: SolverMethod) -> None: + """Implicit/Rosenbrock method without system jacobian → missing_jacobian.""" + engine = OpEngineFlepimop2Engine( + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=method), + ) + system = _GoodSystem() + # no "jacobian" in system.options + + issues = engine.validate_system(system) + assert issues is not None + kinds = [i.kind for i in issues] + assert "missing_jacobian" in kinds + + +@pytest.mark.parametrize( + "method", + [ + SolverMethod.IMPLICIT_EULER, + SolverMethod.TRAPEZOIDAL, + SolverMethod.BDF2, + SolverMethod.ROS2, + ], +) +def test_validate_implicit_with_jacobian(method: SolverMethod) -> None: + """Implicit/Rosenbrock method + system provides jacobian → no warning.""" + engine = OpEngineFlepimop2Engine( + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=method), + ) + system = _GoodSystem() + system.options = { + **(system.options or {}), + "jacobian": lambda _t, y: -np.eye(len(y)), + } + + issues = engine.validate_system(system) + assert issues is None + + +def test_validate_explicit_no_extra_issues() -> None: + """Explicit methods do not trigger operator or jacobian warnings.""" + for method in (SolverMethod.EULER, SolverMethod.HEUN): + engine = OpEngineFlepimop2Engine( + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=method), + ) + system = _GoodSystem() + system.options = {} + assert engine.validate_system(system) is None + + # ----------------------------------------------------------------------------- # Bind API integration # ----------------------------------------------------------------------------- @@ -169,7 +291,7 @@ def test_validate_system_checks_state_change() -> None: def test_engine_uses_bind_not_stepper() -> None: """Engine calls system.bind() rather than accessing system._stepper.""" - engine = OpEngineFlepimop2Engine(state_change="flow") + engine = OpEngineFlepimop2Engine(state_change=StateChangeEnum.FLOW) system = _GoodSystem() bind_called = False original_bind = system.bind @@ -181,10 +303,37 @@ def tracking_bind( bind_called = True return original_bind(params, **kwargs) - system.bind = tracking_bind # type: ignore[assignment] + system.bind = tracking_bind # type: ignore[method-assign] times = np.array([0.0, 0.1], dtype=np.float64) y0 = np.array([1.0], dtype=np.float64) engine.run(system, times, y0, {}) assert bind_called, "Engine should call system.bind()" + + +# ----------------------------------------------------------------------------- +# Jacobian wiring for implicit methods +# ----------------------------------------------------------------------------- + + +def test_run_implicit_method_uses_system_jacobian() -> None: + """Implicit method retrieves jacobian from system.option and runs.""" + engine = OpEngineFlepimop2Engine( + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=SolverMethod.IMPLICIT_EULER), + ) + system = _GoodSystem() + + def neg_identity_jac(_t: float, y: np.ndarray) -> np.ndarray: + return -np.eye(len(y), dtype=np.float64) + + system.options = {**(system.options or {}), "jacobian": neg_identity_jac} + + times = np.array([0.0, 0.1, 0.2], dtype=np.float64) + y0 = np.array([1.0], dtype=np.float64) + + out = engine.run(system, times, y0, {}) + + assert out.shape == (3, 2) + assert out.dtype == np.float64