From c1a02e97913134784e179fcc427edffbe4e60137 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Mon, 9 Feb 2026 14:48:21 -0500 Subject: [PATCH 1/9] edit connector class --- .../src/flepimop2/engine/op_engine/config.py | 8 ++- .../src/flepimop2/engine/op_engine/engine.py | 18 +++-- flepimop2-op_engine/tests/test_engine.py | 69 +++++++++++++++++++ justfile | 42 +++++++---- src/op_engine/model_core.py | 5 -- 5 files changed, 115 insertions(+), 27 deletions(-) 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 e83ebf9..3bd0270 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py @@ -86,14 +86,16 @@ def to_run_config(self) -> RunConfig: fac_max=self.fac_max, ) - # Today: keep operators empty/default; later you can translate self.operators - # into OperatorSpecs() here. + # Operators are validated for presence (for IMEX) but intentionally not + # translated yet; keep RunConfig operators empty until wiring is added. + op_specs: OperatorSpecs | None = None + return RunConfig( method=self.method, adaptive=self.adaptive, strict=self.strict, adaptive_cfg=adaptive_cfg, dt_controller=dt_controller, - operators=OperatorSpecs(), + operators=op_specs or OperatorSpecs(), gamma=self.gamma, ) diff --git a/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py b/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py index 4173dca..addd2f7 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py @@ -226,14 +226,22 @@ def run( msg = "system does not expose a valid flepimop2 SystemProtocol stepper" raise TypeError(msg) - rhs = _rhs_from_stepper(stepper, params=params, n_state=n_state) + # Merge any precomputed mixing kernels exposed by op_system connector + merged_params: dict[IdentifierString, object] = {} + kernels = getattr(system, "mixing_kernels", None) + if isinstance(kernels, dict): + merged_params.update(kernels) + merged_params.update(params) + + rhs = _rhs_from_stepper(stepper, params=merged_params, n_state=n_state) core = _make_core(times, y0) - # Today: operators are not yet wired through the adapter. - # Future: translate self.config.operators into op_engine OperatorSpecs and - # pass them here (and/or via run_cfg) as appropriate. - solver = CoreSolver(core, operators=None, operator_axis="state") + op_default = None + if self.config.operators: + op_default = self.config.operators.get("default") + + solver = CoreSolver(core, operators=op_default, operator_axis="state") run_cfg = self.config.to_run_config() solver.run(rhs, config=run_cfg) diff --git a/flepimop2-op_engine/tests/test_engine.py b/flepimop2-op_engine/tests/test_engine.py index bfa319d..a0dddc1 100644 --- a/flepimop2-op_engine/tests/test_engine.py +++ b/flepimop2-op_engine/tests/test_engine.py @@ -50,6 +50,34 @@ def __init__(self) -> None: self._stepper: object = object() +class _KernelStepper: + """Stepper that uses a kernel parameter to scale a constant RHS.""" + + def __call__( + self, time: np.float64, state: np.ndarray, **params: object + ) -> np.ndarray: + _ = time + _ = state + k = float(params.get("K", 0.0)) + return np.asarray([k], dtype=np.float64) + + +class _KernelSystem: + """System exposing a stepper and precomputed mixing_kernels.""" + + def __init__(self, k: float) -> None: + self._stepper: SystemProtocol = _KernelStepper() + self.mixing_kernels = {"K": k} + + +class _ImexSystem: + """System exposing an identity stepper for IMEX tests.""" + + def __init__(self, n: int) -> None: + self._stepper: SystemProtocol = _GoodStepper() + self.n = n + + # ----------------------------------------------------------------------------- # Engine construction # ----------------------------------------------------------------------------- @@ -107,6 +135,47 @@ def test_engine_run_identity_rhs_behavior() -> None: assert state_values[2] >= state_values[1] +def test_engine_passes_mixing_kernels_into_params() -> None: + """mixing_kernels from the system are merged into RHS params.""" + engine = _OpEngineFlepimop2EngineImpl() + system = cast("SystemABC", _KernelSystem(k=2.5)) + + times = np.array([0.0, 1.0], dtype=np.float64) + y0 = np.array([1.0], dtype=np.float64) + + params: dict[IdentifierString, object] = {} + + out = engine.run(system, times, y0, params) + + # dy/dt = K = 2.5, Heun with dt=1.0 gives y1 = 1 + 0.5*(K+K) = 3.5 + np.testing.assert_allclose(out[-1, 1], 3.5, rtol=1e-12, atol=0.0) + + +def test_engine_imex_identity_with_identity_ops() -> None: + """IMEX path accepts operator specs and runs with identity operators.""" + engine = _OpEngineFlepimop2EngineImpl( + config={ + "method": "imex-euler", + "operators": { + "default": (np.eye(1, dtype=np.float64), np.eye(1, dtype=np.float64)), + }, + "adaptive": False, + } + ) + system = cast("SystemABC", _ImexSystem(n=1)) + + times = np.array([0.0, 0.5, 1.0], dtype=np.float64) + y0 = np.array([1.0], dtype=np.float64) + + params: dict[IdentifierString, object] = {} + + out = engine.run(system, times, y0, params) + + assert out.shape == (3, 2) + # Identity RHS dy/dt = y; implicit Euler with identity L/R behaves like explicit. + assert np.all(np.isfinite(out)) + + # ----------------------------------------------------------------------------- # Error handling # ----------------------------------------------------------------------------- diff --git a/justfile b/justfile index 400cad3..51dc959 100644 --- a/justfile +++ b/justfile @@ -1,16 +1,17 @@ # Run all default tasks for local development -default: format check pytest mypy docs +default: format check provider-sync pytest-core pytest-provider-nosync mypy-core mypy-provider-nosync # ------------------------------------------------- -# Formatting +# Formatting / linting # ------------------------------------------------- +# Tool-only execution (does not require resolving project deps). format: - uv run ruff format --preview + uvx ruff format --preview check: - uv run ruff check --preview --fix + uvx ruff check --preview --fix # ------------------------------------------------- @@ -18,7 +19,6 @@ check: # ------------------------------------------------- # Create/refresh the provider venv and install all deps (including dev group). -# Run this once before running provider pytest/mypy if you aren't using `ci`. provider-sync: cd flepimop2-op_engine && uv venv --clear cd flepimop2-op_engine && uv sync --dev @@ -31,11 +31,18 @@ provider-sync: pytest-core: uv run pytest --doctest-modules -# Assumes `flepimop2-op_engine/.venv` already exists (run `just provider-sync` or `just ci` first). +# Ensure provider venv exists before running provider tests. pytest-provider: provider-sync cd flepimop2-op_engine && .venv/bin/python -m pytest --doctest-modules -pytest: pytest-core pytest-provider +# Provider tests assuming the venv already exists (used by ci for speed). +pytest-provider-nosync: + cd flepimop2-op_engine && .venv/bin/python -m pytest --doctest-modules + +pytest: + just provider-sync + just pytest-core + just pytest-provider-nosync # ------------------------------------------------- @@ -45,13 +52,18 @@ pytest: pytest-core pytest-provider mypy-core: uv run mypy --strict src/op_engine -# Assumes `flepimop2-op_engine/.venv` already exists (run `just provider-sync` or `just ci` first). +# Ensure provider venv exists before running provider mypy. mypy-provider: provider-sync cd flepimop2-op_engine && .venv/bin/python -m mypy --strict src/flepimop2 +# Provider mypy assuming the venv already exists (used by ci for speed). +mypy-provider-nosync: + cd flepimop2-op_engine && .venv/bin/python -m mypy --strict src/flepimop2 + mypy: + just provider-sync just mypy-core - just mypy-provider + just mypy-provider-nosync # ------------------------------------------------- @@ -59,11 +71,13 @@ mypy: # ------------------------------------------------- ci: - uv run ruff format --preview --check - uv run ruff check --preview --no-fix + uvx ruff format --preview --check + uvx ruff check --preview --no-fix just provider-sync - just pytest - just mypy + just pytest-core + just pytest-provider-nosync + just mypy-core + just mypy-provider-nosync # ------------------------------------------------- @@ -81,7 +95,7 @@ clean: # Build API reference for the documentation using `mkdocstrings` api-reference: - uv run scripts/api-reference.py + uv run scripts/api-reference.py # Build the documentation using `mkdocs` docs: api-reference diff --git a/src/op_engine/model_core.py b/src/op_engine/model_core.py index 9e199ad..10e3078 100644 --- a/src/op_engine/model_core.py +++ b/src/op_engine/model_core.py @@ -288,11 +288,6 @@ def reshape_for_axis_solve( Raises: ValueError: if x does not have shape state_shape. - - Contract: - - x2d has shape (axis_len, batch) - - batch is the product of all non-axis dimensions - - unreshape_from_axis_solve(inverse) reconstructs exactly. """ x_arr = np.asarray(x, dtype=self.dtype) if x_arr.shape != self.state_shape: From c37c24cb6d5586a3b1257d3e964453b0d7b044b2 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Mon, 9 Feb 2026 14:57:30 -0500 Subject: [PATCH 2/9] add uv to pyproject --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 5e5adcf..6e31efa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,3 +91,6 @@ exclude = ["examples/"] [tool.hatch.build.targets.wheel] packages = ["src/op_engine"] include = ["src/op_engine/py.typed"] + +[tool.uv] +version = "0.4.17" From 61aedc83b1c8e271bc7b0b2a655476827947d340 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Mon, 9 Feb 2026 15:01:24 -0500 Subject: [PATCH 3/9] Create uv.toml --- uv.toml | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 uv.toml diff --git a/uv.toml b/uv.toml new file mode 100644 index 0000000..71bb649 --- /dev/null +++ b/uv.toml @@ -0,0 +1,2 @@ +[tool.uv] +version = "0.4.17" From 91f3826e54b42e33f77c228dd70a36351c2b9ea1 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" <72512262+MacdonaldJoshuaCaleb@users.noreply.github.com> Date: Mon, 9 Feb 2026 15:03:08 -0500 Subject: [PATCH 4/9] Change 'version' to 'required-version' in uv.toml --- uv.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/uv.toml b/uv.toml index 71bb649..0863f81 100644 --- a/uv.toml +++ b/uv.toml @@ -1,2 +1 @@ -[tool.uv] -version = "0.4.17" +required-version = "0.4.17" From 36a6a94fee548f00a4a384ce755bb2286ce402b3 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Mon, 9 Feb 2026 15:21:38 -0500 Subject: [PATCH 5/9] Update README.md --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 2be121b..e2538cc 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ Operator-Partitioned Engine (OP Engine) is a lightweight multiphysics solver cor - Strong typing, minimal dependencies (NumPy + SciPy for implicit paths). - Separates state/time management (`ModelCore`) from stepping logic (`CoreSolver`). - Optional adapters (e.g., flepimop2) without affecting the core API. +- IMEX paths accept externally supplied operator tuples; defaults remain explicit-only. ## Installation @@ -46,7 +47,7 @@ solution = core.state_array # shape (n_timesteps, state, subgroup) - `ModelCore`: state tensor + time grid manager; supports extra axes and optional history. - `CoreSolver`: explicit and IMEX stepping; methods: `euler`, `heun`, `imex-euler`, `imex-heun-tr`, `imex-trbdf2`. - Operator utilities (`matrix_ops`): Laplacian builders, Crank–Nicolson/implicit Euler/trapezoidal operators, predictor-corrector builders, implicit solve cache, Kronecker helpers, grouped aggregation utilities. -- Adapters: optional flepimop2 integration (extra dependency) via entrypoints in the adapter package. +- Adapters: optional flepimop2 integration (extra dependency) via entrypoints in the adapter package. Adapter merges any `mixing_kernels` exposed by op_system and will consume config-supplied IMEX operator tuples when provided. ## Development From 159d2fe1bdf723e3031cb998e027f53cbc640cde Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Mon, 9 Feb 2026 15:33:41 -0500 Subject: [PATCH 6/9] Update README.md --- README.md | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index e2538cc..9f67de8 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,12 @@ Operator-Partitioned Engine (OP Engine) is a lightweight multiphysics solver cor - Optional adapters (e.g., flepimop2) without affecting the core API. - IMEX paths accept externally supplied operator tuples; defaults remain explicit-only. +## Core surface +- `ModelCore`: state/time manager; configure axes, dtype, and optional history. +- `CoreSolver`: explicit + IMEX methods (`euler`, `heun`, `imex-euler`, `imex-heun-tr`, `imex-trbdf2`); accepts `RunConfig` with `AdaptiveConfig`, `DtControllerConfig`, and `OperatorSpecs`. +- `matrix_ops`: Laplacian/Crank–Nicolson/implicit Euler/trapezoidal builders, predictor–corrector, implicit solve cache, Kronecker helpers, grouped aggregations. +- Extras: `OperatorSpecs`, `RunConfig`, `AdaptiveConfig`, `DtControllerConfig`, `Operator`, `GridGeometry`, `DiffusionConfig`. + ## Installation ```bash @@ -43,11 +49,38 @@ solver.run(rhs) # defaults to Heun/RK2 solution = core.state_array # shape (n_timesteps, state, subgroup) ``` +### IMEX with operators (tuple form) + +```python +import numpy as np +from op_engine import CoreSolver, ModelCore, OperatorSpecs + +n = 4 +times = np.linspace(0.0, 1.0, 11) +core = ModelCore(n_states=n, n_subgroups=1, time_grid=times) +core.set_initial_state(np.ones((n, 1))) + +# Identity implicit operator along state axis +L = np.eye(n) +R = np.eye(n) +ops = OperatorSpecs(default=(L, R)) + +def rhs(t, y): + return -0.1 * y + +solver = CoreSolver(core, operators=ops.default, operator_axis="state") +solver.run(rhs, config=None) # defaults: method="heun" (explicit) + +# For IMEX methods set method and operators via RunConfig: +# from op_engine.core_solver import RunConfig, AdaptiveConfig, DtControllerConfig +``` + ## Public API - `ModelCore`: state tensor + time grid manager; supports extra axes and optional history. - `CoreSolver`: explicit and IMEX stepping; methods: `euler`, `heun`, `imex-euler`, `imex-heun-tr`, `imex-trbdf2`. - Operator utilities (`matrix_ops`): Laplacian builders, Crank–Nicolson/implicit Euler/trapezoidal operators, predictor-corrector builders, implicit solve cache, Kronecker helpers, grouped aggregation utilities. -- Adapters: optional flepimop2 integration (extra dependency) via entrypoints in the adapter package. Adapter merges any `mixing_kernels` exposed by op_system and will consume config-supplied IMEX operator tuples when provided. +- Configuration helpers: `RunConfig`, `OperatorSpecs`, `AdaptiveConfig`, `DtControllerConfig` for method/IMEX/adaptive control. +- Adapters: optional flepimop2 integration (extra dependency) via entrypoints in the adapter package. Adapter merges `mixing_kernels` exposed by op_system and consumes config-supplied IMEX operator tuples when provided; operator metadata in op_system is not auto-translated yet. ## Development From 57801e5246332e36cb911df432cf68c2ccbcf8a8 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Tue, 10 Feb 2026 12:06:52 -0500 Subject: [PATCH 7/9] revert just file --- justfile | 42 ++++++++++++++---------------------------- 1 file changed, 14 insertions(+), 28 deletions(-) diff --git a/justfile b/justfile index 51dc959..400cad3 100644 --- a/justfile +++ b/justfile @@ -1,17 +1,16 @@ # Run all default tasks for local development -default: format check provider-sync pytest-core pytest-provider-nosync mypy-core mypy-provider-nosync +default: format check pytest mypy docs # ------------------------------------------------- -# Formatting / linting +# Formatting # ------------------------------------------------- -# Tool-only execution (does not require resolving project deps). format: - uvx ruff format --preview + uv run ruff format --preview check: - uvx ruff check --preview --fix + uv run ruff check --preview --fix # ------------------------------------------------- @@ -19,6 +18,7 @@ check: # ------------------------------------------------- # Create/refresh the provider venv and install all deps (including dev group). +# Run this once before running provider pytest/mypy if you aren't using `ci`. provider-sync: cd flepimop2-op_engine && uv venv --clear cd flepimop2-op_engine && uv sync --dev @@ -31,18 +31,11 @@ provider-sync: pytest-core: uv run pytest --doctest-modules -# Ensure provider venv exists before running provider tests. +# Assumes `flepimop2-op_engine/.venv` already exists (run `just provider-sync` or `just ci` first). pytest-provider: provider-sync cd flepimop2-op_engine && .venv/bin/python -m pytest --doctest-modules -# Provider tests assuming the venv already exists (used by ci for speed). -pytest-provider-nosync: - cd flepimop2-op_engine && .venv/bin/python -m pytest --doctest-modules - -pytest: - just provider-sync - just pytest-core - just pytest-provider-nosync +pytest: pytest-core pytest-provider # ------------------------------------------------- @@ -52,18 +45,13 @@ pytest: mypy-core: uv run mypy --strict src/op_engine -# Ensure provider venv exists before running provider mypy. +# Assumes `flepimop2-op_engine/.venv` already exists (run `just provider-sync` or `just ci` first). mypy-provider: provider-sync cd flepimop2-op_engine && .venv/bin/python -m mypy --strict src/flepimop2 -# Provider mypy assuming the venv already exists (used by ci for speed). -mypy-provider-nosync: - cd flepimop2-op_engine && .venv/bin/python -m mypy --strict src/flepimop2 - mypy: - just provider-sync just mypy-core - just mypy-provider-nosync + just mypy-provider # ------------------------------------------------- @@ -71,13 +59,11 @@ mypy: # ------------------------------------------------- ci: - uvx ruff format --preview --check - uvx ruff check --preview --no-fix + uv run ruff format --preview --check + uv run ruff check --preview --no-fix just provider-sync - just pytest-core - just pytest-provider-nosync - just mypy-core - just mypy-provider-nosync + just pytest + just mypy # ------------------------------------------------- @@ -95,7 +81,7 @@ clean: # Build API reference for the documentation using `mkdocstrings` api-reference: - uv run scripts/api-reference.py + uv run scripts/api-reference.py # Build the documentation using `mkdocs` docs: api-reference From e1637e92ac8748466b362303bcbc1215b004bd37 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Tue, 10 Feb 2026 12:08:36 -0500 Subject: [PATCH 8/9] remove uv.toml --- uv.toml | 1 - 1 file changed, 1 deletion(-) delete mode 100644 uv.toml diff --git a/uv.toml b/uv.toml deleted file mode 100644 index 0863f81..0000000 --- a/uv.toml +++ /dev/null @@ -1 +0,0 @@ -required-version = "0.4.17" From 9a93a81fbfb5ad0935bd181d63335ba61e2a53ef Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Tue, 10 Feb 2026 12:35:07 -0500 Subject: [PATCH 9/9] address reviewer comments --- README.md | 2 +- .../src/flepimop2/engine/op_engine/config.py | 43 ++++++++++++++++--- .../src/flepimop2/engine/op_engine/engine.py | 22 +++++----- flepimop2-op_engine/tests/test_config.py | 18 +++++++- 4 files changed, 64 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 9f67de8..4ce5ea3 100644 --- a/README.md +++ b/README.md @@ -80,7 +80,7 @@ solver.run(rhs, config=None) # defaults: method="heun" (explicit) - `CoreSolver`: explicit and IMEX stepping; methods: `euler`, `heun`, `imex-euler`, `imex-heun-tr`, `imex-trbdf2`. - Operator utilities (`matrix_ops`): Laplacian builders, Crank–Nicolson/implicit Euler/trapezoidal operators, predictor-corrector builders, implicit solve cache, Kronecker helpers, grouped aggregation utilities. - Configuration helpers: `RunConfig`, `OperatorSpecs`, `AdaptiveConfig`, `DtControllerConfig` for method/IMEX/adaptive control. -- Adapters: optional flepimop2 integration (extra dependency) via entrypoints in the adapter package. Adapter merges `mixing_kernels` exposed by op_system and consumes config-supplied IMEX operator tuples when provided; operator metadata in op_system is not auto-translated yet. +- Adapters: optional flepimop2 integration (extra dependency) via entrypoints in the adapter package. The adapter merges any `mixing_kernels` already computed by op_system (no automatic generation) and consumes config-supplied IMEX operator specs (dict or `OperatorSpecs`), forwarding the chosen `operator_axis` to `CoreSolver`. ## Development 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 3bd0270..af15fb8 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/config.py @@ -49,8 +49,7 @@ class OpEngineEngineConfig(BaseModel): gamma: float | None = Field(default=None, gt=0.0, lt=1.0) - # Future-facing: allow operator specs to be supplied via YAML, even if the - # adapter doesn't fully exploit them yet. + # Operator specs (default/tr/bdf2) for IMEX methods. operators: dict[str, Any] | None = Field( default=None, description=( @@ -59,10 +58,17 @@ class OpEngineEngineConfig(BaseModel): ), ) + operator_axis: str | int = Field( + default="state", + description="Axis along which implicit operators act (name or index).", + ) + @model_validator(mode="after") def _validate_imex_requirements(self) -> OpEngineEngineConfig: method = str(self.method) - if method.startswith("imex-") and not self.operators: + if method.startswith("imex-") and not self._has_any_operator_specs( + self.operators + ): msg = ( f"IMEX method '{method}' requires operator specifications, " "but no operators were provided in the engine config." @@ -70,6 +76,15 @@ def _validate_imex_requirements(self) -> OpEngineEngineConfig: raise ValueError(msg) return self + @staticmethod + def _has_any_operator_specs(operators: dict[str, Any] | None) -> bool: + """Return True if any operator spec is provided.""" + if operators is None: + return False + return any( + operators.get(name) is not None for name in ("default", "tr", "bdf2") + ) + def to_run_config(self) -> RunConfig: """ Convert to op_engine RunConfig. @@ -86,9 +101,7 @@ def to_run_config(self) -> RunConfig: fac_max=self.fac_max, ) - # Operators are validated for presence (for IMEX) but intentionally not - # translated yet; keep RunConfig operators empty until wiring is added. - op_specs: OperatorSpecs | None = None + op_specs = self._coerce_operator_specs(self.operators) return RunConfig( method=self.method, @@ -96,6 +109,22 @@ def to_run_config(self) -> RunConfig: strict=self.strict, adaptive_cfg=adaptive_cfg, dt_controller=dt_controller, - operators=op_specs or OperatorSpecs(), + operators=op_specs, gamma=self.gamma, ) + + @staticmethod + def _coerce_operator_specs(operators: dict[str, Any] | None) -> OperatorSpecs: + """ + Normalize operator inputs into OperatorSpecs. + + Returns: + OperatorSpecs with default/tr/bdf2 fields populated when provided. + """ + if operators is None: + return OperatorSpecs() + return OperatorSpecs( + default=operators.get("default"), + tr=operators.get("tr"), + bdf2=operators.get("bdf2"), + ) diff --git a/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py b/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py index addd2f7..81cddaa 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py @@ -226,24 +226,24 @@ def run( msg = "system does not expose a valid flepimop2 SystemProtocol stepper" raise TypeError(msg) + run_cfg = self.config.to_run_config() + # Merge any precomputed mixing kernels exposed by op_system connector - merged_params: dict[IdentifierString, object] = {} - kernels = getattr(system, "mixing_kernels", None) - if isinstance(kernels, dict): - merged_params.update(kernels) - merged_params.update(params) + kernels = getattr(system, "mixing_kernels", {}) + kernel_params = kernels if isinstance(kernels, dict) else {} + merged_params: dict[IdentifierString, object] = {**kernel_params, **params} rhs = _rhs_from_stepper(stepper, params=merged_params, n_state=n_state) core = _make_core(times, y0) - op_default = None - if self.config.operators: - op_default = self.config.operators.get("default") + op_default = run_cfg.operators.default - solver = CoreSolver(core, operators=op_default, operator_axis="state") - - run_cfg = self.config.to_run_config() + solver = CoreSolver( + core, + operators=op_default, + operator_axis=self.config.operator_axis, + ) solver.run(rhs, config=run_cfg) states = _extract_states_2d(core, n_state=n_state) diff --git a/flepimop2-op_engine/tests/test_config.py b/flepimop2-op_engine/tests/test_config.py index 80bc993..efa9cc3 100644 --- a/flepimop2-op_engine/tests/test_config.py +++ b/flepimop2-op_engine/tests/test_config.py @@ -154,7 +154,21 @@ def test_engine_config_imex_requires_operators() -> None: ) run = cfg.to_run_config() assert run.method == "imex-euler" + assert isinstance(run.operators, OperatorSpecs) + assert _has_any_operator_specs(run.operators) + assert run.operators.default == "sentinel" + + +def test_engine_config_operator_dict_coerces_to_specs() -> None: + """Operator dict input is coerced into OperatorSpecs with stage keys.""" + cfg = OpEngineEngineConfig( + method="imex-trbdf2", + operators={"default": 1, "tr": 2, "bdf2": 3}, + gamma=0.6, + ) + run = cfg.to_run_config() - # Current implementation does not translate self.operators into OperatorSpecs yet. assert isinstance(run.operators, OperatorSpecs) - assert not _has_any_operator_specs(run.operators) + assert run.operators.default == 1 + assert run.operators.tr == 2 + assert run.operators.bdf2 == 3