from typing import Any, Callable, Dict, List, Optional, 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
from .stochsim import OnStepLike, RateLike, Stochsim, TimeDrawLike
_class_ref = TSObject.get_class_ref("MultiPhaseStochsim")
# A channel rate: a constant (float/int) or a callback given the transferred
# species graph -- ``(Graph) -> float | (float, bool)`` -- re-evaluated each step
# (the engine applies mass action on top of the returned rate constant).
ChannelRateLike = Union[float, int, Callable[[Graph], Any]]
# A coupled-channel rate: a constant or a no-argument callback
# ``() -> float | (float, bool)``.
CoupledRateLike = Union[float, int, Callable[[], Any]]
# noinspection PyPep8Naming
[docs]
class MultiPhaseStochsim(TSObject):
"""Exact multiphase (multi-compartment) stochastic simulation.
``n`` phases each run the rule-based Gillespie SSA with lazy network
expansion; the whole system is evolved as one combined continuous-time Markov
chain by a single shared simulator (one clock, one RNG, one global
propensity) -- the exact Direct-Method realization.
Stage A: phases only, no interphase channels yet. A phase may still exchange
with the implicit outside via its own ``inputRate``/``outputRate`` flows, and
a single channel-free phase reproduces :class:`Stochsim` for the same seed.
All phases share one :class:`LabelSettings` so species identity is comparable
across phases by canonical key.
"""
def __init__(
self,
*,
labelSettings: Optional[LabelSettings] = None,
retainHistory: bool = True,
onStep: Optional[OnStepLike] = None,
seed: int = 42,
) -> None:
"""Create a multiphase stochastic simulation.
``retainHistory`` and ``onStep`` have the same meaning as for
:class:`~molfoundry.stochsim.Stochsim`: with ``retainHistory=False`` the
engine holds only O(species) state and the full trajectory is streamed to a
history file under ``./out`` so :meth:`print` can still rebuild the playback
(while :meth:`trajectory` / :meth:`series` / :meth:`times`, which need the
in-memory log, are unavailable). ``onStep`` fires once per step with a
:class:`~molfoundry.stochsim.Stochsim.StepInfo`.
"""
if labelSettings is None:
labelSettings = LabelSettings(LabelType.String, LabelRelation.Isomorphism)
super().__init__(_class_ref, labelSettings._ref, int(seed))
self._phases: List["MultiPhaseStochsim.Phase"] = []
#: History recorder streaming the trajectory to disk 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("multiphase")
self._ref.enableHistoryRecording(self._history.sink)
if onStep is not None:
self.setOnStep(onStep)
[docs]
def setOnStep(self, onStep: Optional[OnStepLike]) -> None:
"""Set a per-step callback, or ``None`` to clear it.
See :meth:`Stochsim.setOnStep`. The callback receives a
:class:`~molfoundry.stochsim.Stochsim.StepInfo` once for the start step and
then for every step across the whole (shared-clock) system.
"""
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)
[docs]
def addPhase(
self,
name: str,
*,
graphDatabase: List[Graph],
expandStrategy: Any,
volume: float = 1.0,
initialState: Optional[Dict[Graph, int]] = None,
reactionRate: Optional[RateLike] = None,
inputRate: Optional[RateLike] = None,
outputRate: Optional[RateLike] = None,
labelSettings: Optional[LabelSettings] = None,
) -> "MultiPhaseStochsim.Phase":
"""Add a phase (compartment) with its own chemistry, volume, and rates.
``graphDatabase``, ``expandStrategy``, and the rate arguments have the
same meaning as for :class:`Stochsim`; ``volume`` is stored for later
partition channels.
``labelSettings`` overrides the constructor's :class:`LabelSettings` for
this phase only. Species identity stays comparable across phases (they
are matched by the graph's intrinsic canonical key), so, e.g., a
wildcard-free phase can use cheap string/isomorphism matching while a
phase with ``_X``-style term rules uses term/specialisation.
"""
strat = dgstrat.DGStrat._dg_strat(expandStrategy)
ls_ref = None if labelSettings is None else labelSettings._ref
ref = TSObject.wrap_exception(
lambda: self._ref.addPhase(
str(name), [g._ref for g in graphDatabase], strat._ref,
float(volume), ls_ref,
)
)
phase = MultiPhaseStochsim.Phase(self, len(self._phases), ref)
self._phases.append(phase)
if initialState:
for g, c in initialState.items():
phase.addInitial(g, c)
if reactionRate is not None:
phase.setReactionRate(reactionRate)
if inputRate is not None:
phase.setInputRate(inputRate)
if outputRate is not None:
phase.setOutputRate(outputRate)
return phase
[docs]
def addChannel(
self,
*,
source: Any,
target: Any,
rate: ChannelRateLike,
species: Optional[Graph] = None,
select: Optional[Any] = None,
) -> "MultiPhaseStochsim.Channel":
"""Add a constant first-order transport channel ``source -> target``.
Selectivity (which species may cross) is one of: ``species`` (a single
graph), ``select`` (a predicate ``Graph -> bool`` or an iterable of
graphs), or neither (all species). ``rate`` is the per-molecule transfer
rate *constant* (the engine applies mass action, ``rate * count(species in
source)``, on top), either a number or a callback ``(Graph) -> float |
(float, bool)`` re-evaluated each step for saturable / gated /
state-dependent transport. Bidirectional exchange is two channels (see
``addPartition`` for the thermodynamically-consistent pair).
"""
si = self._phaseIndex(source)
ti = self._phaseIndex(target)
predicate = MultiPhaseStochsim._buildSelect(species, select)
rate_arg = MultiPhaseStochsim._buildChannelRate(rate)
ref = TSObject.wrap_exception(
lambda: self._ref.addChannel(si, ti, predicate, rate_arg)
)
return MultiPhaseStochsim.Channel(self, ref)
[docs]
def addCoupledChannel(
self, *, reactants: Any, products: Any, rate: CoupledRateLike
) -> "MultiPhaseStochsim.Channel":
"""Add a coupled cross-phase reaction (symport / antiport / reactive transport).
``reactants`` and ``products`` are lists of ``(species_graph, phase)``
pairs (``phase`` a Phase or index). Mass action requires all reactants
present, so e.g. a symport fires only when every co-transported species is
available. ``rate`` is a constant or a no-argument callback ``() -> float
| (float, bool)``. All named species must be representable in their phases
(in the ``graphDatabase`` or derivable there).
"""
rp = [self._phaseIndex(p) for _, p in reactants]
rg = [g._ref for g, _ in reactants]
pp = [self._phaseIndex(p) for _, p in products]
pg = [g._ref for g, _ in products]
rate_arg = MultiPhaseStochsim._buildCoupledRate(rate)
ref = TSObject.wrap_exception(
lambda: self._ref.addCoupledChannel(rp, rg, pp, pg, rate_arg)
)
return MultiPhaseStochsim.Channel(self, ref)
@staticmethod
def _buildChannelRate(rate: ChannelRateLike):
"""Build a channel rate: a JS ``(Graph) -> {rate, cache}`` callback or a float."""
if callable(rate):
return lambda g: Stochsim._rateResult(rate(Graph.wrap(g)))
return float(rate)
@staticmethod
def _buildCoupledRate(rate: CoupledRateLike):
"""Build a coupled-channel rate: a JS ``() -> {rate, cache}`` callback or a float."""
if callable(rate):
return lambda: Stochsim._rateResult(rate())
return float(rate)
[docs]
def addPartition(
self,
*,
source: Any,
target: Any,
Kd: float,
rate: float,
species: Optional[Graph] = None,
select: Optional[Any] = None,
) -> "tuple[MultiPhaseStochsim.Channel, MultiPhaseStochsim.Channel]":
"""Add a thermodynamically-consistent partition (a bidirectional pair).
``Kd`` is the equilibrium distribution coefficient ``[target]/[source]``;
``rate`` is the forward (source -> target) per-molecule rate constant. The
backward constant is derived to enforce detailed balance across the
(possibly unequal) phase volumes: ``k_f/k_b = Kd * (V_target/V_source)``,
so the pair on its own relaxes to ``[target]/[source] = Kd``. Selectivity
is specified as for :meth:`addChannel`. Returns ``(forward, backward)``.
"""
si = self._phaseIndex(source)
ti = self._phaseIndex(target)
predicate = MultiPhaseStochsim._buildSelect(species, select)
refs = TSObject.wrap_exception(
lambda: self._ref.addPartition(si, ti, predicate, float(Kd), float(rate))
)
pair = [MultiPhaseStochsim.Channel(self, r) for r in refs]
return pair[0], pair[1]
@staticmethod
def _buildSelect(species: Optional[Graph], select: Optional[Any]):
"""Build a JS-callable ``(Graph) -> bool`` species predicate."""
if species is not None:
key = species.canonKey
return lambda g: Graph.wrap(g).canonKey == key
if select is None:
return lambda g: True
if callable(select):
return lambda g: bool(select(Graph.wrap(g)))
keys = {x.canonKey for x in select}
return lambda g: Graph.wrap(g).canonKey in keys
[docs]
def setDrawTime(self, drawTime: Optional[TimeDrawLike]) -> None:
"""Set the global waiting-time strategy; see :meth:`Stochsim.setDrawTime`."""
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 simulate(
self,
time: Optional[float] = None,
iterations: Optional[int] = None,
advanceToEndTime: bool = False,
) -> None:
"""Advance the whole system up to a time and/or iteration bound.
Both bounds are relative to the current state, so repeated calls continue;
with neither bound the system runs until a global deadlock.
With a ``time`` bound, ``advanceToEndTime`` parks the shared clock at exactly
that time when no reaction fires before it (the drawn event overshot, or the
whole system is momentarily dead), so a subsequent segment re-evaluates
time-dependent rates there (reviving a dead system) or a scheduled
:meth:`transfer` lands at exactly that time. Off by default.
"""
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),
)
)
[docs]
def transfer(
self,
source: Any,
target: Any,
*,
species: Optional[Graph] = None,
fraction: Optional[float] = None,
amount: Optional[int] = None,
) -> int:
"""Discrete scheduled transfer between two phases, applied at the current time.
This is an operator step *between* :meth:`simulate` segments -- serial
passage / dilution, pipetting an aliquot, decanting a layer, or a
deterministic spike-in -- not an intrinsic reaction, so it does not advance
the iteration count. It is applied atomically (one injection) and the target
phase's network is expanded around any species new to it, so the transferred
molecules react and flow onward on subsequent steps. Provide exactly one of:
* ``fraction`` -- move a binomially-sampled aliquot (each molecule crosses
independently with probability ``fraction``, drawn on the shared RNG). With
``species=None`` every species currently present in ``source`` is aliquoted
(a whole-compartment dilution); with a ``species`` graph, only that one.
* ``amount`` -- move exactly that many molecules of a specific ``species`` (a
deterministic spike-in / decant); raises if fewer are present.
Returns the total number of molecules moved (``source`` and ``target`` are a
Phase or an index).
"""
si = self._phaseIndex(source)
ti = self._phaseIndex(target)
if (fraction is None) == (amount is None):
raise ValueError("Provide exactly one of 'fraction' or 'amount'.")
if amount is not None:
if species is None:
raise ValueError("An exact 'amount' transfer requires a 'species'.")
return int(
TSObject.wrap_exception(
lambda: self._ref.transferAmount(
si, ti, species._ref, int(amount)
)
)
)
sref = None if species is None else species._ref
return int(
TSObject.wrap_exception(
lambda: self._ref.transferFraction(si, ti, sref, float(fraction))
)
)
@property
def phases(self) -> List["MultiPhaseStochsim.Phase"]:
return list(self._phases)
@property
def iteration(self) -> int:
return int(self._ref.getIteration())
@property
def time(self) -> float:
return float(self._ref.getTime())
def _phaseIndex(self, phase: Any) -> int:
if isinstance(phase, MultiPhaseStochsim.Phase):
return phase.index
return int(phase)
[docs]
def state(self, phase: Any, graph: Graph) -> int:
"""Current molecule count of ``graph`` in ``phase`` (a Phase or index)."""
idx = self._phaseIndex(phase)
return int(TSObject.wrap_exception(lambda: self._ref.state(idx, graph._ref)))
[docs]
def concentration(self, phase: Any, graph: Graph) -> float:
"""Current concentration (count / volume) of ``graph`` in ``phase``."""
idx = self._phaseIndex(phase)
return float(
TSObject.wrap_exception(lambda: self._ref.concentration(idx, graph._ref))
)
[docs]
def total(self, graph: Graph) -> int:
"""Current total count of ``graph`` summed across all phases."""
return int(TSObject.wrap_exception(lambda: self._ref.total(graph._ref)))
[docs]
def times(self) -> List[float]:
"""The time stamp of each recorded step (shared across all phases)."""
return [float(x) for x in TSObject.wrap_exception(lambda: self._ref.getTimes())]
[docs]
def totalSeries(self, graph: Graph) -> List[int]:
"""Per-step total count of ``graph`` across all phases, aligned with times."""
traj = self.trajectory()
totals = [0 for _ in traj.times]
for ph in self._phases:
v = ph.dg.findVertex(graph)
if not v:
continue
s = traj.series(ph.index, v.id)
for i in range(len(totals)):
totals[i] += s[i]
return totals
[docs]
def trajectory(self) -> "MultiPhaseStochsim.Trajectory":
"""Full event-by-event trajectory across all phases."""
t = TSObject.wrap_exception(lambda: self._ref.getTrajectory())
times = [float(x) for x in t.times]
columnPhases = [int(x) for x in t.columnPhases]
columnVertexIds = [int(x) for x in t.columnVertexIds]
counts = [[int(c) for c in row] for row in t.counts]
return MultiPhaseStochsim.Trajectory(
times, columnPhases, columnVertexIds, counts
)
[docs]
def series(self, phase: Any, graph: Graph) -> List[int]:
"""Per-step count series of ``graph`` in ``phase``, aligned with times."""
idx = self._phaseIndex(phase)
ph = self._phases[idx]
traj = self.trajectory()
v = ph.dg.findVertex(graph)
if not v:
return [0 for _ in traj.times]
return traj.series(idx, 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 multiphase run to the summary folder.
Produces one standalone, seekable page (``<name>.html`` next to
``summary.html``) with one box per phase (name at the top-left), the final
per-phase derivation graphs as oval-framed molecule depictions with live
counts below them, dotted bent arrows for the inter-phase channels, and
input/output flows drawn to the outside of each box. Nodes and hyperedges
are revealed 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 (across all phases)
so its ``drawMode`` / ``collapseHydrogens`` settings control how the
molecules are drawn. ``name`` sets the filename (default ``multiphase``).
Flushed at process exit; 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, "multiphase", name, printer, dark, trajectory
)
[docs]
class Phase:
"""A single phase (compartment). Obtained from :meth:`addPhase`."""
def __init__(
self, parent: "MultiPhaseStochsim", index: int, ref: Any
) -> None:
self._parent = parent
#: Index of this phase within its owning simulation.
self.index = index
self._ref = ref
@property
def name(self) -> str:
return str(self._ref.getName())
@property
def volume(self) -> float:
return float(self._ref.getVolume())
@property
def dg(self) -> DG:
"""The derivation graph underlying this phase's network."""
return DG.wrap(self._ref.getDG())
[docs]
def addInitial(self, graph: Graph, count: int) -> None:
TSObject.wrap_exception(
lambda: self._ref.addInitial(graph._ref, int(count))
)
[docs]
def setReactionRate(self, rate: RateLike) -> None:
"""See :meth:`Stochsim.setReactionRate` (scoped to this phase)."""
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:
"""See :meth:`Stochsim.setOutputRate` (this phase's flow to the outside)."""
if callable(rate):
self._ref.setOutputRate(
lambda v: Stochsim._rateResult(rate(DG.Vertex.wrap(v)))
)
else:
self._ref.setOutputRate(Stochsim._rateValue(rate))
[docs]
def state(self, graph: Graph) -> int:
"""Current molecule count of ``graph`` in this phase."""
return self._parent.state(self, graph)
[docs]
def concentration(self, graph: Graph) -> float:
"""Current concentration (count / volume) of ``graph`` in this phase."""
return self._parent.concentration(self, graph)
[docs]
def series(self, graph: Graph) -> List[int]:
"""Per-step count series of ``graph`` in this phase."""
return self._parent.series(self, graph)
[docs]
class Channel:
"""A transport channel. Obtained from :meth:`addChannel`."""
def __init__(self, parent: "MultiPhaseStochsim", ref: Any) -> None:
self._parent = parent
self._ref = ref
@property
def index(self) -> int:
return int(self._ref.index)
@property
def rate(self) -> Optional[float]:
"""The constant rate, or ``None`` for a dynamic (callback) channel."""
try:
return float(self._ref.rate)
except (TypeError, ValueError):
return None
[docs]
class Trajectory:
def __init__(
self,
times: List[float],
columnPhases: List[int],
columnVertexIds: List[int],
counts: List[List[int]],
) -> None:
#: Time stamp of each recorded step.
self.times = times
#: Phase index of each column of ``counts``.
self.columnPhases = columnPhases
#: DG vertex index (within its phase) of each column of ``counts``.
self.columnVertexIds = columnVertexIds
#: ``counts[step][column]`` molecule counts.
self.counts = counts
[docs]
def series(self, phaseIndex: int, vertexId: int) -> List[int]:
for col in range(len(self.columnPhases)):
if (
self.columnPhases[col] == phaseIndex
and self.columnVertexIds[col] == vertexId
):
return [row[col] for row in self.counts]
return [0 for _ in self.times]