Core Solver

core_solver

Core semi-implicit solver for time-evolving models (ODE + IMEX multiphysics).

This solver advances a :class:op_engine.model_core.ModelCore instance over its configured time grid. In the updated semantics, ModelCore.time_grid is treated as output times: the times at which the user wants a stored solution state.

Between consecutive output times, the solver may take either: - exactly one step of size dt = t_{i+1} - t_i (adaptive=False), or - multiple internal adaptive substeps that land exactly on t_{i+1} (adaptive=True).

Supported methods (keyword method=): - "euler": Explicit Euler (order 1), adaptive via step-doubling. - "heun": Explicit Heun / RK2 (order 2), embedded Euler estimator. - "imex-euler": IMEX Euler: explicit Euler on F(t,y), implicit Euler on A. Adaptive via step-doubling (IMEX step-doubling). - "imex-heun-tr": IMEX Heun-Trapezoidal: Heun on F, trapezoidal/CN on A. Adaptive via embedded low/high (Euler vs Heun) mapped by the same implicit operator solve. - "imex-trbdf2": IMEX TR-BDF2 (order 2), adaptive via step-doubling.

IMEX structure

We assume a split system: y' = A(t,y) y + F(t,y) where F is provided by rhs_func(t, y), and A is represented by linear operators applied along a single tensor axis. Operators may be: - None (ODE-only / explicit-only behavior), or - provided as tuples (predictor?, L, R), or - provided as factories depending on dt, stage-scale, and context.

Operator application

Operators act along a configured axis (default "state"). All other axes are batched. The solve form is: L @ y_next = R @ x optionally with a preprocessing predictor: x_tilde = predictor @ x

Non-uniform dt

• Explicit methods naturally support non-uniform dt.
• Implicit/IMEX methods require operator factories whenever dt varies across steps (non-uniform output grid or adaptive stepping), because L/R depend on dt.

Performance hygiene

• All major scratch arrays are preallocated.
• Inner loops use in-place NumPy ops and np.copyto.
• Implicit solves are delegated to matrix_ops.implicit_solve (cached factorization).

AdaptiveAdvanceParams(plan, t0, t1, y0, adaptive_cfg, dt_ctrl) dataclass

Bundle of parameters for adaptive advancement to an output time.

Attributes:

Name	Type	Description
plan	RunPlan	Resolved run plan.
t0	float	Start time.
t1	float	End/output time.
y0	NDArray[floating]	Initial state at t0.
adaptive_cfg	AdaptiveConfig	Adaptive stepping configuration.
dt_ctrl	DtControllerConfig	dt controller configuration.

AdaptiveConfig(rtol=1e-06, atol=1e-09, dt_init=None, max_reject=25, max_steps=1000000) dataclass

Configuration for adaptive stepping.

Attributes:

Name	Type	Description
rtol	float	Relative tolerance.
atol	float | NDArray[floating]	Absolute tolerance (scalar or array-like).
dt_init	float | None	Optional initial dt guess; if None, use output dt.
max_reject	int	Maximum number of rejected attempts per accepted step.
max_steps	int	Maximum number of internal substeps per output interval.

CoreSolver(core, operators=None, *, operator_axis='state')

Semi-implicit solver operating on a ModelCore time/state grid.

Initialize CoreSolver.

Parameters:

Name	Type	Description	Default
core	ModelCore	ModelCore instance to solve.	required
operators	CoreOperators | StageOperatorFactory | None	Default operator spec (tuple or factory) for implicit stages.	None
operator_axis	str | int	Axis along which operators act (name or index).	'state'

Source code in src/op_engine/core_solver.py

