Systems

For convenience stochastix defines systems as a bundle of a reaction network and a solution method (simulation routine + solver). They are based on the assumption that most modelling aim to represent the behavior of a specific system (a ReactionNetwork) over a specified time period (T).

In practice, they are high-level convenience wrappers for running simulations with a simple functional interface. stochastix provides two simple templates for this:

• StochasticModel: for stochastic simulations using stochastix.stochsimsolve
• MeanFieldModel: for deterministic ODE simulations using diffrax.diffeqsolve

A whole simulation can be now carried out with a simple function call:

model = stochastix.StochasticModel(network, solver, T=10.0)
result = model(x0)

Both implementations have the same backbone:

class System(equinox.Module):
    network: ReactionNetwork
    solver: Solver
    ...

    def __init__(self, network: ReactionNetwork, solver: Solver):
        ...

    def __call__(self, x0: Array, T: float, **kwargs) -> Array:
        ...

For more complex use cases, you can fall back to the standard solution API or easily create your own custom system by appropriately subclassing equinox.Module.

stochastix.StochasticModel

StochasticModel(
    network: ReactionNetwork,
    solver: AbstractStochasticSolver,
    T: float | ndarray | None = None,
    max_steps: int = 10000,
)

A stochastic model combining a reaction network and a solver.

This class serves as a high-level convenience wrapper for running simulations. It composes a ReactionNetwork with a solver to define a complete executable simulation system.

Attributes:

• network –

  The reaction network defining the system's structure.

• solver –

  The solver for the simulation.

• T –

  The final time of the simulation.

• max_steps –

  The maximum number of simulation steps.

Parameters:

• network (ReactionNetwork) –

  The reaction network defining the system's structure.

• solver (AbstractStochasticSolver) –

  The solver for the simulation.

• T (float | ndarray | None, default: None ) –

  The final time of the simulation. If not provided, it must be specified during the call.

• max_steps (int, default: 10000 ) –

  The maximum number of simulation steps.

__call__

__call__(
    key: ndarray,
    x0: Any,
    T: float | ndarray | None = None,
    **kwargs: Any,
) -> SimulationResults

Execute the simulation.

This method internally calls the stochsimsolve function to run the simulation and returns the simulation results.

Parameters:

• key (ndarray) –

  The JAX random key.

• x0 (Any) –

  The initial state of the system.

• T (float | ndarray | None, default: None ) –

  The final time of the simulation. If provided, it overrides the T attribute of the model. If not provided, the model's T is used.

• **kwargs (Any) –

  Additional keyword arguments to pass to stochsimsolve.

Returns:

• solution ( SimulationResults ) –

  The results of the simulation.

stochastix.MeanFieldModel

MeanFieldModel(
    network: ReactionNetwork,
    solver: AbstractSolver = Dopri5(),
    T: float | ndarray | None = None,
    saveat_steps: int | list[float] | ndarray = -1,
    max_steps: int = 10000,
)

A mean field model combining a reaction network and a diffrax solver.

This class serves as a high-level convenience wrapper for running deterministic simulations. It composes a ReactionNetwork with a diffrax solver to define a complete executable simulation system for ordinary differential equations (ODEs).

Note

This class accommodates only basic use cases and is not very flexible. If you need more control, use diffrax explicitly or create your own callable model.

Attributes:

• network –

  The reaction network defining the system's structure.

• solver –

  The diffrax solver for the ODE simulation.

• T –

  The final time of the simulation.

• saveat –

  A diffrax.SaveAt object specifying when to save the solution.

• max_steps –

  The maximum number of steps for the solver.

Parameters:

• network (ReactionNetwork) –

  The reaction network defining the system's structure.

• solver (AbstractSolver, default: Dopri5() ) –

  The diffrax solver for the ODE simulation. Defaults to diffrax.Dopri5().

• T (float | ndarray | None, default: None ) –

  The final time of the simulation. If not provided, it must be specified during the call.

• saveat_steps (int | list[float] | ndarray, default: -1 ) –

  The time points at which to save the solution. Defaults to -1 (save only at the final time).

• max_steps (int, default: 10000 ) –

  The maximum number of steps to take.

__call__

__call__(
    key: ndarray | None,
    x0: ndarray,
    T: float | ndarray | None = None,
    max_steps: int | None = None,
    saveat_steps: int | list[float] | ndarray = None,
    **kwargs: Any,
) -> SimulationResults

Execute the ODE simulation using diffrax.

Parameters:

• key (ndarray | None) –

  For compatibility with stochastic models. This is ignored and can be safely set to None.

• x0 (ndarray) –

  The initial state of the system (species concentrations/counts).

• T (float | ndarray | None, default: None ) –

  The final time of the simulation. Overrides the model's T attribute.

• max_steps (int | None, default: None ) –

  Maximum number of steps to take. Overrides the model's max_steps attribute.

• saveat_steps (int | list[float] | ndarray, default: None ) –

  The time points to save the solution at. Overrides the model's saveat attribute.

• **kwargs (Any) –

  Additional keyword arguments to pass to diffrax.diffeqsolve.

Returns:

• solution ( SimulationResults ) –

  The results of the simulation, adapted for ODE output.