abstochkin.simulation

  1"""
  2Perform an Agent-based Kinetics simulation.
  3
  4This module contains the code for the class `Simulation`,
  5which, along with the `SimulationMethodsMixin` class,
  6does everything that is needed to run an
  7Agent-based Kinetics simulation and store its results.
  8
  9The class `AgentStateData` is used by a `Simulation`
 10object to store and handle some of the necessary runtime data.
 11"""
 12
 13#  Copyright (c) 2024-2025, Alex Plakantonakis.
 14#
 15#  This program is free software: you can redistribute it and/or modify
 16#  it under the terms of the GNU General Public License as published by
 17#  the Free Software Foundation, either version 3 of the License, or
 18#  (at your option) any later version.
 19#
 20#  This program is distributed in the hope that it will be useful,
 21#  but WITHOUT ANY WARRANTY; without even the implied warranty of
 22#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 23#  GNU General Public License for more details.
 24#
 25#  You should have received a copy of the GNU General Public License
 26#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 27
 28import contextlib
 29from typing import Literal
 30
 31import numpy as np
 32from tqdm import tqdm
 33
 34from ._simulation_methods import SimulationMethodsMixin
 35from .de_calcs import DEcalcs
 36from .graphing import Graph
 37from .het_calcs import get_het_processes
 38from .process import update_all_species, MichaelisMentenProcess, \
 39    RegulatedProcess, RegulatedMichaelisMentenProcess
 40from .utils import rng_streams
 41
 42
 43class Simulation(SimulationMethodsMixin):
 44    """
 45    Run an Agent-based Kinetics simulation.
 46
 47    Attributes
 48    ----------
 49    p0 : dict[str: int]
 50        Dictionary specifying the initial population sizes of all
 51        species in the given processes.
 52    t_max : float or int
 53        Numerical value of the end of simulated time in the units
 54        specified in the class attribute `AbStochKin.time_unit`.
 55    dt : float
 56        The duration of the time interval that the simulation's
 57        algorithm considers. The current implementation only
 58        supports a fixed time step interval whose value is `dt`.
 59    n : int
 60        The number of repetitions of the simulation to be
 61        performed.
 62    random_state : float or int
 63        A number used to seed the random number generator.
 64    use_multithreading : bool
 65        Specify whether to parallelize the simulation
 66        using multithreading. If `False`, the ensemble
 67        of simulations is run sequentially.
 68    max_agents : dict
 69        Specification of the maximum number of agents that each
 70        species should have when running the simulation. An
 71        empty dictionary signifies that a default approach will
 72        be taken and the number for each species will be
 73        automatically determined (see method `_setup_runtime_data()`
 74        for details). The entries in the dictionary should be
 75        `species name (string): number (int)`.
 76    max_agents_multiplier : float or int
 77        Set the number of possible agents for each species to be the
 78        maximum value of the ODE solution for the species
 79        times `max_agents_multiplier`.
 80    time_unit : str
 81        A string of the time unit to be used for describing the
 82        kinetics of the given processes.
 83    """
 84
 85    def __init__(self,
 86                 /,
 87                 p0: dict,
 88                 t_max: float,
 89                 dt: float,
 90                 n: int,
 91                 processes: list,
 92                 *,
 93                 random_state: int,
 94                 do_solve_ODEs: bool,
 95                 ODE_method: str,
 96                 do_run: bool,
 97                 show_graphs: bool,
 98                 graph_backend: Literal['matplotlib', 'plotly'],
 99                 use_multithreading: bool,
100                 max_agents: dict,
101                 max_agents_multiplier: float | int,
102                 time_unit: str):
103        """
104        The parameters below are not class attributes, but are part of a
105        `Simulation` object's initialization to trigger specific actions
106        to be automatically performed. Note that these actions can also
107        be performed manually by calling the appropriate methods once a
108        class object has been instantiated.
109
110        Other Parameters
111        ----------------
112        do_solve_ODEs : bool
113            If `True`, attempt to numerically solve the system
114            of ODEs defined from the given set of processes.
115            If `False`, do not attempt to solve the ODEs and
116            do not run the simulation.
117        ODE_method : str
118            Method to use when attempting to solve the system
119            of ODEs (if `do_solve_ODEs` is `True`).
120            Available ODE methods: RK45, RK23, DOP853, Radau,
121            BDF, LSODA.
122        do_run : bool
123            Specify whether to run the AbStochKin simulation.
124            If `False`, then a `Simulation` object is created but
125            the simulation is not run. A user can then manually
126            run it by calling the class method `run_simulation()`.
127        show_graphs : bool
128            Specify whether to show graphs of the results.
129        graph_backend : str
130            `Matplotlib` and `Plotly` are currently supported.
131        """
132
133        self.p0 = p0  # dictionary of initial population sizes
134        self.t_min = 0  # start simulation at time 0 (by assumption)
135        self.t_max = t_max  # end simulation at time t_max
136        self.dt = dt  # fixed time interval
137        self.n = n  # repeat the simulation this many times
138        self.random_state = random_state
139        self.use_multithreading = use_multithreading
140        self.max_agents = max_agents
141        self.max_agents_multiplier = max_agents_multiplier
142        self.time_unit = time_unit
143
144        self.all_species, self._procs_by_reactant, self._procs_by_product = \
145            update_all_species(tuple(processes))
146
147        """ Generate the list of processes to be used by the algorithm. 
148        This is done because some processes (e.g., reversible processes) 
149        need to be split up into the forward and reverse processes. """
150        self._algo_processes = list()
151        self._gen_algo_processes(processes)
152        self._het_processes = get_het_processes(self._algo_processes)
153        self._het_processes_num = len(self._het_processes)
154
155        self._validate_p0()  # validate initial population sizes
156
157        # Generate streams of random numbers given a seed
158        self.streams = rng_streams(self.n, self.random_state)
159
160        # ******************** Deterministic calculations ********************
161        self.de_calcs = DEcalcs(self.p0, self.t_min, self.t_max, processes,
162                                ode_method=ODE_method, time_unit=self.time_unit)
163        if do_solve_ODEs:
164            try:
165                self.de_calcs.solve_ODEs()
166            except Exception as exc:
167                print(f"ODE solver exception:\n{exc}")
168        else:
169            if do_run:
170                print("Warning: Must specify the maximum number of agents for "
171                      "each species when not solving the system ODEs.")
172        # ********************************************************************
173
174        self.graph_backend = graph_backend
175
176        self.total_time = self.t_max - self.t_min
177        self.t_steps = int(self.total_time / self.dt)
178        self.time = np.linspace(0, self.total_time, self.t_steps + 1)
179
180        # Initialize data structures for results, runtime data (rtd),
181        # process-specific k values, transition probabilities
182        # for 0th and 1st order processes, metrics of population
183        # heterogeneity, sequence of functions th algorithm will
184        # execute, and progress bar.
185        self.results, self.rtd = dict(), dict()
186
187        self.trans_p = dict()  # Process-specific transition probabilities
188
189        # Process-specific parameters that can exhibit heterogeneity
190        self.k_vals = dict()
191        self.Km_vals = dict()
192        self.K50_vals = dict()
193
194        self.k_het_metrics = dict()
195        self.Km_het_metrics = dict()
196        self.K50_het_metrics = dict()
197
198        # Runtime-specific attributes
199        self.algo_sequence = list()
200        self.progress_bar = None
201
202        self.graphs = None  # For storing any graphs that are shown
203
204        if do_run:
205            self.run_simulation()
206            if show_graphs:  # Simulation must first be run before plotting
207                self.graphs = self.graph_results()
208
209    def run_simulation(self):
210        """
211        Run an ensemble of simulations and compute statistics
212        of simulation data.
213        """
214        bar_fmt = f"{{desc}}: {{percentage:3.0f}}% |{{bar}}| " \
215                  f"n={{n_fmt}}/{{total_fmt}} " \
216                  f"[{{elapsed}}{'' if self.use_multithreading else '<{remaining}'}, " \
217                  f"{{rate_fmt}}{{postfix}}]"
218        self.progress_bar = tqdm(total=self.n,
219                                 ncols=65 if self.use_multithreading else 71,
220                                 desc="Progress",
221                                 bar_format=bar_fmt,
222                                 colour='green')
223        # Note on progress bar info: Multi-threading makes calculating
224        # the remaining time of the simulation unreliable, so only
225        # the elapsed time is shown.
226
227        self._setup_data()  # initialize data
228        self._gen_algo_sequence()  # generate sequence of processes for algorithm
229
230        self._parallel_run() if self.use_multithreading else self._sequential_run()
231
232        self._compute_trajectory_stats()  # Get statistics on simulation data
233
234        # self._compute_k_het_stats()  # Get statistics on heterogeneity data (k)
235        # self._compute_Km_het_stats()  # Get statistics on heterogeneity data (Km)
236        # self._compute_K50_het_stats()  # Get statistics on heterogeneity data (Km)
237        self._compute_het_stats()  # Get statistics on heterogeneity data
238
239        self._post_run_cleanup()  # free up some memory
240        self.progress_bar.close()
241
242    def graph_results(self,
243                      /,
244                      graphs_to_show=None,
245                      species_to_show=None):
246        """
247        Make graphs of the results.
248
249        Parameters
250        ----------
251        species_to_show : `None` or list of string(s), default: `None`
252            If `None`, data for all species are plotted.
253
254        graphs_to_show : `None` or string or list of string(s), default: `None`
255            If `None`, all graphs are shown. If a string is given then the
256            graph that matches the string is shown. A list of strings shows
257            all the graphs specified in the list.
258            Graph specifications:
259
260            'avg' : Plot the average trajectories and their
261                one-standard deviation envelopes. The ODE
262                trajectories are also shown.
263            'traj' : Plot the species trajectories of individual
264                simulations.
265            'ode' : Plot the deterministic species trajectories,
266                obtained by numerically solving the ODEs.
267            'eta' : Plot the coefficient of variation (CoV) and
268                the CoV assuming that all processes a species
269                participates in obey Poisson statistics.
270            'het' : Plot species- and process-specific metrics
271                of heterogeneity (`k` and `ψ`) and their
272                one-standard deviation envelopes.
273        """
274        graphs_to_return = list()
275
276        if graphs_to_show is None:
277            graphs_to_show = ['avg', 'het']
278        if species_to_show is None:
279            species_to_show = self.all_species
280
281        if 'avg' in graphs_to_show:
282            graph_avg = Graph(backend=self.graph_backend)
283            with contextlib.suppress(AttributeError):
284                graph_avg.plot_ODEs(self.de_calcs, species=species_to_show, show_plot=False)
285            graph_avg.plot_avg_std(self.time, self.results, species=species_to_show)
286            graphs_to_return.append(graph_avg)
287
288        if 'traj' in graphs_to_show:
289            graph_traj = Graph(backend=self.graph_backend)
290            graph_traj.plot_trajectories(self.time, self.results, species=species_to_show)
291            graphs_to_return.append(graph_traj)
292
293        if 'ode' in graphs_to_show:
294            graph_ode = Graph(backend=self.graph_backend)
295            graph_ode.plot_ODEs(self.de_calcs, species=species_to_show)
296            graphs_to_return.append(graph_ode)
297
298        if 'eta' in graphs_to_show:
299            graph_eta = Graph(backend=self.graph_backend)
300            graph_eta.plot_eta(self.time, self.results, species=species_to_show)
301            graphs_to_return.append(graph_eta)
302
303        if 'het' in graphs_to_show:
304            graph_het = None
305
306            for proc in self._algo_processes:
307                if proc.is_heterogeneous:
308                    graph_het = Graph(backend=self.graph_backend)
309                    graph_het.plot_het_metrics(self.time,
310                                               (str(proc), ''),
311                                               self.k_het_metrics[proc])
312
313                if isinstance(proc, (MichaelisMentenProcess, RegulatedMichaelisMentenProcess)):
314                    if proc.is_heterogeneous_Km:
315                        graph_het = Graph(backend=self.graph_backend)
316                        graph_het.plot_het_metrics(self.time,
317                                                   (str(proc), f"K_m={proc.Km}"),
318                                                   self.Km_het_metrics[proc],
319                                                   het_attr='Km')
320
321                if isinstance(proc, (RegulatedProcess, RegulatedMichaelisMentenProcess)):
322                    if isinstance(proc.is_heterogeneous_K50, list):  # multiple regulators
323                        for i in range(len(proc.regulating_species)):
324                            if proc.is_heterogeneous_K50[i]:
325                                graph_het = Graph(backend=self.graph_backend)
326                                extra_str = f"\\textrm{{{proc.regulation_type[i]} by }}" \
327                                            f"{proc.regulating_species[i]}, " \
328                                            f"K_{{50}}={proc.K50[i]}"
329                                graph_het.plot_het_metrics(self.time,
330                                                           (str(proc), extra_str),
331                                                           self.K50_het_metrics[proc][i],
332                                                           het_attr='K50')
333                    else:  # one regulator
334                        if proc.is_heterogeneous_K50:
335                            graph_het = Graph(backend=self.graph_backend)
336                            graph_het.plot_het_metrics(self.time,
337                                                       (str(proc), f"K_{{50}}={proc.K50}"),
338                                                       self.K50_het_metrics[proc],
339                                                       het_attr='K50')
340
341            if graph_het is not None:
342                graphs_to_return.append(graph_het)
343
344        return graphs_to_return
345
346    def __repr__(self):
347        return f"AbStochKin Simulation(p0={self.p0},\n" \
348               f"                 t_min={self.t_min},\n" \
349               f"                 t_max={self.t_max},\n" \
350               f"                 dt={self.dt},\n" \
351               f"                 n={self.n},\n" \
352               f"                 random_state={self.random_state})"
353
354    def __str__(self):
355        descr = f"AbStochKin Simulation object with\n" \
356                f"Processes: {self._algo_processes}\n" \
357                f"Number of heterogeneous processes: {len(self._het_processes)}\n" \
358                f"Initial population sizes: {self.p0}\n" \
359                f"Simulated time interval: {self.t_min} - {self.t_max} {self.time_unit} " \
360                f"with fixed time step increment {self.dt} {self.time_unit}\n" \
361                f"Simulation runs: {self.n}\n" \
362                f"Use multi-threading: {self.use_multithreading}\n" \
363                f"Random state seed: {self.random_state}\n"
364
365        if self.de_calcs.odes_sol is not None:
366            descr += f"Attempt to solve ODEs with method " \
367                     f"{self.de_calcs.ode_method}: " \
368                     f"{'Successful' if self.de_calcs.odes_sol.success else 'Failed'}\n"
369
370        return descr