AdEx_params_space.py

# importing libraries 
import os
import numpy as np
import json
from scipy.stats import qmc
import math
import matplotlib.pyplot as plt
from exploration_helper import *

##########################
# 0. Simulation settings #
##########################

fitting_folder_name = 'test_cortex_control_v0'
num_samples = 2**12
num_best_model = 2**3

parallel_running = False
#parallel_running = True

# If parallel_running = True, specify the number of threads

#num_threads = os.cpu_count()
#num_threads = 4

# Set the seed for reproducibility
seed = 12345

reading_from_file_AdEx = False
reading_from_file_Err  = False

if not os.path.exists(fitting_folder_name):
    # Create the folder if it doesn't exist
    os.makedirs(fitting_folder_name)
    print(f"Folder '{fitting_folder_name}' created.")
else:
    print(f"Folder '{fitting_folder_name}' already exists.")

#####################
# 1. Setting bounds #
#####################

def params_bound(cm_bound = None,
                 gl_bound = None,
                 v_rest_bound = None,
                 i_offset_bound = None,
                 a_bound = None,
                 b_bound = None,
                 tau_w_bound = None,
                 v_thresh_bound = None,
                 delta_T_bound = None,
                 v_reset_bound = None,
                 v_spike_bound = None,
                 tau_refrac_bound = None,
                 v_0_bound = None,
                 w_0_bound = None):

    args = locals()
    
    intervals = []
    params_to_opt = []
    fixed_params_name = []
    fixed_params_value = []

    for arg, value in args.items():
        if type(value) == tuple:
            intervals.append(value)
            params_to_opt.append(str(arg[:-6]))
        else: 
            fixed_params_name.append(str(arg[:-6]))
            fixed_params_value.append(value)
    return params_to_opt, intervals, fixed_params_name, fixed_params_value

### The units are the standars used in PyNN ###
cm_bound         = (0.1,0.2)       # nF - Capacity of the membrane
gl_bound         = (0.02, 0.05)    # uS - Leak conductance
v_rest_bound     = -60             # mV - Leak reversal potential
i_offset_bound   =  0              # nA - Constant external input current
a_bound          = (0, 10)         # nS - Subthreshold adaptation
b_bound          = (0, 0.3)        # nA - Spike-triggered adaptation
tau_w_bound      = (0, 200)        # ms - Adaptation time constant
v_thresh_bound   = (-45, -43)      # mV - Spike initiation threshold
delta_T_bound    = (2, 8)          # mV - Slope factor
v_reset_bound    = (-53, -50)      # mV - Reset value for V_m after a spike
v_spike_bound    = 17              # mV - Spike detection threshold
tau_refrac_bound = (1, 3)          # ms - Duration of refractory period
v_0_bound        = (-60)           # mV - Initial membrane potential value
w_0_bound        = 0               # nA - Initial spike adaptation current 

params_to_opt, intervals, fixed_params_name, fixed_params_value = params_bound(
                                                                        cm_bound         = cm_bound,
                                                                        gl_bound         = gl_bound,
                                                                        v_rest_bound     = v_rest_bound,
                                                                        i_offset_bound   = i_offset_bound,
                                                                        a_bound          = a_bound,
                                                                        b_bound          = b_bound,
                                                                        tau_w_bound      = tau_w_bound,
                                                                        v_thresh_bound   = v_thresh_bound,  
                                                                        delta_T_bound    = delta_T_bound,
                                                                        v_reset_bound    = v_reset_bound,
                                                                        v_spike_bound    = v_spike_bound,
                                                                        tau_refrac_bound = tau_refrac_bound,
                                                                        v_0_bound        = v_0_bound,
                                                                        w_0_bound        = w_0_bound
                                                                        )

print(f'parameters to optimize for: {params_to_opt}, and their respective boundaries {intervals}')
print(f'parameters fixed: {fixed_params_name} and their respectives values: {fixed_params_value}')

######################################################
# 2. Generating distribution of values in the bounds #
######################################################
rng = np.random.default_rng(seed)

# Number of samples and dimensions
n_samples = num_samples
dim = len(intervals)

# Initialize Sobol engine with the specified seed
engine = qmc.Sobol(d=dim, scramble=True, seed=rng)

