from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import molfoundry.dgstrat as dgstrat
from .core import TSObject, coerce_type
from .dg import DG
from .graph import Graph
from .graphprinter import GraphPrinter
from .history import HistoryRecorder
from .labelrelation import LabelRelation
from .labelsettings import LabelSettings
from .labeltype import LabelType
_class_ref = TSObject.get_class_ref("Stochsim")
# A per-step callback: receives a ``Stochsim.StepInfo`` for the start step and then
# every step as it is produced. Read species counts from the owning simulation
# inside the callback (it is invoked synchronously as each step is recorded).
OnStepLike = Callable[["Stochsim.StepInfo"], Any]
# A rate is either a constant or a callback, mirroring MØD's ``DrawMassAction``.
# A constant is a float (or a ``(rate, cache)`` pair; the cache flag is moot for a
# literal constant). A callback receives a DG.HyperEdge (reaction rates) or
# DG.Vertex (input/output flow rates) and returns a float or a ``(rate, cache)``
# pair. Unless it returns ``cache=True`` the callback is re-evaluated every step,
# so rates may depend on the current ``time`` or ``state(...)``.
RateConstant = Union[float, int, Tuple[float, bool]]
RateLike = Union[RateConstant, Callable[..., Any]]
# A waiting-time strategy: given the total propensity and ``uniform`` -- a
# zero-argument callable yielding independent draws in [0, 1) from the
# simulator's own seeded RNG -- return the time increment until the next
# reaction. Mirrors MØD's ``drawTime`` (default: exponential).
TimeDrawLike = Callable[[float, Callable[[], float]], float]
# noinspection PyPep8Naming
[docs]
class Stochsim(TSObject):
"""Rule-based Gillespie stochastic simulation with on-demand network expansion.
Implements Herrera Machado et al., "Rule-Based Gillespie Simulation of
Chemical Systems" (2025): the reaction network is not enumerated up front but
grown lazily as new species appear, while an exact SSA evolves the species
counts.
Rates may be constants (evaluated once) or callbacks. Following MØD's
``DrawMassAction``, a callback returns either a rate or a ``(rate, cache)``
pair and -- unless it caches -- is re-evaluated every step, so rates may
depend on the current :attr:`time` or :meth:`state`. The underlying simulator
applies the Law of Mass Action on top of the returned rate constant. The
waiting-time distribution is customizable via ``drawTime`` (default:
exponential), mirroring MØD's ``DrawTimeExponential``.
"""
def __init__(
self,
*,
graphDatabase: List[Graph],
expandStrategy: Any,
initialState: Dict[Graph, int],
labelSettings: Optional[LabelSettings] = None,
reactionRate: Optional[RateLike] = None,
inputRate: Optional[RateLike] = None,
outputRate: Optional[RateLike] = None,
drawTime: Optional[TimeDrawLike] = None,
retainHistory: bool = True,
onStep: Optional[OnStepLike] = None,
seed: int = 42,
) -> None:
"""Create a rule-based Gillespie simulation.
``retainHistory`` (default ``True``) keeps the whole trajectory in memory,
so :meth:`trajectory` / :meth:`series` and :meth:`print` work directly. Set
it to ``False`` for long / large runs to bound memory to O(species): the
engine then holds only the start and current step. The full trajectory is
instead streamed to a history file under ``./out`` (adjacent to
``./summary``) so :meth:`print` can still rebuild the interactive playback;
:meth:`trajectory` / :meth:`series`, which need the in-memory log, are
unavailable in that mode.
``onStep`` is an optional callback invoked once for the start step and then
for every step as it is produced, receiving a :class:`StepInfo`
(``index``, ``time``, ``reaction``). It fires synchronously as the step is
recorded, so it may read the live state via this simulation (e.g.
:meth:`state`) — the streaming hook that pairs with ``retainHistory=False``.
"""
if labelSettings is None:
labelSettings = LabelSettings(LabelType.String, LabelRelation.Isomorphism)
strat = dgstrat.DGStrat._dg_strat(expandStrategy)
super().__init__(
_class_ref,
labelSettings._ref,
[g._ref for g in graphDatabase],
strat._ref,
int(seed),
)
for g, c in initialState.items():
self._ref.addInitial(g._ref, int(c))
if reactionRate is not None:
self.setReactionRate(reactionRate)
if inputRate is not None:
self.setInputRate(inputRate)
if outputRate is not None:
self.setOutputRate(outputRate)
if drawTime is not None:
self.setDrawTime(drawTime)
#: History recorder streaming the trajectory to disk (only when history is
#: not retained in memory, so the viz can still be rebuilt); else ``None``.
self._history: Optional[HistoryRecorder] = None
if not retainHistory:
self._ref.setRetainHistory(False)
self._history = HistoryRecorder("stochsim")
self._ref.enableHistoryRecording(self._history.sink)
if onStep is not None:
self.setOnStep(onStep)
# --- rate configuration --------------------------------------------------
@staticmethod
def _rateValue(result: Any) -> float:
"""Scalar rate from a float or an MØD-style ``(rate, cache)`` pair."""
if isinstance(result, (tuple, list)):
return float(result[0])
return float(result)
@staticmethod
def _rateResult(result: Any) -> Dict[str, Any]:
"""Normalize a callback return into the ``{rate, cache}`` shape the engine
expects. A bare scalar defaults to ``cache=False`` (re-evaluated each
step); a ``(rate, cache)`` pair carries its flag through."""
if isinstance(result, (tuple, list)):
cache = bool(result[1]) if len(result) > 1 else False
return {"rate": float(result[0]), "cache": cache}
return {"rate": float(result), "cache": False}
[docs]
def setReactionRate(self, rate: RateLike) -> None:
"""Set the rate for each reaction (default: 1 per hyperedge).
``rate`` is a constant (float, or a ``(rate, cache)`` pair) or a callback
``(DG.HyperEdge) -> float | (float, bool)``. A callback is re-evaluated
every step unless it returns ``cache=True``, so reaction rates may depend
on :attr:`time` or :meth:`state`. The engine applies the Law of Mass
Action on top of the returned rate constant.
"""
if callable(rate):
self._ref.setReactionRate(
lambda he: Stochsim._rateResult(rate(DG.HyperEdge.wrap(he)))
)
else:
self._ref.setReactionRate(Stochsim._rateValue(rate))
[docs]
def setOutputRate(self, rate: RateLike) -> None:
"""Set the output-flow (species -> ∅) rate (default: 0, closed system).
``rate`` is a constant or a callback
``(DG.Vertex) -> float | (float, bool)`` with the same re-evaluation
semantics as :meth:`setReactionRate`.
"""
if callable(rate):
self._ref.setOutputRate(
lambda v: Stochsim._rateResult(rate(DG.Vertex.wrap(v)))
)
else:
self._ref.setOutputRate(Stochsim._rateValue(rate))
[docs]
def setDrawTime(self, drawTime: Optional[TimeDrawLike]) -> None:
"""Set the waiting-time strategy, or ``None`` to restore the default
Gillespie exponential draw ``ln(1 / uniform()) / totalPropensity``.
``drawTime`` is called as ``drawTime(totalPropensity, uniform)``, where
``uniform()`` yields independent draws in [0, 1) from the simulator's own
seeded RNG, and must return the time increment until the next reaction.
Must be set before the simulation starts.
"""
if drawTime is None:
cb = None
else:
cb = lambda total, uniform: float(drawTime(float(total), uniform))
TSObject.wrap_exception(lambda: self._ref.setDrawTime(cb))
[docs]
def setOnStep(self, onStep: Optional[OnStepLike]) -> None:
"""Set a per-step callback, or ``None`` to clear it.
``onStep`` receives a :class:`StepInfo` once for the start step and then for
every step produced. Set it before :meth:`simulate` so the start step is
seen; the callback runs synchronously as the step is recorded and may read
the live state through this simulation.
"""
if onStep is None:
self._ref.setOnStep(None)
return
def cb(index: Any, time: Any, reaction: Any) -> None:
name = coerce_type(reaction)
onStep(
Stochsim.StepInfo(
int(index),
float(time),
None if name is None else str(name),
)
)
self._ref.setOnStep(cb)
# --- running -------------------------------------------------------------
[docs]
def simulate(
self,
time: Optional[float] = None,
iterations: Optional[int] = None,
advanceToEndTime: bool = False,
) -> None:
"""Advance the simulation up to a time bound and/or iteration bound.
Both bounds are relative to the current state, so repeated calls
continue. With neither bound the simulation runs until it deadlocks.
With a ``time`` bound, ``advanceToEndTime`` parks the clock at exactly that
time when no reaction fires before it (the drawn event overshot, or the
system is momentarily dead), appending a no-reaction marker step at the
bound. This lets a later ``simulate`` segment re-evaluate a time-dependent
rate at that instant -- reviving an otherwise dead system -- or a scheduled
intervention (e.g. :meth:`state`-dependent rate, or a follow-up action) land
at exactly that time. Off by default, so the clock otherwise stops at the
last event.
"""
TSObject.wrap_exception(
lambda: self._ref.simulate(
None if time is None else float(time),
None if iterations is None else int(iterations),
bool(advanceToEndTime),
)
)
# --- state ---------------------------------------------------------------
@property
def dg(self) -> DG:
return DG.wrap(self._ref.getDG())
@property
def iteration(self) -> int:
return int(self._ref.getIteration())
@property
def time(self) -> float:
return float(self._ref.getTime())
[docs]
def state(self, graph: Graph) -> int:
"""Current molecule count of a species graph."""
return int(TSObject.wrap_exception(lambda: self._ref.state(graph._ref)))
[docs]
def trajectory(self) -> "Stochsim.Trajectory":
"""Full event-by-event trajectory of the simulation."""
t = TSObject.wrap_exception(lambda: self._ref.getTrajectory())
times = [float(x) for x in t.times]
vertexIds = [int(x) for x in t.vertexIds]
counts = [[int(c) for c in row] for row in t.counts]
return Stochsim.Trajectory(times, vertexIds, counts)
[docs]
def series(self, graph: Graph) -> List[int]:
"""Per-step count series of ``graph`` aligned with ``trajectory().times``."""
traj = self.trajectory()
v = self.dg.findVertex(graph)
if not v or v.id not in traj.vertexIds:
return [0 for _ in traj.times]
return traj.series(v.id)
[docs]
def print(
self,
printer: Optional[GraphPrinter] = None,
*,
name: Optional[str] = None,
dark: Any = "dynamic",
) -> str:
"""Write an interactive HTML playback of this simulation to the summary folder.
Produces one standalone, seekable page (``<name>.html`` next to
``summary.html``) that lays out the final derivation graph as oval-framed
molecule depictions, shows each molecule's current count below it, and
reveals nodes/hyperedges at the step they were first expanded. Play / pause
/ stop and a speed control step through the whole simulation.
``printer`` is an optional :class:`GraphPrinter` -- exactly as for
:meth:`DG.print` -- threaded to every molecule depiction so its ``drawMode``
/ ``collapseHydrogens`` settings control how the molecules are drawn.
``name`` sets the filename (default ``stochsim``). Like the rest of the
summary the file is flushed when the process exits; returns the target
filename.
With ``retainHistory=False`` the trajectory is rebuilt from the streamed
``./out`` history file rather than the (absent) in-memory log.
"""
from . import viz
trajectory = self._history.read() if self._history is not None else None
return viz.printViz(
self._ref, "stochsim", name, printer, dark, trajectory
)
[docs]
class StepInfo:
"""One step handed to an ``onStep`` callback."""
def __init__(
self, index: int, time: float, reaction: Optional[str]
) -> None:
#: Absolute step index in the trajectory (0 is the start step).
self.index = index
#: Simulation time at this step.
self.time = time
#: Name of the reaction that fired, or ``None`` for the start step and
#: no-reaction time markers.
self.reaction = reaction
def __repr__(self) -> str:
return (
f"StepInfo(index={self.index}, time={self.time}, "
f"reaction={self.reaction!r})"
)
[docs]
class Trajectory:
def __init__(
self, times: List[float], vertexIds: List[int], counts: List[List[int]]
) -> None:
#: Time stamp of each recorded step.
self.times = times
#: DG vertex index of each column of ``counts``.
self.vertexIds = vertexIds
#: ``counts[step][column]`` molecule counts.
self.counts = counts
[docs]
def series(self, vertexId: int) -> List[int]:
col = self.vertexIds.index(vertexId)
return [row[col] for row in self.counts]