Metadata-Version: 2.5
Name: op_system
Version: 0.7.0
Summary: Config-based RHS specification + compilation utilities for ODE/PDE systems
Project-URL: Documentation, https://accidda.github.io/op_system/
Project-URL: Repository, https://github.com/ACCIDDA/op_system
Project-URL: Issues, https://github.com/ACCIDDA/op_system/issues
Project-URL: Changelog, https://github.com/ACCIDDA/op_system/blob/main/CHANGELOG.md
Author-email: Joshua Macdonald <jmacdo16@jh.edu>, Carl Pearson <cap1024@unc.edu>, Timothy Willard <twillard@unc.edu>
License: MIT License
        
        Copyright (c) 2026 Joshua Macdonald, Carl Pearson, and Timothy Willard
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: <3.14,>=3.11
Requires-Dist: array-api-compat>=1.12
Requires-Dist: numpy>=2.0
Requires-Dist: pydantic<3.0.0,>=2.12.5
Provides-Extra: data
Requires-Dist: pandas>=2.2; extra == 'data'
Requires-Dist: pyarrow>=16.0; extra == 'data'
Provides-Extra: jax
Requires-Dist: jax>=0.4; extra == 'jax'
Provides-Extra: jax-inference
Requires-Dist: blackjax; extra == 'jax-inference'
Requires-Dist: diffrax; extra == 'jax-inference'
Requires-Dist: jax>=0.4; extra == 'jax-inference'
Provides-Extra: torch
Requires-Dist: torch>=2.6; extra == 'torch'
Description-Content-Type: text/markdown

# 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

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

Optional extras:

```bash
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](https://accidda.github.io/op_system/guides/time-indexed-parameters/)
for the schema, examples, and exactness conditions.

## Quick start

```python
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:

```yaml
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](https://accidda.github.io/op_system/guides/operators/)
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

```python
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)

```yaml
# 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
```

```yaml
# 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):

```yaml
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:

```yaml
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](docs/guides/renewal-births.md) for
reaction metadata, retained group axes, and a stationary age-population example.

### Templated states with `apply_along`

```yaml
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`

```yaml
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:

```yaml
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

```yaml
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`

```yaml
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](https://accidda.github.io/op_system/guides/aging-chains/).

### Continuous axis + kernel

```yaml
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

```python
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

```python
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

```bash
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/](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. |
