diff --git a/README.md b/README.md index 4ce5ea3..2e39b9e 100644 --- a/README.md +++ b/README.md @@ -33,11 +33,13 @@ pip install "op_engine[flepimop2]" import numpy as np from op_engine import ModelCore, CoreSolver + # Define RHS def rhs(t, y): s, i, r = y beta, gamma = 0.3, 0.1 - return np.array([-beta*s*i, beta*s*i - gamma*i, gamma*i]) + return np.array([-beta * s * i, beta * s * i - gamma * i, gamma * i]) + # Time grid and state core = ModelCore(n_states=3, n_subgroups=1, time_grid=np.linspace(0, 10, 101)) @@ -65,9 +67,11 @@ 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) diff --git a/docs/guides/getting-started.md b/docs/guides/getting-started.md index 56eb954..af25eef 100644 --- a/docs/guides/getting-started.md +++ b/docs/guides/getting-started.md @@ -4,6 +4,7 @@ A brief tutorial on how to get started with `OP Engine`. ```python import math + foo = 2 * math.pi ``` diff --git a/examples/simple_sir.py b/examples/simple_sir.py index e2bc809..2245a71 100644 --- a/examples/simple_sir.py +++ b/examples/simple_sir.py @@ -129,11 +129,11 @@ def _run_sir( initial_infected: Initial infected fraction. config: Run configuration for CoreSolver. - Raises: - RuntimeError: If the state history is not available after the run. - Returns: State history array of shape (n_steps, 3, 1). + + Raises: + RuntimeError: If the state history is not available after the run. """ opts = ModelCoreOptions(store_history=True, dtype=np.float64) core = ModelCore(n_states=3, n_subgroups=1, time_grid=time_grid, options=opts) 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 3a319b0..d0fd33a 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py @@ -11,11 +11,15 @@ from __future__ import annotations +from typing import Literal + from .engine import _OpEngineFlepimop2EngineImpl class OpEngineFlepimop2Engine(_OpEngineFlepimop2EngineImpl): # noqa: RUF067 """Public op_engine-backed flepimop2 Engine (default-build enabled).""" + module: Literal["flepimop2.engine.op_engine"] = "flepimop2.engine.op_engine" + __all__ = ["OpEngineFlepimop2Engine"] 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 81cddaa..55ed89b 100644 --- a/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py +++ b/flepimop2-op_engine/src/flepimop2/engine/op_engine/engine.py @@ -80,11 +80,11 @@ def rhs(t: float, y: np.ndarray) -> np.ndarray: t: Current time. y: Current state array with shape (n_state, 1). - Raises: - ValueError: If input or output shapes are invalid. - Returns: 2D array of shape (n_state, 1) representing dstate/dt + + Raises: + ValueError: If input or output shapes are invalid. """ y_arr = np.asarray(y, dtype=np.float64) @@ -124,11 +124,11 @@ def _extract_states_2d(core: ModelCore, *, n_state: int) -> Float64Array2D: core: ModelCore instance n_state: Number of state variables - Raises: - RuntimeError: If state_array is missing or has an unexpected shape. - Returns: 2D float64 array of stored states with shape (T, n_state). + + Raises: + RuntimeError: If state_array is missing or has an unexpected shape. """ state_array = getattr(core, "state_array", None) if state_array is None: @@ -204,11 +204,11 @@ def run( params: Mapping of parameter names to values. **kwargs: Additional engine-specific keyword arguments (ignored). - Raises: - TypeError: If system does not expose a valid stepper. - Returns: 2D array of shape (T, 1 + n_states) with time in the first column. + + Raises: + TypeError: If system does not expose a valid stepper. """ del kwargs diff --git a/flepimop2-op_engine/tests/test_engine.py b/flepimop2-op_engine/tests/test_engine.py index a0dddc1..1532e97 100644 --- a/flepimop2-op_engine/tests/test_engine.py +++ b/flepimop2-op_engine/tests/test_engine.py @@ -13,6 +13,7 @@ import numpy as np import pytest +from flepimop2.engine.op_engine import OpEngineFlepimop2Engine from flepimop2.engine.op_engine.engine import ( _OpEngineFlepimop2EngineImpl, # noqa: PLC2701 ) @@ -83,12 +84,11 @@ def __init__(self, n: int) -> None: # ----------------------------------------------------------------------------- -def test_engine_default_config_constructs() -> None: - """Engine can be constructed with defaults.""" - engine = _OpEngineFlepimop2EngineImpl() +def test_public_engine_wrapper_defines_module() -> None: + """Public engine wrapper satisfies flepimop2's concrete module contract.""" + engine = OpEngineFlepimop2Engine() assert isinstance(engine, _OpEngineFlepimop2EngineImpl) - # Default comes from OpEngineEngineConfig; we do not assert its exact value here. assert engine.module == "flepimop2.engine.op_engine" diff --git a/flepimop2-op_engine/tests/test_integration_cli.py b/flepimop2-op_engine/tests/test_integration_cli.py index 6166bce..cd1b383 100644 --- a/flepimop2-op_engine/tests/test_integration_cli.py +++ b/flepimop2-op_engine/tests/test_integration_cli.py @@ -44,7 +44,13 @@ def _write_config(config_path: Path) -> None: config_path.parent.mkdir(parents=True, exist_ok=True) cfg = { "name": "SIR_op_engine", - "system": [{"module": "wrapper", "script": "model_input/plugins/SIR.py"}], + "system": [ + { + "module": "wrapper", + "state_change": "flow", + "script": "model_input/plugins/SIR.py", + } + ], "engine": [ { "module": "op_engine", diff --git a/src/op_engine/core_solver.py b/src/op_engine/core_solver.py index c67f7dc..7d336db 100644 --- a/src/op_engine/core_solver.py +++ b/src/op_engine/core_solver.py @@ -485,11 +485,11 @@ def _as_scipy_operator(op: OperatorLike) -> ScipyOperator: Args: op: Operator-like object. - Raises: - TypeError: If op is not a supported operator type for this backend. - Returns: Operator as a supported SciPy/NumPy operator type. + + Raises: + TypeError: If op is not a supported operator type for this backend. """ if isinstance(op, np.ndarray): return cast("NDArray[np.floating]", op) @@ -581,11 +581,11 @@ def _normalize_ops_tuple( Args: ops: Operator spec tuple. - Raises: - ValueError: If ops is not a 2- or 3-tuple. - Returns: Tuple of (predictor, L, R) operators. + + Raises: + ValueError: If ops is not a 2- or 3-tuple. """ if len(ops) == 2: left_op, right_op = ops @@ -733,11 +733,11 @@ def _normalize_method(method: str) -> MethodName: Args: method: User-provided method string. - Raises: - ValueError: If method is unknown. - Returns: Normalized method literal. + + Raises: + ValueError: If method is unknown. """ method_norm = str(method).strip().lower() allowed: tuple[str, ...] = ( @@ -760,11 +760,11 @@ def _resolve_gamma(method: MethodName, gamma: float | None) -> float | None: method: Method name. gamma: User-provided gamma (or None for default). - Raises: - ValueError: If gamma is out of range for TR-BDF2. - Returns: Resolved gamma for TR-BDF2, or None for non-TR-BDF2 methods. + + Raises: + ValueError: If gamma is out of range for TR-BDF2. """ if method != "imex-trbdf2": return None @@ -827,12 +827,12 @@ def _plan_for_imex_single( strict: If True, invalid configuration raises. adaptive: Whether adaptive stepping is enabled. - Raises: - ValueError: If required operators are missing. - Returns: RunPlan for IMEX method, or an explicit fallback if strict=False and operators are missing. + + Raises: + ValueError: If required operators are missing. """ if op_default is None: if strict: @@ -887,13 +887,13 @@ def _plan_for_trbdf2( strict: If True, invalid configuration raises. adaptive: Whether adaptive stepping is enabled. - Raises: - RuntimeError: If method_in is not "imex-trbdf2". - ValueError: If required operators are missing. - Returns: RunPlan for TR-BDF2, or a Heun fallback if strict=False and operators are missing. + + Raises: + RuntimeError: If method_in is not "imex-trbdf2". + ValueError: If required operators are missing. """ if method_in != "imex-trbdf2": raise RuntimeError(_UNKNOWN_METHOD_ERROR_MSG.format(method=method_in)) @@ -1083,11 +1083,11 @@ def _require_err_out(step: StepIO) -> NDArray[np.floating]: Args: step: Step bundle. - Raises: - RuntimeError: If err_out is None. - Returns: err_out array. + + Raises: + RuntimeError: If err_out is None. """ if step.err_out is None: raise RuntimeError(_INTERNAL_ERROR_ERR_OUT_MSG) @@ -1447,11 +1447,11 @@ def _attempt_step( plan: Run plan. step: Step bundle. - Raises: - RuntimeError: If an unknown method is encountered or internal errors occur. - Returns: Method order. + + Raises: + RuntimeError: If an unknown method is encountered or internal errors occur. """ if plan.method == "euler": return self._step_euler_doubling(rhs_func, step) @@ -1525,11 +1525,11 @@ def _advance_adaptive_to_time( rhs_func: RHS function. params: Adaptive advance parameters. - Raises: - RuntimeError: If step rejection limits or dt bounds are violated. - Returns: State at params.t1. + + Raises: + RuntimeError: If step rejection limits or dt bounds are violated. """ t0 = float(params.t0) t1 = float(params.t1) diff --git a/src/op_engine/matrix_ops.py b/src/op_engine/matrix_ops.py index fc27adc..55b5827 100644 --- a/src/op_engine/matrix_ops.py +++ b/src/op_engine/matrix_ops.py @@ -192,11 +192,11 @@ def build_laplacian_tridiag( dtype: Floating dtype (e.g. np.float64). bc: Boundary condition; either "neumann" or "absorbing". - Raises: - ValueError: If an unknown boundary condition is provided. - Returns: Sparse CSR matrix representing the Laplacian operator. + + Raises: + ValueError: If an unknown boundary condition is provided. """ dtype_obj = np.dtype(dtype) factor = coeff / dx**2 @@ -456,11 +456,11 @@ def build_implicit_euler_operators( base_op: Base linear operator A. dt_scale: Time-step scaling factor (dt * scale). - Raises: - ValueError: If dt_scale is not finite. - Returns: Tuple of (L, R) operators for implicit Euler scheme. + + Raises: + ValueError: If dt_scale is not finite. """ if not np.isfinite(dt_scale): raise ValueError(_OPERATOR_SCALE_ERROR.format(scale=dt_scale)) @@ -491,11 +491,11 @@ def build_trapezoidal_operators( base_op: Base linear operator A. dt_scale: Time-step scaling factor (dt * scale). - Raises: - ValueError: If dt_scale is not finite. - Returns: Tuple of (L, R) operators for trapezoidal scheme. + + Raises: + ValueError: If dt_scale is not finite. """ if not np.isfinite(dt_scale): raise ValueError(_OPERATOR_SCALE_ERROR.format(scale=dt_scale)) @@ -537,11 +537,11 @@ def make_stage_operator_factory( base_builder: Function that builds a base operator given stage context. scheme: Implicit scheme; either "implicit-euler" or "trapezoidal". - Raises: - ValueError: If an unknown scheme is provided. - Returns: A StageOperatorFactory that builds (L, R) operators for the given scheme. + + Raises: + ValueError: If an unknown scheme is provided. """ scheme_norm = str(scheme).strip().lower() @@ -620,12 +620,12 @@ def kron_sum(ops: list[Operator]) -> Operator: Args: ops: List of 2D square operators. + Returns: + The Kronecker sum operator. + Raises: ValueError: If ops is empty or if operators are not square or have incompatible shapes. - - Returns: - The Kronecker sum operator. """ if not ops: raise ValueError(_KRON_EMPTY_ERROR) @@ -1050,13 +1050,13 @@ def grouped_sum_ids_2d( group_ids: 1D array of integer group IDs of length N. n_groups: Total number of groups. - Raises: - ValueError: If values is not 2D or if group_ids length does not match - the number of items in values. - Returns: A 2D array of shape (n_groups, K) where each row contains the sum of values for the corresponding group ID. + + Raises: + ValueError: If values is not 2D or if group_ids length does not match + the number of items in values. """ values_arr = np.asarray(values) group_ids_arr = np.asarray(group_ids, dtype=np.int64) diff --git a/src/op_engine/model_core.py b/src/op_engine/model_core.py index 10e3078..fca58a1 100644 --- a/src/op_engine/model_core.py +++ b/src/op_engine/model_core.py @@ -219,12 +219,12 @@ def axis_index(self, axis: str | int) -> int: Args: axis: Axis name or index. + Returns: + Axis index as an integer. + Raises: IndexError: if axis index is out of bounds. ValueError: if axis name is unknown. - - Returns: - Axis index as an integer. """ if isinstance(axis, int): if not (0 <= axis < self.state_ndim): @@ -351,11 +351,11 @@ def get_time_at(self, step_idx: int) -> float: Args: step_idx: Timestep index in [0, n_timesteps). - Raises: - IndexError: if step_idx is out of bounds. - Returns: Time as a float. + + Raises: + IndexError: if step_idx is out of bounds. """ if not (0 <= step_idx < self.n_timesteps): raise IndexError(_TIME_INDEX_OOB_ERROR.format(idx=step_idx)) @@ -368,11 +368,11 @@ def get_dt(self, step_idx: int) -> float: Args: step_idx: Timestep index in [0, n_timesteps - 1]. - Raises: - IndexError: if step_idx is out of bounds. - Returns: dt as a float. + + Raises: + IndexError: if step_idx is out of bounds. """ if self.n_timesteps <= 1: return 0.0