abstochkin.graphing

  1"""  Graphing for AbStochKin simulations. """
  2
  3#  Copyright (c) 2024-2025, Alex Plakantonakis.
  4#
  5#  This program is free software: you can redistribute it and/or modify
  6#  it under the terms of the GNU General Public License as published by
  7#  the Free Software Foundation, either version 3 of the License, or
  8#  (at your option) any later version.
  9#
 10#  This program is distributed in the hope that it will be useful,
 11#  but WITHOUT ANY WARRANTY; without even the implied warranty of
 12#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 13#  GNU General Public License for more details.
 14#
 15#  You should have received a copy of the GNU General Public License
 16#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 17
 18from pathlib import Path
 19from typing import Literal, Self
 20
 21import numpy as np
 22import matplotlib.pyplot as plt
 23import matplotlib.ticker as ticker
 24import plotly.graph_objects as go
 25
 26
 27class Graph:
 28    """
 29    Graphing class for displaying the results of AbStochKin simulations.
 30
 31    Notes
 32    -----
 33    To successfully use the LaTeX engine for rendering text on Linux,
 34    run the following command in a terminal: `sudo apt install cm-super`.
 35    """
 36
 37    # First, set some global matplotlib settings
 38    plt.rcParams['figure.autolayout'] = True  # tight layout
 39    plt.rcParams["figure.facecolor"] = 'lightgray'
 40    plt.rcParams['text.usetex'] = True
 41    plt.rcParams['axes.titlesize'] = 9
 42    plt.rcParams['axes.labelsize'] = 7
 43    plt.rcParams['xtick.labelsize'] = 6
 44    plt.rcParams['ytick.labelsize'] = 6
 45    plt.rcParams["legend.fontsize"] = 7
 46    plt.rcParams["legend.framealpha"] = 0.65
 47
 48    def __init__(self,
 49                 /,
 50                 nrows=1,
 51                 ncols=1,
 52                 figsize=(5, 5),
 53                 dpi=300,
 54                 *,
 55                 backend: Literal['matplotlib', 'plotly'] = 'matplotlib',
 56                 **kwargs):
 57        self.backend = backend
 58
 59        if self.backend == 'matplotlib':
 60            self.fig, self.ax = plt.subplots(nrows=nrows, ncols=ncols,
 61                                             figsize=figsize, dpi=dpi,
 62                                             **kwargs)
 63        elif self.backend == 'plotly':
 64            self.fig = go.Figure()
 65            self.fig.update_layout(
 66                width=figsize[0] * dpi * 0.6,  # Convert figsize to pixels
 67                height=figsize[1] * dpi * 0.6,  # Convert figsize to pixels
 68            )
 69            self.fig2 = None  # optional second figure
 70        else:
 71            raise ValueError(f"Unknown backend: {self.backend}. "
 72                             f"Please choose from 'matplotlib' (default), 'plotly'.")
 73
 74    def setup_spines_ticks(self, ax_loc):
 75        """
 76        Set up the spines and ticks in a `matplotlib` graph.
 77        Make only the left and bottom spines/axes visible on the graph
 78        and place major ticks on them. Also set the minor ticks.
 79        """
 80        if self.backend == 'matplotlib':
 81            axs = self.ax if len(ax_loc) == 0 else self.ax[ax_loc]
 82            axs.spines[['left']].set_position('zero')
 83            axs.spines[['top', 'right']].set_visible(False)
 84            axs.xaxis.set_ticks_position('bottom')
 85            axs.yaxis.set_ticks_position('left')
 86            axs.xaxis.set_minor_locator(ticker.AutoMinorLocator())
 87            axs.yaxis.set_minor_locator(ticker.AutoMinorLocator())
 88
 89        else:  # self.backend == 'plotly'
 90
 91            self.fig.update_layout(
 92                xaxis=dict(zeroline=True, showline=True, linewidth=2,
 93                           zerolinecolor="black", showticklabels=True,
 94                           tickmode='auto', ticks='outside'),
 95                yaxis=dict(zeroline=True, showline=True, linewidth=2,
 96                           zerolinecolor="black", showticklabels=True,
 97                           tickmode='auto', ticks='outside')
 98            )
 99
