feat: add intermediate results plotting to epidemic demo

src/dynadojo/systems/epidemic/seis.py

@@ -6,7 +6,9 @@
 import networkx as nx
 
 from ..utils.epidemic import EpidemicSystem
+from ...utils.opinion import plot
 
+import numpy as np
 
 class SEISSystem(EpidemicSystem):
     """
@@ -102,4 +104,20 @@ def create_model(self, x0):
         self.model = ep.SEISModel(self.g)
         self.model.set_initial_status(self.config)
         self.model.status = x0
-        self.model.initial_status = x0
+        self.model.initial_status = x0
+
+
+    def save_plotted_trajectories( self, 
+            y_true:np.ndarray, 
+            y_pred: np.ndarray,
+            filepath: str,
+            tag: str = "", 
+        ):
+
+        fig, ax = plot([y_true, y_pred], 
+                       target_dim=min(self._embed_dim, 3), 
+                       labels=["true", "pred"], 
+                       max_lines=10,
+                       title=f"Epidemic l={self.latent_dim}, e={self._embed_dim} - {tag}")
+        fig.savefig(filepath, bbox_inches='tight', dpi=300, transparent=True, format='pdf')
+        return fig, ax

src/dynadojo/systems/epidemic/sir.py

@@ -6,6 +6,9 @@
 import networkx as nx
 
 from ..utils.epidemic import EpidemicSystem
+from ...utils.opinion import plot
+
+import numpy as np
 
 
 class SIRSystem(EpidemicSystem):
@@ -87,4 +90,19 @@ def create_model(self, x0):
         self.model = ep.SIRModel(self.g)
         self.model.set_initial_status(self.config)
         self.model.status = x0
-        self.model.initial_status = x0
+        self.model.initial_status = x0
+
+    def save_plotted_trajectories( self, 
+            y_true:np.ndarray, 
+            y_pred: np.ndarray,
+            filepath: str,
+            tag: str = "", 
+        ):
+
+        fig, ax = plot([y_true, y_pred], 
+                       target_dim=min(self._embed_dim, 3), 
+                       labels=["true", "pred"], 
+                       max_lines=10,
+                       title=f"Epidemic l={self.latent_dim}, e={self._embed_dim} - {tag}")
+        fig.savefig(filepath, bbox_inches='tight', dpi=300, transparent=True, format='pdf')
+        return fig, ax

src/dynadojo/systems/epidemic/sis.py

@@ -6,6 +6,9 @@
 import networkx as nx
 
 from ..utils.epidemic import EpidemicSystem
+from ...utils.opinion import plot
+
+import numpy as np
 
 
 class SISSystem(EpidemicSystem):
@@ -83,4 +86,19 @@ def create_model(self, x0):
         self.model = ep.SISModel(self.g)
         self.model.set_initial_status(self.config)
         self.model.status = x0
-        self.model.initial_status = x0
+        self.model.initial_status = x0
+
+    def save_plotted_trajectories( self, 
+            y_true:np.ndarray, 
+            y_pred: np.ndarray,
+            filepath: str,
+            tag: str = "", 
+        ):
+
+        fig, ax = plot([y_true, y_pred], 
+                       target_dim=min(self._embed_dim, 3), 
+                       labels=["true", "pred"], 
+                       max_lines=10,
+                       title=f"Epidemic l={self.latent_dim}, e={self._embed_dim} - {tag}")
+        fig.savefig(filepath, bbox_inches='tight', dpi=300, transparent=True, format='pdf')
+        return fig, ax

src/dynadojo/utils/opinion.py

@@ -4,21 +4,32 @@
 MAX_LINES = 30
 
 
-def plot(grid: list[np.ndarray], target_dim: int = 3, max_lines=MAX_LINES, datalabels: list[str] = None, gridlabels: list[str] = None):
-    grid = np.array([x[:max_lines] for x in grid])
+def plot(
+        datasets: list[np.ndarray],
+        target_dim: int = 3, 
+        max_lines=MAX_LINES, 
+        datalabels: list[str] = None, 
+        labels: list[str] = None,
+        title: str = None
+    ):
+
+    datasets = np.array([x[:max_lines] for x in datasets]) 
     fig = plt.figure(figsize=(16, 4))
 
     posidx = 1
-    for idx, dataset in enumerate(grid):
-        ax = fig.add_subplot(1, len(grid), posidx)
+    for idx, dataset in enumerate(datasets):
+        ax = fig.add_subplot(1, len(datasets), posidx)
 
         ydata = dataset[0]
         if(datalabels):
             ax.plot(ydata, label=datalabels)
-            ax.legend()
         else:
             ax.plot(ydata)
-        ax.set_title(gridlabels[idx])
+        ax.set_title(labels[idx])
+        ax.legend()
         posidx += 1
 
-    plt.show()
+
+    if title:
+        fig.suptitle(title)
+    return fig, ax