# Generate Sobol samples
sobol_samples = engine.random(n=n_samples)

# Scale samples to the specified intervals
scaled_samples = np.zeros_like(sobol_samples)
for i, (a, b) in enumerate(intervals):
    scaled_samples[:, i] = np.round(a + sobol_samples[:, i] * (b - a), 3)

# Print the samples
print("Samples with seed", seed)

#Saving a json file containing such information
def append_to_json_file(file_path, data):
    # Check if the file exists
    if os.path.exists(file_path):
        # Load existing data from the JSON file
        with open(file_path, 'r') as file:
            existing_data = json.load(file)
    else:
        existing_data = []

    # Append the new data
    existing_data.append(data)

    # Write the updated data back to the JSON file
    with open(file_path, 'w') as file:
        json.dump(existing_data, file, indent=4)

def AdEx_params(cm = None, gl = None, v_rest = None, i_offset = None, a = None, b = None,
                tau_w = None, v_thresh = None, delta_T = None, v_reset = None, v_spike = None, tau_refrac = None,
                v_0 = None, w_0 = None):
            
    params =  { 'C_m': cm ,                         # nF - Capacity of the membrane
                'g_L': gl,                          # uS - Leak conductance
                'E_L': v_rest,                      # mV - Leak reversal potential
                'I_e': i_offset,                    # nA - Constant external input current
                'a': np.round(a * 1e-3, 3),         # uS - Subthreshold adaptation
                'b': b,                             # nA - Spike-triggered adaptation
                'tau_w': tau_w,                     # ms - Adaptation time constant
                'V_th': v_thresh,                   # mV - spike threshold 
                'Delta_T': delta_T,                 # mV - Slope factor
                'V_reset': v_reset,                 # mV - Reset value for V_m after a spike
                'V_peak': v_spike,                  # mV - Spike detection threshold
                't_ref': tau_refrac                 # ms - Duration of refractory period
                }
        
    initial_values = {'v': v_0, 'w': w_0}
    return params, initial_values

json_file_model_name = (os.path.join(fitting_folder_name,
                                     'AdEx_models_testing_' + str(num_samples) + '.json'))

# Prepare fixed parameters outside the loop
fixed_params = dict(zip(fixed_params_name, fixed_params_value))

# List to hold data for bulk appending
params_data_to_append_list = []

for i in range(num_samples):
    
    # Prepare optimized parameters
    opt_params = dict(zip(params_to_opt, scaled_samples[i, :]))
    
    # Combine fixed and optimized parameters
    params = {**opt_params, **fixed_params}
    
    # Generate AdEx parameters and initial values
    temp_params, temp_initial_values = AdEx_params(**params)
    
    # Append data to the list
    params_data_to_append_list.append((temp_params, temp_initial_values))

# Append all data to the JSON file in one go
append_to_json_file(json_file_model_name, params_data_to_append_list)

#########################################################
# 3. Setting frequency data and other selected features #
#########################################################
data_file = os.path.join('../', 'extracting_features', 'extracting_features_test.json')
with open(data_file, 'r') as file:
    exp_data = json.load(file)

data_current_tot_sup = exp_data['current']           # pA
data_freq_tot = exp_data['mean_frequency']                          # Hz 
data_inv_first_ISI_tot = exp_data['inv_first_ISI']                  # Hz  
data_inv_last_ISI_tot = exp_data['inv_last_ISI']                    # Hz 
data_time_to_first_spike_tot = exp_data['time_to_first_spike']      # ms 
data_time_to_second_spike_tot = exp_data['time_to_second_spike']    # ms 
data_time_to_third_spike_tot = exp_data['time_to_third_spike']      # ms 
data_time_to_last_spike_tot = exp_data['time_to_last_spike']        # ms 

data_volt_stimend = None

# suprathreshold features
indices_sup = [0]
data_current_sup = [data_current_tot_sup[i] for i in indices_sup]
data_freq = [data_freq_tot[i] for i in indices_sup]
data_inv_first_ISI = [data_inv_first_ISI_tot[i] for i in indices_sup]
data_inv_last_ISI = [data_inv_last_ISI_tot[i] for i in indices_sup]
data_time_to_first_spike = [data_time_to_first_spike_tot[i] for i in indices_sup]
data_time_to_second_spike = [data_time_to_second_spike_tot[i] for i in indices_sup]
data_time_to_third_spike = [data_time_to_third_spike_tot[i] for i in indices_sup]
data_time_to_last_spike = [data_time_to_last_spike_tot[i] for i in indices_sup]

