Create app.py

app.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+import streamlit as st
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Define the MemoryAugmentedLSTM model
+class MemoryAugmentedLSTM(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size, num_memories=10):
+        super(MemoryAugmentedLSTM, self).__init__()
+        self.lstm = nn.LSTM(input_size, hidden_size)
+        self.fc = nn.Linear(hidden_size, output_size)
+        
+        # Memory matrix (example: 10 slots, each with hidden_size dimensions)
+        self.memory_bank = torch.zeros(num_memories, hidden_size)
+        
+        self.num_memories = num_memories
+        self.hidden_size = hidden_size
+
+    def forward(self, x, pressure):
+        lstm_out, _ = self.lstm(x)
+        
+        # Memory allocation based on pressure
+        memory_idx = self._get_memory_idx(pressure)
+        
+        # Read from memory (get the most relevant memory slot)
+        memory_out = self.memory_bank[memory_idx]
+
+        # Combine the LSTM output and the relevant memory
+        combined_output = lstm_out[-1] + memory_out
+        
+        # Final prediction
+        prediction = self.fc(combined_output)
+        return prediction
+
+    def _get_memory_idx(self, pressure):
+        # Normalize the pressure to determine the most relevant memory slot
+        normalized_pressure = (pressure - 0.1) / (1.0 - 0.1)  # Normalize between 0 and 1
+        memory_idx = int(normalized_pressure * (self.num_memories - 1))
+        return memory_idx
+
+# Load the trained model
+model = MemoryAugmentedLSTM(input_size=10, hidden_size=50, output_size=1, num_memories=10)
+model.load_state_dict(torch.load('memristor_lstm.pth'))
+
+# Streamlit UI components
+st.title("Pressure to Spike Prediction Model")
+st.write("Select a pressure value between 0.1 MPa and 1 MPa:")
+
+# Slider for selecting pressure
+pressure_input = st.slider("Pressure (MPa)", 0.1, 1.0, 0.2, step=0.1)
+
+# Example input data (adjust according to your model input)
+input_data = torch.randn(5, 1, 10)  # Example random input data for the LSTM
+
+# Make prediction using the selected pressure
+prediction = model(input_data, pressure_input).detach().numpy()
+
+# Display the prediction outcome (spike generation)
+st.write("Predicted spike outcome:")
+
+# Plot the spikes (for demonstration, we use a simple plot with noise for now)
+spike_time = np.linspace(0, 10, 100)  # Time from 0 to 10 (for example)
+spike_amplitude = prediction * np.sin(spike_time)  # Example of a spike pattern
+
+# Plot the spike generation
+plt.figure(figsize=(10, 6))
+plt.plot(spike_time, spike_amplitude, label=f"Spike Pattern for {pressure_input} MPa")
+plt.title(f"Spike Generation for Pressure = {pressure_input} MPa")
+plt.xlabel("Time (s)")
+plt.ylabel("Spike Amplitude")
+plt.legend()
+plt.grid(True)
+st.pyplot(plt)
+