op_system

Domain-agnostic specification and compilation of right-hand sides (RHS) for ODE, PDE, and multi-physics / multi-scale compartmental systems. op_system takes a YAML/JSON-friendly spec, validates and normalizes it, then compiles it into a fast, array-API-polymorphic callable whose namespace is selected from the inputs at call time. NumPy, JAX (concrete and traced), and raw PyTorch tensors are covered by the test suite. Other Array-API implementations may work through array-api-compat, but should be qualified before production use.

Docs: https://accidda.github.io/op_system/
License: MIT
Python: 3.11 – 3.13

Why op_system?

Modelers often combine compartment hazards, templated populations, and rich metadata (axes, kernels, operators) that must be validated and preserved for downstream solvers. op_system provides:

Two equivalent surfaces — expr (explicit equations) and transitions (hazard / flow style) — that share the same axis, alias, template, and reducer machinery.
Validated, restricted expression parsing with a small allowlist of NumPy ops and helpers; no arbitrary code execution.
A typed intermediate representation (IR) that handles template expansion, alias inlining, and apply_along/sum_over reductions symbolically before code generation.
Vectorized compilation that operates on shaped state buffers (one tensor expression per template) rather than per-cell scalar code, with template-level common-subexpression elimination.
Backend polymorphism at call time — array-api-compat selects the namespace from each input, so the same compiled artifact serves NumPy, JAX jit/vmap/grad, and raw PyTorch tensors with autograd.
First-class PyTree interface (pytree_eval_fn) for engines that want to keep state as a dict of shaped arrays rather than a flat vector.
Block-axis vmap support (block_pytree_eval_fn) for hierarchical models — declare factorize_axes and the engine can vmap a stripped per-block RHS over a block axis instead of evaluating a monolithic flat state.
Picklable CompiledRhs — round-trips through pickle.dumps/loads by retaining the source spec and recompiling on load.

Installation

pip install op-system
# or, from a checkout, using uv:
uv pip install .

Optional extras:

pip install "op-system[jax]"            # JAX runtime support
pip install "op-system[jax-inference]"  # adds diffrax + blackjax
pip install "op-system[torch]"          # PyTorch runtime support
pip install "op-system[data]"           # pandas + pyarrow helpers

Time-indexed parameters

Time-indexed parameter tables use linear interpolation by default. Set time_interpolation: previous in a specification to hold each table value until the next coordinate, with right-continuous changes and constant endpoint extrapolation. The same policy applies to flat, PyTree, block, and reaction evaluators. Compiled metadata and the flepimop2 provider publish immutable forcing coordinates for numerical engines. See the time-indexed parameter guide for the schema, examples, and exactness conditions.

Quick start

import numpy as np

from op_system import compile_spec

spec = {
    "kind": "expr",
    "state": ["S", "I", "R"],
    "aliases": {"N": "S + I + R"},
    "equations": {
        "S": "-beta * S * I / N",
        "I": "beta * S * I / N - gamma * I",
        "R": "gamma * I",
    },
}

compiled = compile_spec(spec)
dydt = compiled.eval_fn(0.0, np.asarray([999.0, 1.0, 0.0]), beta=0.3, gamma=0.1)

The compiled object exposes:

Attribute	Description
`eval_fn(t, y, **params) -> dydt`	Flat-vector RHS; array namespace inferred from `y`.
`pytree_eval_fn(t, state_dict, **params) -> dict`	PyTree RHS keyed by state template base name (axis-indexed specs).
`template_shapes`	`{base: shape}` for each state template; axis-less states are `()`.
`state_names`, `param_names`	Tuples of expanded state cells and parameter names.
`factorize_axes`, `block_axes`	Axes the IR proved separable for block vmap.
`block_pytree_eval_fn`, `block_template_shapes`	Per-block PyTree RHS with the first factorize axis stripped.
`meta`	Normalized metadata (axes, state_axes, kernels, operators, reserved blocks).
`operators`	Tuple of `OperatorDescriptor` preserving normalized names, state selectors, coefficients, directions, boundary conditions, and kernel metadata.
`reactions`, `reaction_gaps`	One `CompiledReaction` per named transition with a reaction artifact, and one `ReactionGap` per transition without one (empty when every transition is covered).