data_current = data_current_sup

################################
# 4. Setting up the AdEx model #
################################
# If data is read from a json file
if reading_from_file_AdEx == True:
    json_file_name = 'test_data_' + str(num_samples) + '.json'  
    with open(json_file_name, 'r') as file:
        data = json.load(file)
    # the following line is needed to avoid having a list of lists    
    models = [item for sublist in data for item in sublist]     
    print(f"Imported json data from json file: {json_file_name}")
else:
    models = params_data_to_append_list

############################################
# 5. Setting up current injection protocol #
############################################
Time = exp_data["stimulation_protocol"][0][3]  # ms
dt = 0.01 # ms

# explicitly defining the time
time=np.arange(0,Time,dt)

amps = data_current                                         # pA
delay = exp_data["stimulation_protocol"][0][1]              # ms
duration = exp_data["stimulation_protocol"][0][2] - delay   # ms
holding_current = exp_data["stimulation_protocol"][1][0]    #pA

stim_delay = delay / dt     
stim_duration = duration / dt    # ms
current_protocols =[]
for amp in data_current:
    temp_current_protocol = np.zeros(len(time))
    temp_current_protocol[0:int(stim_delay)] = holding_current #pA 
    temp_current_protocol[int(stim_delay):int(stim_delay+stim_duration)] = amp
    temp_current_protocol[int(stim_delay+stim_duration):] = holding_current #pA
    current_protocols.append(temp_current_protocol)

##############################
# 6. Running the simulations #
##############################
model_current = np.zeros((len(models), len(data_current)))
model_volt_stimend = np.zeros((len(models), len(data_current)))
model_freq = np.zeros((len(models), len(data_current)))
model_inv_first_ISI = np.zeros((len(models), len(data_current)))
model_inv_last_ISI = np.zeros((len(models), len(data_current)))
model_time_to_first_spike = np.zeros((len(models), len(data_current)))
model_time_to_second_spike = np.zeros((len(models), len(data_current)))
model_time_to_third_spike = np.zeros((len(models), len(data_current)))
model_time_to_last_spike = np.zeros((len(models), len(data_current)))

if parallel_running == True:
    from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
    with ProcessPoolExecutor(max_workers=num_threads) as executor:
        futures = [executor.submit(run_parallel_Euler_simulation_features, idx, model,
                                   current_protocols, model_current, model_freq,
                                   dt, time, stim_delay, stim_duration) for idx, model in enumerate(models)]
    
        for future in as_completed(futures):
            results = future.result()

            for idx, cc, current, volt_stimend, freq, inv_first_ISI, inv_last_ISI, time_to_first_spike, time_to_second_spike, time_to_third_spike, time_to_last_spike in results:
                model_current[idx, cc] = current
                model_volt_stimend[idx, cc] = volt_stimend
                model_freq[idx, cc] = freq
                model_inv_first_ISI[idx, cc] = inv_first_ISI
                model_inv_last_ISI[idx, cc] = inv_last_ISI
                model_time_to_first_spike[idx, cc] = time_to_first_spike  
                model_time_to_second_spike[idx, cc] = time_to_second_spike  
                model_time_to_third_spike[idx, cc] = time_to_third_spike  
                model_time_to_last_spike[idx, cc] = time_to_last_spike  
else:
   for idx, model in enumerate(models):
    results = run_parallel_Euler_simulation_features(
        idx, model, current_protocols, model_current, model_freq,
        dt, time, stim_delay, stim_duration
    )

    for idx, cc, current, volt_stimend, freq, inv_first_ISI, inv_last_ISI, time_to_first_spike, time_to_second_spike, time_to_third_spike, time_to_last_spike in results:
        model_current[idx, cc] = current
        model_volt_stimend[idx, cc] = volt_stimend
        model_freq[idx, cc] = freq
        model_inv_first_ISI[idx, cc] = inv_first_ISI
        model_inv_last_ISI[idx, cc] = inv_last_ISI
        model_time_to_first_spike[idx, cc] = time_to_first_spike  
        model_time_to_second_spike[idx, cc] = time_to_second_spike  
        model_time_to_third_spike[idx, cc] = time_to_third_spike  
        model_time_to_last_spike[idx, cc] = time_to_last_spike   
        
