Update app.py

app.py

@@ -40,15 +40,14 @@ def load_model(model_path="memristor_lstm.pth"):
         return None
 
 # Visualize the 4x4 Memory Grid
-def visualize_memory_grid(memory_grid, plot_container):
+def visualize_memory_grid(memory_grid):
     fig, ax = plt.subplots(figsize=(6, 6))
     ax.imshow(memory_grid.detach().numpy(), cmap='viridis', interpolation='nearest')
     ax.set_title("4x4 Memory Grid - VRAM")
     ax.set_xlabel("Memory Cells (Columns)")
     ax.set_ylabel("Memory Cells (Rows)")
     
-    plot_container.pyplot(fig)  # Streamlit plot
-    plt.close(fig)  # Close the figure to save memory
+    return fig  # Return the figure
 
 # Generate Spike Train
 def generate_spike_train(length=50, threshold=0.5):
@@ -56,16 +55,15 @@ def generate_spike_train(length=50, threshold=0.5):
     return spike_train.astype(int)
 
 # Visualize the Spike Train
-def visualize_spike_train(spike_train, plot_container):
+def visualize_spike_train(spike_train):
     fig, ax = plt.subplots(figsize=(6, 4))
     ax.plot(spike_train, marker='o', linestyle='-', color='r', markersize=5)
     ax.set_title("Spike Train")
     ax.set_xlabel("Time Steps")
     ax.set_ylabel("Spike Event (1 = Spike, 0 = No Spike)")
     ax.set_yticks([0, 1])
-
-    plot_container.pyplot(fig)
-    plt.close(fig)
+    
+    return fig  # Return the figure
 
 # Generate spikes and VRAM updates
 def generate_spikes_for_pressure(pressure):
@@ -92,18 +90,20 @@ def app():
 
         plot_col1, plot_col2 = st.columns(2)  # Side-by-side layout
 
+        # Create empty containers to update plots in real-time
+        plot_vram = plot_col1.empty()
+        plot_spike_train = plot_col2.empty()
+
         end_time = time.time() + duration * 60  # Run for selected time
         while time.time() < end_time:
             output, memory_grid, spike_train = generate_spikes_for_pressure(pressure)
 
             if output is not None and memory_grid is not None:
-                with plot_col1:
-                    visualize_memory_grid(memory_grid, st)
-
-                with plot_col2:
-                    visualize_spike_train(spike_train, st)
+                # Update plots dynamically in existing containers
+                plot_vram.pyplot(visualize_memory_grid(memory_grid))
+                plot_spike_train.pyplot(visualize_spike_train(spike_train))
 
-            time.sleep(1)  # Update every second
+            time.sleep(0.2)  # Update every 0.2 seconds (FAST UPDATES)
 
 if __name__ == "__main__":
     app()