Advection contract

Advection and transport act along the declared coordinate order. A signed velocity without direction is used directly: positive moves toward increasing indices and negative moves toward decreasing indices. An optional direction makes the orientation explicit while keeping a dynamic coefficient:

operators:
  - kind: advection
    axis: imm
    velocity: waning_rate
    direction: decreasing
    bc: reflecting

Providers multiply an increasing coefficient by +1 and a decreasing coefficient by -1. Coefficients used with explicit direction should therefore be non-negative; producers of traced dynamic values are responsible for that invariant.

Boundary conditions are defined relative to the resolved direction:

absorbing uses zero upstream inflow and permits downstream outflow;
reflecting uses zero upstream inflow and zero downstream flux, so mass accumulates in the terminal cell;
periodic wraps downstream outflow to the upstream cell.

Engines must apply these semantics identically for either velocity sign.

Jump-integral contract

jump_integral metadata defines a conservative row-source, column-target matrix generator along an axis. direction: up|down|both masks destinations; continuous axes use target trapezoidal weights; and the currently supported reflecting boundary truncates out-of-domain jumps without renormalizing or losing mass. See the operator guide for the exact schema, units, and Array-API reference functions.

compile_spec accepts legacy backend= / xp= keyword arguments but they are deprecated and ignored — the compiled callable infers its array namespace from the input y on every call.

JAX usage

import jax, jax.numpy as jnp
from op_system import compile_spec

compiled = compile_spec(spec)
y0 = jnp.asarray([999.0, 1.0, 0.0])

# Native JAX call — eval_fn returns a jnp array.
dydt = compiled.eval_fn(0.0, y0, beta=0.3, gamma=0.1)

# Works inside jit / vmap / grad without recompilation.
solve = jax.jit(lambda y: compiled.eval_fn(0.0, y, beta=0.3, gamma=0.1))

For diffrax-based ODE solves and NUTS / HMC inference, install the jax-inference extra above.

YAML examples

The full guide of YAML patterns — including templates, axis asymmetry, chains, continuous axes with kernels, and block-axis hierarchical models — lives at https://accidda.github.io/op_system/guides/getting-started/. A few highlights:

Baseline SIR (two pathways)

# expr
spec:
  kind: expr
  state: [S, I, R]
  equations:
    S: -beta * S * I / sum_state()
    I:  beta * S * I / sum_state() - gamma * I
    R:  gamma * I

# transitions
spec:
  kind: transitions
  state: [S, I, R]
  transitions:
    - {from: S, to: I, rate: beta * I / sum_state()}
    - {from: I, to: R, rate: gamma}

Source-only tracking transitions are also supported (from: null or omitted):

spec:
  kind: transitions
  state: [I, H_cum]
  transitions:
    - {to: H_cum, rate: k * I}  # equivalent to {from: null, ...}

This pattern is useful for cumulative trackers (e.g., weekly admissions via diff(H_cum)) without introducing a dummy donor compartment.

Named transitions may also declare the molecular reactants needed by stochastic solvers. The list is independent of net source/target stoichiometry, so it must include the consumed source as well as catalysts:

spec:
  kind: transitions
  axes:
    - {name: age, coords: [child, adult]}
    - {name: vax, coords: [u, v]}
  state: [S[age,vax], E[age,vax], I[age]]
  transitions:
    - name: infection
      from: S[age,vax]
      to: E[age,vax]
      rate: beta * I[age]
      reactants:
        - {state: S[age,vax], order: 1}
        - {state: I[age], order: 1}  # catalytic: not consumed

The compiled reaction exposes these entries as array-neutral structural metadata. reactants: auto derives the list from the rate, after aliases are inlined: the consumed source at order one plus each state factor at its integer power. For the transition above it infers exactly the declared list. Check compiled.reactions[i].reactants to confirm what was inferred.

A rate that is not a single product of states, such as the frequency-dependent beta * sum_over(I[age:a], age=a) / N, has no molecular reactants beyond the consumed source. Under reactants: auto it instead publishes what adaptive tau-leaping needs:

