SimulationResults

stochastix.SimulationResults

SimulationResults(
    x: Any,
    t: ndarray,
    propensities: ndarray | None = None,
    reactions: ndarray | None = None,
    time_overflow: bool | None = None,
    species: tuple[str, ...] | None = None,
)

Container for simulation results.

This class is used to store the results of a stochastic simulation. Includes utilities for indexing, cleaning, and interpolating the results.

Attributes:

• x –

  State trajectory over time. Shape (n_timepoints, n_species) for full trajectory, or (2, n_species) for initial/final only mode.

• t –

  Time points corresponding to state changes. Shape (n_timepoints,) for full trajectory, or (2,) for initial/final only mode.

• propensities –

  Reaction propensities at each time step, or None if not saved.

• reactions –

  Index of reactions that occurred at each step, or None if not saved.

• time_overflow –

  True if simulation stopped due to reaching max_steps before time T.

• species –

  Names of the species in the simulation.

Methods:

• clean –

  Remove padding steps from the simulation results (returns a new object).

• interpolate –

  Interpolate the simulation results to new time points (returns a new object).

Indexing into batched simulation results

It is pretty common to generate a batch of simulation results, e.g. by doing:

keys = jax.random.split(key, num_simulations)
results = equinox.filter_vmap(stochsimsolve, in_axes=(0, None, None, None))(keys, network, x0, T)

In this case, results is a SimulationResults object where each field but species has an extra first dimension of size num_simulations. For convenience, you can directly index into the batch of results (this is not standard JAX behavior):

>>> results[0]  # First simulation result
# SimulationResults(x[0, ...], t[0, ...], propensities[0, ...], reactions[0, ...], time_overflow=[0, ...], species=species)
>>> results[1:]  # All but the first simulation result
# SimulationResults(x[1:, ...], t[1:, ...], propensities[1:, ...], reactions[1:, ...], time_overflow=[1, ...], species=species)

Removing padding from simulation results

If simulation time did not overflow, i.e. max_steps was not reached, SimulationResults will contain unused steps at the end of the simulation, with reactions set to -1. You can remove these steps using the clean method:

cleaned_results = results.clean()

Interpolating simulation results

You can interpolate the simulation results to new time points using the interpolate method. This is especially useful for plotting or calculating statistics over the simulation results that require a regular time grid.

t_interp = jnp.linspace(0, T, 100)
interpolated_results = results.interpolate(t_interp)

clean

clean() -> SimulationResults

Cleans simulation results by removing padded, unused steps.

Note

Does not work with batched simulations or initial/final only mode (when x has shape (2, ...)).

Stochastic simulations pre-allocate arrays of a fixed size for JIT compilation. This function removes the trailing steps that were allocated but not used because the simulation terminated early (e.g., reached the final time or ran out of reactants).

It identifies unused steps by checking for negative reaction indices, which are used as padding values by the solvers. It correctly handles results from both exact solvers (where reactions is 1D) and approximate solvers (where reactions is 2D).

Returns:

• SimulationResults ( SimulationResults ) –

  A new container with only the valid steps of the simulation trajectory.

Raises:

• NotImplementedError –

  If called on batched simulation results or initial/final only mode.

interpolate

interpolate(t_interp: ndarray) -> SimulationResults

Interpolates stochastic simulation results to new time points.

This function performs a piecewise-constant (or "forward-fill") interpolation of the simulation's state trajectory x. The state x[i] is considered constant over the time interval [t[i], t[i+1]).

Parameters:

• t_interp (ndarray) –

  A sorted array of time points at which to interpolate the state. Values outside the simulation time range [self.t[0], self.t[-1]] will be automatically clamped to the nearest boundary.

Returns:

• SimulationResults ( SimulationResults ) –

  A new container with the interpolated state trajectory at the specified time points. The reactions and propensities fields are set to None, as they are not meaningful at arbitrary time points. time_overflow is also set to None, since interpolations do not exceed the simulation time range.

Note

This method is JIT-compatible. Time points outside the simulation range are automatically clamped to [self.t[0], self.t[-1]] to prevent out-of-bounds errors. Input validation should be performed before calling this method if strict bounds checking is required.