Update app.py

app.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 import matplotlib.pyplot as plt
 import streamlit as st
 import time
+import numpy as np
 
 # Simple LSTM with a VRAM-like 4x4 Memory Grid, including gates and filters
 class MemristorLSTM(nn.Module):
@@ -11,46 +12,27 @@ class MemristorLSTM(nn.Module):
         self.memory_size = memory_size
         self.lstm = nn.LSTM(input_size=1, hidden_size=50, num_layers=2, batch_first=True)
         self.fc = nn.Linear(50, self.memory_size * self.memory_size)  # Output 16 values (4x4 grid)
-        # Gate to modulate the output
-        self.gate = nn.Sigmoid()
-        # Tanh filter to apply non-linearity to the output
-        self.filter = nn.Tanh()
-
-        # 4x4 memory grid initialized to zeros
-        self.memory = torch.zeros(memory_size, memory_size)  
+        self.gate = nn.Sigmoid()  # Gate to modulate the output
+        self.filter = nn.Tanh()   # Non-linearity filter
+        self.memory = torch.zeros(memory_size, memory_size)  # 4x4 memory grid initialized to zeros
 
     def forward(self, x):
-        # Forward pass through LSTM
         lstm_out, _ = self.lstm(x)
         output = self.fc(lstm_out[:, -1, :])
-
-        # Apply sigmoid gate to control the flow of information
         gated_output = self.gate(output)
-
-        # Apply tanh filter for non-linearity
         filtered_output = self.filter(gated_output)
-
-        # Update memory (4x4 grid), simulating synaptic memory flow
         self.memory = self.memory + filtered_output.view(self.memory_size, self.memory_size)
         return filtered_output, self.memory
 
-# Load Pretrained Model with Weights Compatibility
+# Load Pretrained Model
 def load_model(model_path="memristor_lstm.pth"):
     model = MemristorLSTM(memory_size=4)
     try:
-        # Load the pretrained weights
         pretrained_dict = torch.load(model_path, map_location=torch.device('cpu'))
-        
-        # Extract the state_dict of the LSTM layers only
         model_dict = model.state_dict()
-        
-        # Only update the LSTM layers with the pretrained weights, leaving fc layer to be reinitialized
         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and 'fc' not in k}
-        
-        # Update the model's state_dict
         model_dict.update(pretrained_dict)
         model.load_state_dict(model_dict)
-        
         model.eval()
         return model
     except Exception as e:
@@ -64,75 +46,64 @@ def visualize_memory_grid(memory_grid, plot_container):
     ax.set_title("4x4 Memory Grid - VRAM")
     ax.set_xlabel("Memory Cells (Columns)")
     ax.set_ylabel("Memory Cells (Rows)")
-    plot_container.pyplot(fig)  # Render plot in Streamlit using st.pyplot
+    
+    plot_container.pyplot(fig)  # Streamlit plot
+    plt.close(fig)  # Close the figure to save memory
+
+# Generate Spike Train
+def generate_spike_train(length=50, threshold=0.5):
+    spike_train = np.random.rand(length) > threshold
+    return spike_train.astype(int)
 
-# Visualize the Spike Train Output
+# Visualize the Spike Train
 def visualize_spike_train(spike_train, plot_container):
-    # Plot the spike train (1's and 0's for spike events)
     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])  # Show only 0 and 1 on the y-axis
+    ax.set_yticks([0, 1])
+
     plot_container.pyplot(fig)
+    plt.close(fig)
 
-# Generate spikes for the given pressure value
-def generate_spikes_for_pressure(pressure, memory_plot_container, spike_plot_container, duration=1):
-    # Simulate the pressure as the input (scale to appropriate format for LSTM)
-    # Assuming the pressure range is from 0.1 to 1.0 MPa, and normalize the value
-    pressure_normalized = (pressure - 0.1) / (1.0 - 0.1)  # Normalize between 0 and 1
-    
-    # Reshape the pressure value to match LSTM input format: (batch_size, seq_len, input_size)
+# Generate spikes and VRAM updates
+def generate_spikes_for_pressure(pressure):
+    pressure_normalized = (pressure - 0.1) / (1.0 - 0.1)  
     pressure_input = torch.tensor([[pressure_normalized]], dtype=torch.float32).view(1, 1, 1)
-    
-    # Load the pre-trained model
-    model = load_model()  # Ensure you have 'memristor_lstm.pth' in the correct path
-    
-    if model is not None:
-        # Loop for live plotting over the given duration
-        start_time = time.time()
-        while time.time() - start_time < duration * 60:  # duration in minutes
-            # Forward pass through the model
-            output, memory_grid = model(pressure_input)
-
-            # Visualize the memory grid (4x4 grid) on the left plot container
-            visualize_memory_grid(memory_grid, memory_plot_container)
+    model = load_model()
 
-            # Convert the output into a spike train by thresholding the filtered output
-            spike_train = (output.detach().numpy() > 0).astype(int)  # 1 for spike, 0 for no spike
-
-            # Visualize the spike train on the right plot container
-            visualize_spike_train(spike_train, spike_plot_container)
-
-            # Delay to simulate real-time plotting (e.g., update every second)
-            time.sleep(1)
+    if model is not None:
+        output, memory_grid = model(pressure_input)
+        spike_train = generate_spike_train(length=50)
+        return output, memory_grid, spike_train
+    return None, None, None
 
-# Streamlit UI for interacting with pressure input
+# Streamlit UI
 def app():
-    st.title("Memristor Response Spike Generator with Pressure Input")
-    st.write("Generate spike patterns and visualize memory passing in a 4x4 memory grid.")
+    st.title("Memristor VRAM & Spike Train Live Visualization")
+    st.write("Watch how VRAM and spike trains change over time with pressure input.")
 
     pressure = st.slider("Select Pressure (MPa)", 0.1, 1.0, 0.5, 0.1)
-    duration = st.selectbox("Select Duration", [1, 2], help="Choose duration for the live plot (minutes)")
+    duration = st.radio("Select Duration", [1, 2], index=0)
 
-    if st.button(f"Generate Spikes & Visualize Memory for {pressure} MPa"):
-        st.info(f"Generating spikes for {pressure} MPa...")
+    if st.button("Start Live Visualization"):
+        st.info(f"Running live update for {duration} minute(s)...")
 
-        # Create columns for both plots (left for memory grid, right for spike train)
-        col1, col2 = st.columns(2)
+        plot_col1, plot_col2 = st.columns(2)  # Side-by-side layout
 
-        # Create plot containers for each column
-        with col1:
-            memory_plot_container = st.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)
 
-        with col2:
-            spike_plot_container = st.empty()
+            if output is not None and memory_grid is not None:
+                with plot_col1:
+                    visualize_memory_grid(memory_grid, st)
 
-        # Generate spikes and memory updates for the given pressure
-        generate_spikes_for_pressure(pressure, memory_plot_container, spike_plot_container, duration)
+                with plot_col2:
+                    visualize_spike_train(spike_train, st)
 
-        st.success(f"Spike generation and memory passing complete for {pressure} MPa!")
+            time.sleep(1)  # Update every second
 
 if __name__ == "__main__":
     app()