main.py

# main.py

import numpy as np
import matplotlib.pyplot as plt

from network import Protocol, NetworkManager, BCPNNPerfect  # , BCPNNModular
from plotting_functions import plot_weight_matrix, plot_state_variables_vs_time, plot_winning_pattern
from plotting_functions import plot_network_activity, plot_network_activity_angle
from analysis_functions import calculate_recall_time_quantities, calculate_angle_from_history

# ---- 1. DEFINE SIMULATION PARAMETERS ----

# Network parameters
g_w_ampa = 2.0
g_w = 0.0
g_a = 10.0
tau_a = 0.250
G = 1.0
sigma = 0.0

# Pattern parameters
hypercolumns = 1
minicolumns = 10
n_patterns = 10

# Manager properties
dt = 0.001
values_to_save = ['o']  # Only save 'o' (output activity) for this example

# Protocol parameters
training_time = 0.100
inter_pulse_interval = 0.0
inter_sequence_interval = 1.0
epochs = 15

# ---- 2. BUILD AND TRAIN THE NETWORK ----

# Initialize the BCPNN network with specified parameters
nn = BCPNNPerfect(hypercolumns, minicolumns, g_w_ampa=g_w_ampa, g_w=g_w, g_a=g_a, tau_a=tau_a,
                  sigma=sigma, G=G,
                  z_transfer=False, diagonal_zero=False, strict_maximum=True, perfect=True)

# Initialize the manager to handle the network simulation
manager = NetworkManager(nn=nn, dt=dt, values_to_save=values_to_save)

# Define the training protocol
protocol = Protocol()
patterns_indexes = list(range(n_patterns))
protocol.simple_protocol(patterns_indexes, training_time=training_time,
                         inter_pulse_interval=inter_pulse_interval,
                         inter_sequence_interval=inter_sequence_interval,
                         epochs=epochs)

# Run the training protocol
print("Starting network training...")
manager.run_network_protocol(protocol=protocol, verbose=True)
print("Training complete.")

# Plot the learned weight matrix
plot_weight_matrix(manager.nn, ampa=True)
plt.suptitle("Learned AMPA Weight Matrix")
plt.show()


# ---- 3. RUN RECALL AND ANALYSIS ----

# Recall parameters
T_recall = 2.0
T_cue = 0.100
sequences = [patterns_indexes]
I_cue = 0.0
n_trials = 1  # Number of recall trials

# Calculate recall performance metrics
print("\nStarting recall analysis...")
aux = calculate_recall_time_quantities(manager, T_recall, T_cue, n_trials, sequences)
total_sequence_time, mean_time, std_time, success_rate, timings = aux

# Plot network activity during recall
plot_network_activity_angle(manager)
plt.suptitle("Network Activity During Recall")
plt.show()

print(f"Recall Success Rate: {success_rate}%")
print(f"Total time for sequence recall: {total_sequence_time:.2f}s")