#####################################
# 7. Computing and saving the error #
#####################################

# Computing the relative error. 
def compute_error(data_current, indices_sup, 
                 data_freq, model_freq,
                 data_volt_stimend, model_volt_stimend,
                 data_inv_first_ISI, model_inv_first_ISI,
                 data_inv_last_ISI, model_inv_last_ISI,
                 data_time_to_first_spike, model_time_to_first_spike,
                 data_time_to_second_spike, model_time_to_second_spike,
                 data_time_to_third_spike, model_time_to_third_spike,
                 data_time_to_last_spike, model_time_to_last_spike
                 ):
    err = np.zeros((num_samples, len(data_current)))
    err_tot = np.zeros(num_samples)

    # Definig the weight of each feature
    # Default value is 1
    w_freq = 1
    w_first_spike = 1
    w_second_spike = 1
    w_inv_first_ISI = 1
    w_third_spike = 1
    w_last_spike = 1
    w_inv_last_ISI = 1
    
    for i in range(0, num_samples):
        for j in range(0, len(data_current)):
            if abs(data_freq[j]) != 0:
                ### adding relative error freq
                
                err[i,j] = (abs((abs(data_freq[j]) - abs(model_freq[i,j]))/abs(data_freq[j]))) * w_freq 
                
                if math.isnan(model_time_to_first_spike[i,j]):
                        err[i,j] = err[i,j] + 3
                else:
                    time_to_first_spike_diff = (abs(data_time_to_first_spike[j] - model_time_to_first_spike[i,j])/data_time_to_first_spike[j]) * w_first_spike
                    ### adding relative error time_to_first_spike
                    err[i,j] = err[i,j] + time_to_first_spike_diff
                if abs(data_freq[j]) > 1:
                    if math.isnan(model_inv_first_ISI[i,j]):
                        err[i,j] = err[i,j] + 3
                    else:
                        inv_first_ISI_diff = (abs(abs(data_inv_first_ISI[j]) - abs(model_inv_first_ISI[i,j]))/abs(data_inv_first_ISI[j])) * w_inv_first_ISI
                        ### adding relative error inv_first_ISI 
                        err[i,j] = err[i,j] + inv_first_ISI_diff 
                    if math.isnan(model_time_to_second_spike[i,j]):
                        err[i,j] = err[i,j] + 3    
                    else:    
                        time_to_second_spike_diff = (abs(data_time_to_second_spike[j] - model_time_to_second_spike[i,j])/data_time_to_second_spike[j]) * w_second_spike
                        ### adding relative error time_to_second_spike
                        err[i,j] = err[i,j] + time_to_second_spike_diff
                if abs(data_freq[j]) > 2:
                    if math.isnan(model_inv_last_ISI[i,j]):
                        err[i,j] = err[i,j] + 3
                    else:                    
                        inv_last_ISI_diff = (abs(abs(data_inv_last_ISI[j]) - abs(model_inv_last_ISI[i,j]))/abs(data_inv_last_ISI[j])) * w_inv_last_ISI
                        ### adding relative error inv_last_ISI 
                        err[i,j] = err[i,j] + inv_last_ISI_diff
                    if math.isnan(model_time_to_third_spike[i,j]):
                        err[i,j] = err[i,j] + 3    
                    else:                       
                        time_to_third_spike_diff = (abs(data_time_to_third_spike[j] - model_time_to_third_spike[i,j])/data_time_to_third_spike[j]) * w_third_spike
                        ### adding relative error time_to_third_spike
                        err[i,j] = err[i,j] + time_to_third_spike_diff
                    if math.isnan(model_time_to_last_spike[i,j]):
                        err[i,j] = err[i,j] + 3    
                    else:                      
                        time_to_last_spike_diff = (abs(data_time_to_last_spike[j] - model_time_to_last_spike[i,j])/data_time_to_last_spike[j]) * w_last_spike
                        ### adding relative error time_to_last_spike
                        err[i,j] = err[i,j] + time_to_last_spike_diff
            else:
                ### adding error freq 
                err[i,j] = abs(abs(data_freq[j]) - abs(model_freq[i,j]))
            
        err_tot[i] = np.sum(err[i])
            
    return err, err_tot

