op_engine

Operator-Partitioned Engine (OP Engine) is a lightweight multiphysics solver core for time-dependent systems. It supports explicit ODE solvers, IMEX/operator-based schemes for PDE-like models, and exact or approximate stochastic reaction-network solvers while staying framework-agnostic.

Why use it?

• Shared solver surface for ODEs and operator-split PDEs.
• Strong typing and Array-API explicit and dense-implicit paths, including reusable higher-order Runge--Kutta tableaus.
• JAX autodiff and JIT support for the same portable fixed-step methods.
• 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.

Core surface

• ModelCore: state/time manager; configure axes, dtype, and optional history.
• CoreSolver: portable explicit methods through Dormand--Prince 5(4), plus dense IMEX and linearly implicit methods; accepts RunConfig with AdaptiveConfig, DtControllerConfig, and OperatorSpecs.
• DirectSSASolver: exact Gillespie direct-method trajectories with injected, backend-specific exponential and categorical sampling.
• AdaptiveTauLeapingSolver: bounded adaptive tau-leaping with leap-condition control, exact critical events, and explicit post-leap rejection. Non-mass-action propensities (frequency-dependent infection) are supported through per-reaction dependency incidence and elasticity orders.
• TauLeapingSolver: fixed-step stochastic reaction-network integration with injected, backend-specific Poisson sampling.
• from_compiled_rhs / compile_reaction_network: build the flat stochastic network (stoichiometry, reactant orders, propensity) from an op_system spec's reaction artifacts, with no flepimop2 dependency.
• steady_state: equilibria by pseudo-transient continuation, with absorbing states held fixed, conserved totals enforced exactly, and a Newton-correction convergence test; runs under jax.jit/vmap.
• matrix_ops: portable dense advection/diffusion, sparse Laplacian/Crank–Nicolson, implicit Euler/trapezoidal builders, predictor–corrector, implicit solve cache, Kronecker helpers, and grouped aggregations.
• Extras: OperatorSpecs, RunConfig, AdaptiveConfig, DtControllerConfig, Operator, GridGeometry, DiffusionConfig.

See the documentation guides for solver selection, the biogeochemical splitting tutorial, stochastic simulation from an op_system spec, steady states, and backend boundaries.

Installation

pip install op_engine

With JAX as the array namespace:

pip install "op_engine[jax]"

For the other packaged Array-API backends:

pip install "op_engine[torch]"
pip install "op_engine[cupy]"

NumPy is installed with the core package. Backend support is method-specific; consult the backend guide before assuming that JIT, automatic differentiation, adaptivity, or stochastic sampling have identical semantics across namespaces.

With flepimop2 adapter:

pip install flepimop2-op-engine

To use the op_system provider contract as well:

pip install "flepimop2-op-engine[op-system]"

Quickstart

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])


# Time grid and state
core = ModelCore(n_states=3, n_subgroups=1, time_grid=np.linspace(0, 10, 101))
core.set_initial_state(np.array([0.999, 0.001, 0.0])[..., None])

solver = CoreSolver(core)
solver.run(rhs)  # defaults to Heun/RK2

solution = core.state_array  # shape (n_timesteps, state, subgroup)

Array namespaces

ModelCore infers its numerical namespace from the initial state through array_api_compat.array_namespace(). The state, stored history, solver stages, and dense solver operations stay in that namespace. There is no xp= or backend option. This supports standard Array-API objects and native arrays such as torch.Tensor. Fixed-step methods can be differentiated by JAX; the built-in adaptive controller is eager Python control flow. See the backend guide for the precise contract and the solver guide for method-specific examples.

IMEX with operators (tuple form)

import numpy as np
from op_engine import CoreSolver, ModelCore
from op_engine.core_solver import OperatorSpecs, RunConfig

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, operator_axis="state")
solver.run(
    rhs,
    config=RunConfig(
        method="imex-heun-tr",
        operators=ops,
    ),
)

This compact example uses identity left/right operators to show the API. For a nontrivial implicit term, construct timestep-aware operator factories as in the solver guide.

Public API

• ModelCore: state tensor + time grid manager; supports extra axes and optional history.
• CoreSolver: portable explicit and dense IMEX/implicit stepping selected by the state array namespace.
• DirectSSASolver: exact event-by-event stochastic reaction trajectories with NumPy, JAX, or another Array-API namespace supplying random draws.
• AdaptiveTauLeapingSolver: adaptive non-negative stochastic leaps using explicit reactant stoichiometry and injected Poisson plus exact-event sampling.
• TauLeapingSolver: portable stoichiometric updates with NumPy, JAX, or another Array-API namespace supplying Poisson samples.
• op_engine.reactions: from_compiled_rhs(compiled, params) and compile_reaction_network(reactions, template_shapes=..., axis_sizes=..., params=...) return a CompiledReactionNetwork whose stoichiometry, reactant_stoichiometry, and propensity feed the stochastic solvers; mean_drift checks the network against the deterministic RHS.
• Operator utilities (matrix_ops): portable upwind advection, Laplacian, Crank–Nicolson/implicit Euler/trapezoidal operators, predictor-corrector builders, implicit solve cache, Kronecker helpers, and grouped aggregation utilities.
• Configuration helpers: RunConfig, OperatorSpecs, AdaptiveConfig, and DtControllerConfig from op_engine.core_solver for method, IMEX, and adaptive control.
• 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

uv sync --dev
just ci

License

MIT License