100    def plot_ODEs(self,
101                  de_data,
102                  *,
103                  num_pts: int = 1000,
104                  species: list[str] | tuple[str] = (),
105                  show_plot: bool = True,
106                  ax_loc: tuple = ()
107                  ) -> Self:
108        """
109        Plot the deterministic trajectories of all species obtained
110        by obtaining the solution to a system of ODEs.
111
112        Parameters
113        ----------
114        de_data : DEcalcs object
115                 Data structure containing all the data related to
116                 solving the system of ODEs.
117
118        num_pts : int, default: 1000, optional
119                 Number of points used to calculate DE curves at.
120                 Used to approximate a smooth/continuous curve.
121
122        species : sequence of strings, default: (), optional
123                 An iterable sequence of strings specifying the species
124                 names to plot. If no species are specified (the default),
125                 then all species trajectories are plotted.
126
127        show_plot : bool, default: True, optional
128                 If True, show the plot.
129
130        ax_loc : tuple, optional
131                If the figure is made up of subplots, specify the location
132                of the axis to draw the data at.
133                Ex: for two subplots, the possible values of `ax_loc`
134                are (0, ) and (1, ). That's because the `self.ax` object is
135                a 1-D array. For figures with multiple rows and columns of
136                subplots, a 2-D tuple is needed.
137        """
138        species = list(de_data.odes.keys()) if len(species) == 0 else species
139        # t, y = ode_sol.t, ode_sol.y.T  # values at precomputed time pts
140        t = np.linspace(de_data.odes_sol.t[0], de_data.odes_sol.t[-1],
141                        num_pts)  # time points for obtaining...
142        y = de_data.odes_sol.sol(t).T  # an approximately continuous solution
143
144        self.setup_spines_ticks(ax_loc)
145
146        if self.backend == 'matplotlib':
147            axs = self.ax if len(ax_loc) == 0 else self.ax[ax_loc]
148            axs.set(xlim=(0, de_data.odes_sol.t[-1]))
149
150            for i, sp in enumerate(list(de_data.odes.keys())):
151                if sp in species:
152                    axs.plot(t, y[:, i], label=f"${sp}_{{DE}}$",
153                             linestyle='--', linewidth=0.75, alpha=0.75)
154
155            # axs.set(title="Deterministic trajectories")
156            axs.set(xlabel=f"$t$ ({de_data.time_unit})", ylabel="$N$")
157            axs.legend(loc='upper right')
158            self.fig.tight_layout()
159
160            if show_plot:
161                plt.show()
162
163        else:  # self.backend == 'plotly'
164
165            for i, sp in enumerate(list(de_data.odes.keys())):
166                if sp in species:
167                    self.fig.add_trace(
168                        go.Scatter(
169                            x=t.tolist(),
170                            y=y[:, i].tolist(),
171                            mode='lines',
172                            name=f"${sp}_{{DE}}$",
173                            line=dict(dash='dash', width=1)
174                        )
175                    )
176
177            self.fig.update_layout(
178                xaxis=dict(title=f"$t \\, ({de_data.time_unit})$",
179                           range=[-0.01 * de_data.odes_sol.t[-1], de_data.odes_sol.t[-1]]),
180                yaxis=dict(title="$N$"),
181            )
182
183            if show_plot:
184                self.fig.show()
185
186        return self
187
188    def plot_trajectories(self,
189                          time,
190                          data,
191                          *,
192                          species: list[str] | tuple[str] = (),
193                          show_plot: bool = True,
194                          ax_loc: tuple = ()
195                          ) -> Self:
196        """ Graph simulation time trajectories. """
197        self.setup_spines_ticks(ax_loc)
198        species = list(data.keys()) if len(species) == 0 else species
199
200        if self.backend == 'matplotlib':
201            axs = self.ax if len(ax_loc) == 0 else self.ax[ax_loc]
202            axs.set(xlim=(0, time[-1]))
203
204            for sp, sp_data in data.items():
205                if sp in species:
206                    trajs = sp_data['N'].T
207                    for traj in trajs:
208                        axs.plot(time, traj, linewidth=0.25)
209
210            axs.set(title="ABK trajectories")
211            axs.set(xlabel=f"$t$ (sec)", ylabel="$N$")
212            # axs.legend(loc='best')
213
214            if show_plot:
215                plt.show()
216
217        else:  # self.backend == 'plotly'
218
219            for sp, sp_data in data.items():
220                if sp in species:
221                    trajs = sp_data['N'].T
222                    for i, traj in enumerate(trajs):
223                        self.fig.add_trace(
224                            go.Scatter(
225                                x=time.tolist(),
226                                y=traj.tolist(),
227                                mode='lines',
228                                name=f"${sp} \\; Run \\, {i}$",
229                                line=dict(width=0.5)
230                            )
231                        )
232            self.fig.update_layout(
233                xaxis=dict(title=f"$t \\; (sec)$",
234                           range=[-0.01 * time[-1], time[-1]]),
235                yaxis=dict(title="$N$")
236            )
237
238            if show_plot:
239                self.fig.show()
240
241        return self
242
243    def plot_avg_std(self,
244                     time,
245                     data,
246                     *,
247                     species: list[str] | tuple[str] = (),
248                     show_plot: bool = True,
249                     ax_loc: tuple = ()
250                     ) -> Self:
251        """
252        Graph simulation average trajectories and
253        1-standard-deviation envelopes.
254        """
255        self.setup_spines_ticks(ax_loc)
256        species = list(data.keys()) if len(species) == 0 else species
257
258        if self.backend == 'matplotlib':
259            axs = self.ax if len(ax_loc) == 0 else self.ax[ax_loc]
260            axs.set(xlim=(0, time[-1]))
261
262            for sp, sp_data in data.items():
263                if sp in species:
264                    axs.plot(time, sp_data['N_avg'],
265                             linewidth=1.5, label=f"$<{sp}>$", alpha=0.5)
266                    axs.fill_between(time,
267                                     sp_data['N_avg'] - sp_data['N_std'],
268                                     sp_data['N_avg'] + sp_data['N_std'],
269                                     alpha=0.5, linewidth=0)
270
271            axs.set(xlabel="$t$ (sec)", ylabel="$N$")
272            axs.legend(loc='upper right')
273            self.fig.tight_layout()
274
275            if show_plot:
276                plt.show()
277
278        else:  # self.backend == 'plotly'
279
280            for sp, sp_data in data.items():
281                if sp in species:
282                    self.fig.add_trace(
283                        go.Scatter(
284                            x=time.tolist(),
285                            y=sp_data['N_avg'].tolist(),
286                            mode='lines',
287                            name=f"$<{sp}>$",
288                            line=dict(width=2)
289                        )
290                    )
291
292                    self.fig.add_trace(
293                        go.Scatter(
294                            x=time.tolist(),
295                            y=(sp_data['N_avg'] + sp_data['N_std']).tolist(),
296                            mode='lines',
297                            name=f"$<{sp}> + \\sigma$",
298                            line=dict(width=0),
299                            showlegend=False
300                        )
301                    )
302
303                    self.fig.add_trace(
304                        go.Scatter(
305                            x=time.tolist(),
306                            y=(sp_data['N_avg'] - sp_data['N_std']).tolist(),
307                            mode='lines',
308                            line=dict(width=0),
309                            name=f"$<{sp}> - \\sigma$",
310                            fill='tonexty',
311                            showlegend=False
312                        )
313                    )
314
315            self.fig.update_layout(
316                xaxis=dict(title=f"$t \\; (sec)$",
317                           range=[-0.01 * time[-1], time[-1]]),
318                yaxis=dict(title="$N$")
319            )
320
321            if show_plot:
322                self.fig.show()
323
324        return self
325
326    def plot_eta(self,
327                 time,
328                 data,
329                 *,
330                 species: list[str] | tuple[str] = (),
331                 show_plot: bool = True,
332                 ax_loc: tuple = ()
333                 ) -> Self:
334        """ Graph the coefficient of variation. """
335        self.setup_spines_ticks(ax_loc)
336        species = list(data.keys()) if len(species) == 0 else species
337
338        if self.backend == "matplotlib":
339            axs = self.ax if len(ax_loc) == 0 else self.ax[ax_loc]
340            axs.set(xlim=(0, time[-1]))
341
342            for sp, sp_data in data.items():
343                if sp in species:
344                    axs.plot(time, sp_data['eta'], linewidth=1.5, label=f"${sp}$")
345                    axs.plot(time, sp_data['eta_p'], linewidth=1, linestyle='--',
346                             label=f"${sp}_{{Poisson}}$", color=(0.5, 0.5, 0.5))
347
348            axs.set(title="Coefficient of Variation, $\\eta$")
349            axs.set(xlabel=f"$t$ (sec)", ylabel="$\\eta$")
350            axs.legend(loc='upper right')
351
352            if show_plot:
353                plt.show()
354
355        else:  # self.backend == 'plotly'
356
357            for sp, sp_data in data.items():
358                if sp in species:
359                    self.fig.add_trace(
360                        go.Scatter(
361                            x=time.tolist(),
362                            y=sp_data['eta'].tolist(),
363                            mode='lines',
364                            name=f"${sp}$",
365                            line=dict(width=2)
366                        ))
367
368                    self.fig.add_trace(
369                        go.Scatter(
370                            x=time.tolist(),
371                            y=sp_data['eta_p'].tolist(),
372                            mode='lines',
373                            name=f"${sp}_{{Poisson}}$",
374                            line=dict(width=2, dash="dash")
375                        ))
376
377            self.fig.update_layout(
378                title="Coefficient of Variation",
379                xaxis=dict(title=f"$t \\; (sec)$",
380                           range=[-0.01 * time[-1], time[-1]]),
381                yaxis=dict(title="$\\eta$")
382            )
383
384            if show_plot:
385                self.fig.show()
386
387        return self
388
389    def plot_het_metrics(self,
390                         time,
391                         proc_str: tuple[str, str],
392                         proc_data: dict,
393                         *,
394                         het_attr='k',
395                         show_plot: bool = True,
396                         ax_loc: tuple = ()
397                         ) -> Self:
398        """
399        Graph species- and process-specific metrics of population heterogeneity.
400        """
401        self.setup_spines_ticks(ax_loc)
402        title = f"${proc_str[0].split(';')[0].replace(' ,', chr(92) + 'hspace{10pt} ,').replace('->', chr(92) + 'rightarrow')}$"
403
404        if self.backend == 'matplotlib':
405            axs = self.ax if len(ax_loc) == 0 else self.ax[ax_loc]
406            axs.set(xlim=(0, time[-1]))
407            axs.set(xlim=(0, time[-1]),
408                    ylim=(0, 1.5 * np.max(proc_data[f"<{het_attr}_avg>"] + proc_data[
409                        f"<{het_attr}_std>"])))
410
411            axs.plot(time, proc_data[f'<{het_attr}_avg>'],
412                     linewidth=1.5, label=f"$<{het_attr}>$", alpha=0.5)
413            axs.fill_between(time,
414                             proc_data[f"<{het_attr}_avg>"] - proc_data[f"<{het_attr}_std>"],
415                             proc_data[f"<{het_attr}_avg>"] + proc_data[f"<{het_attr}_std>"],
416                             alpha=0.5, linewidth=0)
417            axs.tick_params(axis='y', labelcolor='blue')
418
419            axs.set(title=title + (f"$, {proc_str[1]}$" if proc_str[1] != "" else ""),
420                    xlabel=f"$t$ (sec)")
421
422            if het_attr == 'Km':
423                axs.set_ylabel("$K_m$", color='blue')
424            elif het_attr == 'K50':
425                axs.set_ylabel("$K_{50}$", color='blue')
426            else:
427                axs.set_ylabel(f"${het_attr}$", color='blue')
428
429            # Second y-axis
430            axs2 = axs.twinx()
431            axs2.spines[['bottom']].set_position('zero')  # x axis
432            axs2.spines[['right']].set_position(('axes', 1))  # y axis
433            axs2.spines[['top', 'left', 'bottom']].set_visible(False)
434            # axs2.spines['right'].set_color('red')
435            axs2.set(ylim=(0, 1))
436            axs2.tick_params(axis='y', labelcolor='red')
437            axs2.yaxis.set_ticks_position('right')
438            axs2.set_yticks([i for i in np.arange(0, 1.1, 0.1)])
439            axs2.yaxis.set_minor_locator(ticker.AutoMinorLocator())
440            axs2.grid(which='major', axis='y', color='r',
441                      linestyle='--', linewidth=0.25, alpha=0.25)
442
443            axs2.plot(time, proc_data['psi_avg'],
444                      linewidth=1.5, label='$<\\psi>$', color='red', alpha=0.5)
445            axs2.fill_between(time,
446                              proc_data['psi_avg'] - proc_data['psi_std'],
447                              proc_data['psi_avg'] + proc_data['psi_std'],
448                              color='red', alpha=0.5, linewidth=0)
449            axs2.set_ylabel("$\\psi$", color='red')
450
451            if show_plot:
452                plt.show()
453
454            return self
455
456        else:  # self.backend == 'plotly'
457
458            self.fig.add_trace(
459                go.Scatter(
460                    x=time.tolist(),
461                    y=proc_data[f'<{het_attr}_avg>'].tolist(),
462                    mode='lines',
463                    name=f"$<{het_attr}>$",
464                    line=dict(width=2, color='blue'),
465                )
466            )
467
468            self.fig.add_trace(
469                go.Scatter(
470                    x=time.tolist(),
471                    y=(proc_data[f"<{het_attr}_avg>"] + proc_data[f"<{het_attr}_std>"]).tolist(),
472                    mode='lines',
473                    name=f"$<{het_attr}> + \\sigma$",
474                    line=dict(width=0),
475                    showlegend=False
476                )
477            )
478
479            self.fig.add_trace(
480                go.Scatter(
481                    x=time.tolist(),
482                    y=(proc_data[f"<{het_attr}_avg>"] - proc_data[f"<{het_attr}_std>"]).tolist(),
483                    mode='lines',
484                    name=f"$<{het_attr}> - \\sigma$",
485                    line=dict(width=0),
486                    fill='tonexty',
487                    showlegend=False
488                )
489            )
490
491            self.fig.update_layout(
492                title=title + (f"$, {proc_str[1]}$" if proc_str[1] != "" else ""),
493                xaxis=dict(title=f"$t \\; (sec)$",
494                           range=[-0.01 * time[-1], time[-1]]),
495                yaxis=dict(title="$K_m$" if het_attr == "Km" else "$K_{50}$" if het_attr == "K50" else f"${het_attr}$",
496                           color="blue",
497                           range=[0, 1.5 * np.max(proc_data[f"<{het_attr}_avg>"] + proc_data[f"<{het_attr}_std>"])])
498            )
499
500            if show_plot:
501                self.fig.show()
502
503            # Plot psi on a separate figure
504            self.fig2 = go.Figure()
505            self.fig2.update_layout(
506                width=self.fig.layout.width,
507                height=self.fig.layout.height,
508                xaxis=dict(zeroline=True, showline=True, linewidth=2,
509                           zerolinecolor="black", showticklabels=True,
510                           tickmode='auto', ticks='outside'),
511                yaxis=dict(zeroline=True, showline=True, linewidth=2,
512                           zerolinecolor="black", showticklabels=True,
513                           tickmode='auto', ticks='outside')
514            )
515
516            self.fig2.add_trace(
517                go.Scatter(
518                    x=time.tolist(),
519                    y=proc_data['psi_avg'],
520                    mode='lines',
521                    name=f"$<\\psi>$",
522                    line=dict(width=2, color='red'),
523                )
524            )
525
526            self.fig2.add_trace(
527                go.Scatter(
528                    x=time.tolist(),
529                    y=proc_data['psi_avg'] + proc_data['psi_std'],
530                    mode='lines',
531                    name=f"$<\\psi> + \\sigma$",
532                    line=dict(width=0),
533                    showlegend=False
534                )
535            )
536
537            self.fig2.add_trace(
538                go.Scatter(
539                    x=time.tolist(),
540                    y=proc_data['psi_avg'] - proc_data['psi_std'],
541                    mode='lines',
542                    name=f"$<\\psi> + \\sigma$",
543                    line=dict(width=0),
544                    fill='tonexty',
545                    showlegend=False
546                )
547            )
548
549            self.fig2.update_layout(
550                title=title + (f"$, {proc_str[1]}$" if proc_str[1] != "" else ""),
551                xaxis=dict(title=f"$t \\; (sec)$",
552                           range=[-0.01 * time[-1], time[-1]]),
553                yaxis=dict(title="$\\psi$",
554                           color="red",
555                           range=[0, 1],
556                           tick0=0,
557                           dtick=0.1)
558            )
559
560            if show_plot:
561                self.fig2.show()
562
563            return self
564
565    def savefig(self, filename: str, image_format: str = 'svg', **kwargs):
566        """ Save the figure as a file. """
567        graph_path = Path('.') / 'plots_output'
568        graph_path.mkdir(exist_ok=True)
569        graph_path_svg = graph_path / f"{filename}.{image_format}"
570        if self.backend == 'matplotlib':
571            self.fig.savefig(graph_path_svg, format=image_format, **kwargs)
572            plt.close(self.fig)
573        else:  # self.backend == 'plotly'
574            self.fig.write_image(graph_path_svg, format=image_format, **kwargs)