diff options
Diffstat (limited to 'src/torch_reservoir.py')
| -rw-r--r-- | src/torch_reservoir.py | 151 |
1 files changed, 0 insertions, 151 deletions
diff --git a/src/torch_reservoir.py b/src/torch_reservoir.py deleted file mode 100644 index a21ed75..0000000 --- a/src/torch_reservoir.py +++ /dev/null @@ -1,151 +0,0 @@ -import numpy as np -import torch -import memtorch -from spires_interface import ( - free_spires_reservoir, - init_spires_reservoir, - send_currents_to_spires, - read_spikes_from_spires, -) -import matplotlib.pyplot as plt -from sklearn.linear_model import Ridge -import time - -# ----- PARAMETERS ----- -NUM_INPUTS = 1 -NUM_OUTPUTS = 1 -NUM_NEURONS = 800 -time_steps = 2000 - -input_layer = torch.nn.Linear(NUM_INPUTS, NUM_NEURONS, bias=False) -#this is a recurrent layer sort of??? -reservoir_layer = torch.nn.Linear(NUM_NEURONS, NUM_NEURONS, bias=False) -readout_layer = torch.nn.Linear(NUM_NEURONS, NUM_OUTPUTS, bias=False) - -spires_reservoir = init_spires_reservoir(NUM_NEURONS) - -#make the reservoir sparse and random -with torch.no_grad(): - reservoir_layer.weight.data.normal_(0.0, 0.5) #random weight - - #mask so 10% of connections exist - mask = (torch.rand(NUM_NEURONS, NUM_NEURONS) < 0.10).float() - reservoir_layer.weight.data *= mask - -# patch layers into memristor crossbars with memTorch -memristor_model = memtorch.bh.memristor.VTEAM -memristor_model_params = { - 'time_series_resolution': 1e-10, - 'r_on': 50, - 'r_off': 1000, -} - -mem_input_layer = memtorch.mn.Module.patch_model( - input_layer, - memristor_model, - memristor_model_params, -) - -mem_reservoir_layer = memtorch.mn.Module.patch_model( - reservoir_layer, - memristor_model, - memristor_model_params, -) - -start_time = time.perf_counter() - -# ----- run spires & collect spikes----- -previous_spikes = torch.zeros(1, NUM_NEURONS) -spike_history = [] - -input_signal = [] -t = np.linspace(0, 20, time_steps) - -input_signal = np.sin(-2.5 * t) + np.sin(5 * t) -target_signal = np.roll(input_signal, -5) #predicting 5 time steps in the future - -for step in range(time_steps): - - input_tensor = torch.tensor([input_signal[step]]).float().unsqueeze(0) - - currents_in = mem_input_layer(input_tensor) - currents_recv = mem_reservoir_layer(previous_spikes) - - noise_multiplier = 0.08 - current_scaler = 0.05 - - total_currents = (currents_in + currents_recv) - total_currents = total_currents - total_currents.mean() # remove DC bias - total_currents = total_currents + (torch.randn_like(total_currents) * noise_multiplier ) - total_currents = total_currents * current_scaler #scale to fit threshhold - - send_currents_to_spires(spires_reservoir, total_currents) - current_spikes = read_spikes_from_spires(spires_reservoir, NUM_NEURONS) - - spike_history.append(current_spikes.copy()) - - # convert spikes back to tensor - previous_spikes = torch.from_numpy(current_spikes).float().unsqueeze(0) - -#low pass filter -spike_matrix = np.array(spike_history) - -decay_rate = 0.85 -filtered_spikes = np.zeros_like(spike_matrix, dtype=float) - -current_trace = np.zeros(NUM_NEURONS) -for i in range(len(spike_matrix)): - current_trace = current_trace * decay_rate + spike_matrix[i] - filtered_spikes[i] = current_trace - -#training readout layer -print("training the readout layer") -x_train = filtered_spikes[100:800] -y_train = target_signal[100:800] - -x_test = filtered_spikes[800:] -y_test = target_signal[800:] - -ridge = Ridge(alpha=5) -ridge.fit(x_train, y_train) - -#predict and plot results -print("Generating predictions...") -predictions = ridge.predict(x_test) - -end_time = time.perf_counter() -execution_time = end_time - start_time -print(f"execution time: {execution_time}") - -test_steps = np.arange(800, time_steps) - -plt.figure(figsize=(10, 5)) -plt.plot(test_steps, y_test, label="True Future Wave", color="black", linestyle="dashed") -plt.plot(test_steps, predictions, label="Reservoir Prediction", color="blue", alpha=0.8) -plt.title("spires memristor Time-Series prediction") -plt.xlabel("time steps") -plt.ylabel("Amplitude") -plt.legend() -plt.tight_layout() -plt.savefig("memristor_spires_proof.png") - -free_spires_reservoir(spires_reservoir) - -# # ---------- Plotting ---------- -# #This block is all AI generated to be transparent -# #plotting to verify neurons are firing randomly -# -# print("Simulation complete, plotting results") #why? cuz its fun and almost 5 -# spike_matrix = np.array(spike_history) -# # Plotting -# plt.figure(figsize=(12, 6)) -# # Transpose (.T) so Time is the X-axis and Neurons are the Y-axis -# plt.imshow(spike_matrix.T, aspect='auto', cmap='binary', interpolation='nearest') -# -# plt.title("Reservoir Spiking Activity (Raster Plot)") -# plt.xlabel("Time Step") -# plt.ylabel("Neuron ID (0 to 799)") -# plt.colorbar(label="Spike (0 or 1)") -# plt.tight_layout() -# plt.show() -# |
