From cc0e47b7ace9c1aaf32beb3514404de1b1530c21 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" <72512262+jc-macdonald@users.noreply.github.com> Date: Thu, 9 Apr 2026 08:39:41 -0400 Subject: [PATCH 1/5] feat: widen MethodName to all 9 solvers, extend validate_system() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - config.py: MethodName now includes implicit-euler, trapezoidal, bdf2, ros2 - validate_system(): IMEX + missing operators → missing_operators warning - validate_system(): implicit/Rosenbrock + missing jacobian → missing_jacobian warning - Existing state_change check preserved Refs #41, #44 --- .../flepimop2/engine/op_engine/__init__.py | 48 +++++++++++++++++-- .../src/flepimop2/engine/op_engine/config.py | 14 ++++-- 2 files changed, 55 insertions(+), 7 deletions(-) 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..1a19764 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py @@ -105,10 +105,16 @@ class OpEngineFlepimop2Engine(ModuleModel, EngineABC): state_change: StateChangeEnum config: OpEngineEngineConfig = Field(default_factory=OpEngineEngineConfig) + _IMPLICIT_METHODS: frozenset[str] = frozenset( + {"implicit-euler", "trapezoidal", "bdf2", "ros2"}, + ) + 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 +122,43 @@ def validate_system(self, system: SystemABC) -> list[ValidationIssue] | None: f"'{system.state_change}'." ), kind="incompatible_system", + ), + ) + + method = self.config.method + is_imex = method.startswith("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 in self._IMPLICIT_METHODS: + 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, 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..679e547 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py @@ -37,9 +37,17 @@ class OpEngineEngineConfig(BaseModel): model_config = ConfigDict(extra="allow") - method: Literal["euler", "heun", "imex-euler", "imex-heun-tr", "imex-trbdf2"] = ( - "heun" - ) + method: Literal[ + "euler", + "heun", + "imex-euler", + "imex-heun-tr", + "imex-trbdf2", + "implicit-euler", + "trapezoidal", + "bdf2", + "ros2", + ] = "heun" adaptive: bool = False strict: bool = True rtol: float = Field(default=1e-6, ge=0.0) From c5ed33ac4e22f74dad345ad9ea7879f206f47a96 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" <72512262+jc-macdonald@users.noreply.github.com> Date: Thu, 9 Apr 2026 08:41:00 -0400 Subject: [PATCH 2/5] feat: wire system.option('jacobian') into RunConfig for implicit methods For implicit-euler, trapezoidal, bdf2, and ros2, the provider now retrieves the Jacobian callable from system.option('jacobian') and passes it through to CoreSolver via RunConfig. Refs #41, #44 --- .../src/flepimop2/engine/op_engine/__init__.py | 5 +++++ 1 file changed, 5 insertions(+) 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 1a19764..fa8c2eb 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py @@ -193,6 +193,11 @@ def run( ) raise ValueError(msg) + if run_cfg.method in self._IMPLICIT_METHODS: + 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) From 5e08cc12c367dc224e4a4424f2d5b8ae186876c8 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" <72512262+jc-macdonald@users.noreply.github.com> Date: Thu, 9 Apr 2026 08:43:18 -0400 Subject: [PATCH 3/5] test: validation + jacobian wiring tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - IMEX + missing operators → missing_operators (3 tests) - implicit/Rosenbrock + missing jacobian → missing_jacobian (4 parametrized) - implicit + jacobian provided → no warning (4 parametrized) - explicit methods → no extra issues - implicit-euler run with system jacobian verifies end-to-end wiring Refs #41, #44 --- flepimop2-op_engine/tests/test_engine.py | 134 +++++++++++++++++++++++ 1 file changed, 134 insertions(+) diff --git a/flepimop2-op_engine/tests/test_engine.py b/flepimop2-op_engine/tests/test_engine.py index ddb3a48..ad71f1f 100644 --- a/flepimop2-op_engine/tests/test_engine.py +++ b/flepimop2-op_engine/tests/test_engine.py @@ -162,6 +162,113 @@ 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="flow", + config={"method": "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="flow", + config={"method": "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="flow", + config={ + "method": "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", + ["implicit-euler", "trapezoidal", "bdf2", "ros2"], +) +def test_validate_implicit_missing_jacobian(method: str) -> None: + """Implicit/Rosenbrock method without system jacobian → missing_jacobian.""" + engine = OpEngineFlepimop2Engine( + state_change="flow", + config={"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", + ["implicit-euler", "trapezoidal", "bdf2", "ros2"], +) +def test_validate_implicit_with_jacobian(method: str) -> None: + """Implicit/Rosenbrock method + system provides jacobian → no warning.""" + engine = OpEngineFlepimop2Engine( + state_change="flow", + config={"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 ("euler", "heun"): + engine = OpEngineFlepimop2Engine( + state_change="flow", + config={"method": method}, + ) + system = _GoodSystem() + system.options = {} + assert engine.validate_system(system) is None + + # ----------------------------------------------------------------------------- # Bind API integration # ----------------------------------------------------------------------------- @@ -188,3 +295,30 @@ def tracking_bind( 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="flow", + config={"method": "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 From 25d3dbcb602f84e83b4de0808fc4fb521c470678 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" <72512262+jc-macdonald@users.noreply.github.com> Date: Thu, 9 Apr 2026 08:46:25 -0400 Subject: [PATCH 4/5] fix: resolve mypy --strict errors in test_engine.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add py.typed marker to provider package - Use StateChangeEnum.FLOW/DELTA instead of string literals - Use OpEngineEngineConfig() instead of raw dicts - Narrow type: ignore[assignment] → type: ignore[method-assign] --- .../src/flepimop2/engine/op_engine/py.typed | 0 flepimop2-op_engine/tests/test_engine.py | 57 ++++++++++--------- 2 files changed, 29 insertions(+), 28 deletions(-) create mode 100644 flepimop2-op_engine/src/flepimop2/engine/op_engine/py.typed 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_engine.py b/flepimop2-op_engine/tests/test_engine.py index ad71f1f..a62fcc1 100644 --- a/flepimop2-op_engine/tests/test_engine.py +++ b/flepimop2-op_engine/tests/test_engine.py @@ -8,8 +8,9 @@ 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 if TYPE_CHECKING: from flepimop2.typing import IdentifierString, SystemProtocol @@ -37,7 +38,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 +59,7 @@ class _DeltaSystem(_GoodSystem): """SystemABC implementation with incompatible state_change.""" module = "flepimop2.system.test_delta" - state_change = "delta" + state_change = StateChangeEnum.DELTA # ----------------------------------------------------------------------------- @@ -68,7 +69,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 +82,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 +103,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 +125,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 +139,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 +153,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 @@ -170,8 +171,8 @@ def test_validate_system_checks_state_change() -> None: def test_validate_imex_missing_operators() -> None: """IMEX method without operators in config or system → missing_operators.""" engine = OpEngineFlepimop2Engine( - state_change="flow", - config={"method": "imex-euler"}, + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method="imex-euler"), ) system = _GoodSystem() system.options = {} @@ -185,8 +186,8 @@ def test_validate_imex_missing_operators() -> None: def test_validate_imex_system_provides_operators() -> None: """IMEX method + system.option('operators') provided → no operator warning.""" engine = OpEngineFlepimop2Engine( - state_change="flow", - config={"method": "imex-euler"}, + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method="imex-euler"), ) system = _GoodSystem() # _GoodSystem already has operators in options @@ -198,13 +199,13 @@ def test_validate_imex_system_provides_operators() -> None: def test_validate_imex_config_provides_operators() -> None: """IMEX method + operators in engine config → no operator warning.""" engine = OpEngineFlepimop2Engine( - state_change="flow", - config={ - "method": "imex-euler", - "operators": { + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig( + method="imex-euler", + operators={ "default": [np.eye(1).tolist(), np.eye(1).tolist()], }, - }, + ), ) system = _GoodSystem() system.options = {} @@ -225,8 +226,8 @@ def test_validate_imex_config_provides_operators() -> None: def test_validate_implicit_missing_jacobian(method: str) -> None: """Implicit/Rosenbrock method without system jacobian → missing_jacobian.""" engine = OpEngineFlepimop2Engine( - state_change="flow", - config={"method": method}, + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=method), # type: ignore[arg-type] ) system = _GoodSystem() # no "jacobian" in system.options @@ -244,8 +245,8 @@ def test_validate_implicit_missing_jacobian(method: str) -> None: def test_validate_implicit_with_jacobian(method: str) -> None: """Implicit/Rosenbrock method + system provides jacobian → no warning.""" engine = OpEngineFlepimop2Engine( - state_change="flow", - config={"method": method}, + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=method), # type: ignore[arg-type] ) system = _GoodSystem() system.options = { @@ -261,8 +262,8 @@ def test_validate_explicit_no_extra_issues() -> None: """Explicit methods do not trigger operator or jacobian warnings.""" for method in ("euler", "heun"): engine = OpEngineFlepimop2Engine( - state_change="flow", - config={"method": method}, + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method=method), ) system = _GoodSystem() system.options = {} @@ -276,7 +277,7 @@ def test_validate_explicit_no_extra_issues() -> 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 @@ -288,7 +289,7 @@ 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) @@ -305,8 +306,8 @@ def tracking_bind( def test_run_implicit_method_uses_system_jacobian() -> None: """Implicit method retrieves jacobian from system.option and runs.""" engine = OpEngineFlepimop2Engine( - state_change="flow", - config={"method": "implicit-euler"}, + state_change=StateChangeEnum.FLOW, + config=OpEngineEngineConfig(method="implicit-euler"), ) system = _GoodSystem() From d3884fe9f87e9802f5ba69c368ba13715186e06b Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" <72512262+jc-macdonald@users.noreply.github.com> Date: Thu, 9 Apr 2026 12:32:01 -0400 Subject: [PATCH 5/5] refactor: replace MethodName literal with SolverMethod StrEnum Address review feedback from @pearsonca: method names are now a proper StrEnum (SolverMethod) with is_imex/is_implicit properties, replacing string literals and the _IMPLICIT_METHODS frozenset. - Add SolverMethod(StrEnum) to config.py with all 9 methods - Add is_imex and is_implicit properties on enum - Move _IMPLICIT_METHODS frozenset to config.py module level - Remove _IMPLICIT_METHODS class variable from engine - Update all tests to use SolverMethod enum members - Drop type: ignore[arg-type] from parametrized tests --- .../flepimop2/engine/op_engine/__init__.py | 22 ++++---- .../src/flepimop2/engine/op_engine/config.py | 56 ++++++++++++++----- flepimop2-op_engine/tests/test_config.py | 22 ++++---- flepimop2-op_engine/tests/test_engine.py | 38 +++++++++---- 4 files changed, 92 insertions(+), 46 deletions(-) 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 fa8c2eb..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 @@ -105,10 +110,6 @@ class OpEngineFlepimop2Engine(ModuleModel, EngineABC): state_change: StateChangeEnum config: OpEngineEngineConfig = Field(default_factory=OpEngineEngineConfig) - _IMPLICIT_METHODS: frozenset[str] = frozenset( - {"implicit-euler", "trapezoidal", "bdf2", "ros2"}, - ) - def validate_system(self, system: SystemABC) -> list[ValidationIssue] | None: """Validate system compatibility with engine config.""" issues: list[ValidationIssue] = [] @@ -126,7 +127,7 @@ def validate_system(self, system: SystemABC) -> list[ValidationIssue] | None: ) method = self.config.method - is_imex = method.startswith("imex-") + is_imex = method.is_imex if is_imex and not _has_operator_specs( _coerce_operator_specs(self.config.operators), @@ -144,7 +145,7 @@ def validate_system(self, system: SystemABC) -> list[ValidationIssue] | None: ), ) - if method in self._IMPLICIT_METHODS: + if method.is_implicit: jac = system.option("jacobian", None) if jac is None: issues.append( @@ -177,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): @@ -193,7 +195,7 @@ def run( ) raise ValueError(msg) - if run_cfg.method in self._IMPLICIT_METHODS: + if method.is_implicit: jacobian = system.option("jacobian", None) if callable(jacobian): run_cfg = replace(run_cfg, jacobian=jacobian) @@ -225,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 679e547..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,22 +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", - "implicit-euler", - "trapezoidal", - "bdf2", - "ros2", - ] = "heun" + method: SolverMethod = SolverMethod.HEUN adaptive: bool = False strict: bool = True rtol: float = Field(default=1e-6, ge=0.0) @@ -99,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/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 a62fcc1..a9cb543 100644 --- a/flepimop2-op_engine/tests/test_engine.py +++ b/flepimop2-op_engine/tests/test_engine.py @@ -10,7 +10,11 @@ from flepimop2.system.abc import SystemABC from flepimop2.typing import StateChangeEnum -from flepimop2.engine.op_engine import OpEngineEngineConfig, OpEngineFlepimop2Engine +from flepimop2.engine.op_engine import ( + OpEngineEngineConfig, + OpEngineFlepimop2Engine, + SolverMethod, +) if TYPE_CHECKING: from flepimop2.typing import IdentifierString, SystemProtocol @@ -172,7 +176,7 @@ 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="imex-euler"), + config=OpEngineEngineConfig(method=SolverMethod.IMEX_EULER), ) system = _GoodSystem() system.options = {} @@ -187,7 +191,7 @@ 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="imex-euler"), + config=OpEngineEngineConfig(method=SolverMethod.IMEX_EULER), ) system = _GoodSystem() # _GoodSystem already has operators in options @@ -201,7 +205,7 @@ def test_validate_imex_config_provides_operators() -> None: engine = OpEngineFlepimop2Engine( state_change=StateChangeEnum.FLOW, config=OpEngineEngineConfig( - method="imex-euler", + method=SolverMethod.IMEX_EULER, operators={ "default": [np.eye(1).tolist(), np.eye(1).tolist()], }, @@ -221,13 +225,18 @@ def test_validate_imex_config_provides_operators() -> None: @pytest.mark.parametrize( "method", - ["implicit-euler", "trapezoidal", "bdf2", "ros2"], + [ + SolverMethod.IMPLICIT_EULER, + SolverMethod.TRAPEZOIDAL, + SolverMethod.BDF2, + SolverMethod.ROS2, + ], ) -def test_validate_implicit_missing_jacobian(method: str) -> None: +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), # type: ignore[arg-type] + config=OpEngineEngineConfig(method=method), ) system = _GoodSystem() # no "jacobian" in system.options @@ -240,13 +249,18 @@ def test_validate_implicit_missing_jacobian(method: str) -> None: @pytest.mark.parametrize( "method", - ["implicit-euler", "trapezoidal", "bdf2", "ros2"], + [ + SolverMethod.IMPLICIT_EULER, + SolverMethod.TRAPEZOIDAL, + SolverMethod.BDF2, + SolverMethod.ROS2, + ], ) -def test_validate_implicit_with_jacobian(method: str) -> None: +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), # type: ignore[arg-type] + config=OpEngineEngineConfig(method=method), ) system = _GoodSystem() system.options = { @@ -260,7 +274,7 @@ def test_validate_implicit_with_jacobian(method: str) -> None: def test_validate_explicit_no_extra_issues() -> None: """Explicit methods do not trigger operator or jacobian warnings.""" - for method in ("euler", "heun"): + for method in (SolverMethod.EULER, SolverMethod.HEUN): engine = OpEngineFlepimop2Engine( state_change=StateChangeEnum.FLOW, config=OpEngineEngineConfig(method=method), @@ -307,7 +321,7 @@ 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="implicit-euler"), + config=OpEngineEngineConfig(method=SolverMethod.IMPLICIT_EULER), ) system = _GoodSystem()