Adds snapshots slider for sim engine plots

pdr_backend/sim/sim_plotter.py

@@ -3,6 +3,7 @@
 import pickle
 import time
 from datetime import datetime
+from typing import Optional
 
 import numpy as np
 import pandas as pd
@@ -24,7 +25,20 @@ def __init__(
         self.st = None
         self.aimodel_plotdata = None
 
-    def load_state(self):
+    @staticmethod
+    def available_snapshots():
+        all_state_files = glob.glob("sim_state/st_*.pkl")
+
+        all_timestamps = [
+            f.replace("sim_state/st_", "").replace(".pkl", "")
+            for f in all_state_files
+            if "final" not in f
+        ]
+        all_timestamps.sort()
+
+        return all_timestamps + ["final"]
+
+    def load_state(self, timestamp: Optional[str] = None):
         if not os.path.exists("sim_state"):
             raise Exception(
                 "sim_state folder does not exist. Please run the simulation first."
@@ -34,6 +48,15 @@ def load_state(self):
         if not all_state_files:
             raise Exception("No state files found. Please run the simulation first.")
 
+        if timestamp:
+            with open(f"sim_state/st_{timestamp}.pkl", "rb") as f:
+                self.st = pickle.load(f)
+
+            with open(f"sim_state/aimodel_plotdata_{timestamp}.pkl", "rb") as f:
+                self.aimodel_plotdata = pickle.load(f)
+
+            return self.st, "final"
+
         if not os.path.exists("sim_state/st_final.pkl"):
             # plot previous state to avoid using a pickle that hasn't finished
             all_state_files = glob.glob("sim_state/st_*.pkl")

sim_plots.py

@@ -8,6 +8,7 @@
 streamlit.set_page_config(layout="wide")
 
 title = streamlit.empty()
+inputs = streamlit.empty()
 c1, c2, c3 = streamlit.columns((1, 1, 2))
 c4, c5 = streamlit.columns((1, 1))
 c6, c7 = streamlit.columns((1, 1))
@@ -27,6 +28,21 @@
 last_ts = None
 sim_plotter = SimPlotter()
 
+
+def load_canvas_on_state(ts):
+    titletext = f"Iter #{st.iter_number} ({ts})" if ts != "final" else "Final sim state"
+    title.title(titletext)
+
+    for key in canvas:
+        if not key.startswith("aimodel"):
+            fig = getattr(sim_plotter, f"plot_{key}")()
+        else:
+            func_name = getattr(aimodel_plotter, f"plot_{key}")
+            fig = func_name(sim_plotter.aimodel_plotdata)
+
+        canvas[key].plotly_chart(fig, use_container_width=True, theme="streamlit")
+
+
 while True:
     try:
         sim_plotter.load_state()
@@ -43,23 +59,16 @@
         time.sleep(1)
         continue
 
-    title.title(f"Iter #{st.iter_number} ({new_ts})")
-
     if new_ts == last_ts:
         time.sleep(1)
         continue
 
-    for key in canvas:
-        if not key.startswith("aimodel"):
-            fig = getattr(sim_plotter, f"plot_{key}")()
-        else:
-            func_name = getattr(aimodel_plotter, f"plot_{key}")
-            fig = func_name(sim_plotter.aimodel_plotdata)
-
-        canvas[key].plotly_chart(fig, use_container_width=True, theme="streamlit")
-
+    load_canvas_on_state(new_ts)
     last_ts = new_ts
 
     if last_ts == "final":
-        title.title("Final sim state")
+        snapshots = SimPlotter.available_snapshots()
+        timestamp = inputs.select_slider("Go to snapshot", snapshots, value="final")
+        st, new_ts = sim_plotter.load_state(timestamp)
+        load_canvas_on_state(timestamp)
         break