def __init__(
    self,
    core: ModelCore,
    operators: CoreOperators | StageOperatorFactory | None = None,
    *,
    operator_axis: str | int = "state",
) -> None:
    """Initialize CoreSolver.

    Args:
        core: ModelCore instance to solve.
        operators: Default operator spec (tuple or factory) for implicit stages.
        operator_axis: Axis along which operators act (name or index).
    """
    self.core = core
    self.dtype = core.dtype
    self.state_shape = core.state_shape
    self.state_ndim = len(self.state_shape)

    # Operator axis resolution
    self._op_axis = operator_axis
    self._op_axis_idx: int | None = None
    self._op_axis_len: int | None = None

    # Default operator spec (tuple or factory or None)
    self._default_operator_spec: CoreOperators | StageOperatorFactory | None = (
        operators
    )

    # Preallocate buffers (full tensor shape)
    self._rhs_buffer: NDArray[np.floating] = np.zeros(
        self.state_shape,
        dtype=self.dtype,
    )
    self._next_state_buffer: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # Shared stepping buffers
    self._f_n: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._f_pred: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._state_pred: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # Adaptive buffers
    self._y_full: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._y_half: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._y_two_half: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._y_low: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._err: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._scale: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._ratio: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # Working state buffers (avoid allocating per substep)
    self._y_curr: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._y_try: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # TR-BDF2 additional buffers
    self._y_stage1: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._f_stage1: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._f_extrap: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # Validate operator sizes if default spec is a static tuple
    if operators is not None and not callable(operators):
        _predictor, left_op, right_op = self._normalize_ops_tuple(operators)
        if left_op is not None and right_op is not None:
            self._resolve_operator_axis()
            self._validate_operator_sizes(left_op, right_op)

run(rhs_func, *, config=None)

Advance the ModelCore state through its time grid.

Parameters:

Name	Type	Description	Default
rhs_func	RHSFunction	Function computing the explicit RHS F(t, y).	required
config	RunConfig | None	Optional run configuration. If None, defaults are used.	None

Raises:

Type	Description
ValueError	If invalid parameters are provided.

Source code in src/op_engine/core_solver.py

def run(self, rhs_func: RHSFunction, *, config: RunConfig | None = None) -> None:
    """Advance the ModelCore state through its time grid.

    Args:
        rhs_func: Function computing the explicit RHS F(t, y).
        config: Optional run configuration. If None, defaults are used.

    Raises:
        ValueError: If invalid parameters are provided.
    """
    cfg = config or RunConfig()
    plan = self._resolve_run_plan(cfg)

    time_grid = np.asarray(self.core.time_grid, dtype=float)
    n_steps = int(self.core.n_timesteps)

    for idx in range(n_steps - 1):
        t0 = float(time_grid[idx])
        t1 = float(time_grid[idx + 1])
        if t1 <= t0:
            raise ValueError(_TIME_GRID_INCREASING_ERROR_MSG)
        dt_out = t1 - t0

        np.copyto(self._y_curr, self.core.get_current_state())

        if not cfg.adaptive:
            y_next = self._advance_nonadaptive_to_time(
                rhs_func,
                plan=plan,
                t0=t0,
                dt_out=dt_out,
                y0=self._y_curr,
            )
            self.core.advance_timestep(y_next)
            continue

        y_end = self._advance_adaptive_to_time(
            rhs_func,
            AdaptiveAdvanceParams(
                plan=plan,
                t0=t0,
                t1=t1,
                y0=self._y_curr,
                adaptive_cfg=cfg.adaptive_cfg,
                dt_ctrl=cfg.dt_controller,
            ),
        )
        self.core.advance_timestep(y_end)

DtControllerConfig(dt_min=0.0, dt_max=float('inf'), safety=0.9, fac_min=0.2, fac_max=5.0) dataclass

Configuration for adaptive timestep control.

Attributes:

Name	Type	Description
dt_min	float	Minimum allowed dt.
dt_max	float	Maximum allowed dt.
safety	float	Safety factor applied to dt updates.
fac_min	float	Minimum multiplicative change factor.
fac_max	float	Maximum multiplicative change factor.

ImexEulerOnceParams(t, y, dt, op_spec, out) dataclass

Bundle of parameters for one IMEX Euler step (non-doubling).

Attributes:

Name	Type	Description
t	float	Current time.
y	NDArray[floating]	Current state.
dt	float	Step size.
op_spec	CoreOperators | StageOperatorFactory | None	Operator spec for implicit stage.
out	NDArray[floating]	Output state array (written in-place).

ImplicitStageParams(spec, dt, scale, t_stage, y_stage, stage, x, out) dataclass

Bundle of parameters for one implicit operator application.

Attributes:

Name	Type	Description
spec	CoreOperators | StageOperatorFactory | None	Operator spec (tuple or factory) or None for identity.
dt	float	Full-step dt for operator factory context.
scale	float	Stage scaling factor for dt-dependent operators.
t_stage	float	Stage time.
y_stage	NDArray[floating]	Stage state proxy for operator factories.
stage	str	Stage label (e.g., "be", "tr", "bdf2").
x	NDArray[floating]	Input array to map.
out	NDArray[floating]	Output array (written in-place).

OperatorLike

Bases: Protocol

Minimal operator interface required by CoreSolver.

Implementations are expected to behave like 2D linear operators suitable for implicit_solve(L, R, rhs2d). Only shape is required for validation.

shape property

Return the operator shape.

OperatorSpecs(default=None, tr=None, bdf2=None) dataclass

Operator specifications for implicit/IMEX methods.

Attributes:

Name	Type	Description
default	CoreOperators | StageOperatorFactory | None	Default operator spec (tuple or factory) used by IMEX Euler/Heun-TR and as a fallback for TR/BDF2 stages.
tr	CoreOperators | StageOperatorFactory | None	Operator spec for trapezoidal stage of TR-BDF2 (optional).
bdf2	CoreOperators | StageOperatorFactory | None	Operator spec for BDF2 stage of TR-BDF2 (optional).

PredictorLike

Bases: Protocol

Minimal predictor interface required by CoreSolver.

The predictor is an optional preprocessing operator applied as

rhs2d = predictor @ rhs2d

__matmul__(other)

Apply the predictor to a 2D array.

Source code in src/op_engine/core_solver.py

def __matmul__(self, other: NDArray[np.floating]) -> NDArray[np.floating]:
    """Apply the predictor to a 2D array."""
    ...

RunConfig(method='heun', adaptive=False, strict=True, dt_controller=DtControllerConfig(), adaptive_cfg=AdaptiveConfig(), operators=OperatorSpecs(), gamma=None) dataclass

Configuration for CoreSolver.run.

Attributes:

Name	Type	Description
method	str	Method name.
adaptive	bool	Whether to use adaptive substepping between output times.
strict	bool	If True, invalid configurations raise; otherwise warnings and method downshifts may occur.
dt_controller	DtControllerConfig	Parameters for dt controller when adaptive=True.
adaptive_cfg	AdaptiveConfig	Parameters controlling error tolerances and limits.
operators	OperatorSpecs	Operator specifications for implicit/IMEX methods.
gamma	float | None	Optional TR-BDF2 gamma (if None, uses default).

RunPlan(method, gamma, op_default, op_tr, op_bdf2) dataclass

Resolved execution plan derived from RunConfig.

This is the internal, validated form used by the stepping loops.

Attributes:

Name	Type	Description
method	MethodName	Final method after any strict=False downshifts.
gamma	float | None	TR-BDF2 gamma, or None for non-TR-BDF2 methods.
op_default	CoreOperators | StageOperatorFactory | None	Operator spec for IMEX Euler/Heun-TR.
op_tr	CoreOperators | StageOperatorFactory | None	TR-stage operator spec for TR-BDF2.
op_bdf2	CoreOperators | StageOperatorFactory | None	BDF2-stage operator spec for TR-BDF2.

StepIO(t, dt, y, out, err_out=None) dataclass

Bundle of per-step state for stepping kernels.

Attributes:

Name	Type	Description
t	float	Current time.
dt	float	Step size.
y	NDArray[floating]	Current state array (input).
out	NDArray[floating]	Output state array (written in-place).
err_out	NDArray[floating] | None	Error estimate array (written in-place) for adaptive methods.

Trbdf2OnceParams(t, y, dt, operators_tr, operators_bdf2, gamma, out) dataclass

Bundle of parameters for one TR-BDF2 step (non-doubling).

Attributes:

Name	Type	Description
t	float	Current time.
y	NDArray[floating]	Current state.
dt	float	Step size.
operators_tr	CoreOperators | StageOperatorFactory | None	TR stage operator spec.
operators_bdf2	CoreOperators | StageOperatorFactory | None	BDF2 stage operator spec.
gamma	float	TR-BDF2 gamma.
out	NDArray[floating]	Output state array (written in-place).