err, err_tot = compute_error(data_current, indices_sup, 
                 data_freq, model_freq,
                 data_volt_stimend, model_volt_stimend,
                 data_inv_first_ISI, model_inv_first_ISI,
                 data_inv_last_ISI, model_inv_last_ISI,
                 data_time_to_first_spike, model_time_to_first_spike,
                 data_time_to_second_spike, model_time_to_second_spike,
                 data_time_to_third_spike, model_time_to_third_spike,
                 data_time_to_last_spike, model_time_to_last_spike
                 )


json_file_error_name = json_file_model_name[:-5] + '_err.json'

err_data_to_append_list = []
for i in range(0,len(model_current)):
    err_data_to_append = {'err_tot': err_tot[i].tolist(), 'err_mismatch': err[i].tolist()}
    err_data_to_append_list.append(err_data_to_append)
append_to_json_file(json_file_error_name, err_data_to_append_list)

###########################
# 8. Plotting best models #
###########################

# If data is read from a json file
if reading_from_file_Err == True:
    json_file_name = 'euler_test_data_' + str(num_samples) + '_err.json'  
    with open(json_file_name, 'r') as file:
        models = json.load(file)
    print(f"Data containing frequency mistmatch imported from json file: {json_file_name}")
else:
    errors = err_data_to_append_list

arr_err_tot = np.zeros(len(errors))
for i in range(0, len(errors)):
    arr_err_tot[i] = errors[i]['err_tot']

idx_arr_err_tot_sorted = np.argsort(arr_err_tot)
increasing_err_tot = arr_err_tot[idx_arr_err_tot_sorted]


smallest_error = increasing_err_tot[:num_best_model]
smallest_error_idx = idx_arr_err_tot_sorted[:num_best_model] 

# Create a figure with an extra subplot for the legend
fig, axs = plt.subplots(1, 3, figsize=(15, 5),
                        gridspec_kw={'width_ratios': [1, 1,0.2]})  # Adjust width_ratios as needed

# Plot the data in the first three subplots
for idx in smallest_error_idx:

    axs[0].plot(model_current[idx, 0:], model_freq[idx, 0:], 
                's-', alpha=0.5, label=idx)

    axs[1].plot(model_current[idx, 0:], model_inv_first_ISI[idx, 0:], 
                's-', alpha=0.5, label=idx)

# Plot the reference data


axs[0].plot(data_current[0:], data_freq, 
            'o-', color='k', label='data')

axs[1].plot(data_current[0:], data_inv_first_ISI, 
            'o-', color='k', label='data')

# Set the titles and labels
plt.suptitle(f'{num_best_model} best models out of {num_samples}')

axs[0].set_xlabel("current (pA)")
axs[0].set_ylabel("frequency (Hz)")
axs[0].set_title('Suprathreshold response')

axs[1].set_xlabel("current (pA)")
axs[1].set_ylabel("frequency (Hz)")
axs[1].set_title('Inv first ISI')

# Collect handles and labels from one of the subplots
handles, labels = axs[0].get_legend_handles_labels()

# Create the legend in the fourth subplot
axs[2].legend(handles, labels, loc='center')
axs[2].axis('off')  # Hide the axes of the legend subplot

plt.tight_layout(rect=[0, 0, 0.9, 1])  # Adjust the layout to make space for the legend
plt.savefig(os.path.join(fitting_folder_name, 'summary_' + str(num_samples) + '.pdf'))


#########################
# 9. Saving best models #
#########################
json_file_best_models = os.path.join(fitting_folder_name,'best_models.json')
best_data_to_append_list = []
for i, idx in enumerate(smallest_error_idx):
    best_data_to_append = {'model': models[idx][0],
                           'init' : models[idx][1],
                           'err_tot': err_tot[idx].tolist(),
                           'err_mismatch': err[idx].tolist()}
    best_data_to_append_list.append(best_data_to_append)
append_to_json_file(json_file_best_models, best_data_to_append_list)