dependencies: every state selection the propensity reads, aligned to the channels like reactants. A reduction contributes one pinned entry per coordinate it visits.
propensity_order: a whole-number bound on the propensity's total elasticity, sum_i |d log a / d log x_i|. It is derived from the expression: products and quotients add their operands' bounds, a literal power p multiplies by |p|, and a sum or reduction of non-negative terms takes the largest term's bound. beta * S * sum(I) / N has order 3.
dependencies_complete=True.

These reactions keep reactants_complete=false, so a consumer that only understands reactants refuses adaptive tau-leaping rather than misreading them. Subtraction, negation, other functions of a state (exp, min, ...), symbolic powers, and history operators have no such bound: reactants: auto rejects them at compile time, naming the construct. Parameters are assumed non-negative.

If reactants is omitted, op_system publishes the consumed source at order one. That is complete (reactants_complete=true) when the rate reads no state, because nothing else can then be a reactant. Otherwise it is reactants_complete=false, and adaptive stochastic consumers should require complete metadata. Parameters, including time-varying ones, are assumed not to depend on the state. An explicit empty list marks a source-only zero-order reaction as complete.

Not every transition publishes a reaction. CompiledRhs.reaction_gaps (and the provider's reaction_gaps option) lists each one that does not, with its spec origin (transitions[1], chain[0].forward[0]), selectors, and a reason such as unnamed, rate_axis_out_of_scope, or unsupported_layout. An expr spec reports a single expr_spec gap. A consumer that executes only the reactions, such as a pure stochastic simulation, should reject a non-empty value rather than silently drop those dynamics.

A rate may name aliases, either bracketed (foi[age]) or, for an axis-less alias, by bare name (lam). Their bodies are inlined into the propensity, following chains of aliases. A rate that still names an alias afterwards, for example one on a reference cycle or a templated alias referenced without its axes, gets an unresolved_alias gap.

Axis-less states take part like any other: a scalar S→I→R model publishes 0-d reactions and template_shapes of (). When a reaction's source and target templates have different axes, for example an axis-less source depositing into a pinned cell, to_full_axes gives the target's axis order for indexing it.

Source-only rates may also depend on population through a bound reduction, such as sum_over(B[age:a] * N[age:a], age=a), while their destination pins age=a0. This produces one total birth hazard into that cell, without donor depletion. See the renewal births guide for reaction metadata, retained group axes, and a stationary age-population example.

Templated states with `apply_along`

spec:
  kind: expr
  axes:
    - {name: age,  coords: [child, adult]}
    - {name: vax,  coords: [u, v]}
  state: [S[age,vax], I[age,vax], R[age,vax]]
  aliases:
    lambda[age]: beta * apply_along(vax=j, I[age,vax=j]) / sum_state()
  equations:
    S[age,vax]: -lambda[age] * S[age,vax]
    I[age,vax]:  lambda[age] * S[age,vax] - gamma * I[age,vax]
    R[age,vax]:  gamma * I[age,vax]

apply_along(axis=var, expr) contracts expr along one or more axes in a single call. Categorical / ordinal axes use uniform weights of 1; continuous axes use trapezoidal weights derived from axis spacing (non-uniform supported). Bindings can be restricted with axis=var in [...] for sub-range integration.

Routing transitions with `axis:alias`

spec:
  kind: transitions
  axes:
    - {name: vax, coords: [u, v]}
    - {name: imm, type: ordinal, coords: [x0, x1, x2, x3]}
  state: [X[vax, imm]]
  transitions:
    - from: X[vax, imm:i]            # waning along a generator G
      to:   X[vax, imm:j]
      rate: waning_rate * G[imm:i, imm:j]
    - from: X[vax=u, imm:i]          # vaccination with routing weights eta
      to:   X[vax=v, imm:j]
      rate: nu * eta[time, imm:i, imm:j]

Binding the same axis under one alias in from and another in to moves mass along that axis with a matrix-valued per-capita rate: dX_from[i] -= r X_from[i] sum_j K[i, j] and dX_to[j] += r sum_i K[i, j] X_from[i]. The rate must reference both aliases on that axis; other axes are shared or pinned as usual. When from and to are otherwise the same slice, the diagonal K[i, i] is a no-op. One routed axis per transition; it cannot be a factorize_axes block axis. Routing is lowered once per template, so its cost does not grow with the number of matrix entries. A named routing transition publishes one reaction whose propensity is shaped like the source plus the routed target axis (routed_axes): the channel for source i and target j has hazard R[i, j] X[i]. When source and target are otherwise the same slice, its no-op diagonal channels have zero propensity, so a generator's negative diagonal never becomes a hazard.

A target-only alias fans one source cell into a target axis the source does not own:

spec:
  kind: transitions
  axes:
    - {name: age, coords: [child, adult]}
    - {name: imm, type: ordinal, coords: [x0, x1, x2]}
  state: [I3[age], X[age,imm]]
  transitions:
    - from: I3[age]
      to: X[age,imm:j]
      rate: reset_rate * reset_kernel[imm:j]

This compiles as one lazy transition. Each target receives reset_rate * reset_kernel[j] * I3, while the source loses reset_rate * sum_j(reset_kernel[j]) * I3 exactly once. The weights are arbitrary per-target rates; op_system does not force normalization. When they sum to one, reset_rate is the total departure hazard. In every case the generated source loss equals the summed target inflow, so the transition is mass-conserving algebraically. Physical rate non-negativity remains a model input responsibility, consistent with other transition rates.

Chain helper

spec:
  kind: transitions
  state: [S, I, R]
  chain:
    - name: I
      length: 3
      entry:   {from: S, rate: beta * S / sum_state()}
      forward: [gamma12, gamma23]
      exit:    {to: R, rate: gamma3r}
  transitions: []

chain synthesizes the staged compartments (I1..I3) and the internal forward / exit transitions; declare only the base I in state. The generated transitions publish reactions named I_entry, I_advance_1, I_advance_2, and I_exit. A stage whose rate reads no state is already complete for adaptive stochastic solvers. For the others, list the reactants beyond each consumed stage: entry.catalysts for the entry rate (here [{state: I1, order: 1}, ...] for every infectious stage it reads) and the chain's catalysts for the forward and exit rates. catalysts: auto infers them instead when a rate is a single product of states, like reactants: auto.

Axis-wide aging with `coord_shift`

spec:
  kind: transitions
  axes:
    - {name: age, type: ordinal, coords: [a0, a1, a2, a3]}
  state: [S[age], I[age]]
  transitions:
    - name: aging
      coord_shift: {axis: age, step: 1, rate: "aging_rate[age]", boundary: absorb}
      apply_to: [S, I]

Every bin k moves to k + step at the source bin's rate. boundary: absorb removes mass shifted off the axis, and stay keeps it in the terminal bin. The entry lowers once per state, and named entries publish one templated reaction per state with an offsets field. See the aging-chain guide.

Continuous axis + kernel

spec:
  kind: expr
  axes:
    - name: x
      type: continuous
      domain: {lb: 0.0, ub: 10.0}
      size: 5
      spacing: linear
  state: [u[x]]
  state_axes: {u: [x]}
  kernels:
    - {name: K, axes: [x], form: gaussian, params: {scale: 1.0, sigma: 0.5}}
  equations:
    u[x]: apply_along(x=xi, K[x=xi] * u[x=xi]) - decay * u[x]

Public API

from op_system import (
    compile_spec,  # validate + normalize + compile
    compile_rhs,  # compile a pre-normalized NormalizedRhs
    normalize_rhs,  # validate + normalize only
    normalize_expr_rhs,
    normalize_transitions_rhs,
    CompiledRhs,
    NormalizedRhs,
    ExprRhs,
    TransitionsRhs,
    BodyEvalFn,
    EvalFn,
    PytreeEvalFn,
    StateDict,
    OperatorDescriptor,
    BlockAxisInfo,
)

NormalizedRhs is a discriminated union of ExprRhs | TransitionsRhs; use isinstance to dispatch.

Expression guardrails

Expressions are parsed with ast and restricted to:

Arithmetic, comparisons, ternary, boolean ops, names and constants.
A NumPy allowlist under the np. root: abs, exp, expm1, log, log1p, log2, log10, sqrt, maximum, minimum, clip, where, sin, cos, tan, sinh, cosh, tanh, hypot, arctan2.
Helpers: sum_state(), sum_prefix(prefix), apply_along(...), sum_over(...).

convolve_history(...) is available via the history-provider runtime hook (CompiledRhs.history_eval_fn and OpSystemSystem's options["history_stepper_fn"]). history(...) and delay(...) remain reserved for issue #173 and still raise a targeted unsupported-feature error with history_requirements=... payloads.

For adaptive ring-buffer engines, use CompiledRhs.body_eval_fn (or OpSystemSystem's options["body_eval_fn"]) to evaluate each history signal body exactly once at a known outer-step boundary. This complements history_eval_fn, which is still responsible for in-RHS history queries.

Each history requirement record currently includes: scope, kind, signal_expr, options, required_options, missing_required_options, and unknown_options.

Runnable convolve_history example

import numpy as np

from op_system import compile_spec

spec = {
    "kind": "expr",
    "axes": [{"name": "loc", "coords": ["a", "b"]}],
    "state": ["x[loc]"],
    "equations": {"x[loc]": "convolve_history(inflow[loc], kernel=gamma, window=14)"},
}
compiled = compile_spec(spec)

# history_eval_fn is available for axis-indexed convolve_history specs.
assert compiled.history_eval_fn is not None
print(compiled.history_requirements)


class ZeroHistoryProvider:
    def query(self, signal_id: int, body: object, **options: object) -> object:
        # Runtime contract from lowering: __hist_query(signal_id, body, **options)
        return np.zeros_like(body)


state = {"x": np.array([1.0, 2.0], dtype=np.float64)}
out = compiled.history_eval_fn(
    0.0,
    state,
    history_provider=ZeroHistoryProvider(),
    inflow=np.array([0.2, 0.4], dtype=np.float64),
)
print(out["x"])  # [0. 0.]

Anything else — non-np attribute access, imports, lambdas, comprehensions, other AST nodes — raises ValueError / TypeError / UnsupportedFeatureError at normalize time.

Development

just ci      # ruff + pytest + mypy (core + flepimop2-op_system mirror) + docs
just test    # pytest only
just ruff
just mypy
just docs    # mkdocs build

See docs/development/ for the IR architecture, block axis plan, and code-style guide.

Repository layout

Path	Purpose
`src/op_system/`	Library source (specs, IR, normalize, vectorize, compile).
`flepimop2-op_system/`	Thin adapter package exposing `op_system` to flepimop2.
`tests/op_system/`	Pytest suite (~430 tests).
`docs/`	mkdocs sources; built site published to GitHub Pages.
`scripts/`	Release validation and API-reference generation helpers.

Name		Name	Last commit message	Last commit date
.github/workflows		.github/workflows
docs		docs
flepimop2-op_system		flepimop2-op_system
overrides		overrides
scripts		scripts
src/op_system		src/op_system
tests		tests
.gitignore		.gitignore
CHANGELOG.md		CHANGELOG.md
LICENSE		LICENSE
README.md		README.md
justfile		justfile
mkdocs.yml		mkdocs.yml
pyproject.toml		pyproject.toml

Navigation Menu

Repository files navigation

op_system

Why op_system?

Installation

Time-indexed parameters

Quick start

Advection contract

Jump-integral contract

JAX usage

YAML examples

Baseline SIR (two pathways)

Templated states with `apply_along`

Routing transitions with `axis:alias`

Chain helper

Axis-wide aging with `coord_shift`

Continuous axis + kernel

Public API

Expression guardrails

Runnable convolve_history example

Development

Repository layout

About

Releases

Packages

Used by

Contributors

Languages

Latest commit

History

Folders and files

Repository files navigation

op_system

Why op_system?

Installation

Time-indexed parameters

Quick start

Advection contract

Jump-integral contract

JAX usage

YAML examples

Baseline SIR (two pathways)

Templated states with apply_along

Routing transitions with axis:alias

Chain helper

Axis-wide aging with coord_shift

Continuous axis + kernel

Public API

Expression guardrails

Runnable convolve_history example

Development

Repository layout

About

Resources

Stars

Watchers

Forks

Releases

Packages

Used by

Contributors

Languages

Templated states with `apply_along`

Routing transitions with `axis:alias`

Axis-wide aging with `coord_shift`