op-engine 0.3.0__tar.gz → 0.4.0__tar.gz

{op_engine-0.3.0 → op_engine-0.4.0}/CHANGELOG.md

@@ -5,6 +5,65 @@ versioning while the project remains pre-1.0.
 
 ## [Unreleased]
 
+## [0.4.0] - 2026-10-02
+
+### Added
+
+- Prepared Flepimop2 execution plans now support explicit PyTree and JAX block
+  layouts for fixed steps and frozen adaptive replay, including structured typed
+  operators, while preserving dynamic-value cache contracts (#168).
+- Direct SSA accepts explicit forcing breakpoints for exact piecewise-constant
+  rates, including dormant zero-rate intervals, while retaining pending events
+  across observation times. The provider exposes this for pure direct SSA (#171).
+- Fixed and adaptive tau-leaping stop at declared forcing changes, so each
+  leap segment uses its own starting rate; forcing boundaries take precedence
+  over tied critical or exact-fallback events (#174).
+- Pure stochastic runs consume the producer's declared `forcing_breakpoints`
+  (op_system `time_interpolation: previous`), combined with explicit engine
+  boundaries. Hybrid mode reports producer forcing as unsupported (#175).
+- `stochastic_method: thinning-ssa` gives exact bounded thinning for smoothly
+  time-varying rates. Supply a total-rate bound as `thinning_rate_bound` or a
+  `RateBoundFunction` through `run(..., rate_bound=...)`; `ThinningSSASolver`
+  is available in the core package (#173, #176, #177).
+- Stochastic execution consumes axis-wide `coord_shift` reactions from
+  op_system: each source bin of an `offsets` axis becomes a channel that moves
+  one unit by the published step, and off-axis destinations only deplete the
+  donor. Direct SSA reproduces Poisson bin occupancy and Erlang exit times for
+  pure aging chains (#180).
+- Pure stochastic validation and execution reject models with dynamics that
+  have no reaction artifact: transitions listed in the producer's
+  `reaction_gaps` coverage records, and typed operators. Each uncovered
+  transition is named by origin, selectors, and reason. Hybrid mode and
+  producers without coverage records are unchanged (#182).
+- Stochastic execution supports reactions between templates with different
+  axes, including axis-less states: scalar S→I→R models, axis-less donors
+  depositing into a pinned cell, and templated sources collapsing into an
+  axis-less state. Targets are indexed with op_system's `to_full_axes`,
+  falling back to `full_axes` for older producers (#184).
+- Stochastic execution runs routing and target-only fan-out reactions: each
+  source cell and routed target coordinate (op_system's `routed_axes`) is one
+  channel, so routing generators move discrete units between bins (#186).
+- Adaptive tau-leaping validation reports incomplete reactant metadata before
+  running, and both validation and run name each incomplete reaction and the
+  remedy: `reactants` on ordinary op_system transitions, or `catalysts` on
+  `chain:` and `coord_shift` entries (#188).
+
+### Changed
+
+- The provider's optional `op_system` integration requires
+  `flepimop2-op-system>=0.6.0` and `op-system>=0.6.0`.
+- Pure stochastic runs that previously dropped uncovered transitions or typed
+  operators silently now fail validation; use `mode: hybrid` or give each
+  transition a reaction artifact (#182).
+
+### Fixed
+
+- `PreparedExecution` can be passed directly to JAX transformations, including
+  mixed NumPy/JAX samples whose effective namespaces and dtypes are canonicalized
+  by JAX during tracing (#168).
+- Pinned source-only renewal births are accepted: donor-axis coverage is
+  validated only for reactions with a donor (#178, #179).
+
 ## [0.3.0] - 2026-09-27
 
 ### Added

{op_engine-0.3.0 → op_engine-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: op_engine
-Version: 0.3.0
+Version: 0.4.0
 Summary: Multiphysics integration engine for ODEs and PDEs using implicit-explicit operator splitting
 Project-URL: Repository, https://github.com/ACCIDDA/op_engine
 Project-URL: Documentation, https://accidda.github.io/op_engine/
@@ -88,6 +88,18 @@ With JAX as the array namespace:
 pip install "op_engine[jax]"
 ```
 
+For the other packaged Array-API backends:
+
+```bash
+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](docs/guides/backends.md) before assuming that JIT,
+automatic differentiation, adaptivity, or stochastic sampling have identical
+semantics across namespaces.
+
 With flepimop2 adapter:
 
 ```bash
@@ -139,7 +151,8 @@ adaptive controller is eager Python control flow. See the
 
 ```python
 import numpy as np
-from op_engine import CoreSolver, ModelCore, OperatorSpecs
+from op_engine import CoreSolver, ModelCore
+from op_engine.core_solver import OperatorSpecs, RunConfig
 
 n = 4
 times = np.linspace(0.0, 1.0, 11)
@@ -156,13 +169,20 @@ 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)
-
-# For IMEX methods set method and operators via RunConfig:
-# from op_engine.core_solver import RunConfig, AdaptiveConfig, DtControllerConfig
+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](docs/guides/solver-methods.md#imex-methods-and-operator-factories).
+
 ## 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.
@@ -170,7 +190,9 @@ solver.run(rhs, config=None)  # defaults: method="heun" (explicit)
 - `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.
 - 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`, `DtControllerConfig` for method/IMEX/adaptive control.
+- 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

{op_engine-0.3.0 → op_engine-0.4.0}/README.md

@@ -36,6 +36,18 @@ With JAX as the array namespace:
 pip install "op_engine[jax]"
 ```
 
+For the other packaged Array-API backends:
+
+```bash
+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](docs/guides/backends.md) before assuming that JIT,
+automatic differentiation, adaptivity, or stochastic sampling have identical
+semantics across namespaces.
+
 With flepimop2 adapter:
 
 ```bash
@@ -87,7 +99,8 @@ adaptive controller is eager Python control flow. See the
 
 ```python
 import numpy as np
-from op_engine import CoreSolver, ModelCore, OperatorSpecs
+from op_engine import CoreSolver, ModelCore
+from op_engine.core_solver import OperatorSpecs, RunConfig
 
 n = 4
 times = np.linspace(0.0, 1.0, 11)
@@ -104,13 +117,20 @@ 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)
-
-# For IMEX methods set method and operators via RunConfig:
-# from op_engine.core_solver import RunConfig, AdaptiveConfig, DtControllerConfig
+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](docs/guides/solver-methods.md#imex-methods-and-operator-factories).
+
 ## 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.
@@ -118,7 +138,9 @@ solver.run(rhs, config=None)  # defaults: method="heun" (explicit)
 - `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.
 - 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`, `DtControllerConfig` for method/IMEX/adaptive control.
+- 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

{op_engine-0.3.0 → op_engine-0.4.0}/docs/guides/backends.md

@@ -82,6 +82,9 @@ semantics.
 `NumpySSASampler` and `NumpyPoissonSampler` are provided as conveniences.
 JAX users can construct fresh keys from the stable direct-SSA `draw_index` or
 tau-leaping `step_index`.
+Direct SSA's optional forcing schedule is static Python configuration. It
+preserves the state namespace and uses fresh draw indices when a forcing
+boundary invalidates a pending event; observation times retain that event.
 
 The current safety checks and stochastic event loops are eager, so these
 stochastic solvers are not JIT-compatible paths. Adaptive selection also reads

op_engine-0.4.0/docs/guides/bounded-thinning.md

@@ -0,0 +1,112 @@
+# Bounded thinning SSA
+
+`ThinningSSASolver` provides an exact waiting-time method for smoothly varying
+propensities with user-certified bounds. This design addresses
+[#173](https://github.com/ACCIDDA/op_engine/issues/173), including the flepimop2
+provider's explicit `thinning-ssa` method.
+
+## Bound contract
+
+`ThinningSSASolver.run(propensity_func, thinning_sampler, *, rate_bound, config=None)`
+uses the same stoichiometry, reaction axis, and batched propensity contract as
+direct SSA. It permits arbitrary time dependence between events when a valid
+upper bound is supplied.
+
+`rate_bound` is either a finite non-negative real constant, or a
+`RateBoundFunction(t, state, limit)` returning `TotalRateBound(rate, valid_until)`.
+The bound covers the sum of **all** reaction and batch-cell propensities with
+the supplied state held fixed, over `[t, valid_until)`. The callback chooses a
+finite endpoint satisfying `t < valid_until <= limit`. `limit` is the earlier
+of the next declared forcing change and the solve's final time; observations
+do not shorten it. A constant bound applies to every state reached during the
+solve, until the next forcing change or final time.
+
+After an accepted event, the state changes and the callback is queried again.
+A rejected candidate leaves the state and current bound unchanged. At expiry,
+the solver discards a tied or later candidate, advances without firing, and
+requests a new bound. The same rule applies at forcing changes, with
+right-continuous propensities. No candidate fires exactly at an interval's
+right endpoint, including the solve's final endpoint.
+
+A zero bound declares the entire interval inactive. The solver checks rates
+at the interval's start, advances to its endpoint without sampling, and can
+resume with a positive bound. Zero instantaneous propensity with a positive
+bound still produces candidates, allowing smooth rates to become active.
+The caller must prove the bound over the whole interval. Checks at interval
+starts and candidate times detect encountered violations but cannot certify
+unsampled times. A violation fails visibly rather than clipping acceptance.
+Bounds and cumulative rates use the state's floating dtype. Bounds that
+overflow or become zero on conversion fail; tight bounds should allow for
+floating-point rounding when summing many channels. Callback inputs must be
+treated as read-only, and callbacks must not depend on observation storage.
+
+## Sampling contract
+
+`ThinningSampler(bound_rate, draw_index)` returns
+`ThinningSample(waiting_time, uniform)`, both scalar arrays in the state array
+namespace. The wait is finite and strictly positive, sampled exponentially at
+`bound_rate`. The independent uniform lies in `[0, 1)`. Draw indices start at
+zero and increase globally, including rejected and expired candidates.
+`NumpyThinningSampler(seed)` supplies the NumPy implementation; other backends
+inject their own sampler.
+
+At a candidate time, the solver evaluates the actual channel rates. If
+`uniform * bound_rate` is at least the actual total rate, it rejects the
+candidate. Otherwise it selects the flattened channel/batch event whose
+cumulative propensity contains that threshold. This uses one uniform for
+acceptance and channel selection: each channel has probability
+`channel_rate / bound_rate`, and the remaining probability rejects.
+
+Candidates beyond an observation are retained with their sampled uniform;
+their propensities are evaluated only when the solver reaches them. Inserting
+observations with the same solve endpoints therefore preserves the seeded path
+and callback/sampler history. `ThinningSSAConfig.max_candidates` bounds all
+candidate draws during one run, including rejection and expiry. Invalid
+sampler values, non-advancing clocks, invalid bounds, and impossible population
+updates fail explicitly.
+
+The construction follows
+[Lewis and Shedler's thinning method](https://doi.org/10.1002/nav.3800260304).
+The bound and expiry contract extends it to state-dependent reaction networks;
+exactness remains conditional on the caller's bound and random sampling laws.
+Existing direct SSA and tau-leaping behavior and samplers are unchanged.
+
+## Example
+
+For an independent birth channel in each batch cell with rate `2 * time`, the
+integrated intensity over `[0, 1]` is one per cell. The total upper bound must
+include every batch cell:
+
+```python
+import numpy as np
+
+from op_engine import ModelCore, NumpyThinningSampler, ThinningSSASolver
+
+core = ModelCore(1, 100, np.linspace(0, 1, 11))
+core.set_initial_state(np.zeros((1, 100)))
+
+
+def smooth_birth(time, state):
+    return np.full(state.shape, 2 * time, dtype=state.dtype)
+
+
+ThinningSSASolver(core, np.asarray([[1]])).run(
+    smooth_birth, NumpyThinningSampler(seed=173), rate_bound=200,
+)
+```
+
+A callback can instead return `TotalRateBound(2 * limit * state.shape[1], limit)`
+for this example. For consuming reactions, a bound that depends on the current
+state can be refreshed after each accepted event. A loose bound remains valid
+but generates more rejected candidates.
+
+The flepimop2 provider selects this method through
+`mode: stochastic` and `stochastic_method: thinning-ssa`. Supply a constant
+`thinning_rate_bound` in configuration, or leave it unset and pass a scalar or
+`RateBoundFunction` through `run(..., rate_bound=...)`. Both interfaces cannot
+be supplied together. `thinning_max_candidates` sets the per-run candidate
+guard. NumPy has seeded sampling through `random_seed`; other backends inject
+`thinning_sampler=`. Hybrid thinning is unsupported. See the
+[provider examples](https://github.com/ACCIDDA/op_engine/tree/main/flepimop2-op_engine#bounded-thinning-ssa)
+for configuration and callback usage. Existing direct SSA and tau-leaping
+defaults remain unchanged.

{op_engine-0.3.0 → op_engine-0.4.0}/docs/guides/solver-methods.md

@@ -2,8 +2,9 @@
 
 `ModelCore` owns output times and state storage. Deterministic systems use
 `CoreSolver` and a method selected by `RunConfig`. Stochastic reaction networks
-use `DirectSSASolver` or `TauLeapingSolver`, because propensities and
-stoichiometry have different semantics from an ODE right-hand side. In every
+use `DirectSSASolver`, `ThinningSSASolver`, `TauLeapingSolver`, or
+`AdaptiveTauLeapingSolver`, because
+propensities and stoichiometry have different semantics from an ODE right-hand side. In every
 case the state array selects the Array-API namespace; choosing JAX instead of
 NumPy does not select a different numerical method.
 
@@ -48,7 +49,8 @@ reaction axis.
 exponential waiting time and one categorical reaction/batch event, applies
 exactly one stoichiometric update, and repeats. A draw beyond an output boundary
 is retained rather than discarded, so adding observation times does not alter
-the simulated path. Zero total propensity is absorbing.
+the simulated path. Zero total propensity is absorbing when no declared future
+forcing boundary remains.
 
 Use direct SSA as the reference method for low-copy-number networks, rare-event
 questions, or checking a tau-leaping approximation. Its cost is proportional to
@@ -78,13 +80,111 @@ DirectSSASolver(core, stoichiometry).run(
 )
 ```
 
-The direct method assumes propensities remain constant between reaction events,
-as in a time-homogeneous continuous-time Markov chain. Across batch cells it
-samples from the superposed event process; the result is equivalent to
+The direct method assumes propensities remain constant in time while the state
+is unchanged, except at explicitly declared forcing boundaries. Across batch
+cells it samples from the superposed event process; the result is equivalent to
 independent direct-method trajectories for independent batch cells.
 
+#### Piecewise-constant forcing
+
+Declare every forcing change through `DirectSSAConfig.forcing_breakpoints`:
+
+```python
+from op_engine import DirectSSAConfig, array_namespace
+
+core = ModelCore(n_states=1, n_subgroups=1, time_grid=np.linspace(0.0, 3.0, 13))
+core.set_initial_state(np.zeros((1, 1)))
+
+
+def forced_birth_rate(time, state):
+    xp = array_namespace(state)
+    rate = 0.0 if time < 1.0 else (3.0 if time < 2.0 else 1.0)
+    return xp.full(state.shape, rate, dtype=state.dtype)
+
+
+DirectSSASolver(core, np.asarray([[1]])).run(
+    forced_birth_rate,
+    NumpySSASampler(seed=2026),
+    config=DirectSSAConfig(forcing_breakpoints=(1.0, 2.0)),
+)
+```
+
+If an event would reach or cross a forcing boundary, the solver advances to
+that boundary without firing, reevaluates the rates, and redraws. Exponential
+memorylessness makes this exact for piecewise-constant forcing. Zero-rate
+intervals advance to the next forcing boundary without drawing randomness.
+
+Forcing is right-continuous: the callback must return the new rates at a
+breakpoint. In an exact event/breakpoint tie, the forcing boundary wins and the
+old event is discarded. Observation times remain independent of the forcing
+schedule, even when an observation coincides with a breakpoint.
+
+Breakpoints must be finite real scalars in strictly increasing order. They are
+snapshotted as an immutable tuple; a global schedule may include times outside
+the solve interval. Boundaries at the initial time are already reflected in
+the first evaluation, and the final boundary causes no extra random draw.
+An omitted or empty schedule preserves ordinary time-homogeneous execution.
+
+The schedule does not freeze or interpolate callback values. All time
+dependence between boundaries must be constant for this exactness claim;
+smoothly varying rates require a different waiting-time method. In particular,
+`op_system` defaults to linear interpolation of time-indexed parameters, so
+listing their grid points as breakpoints does not make direct SSA exact for
+those interpolated rates. Producers containing
+[op_system PR #241](https://github.com/ACCIDDA/op_system/pull/241) can select
+`time_interpolation: previous` to hold table values until the next coordinate.
+The flepimop2 engine provider consumes their published `forcing_breakpoints`
+automatically in pure stochastic mode and combines them with any extra
+boundaries in the engine configuration. This also supports fixed and adaptive
+tau-leaping. Any additional direct time dependence must still satisfy the
+chosen solver's exactness conditions.
+
+### Exact SSA for smooth forcing
+
+`ThinningSSASolver` accepts time-dependent propensities and a user-certified
+total-rate bound. It samples candidates at the bound rate, then evaluates
+channel rates at the candidate time to reject or apply one event. See the
+[bounded thinning guide](bounded-thinning.md) for constant and callback bounds,
+sampling injection, interval expiry, and a runnable smooth-birth example.
+
 ### Bounded adaptive tau-leaping
 
+Both tau-leaping configurations accept the same `forcing_breakpoints` schedule
+as direct SSA. Each leap is capped at the next observation or forcing time,
+as well as the method's usual step-size limit. The next leap evaluates the
+right-continuous rates at that boundary. Boundaries count toward the fixed or
+adaptive method's `max_steps` guard for the entire output interval; they do
+not reset it.
+
+For example, the piecewise birth callback above can also use fixed tau-leaping:
+
+```python
+from op_engine import NumpyPoissonSampler, TauLeapingConfig, TauLeapingSolver
+
+core = ModelCore(n_states=1, n_subgroups=1, time_grid=np.linspace(0.0, 3.0, 13))
+core.set_initial_state(np.zeros((1, 1)))
+
+TauLeapingSolver(core, np.asarray([[1]])).run(
+    forced_birth_rate,
+    NumpyPoissonSampler(seed=2026),
+    config=TauLeapingConfig(max_step=0.1, forcing_breakpoints=(1.0, 2.0)),
+)
+```
+
+Adaptive tau-leaping uses these caps for both accepted leaps and rejected
+proposals. A critical event tied with a forcing boundary is discarded, while
+noncritical Poisson counts cover the segment ending there. During exact
+fallback, a pending event survives observation times but is discarded at or
+beyond the next forcing change. Zero total propensity waits for that change
+without drawing randomness; without a future boundary it remains absorbing.
+
+Boundary caps keep a leap from using the old rate across a known change. The
+usual tau-leaping approximation still freezes state-dependent propensities
+within each leap. Smooth time variation between boundaries also remains an
+approximation for leaps and is unsupported by adaptive exact fallback. For
+that fallback, use rates constant in time between declared forcing changes
+while the state is unchanged.
+
 `AdaptiveTauLeapingSolver` implements the species-based pre-leap selector of
 [Cao, Gillespie, and Petzold](https://doi.org/10.1063/1.2159468) together with
 their [non-negative modified Poisson