Update app.py

app.py

@@ -78,7 +78,7 @@ def visualize_spike_train(spike_train, plot_container):
     plot_container.pyplot(fig)
 
 # Generate spikes for the given pressure value
-def generate_spikes_for_pressure(pressure, plot_container, duration=1):
+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
@@ -96,14 +96,14 @@ def generate_spikes_for_pressure(pressure, plot_container, duration=1):
             # Forward pass through the model
             output, memory_grid = model(pressure_input)
 
-            # Visualize the memory grid (4x4 grid)
-            visualize_memory_grid(memory_grid, plot_container)
+            # Visualize the memory grid (4x4 grid) on the left plot container
+            visualize_memory_grid(memory_grid, memory_plot_container)
 
             # 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
-            visualize_spike_train(spike_train, plot_container)
+            # 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)
@@ -119,12 +119,18 @@ def app():
     if st.button(f"Generate Spikes & Visualize Memory for {pressure} MPa"):
         st.info(f"Generating spikes for {pressure} MPa...")
 
-        # Create containers for both plots
-        memory_plot_container = st.empty()
-        spike_train_plot_container = st.empty()
+        # Create columns for both plots (left for memory grid, right for spike train)
+        col1, col2 = st.columns(2)
+
+        # Create plot containers for each column
+        with col1:
+            memory_plot_container = st.empty()
+
+        with col2:
+            spike_plot_container = st.empty()
 
         # Generate spikes and memory updates for the given pressure
-        generate_spikes_for_pressure(pressure, memory_plot_container, duration)
+        generate_spikes_for_pressure(pressure, memory_plot_container, spike_plot_container, duration)
 
         st.success(f"Spike generation and memory passing complete for {pressure} MPa!")