import numpy as np import pickle from multiprocessing import Pool, cpu_count import datetime, os import random import matplotlib.pyplot as plt ####### from source.utils import parameters_linear_track, run_td_lambda_new_continuousTime,run_MC from schema import Schema, Optional, And from types import SimpleNamespace def fun_to_run(traj): traj,nr,ini,path, tlist, mode, store_data = traj print(nr) params, gamma, eta, lambda_var = parameters_linear_track(no_states=3) params['trajectories'] = traj params = update_params(mode, params, tlist, gamma, eta, state_3_rate=1) if mode == 'TD': M_stdp = run_td_lambda_new_continuousTime(**params) if mode == 'MC': M_stdp = run_MC(**params) if mode == 'MX': Trials = len(traj) tau = 6 # ratios = np.exp(-1*np.arange(Trials)/tau) M_stdp = np.zeros((Trials, params['no_states'], params['no_states'])) w = [] for trial in range(Trials): params, _, _, _ = parameters_linear_track(no_states=3) params['trajectories'] = [traj[trial]] params['w'] = w if np.random.rand() -1 and state < 3: next_state = state + random.choice([-1,1]) traj.append(next_state) state = next_state return traj[:-1] def root_mse(M_est, M_theor, cutoff=0): # trials, epochs, state, state mse = ((((M_est-M_theor)**2).mean(axis=0))**0.5).mean(axis=1).mean(axis=1) stderr = ((((M_est-M_theor)**2).mean(axis=2).mean(axis=2))**0.5).std(axis=0)/np.sqrt(len(M_est)) return mse[cutoff:], stderr[cutoff:] def create_path(conf: dict): """ Create experiment path. Parameters ---------- conf: dict Dict of the yaml configuration file. logg: LOGG Logger object. Returns ------- newpath: str Path to experiment folder. """ # Create experiment folder date_time = str(datetime.datetime.now()).split('.')[0].replace('-','_').replace(':','_').replace(' ','_') newpath = os.path.join(conf.experiment_folder, date_time) newpath = f'{conf.mode}_runs/'+newpath # Create storage folder if conf.store_data: os.makedirs(newpath) print('Folder {} created!'.format(newpath)) return newpath def validate_conf(conf: dict) -> SimpleNamespace : """ Validate the hyperparameters in the dictionary. Parameters ---------- conf : dict Dictionary containing hyperparemeters to run the simulation. Returns ------- SimpleNamespace Namespace containing the validated hyperparameters. """ valid_conf = Schema({ Optional('state_3_T', default=100): And(int, lambda s: s in [100]), Optional('experiment_folder', default="simulation"): str, Optional('initial_trial_nr', default=1): And(int, lambda s: s>0), Optional('final_trial_nr', default=2): And(int, lambda s: s>0), Optional('nr_epochs', default=2): And(int, lambda s: s>0), # Optional('nr_states', default=10): And(int, lambda s: s>2), Optional('cpu_percentage', default=0.9): And(float, lambda s: s>0, lambda s: s<=1), Optional('multiprocessing', default=False): bool, Optional('store_data', default=False): bool, Optional('mode', default="MC"): And(str, lambda s: s in ['TD', 'MC', 'MX']), Optional('state_3_rate', default=1): And(int, lambda s: s in [1]), }).validate(conf) assert valid_conf['final_trial_nr'] > valid_conf['initial_trial_nr'], 'ERROR: Number of final trial should be larger than number of initial trial!' return SimpleNamespace(**valid_conf) def run_simulation(conf): """ Run a simulation on the linear track. Parameters ---------- """ # validate cofiguration conf = validate_conf(conf) # Define total number of trials tot_trials = conf.final_trial_nr-conf.initial_trial_nr # Define total number of trials tot_trials = conf.final_trial_nr-conf.initial_trial_nr # Define trajectory and time per state nr_states = 3 traj = [ [generate_traj() for _ in range(conf.nr_epochs)] for _ in range(tot_trials)] T_list = [[[100]*len(t) for t in tt] for tt in traj] # Define function to run and storage path newpath = create_path(conf) # Run in parallel if set if conf.multiprocessing: cc = int(conf.cpu_percentage*cpu_count()) print(f'Using max {cc} cpus...') pool = Pool(cc) output = pool.map(lambda x: fun_to_run(x), zip(traj,range(len(traj)),[conf.initial_trial_nr]*len(traj), [newpath]*len(traj), T_list, [conf.mode]*len(traj), [conf.store_data]*len(traj), ) ) outputs=np.array(output) else: output = [] for idx,tt in enumerate(traj): output.append(fun_to_run((tt,idx,conf.initial_trial_nr, newpath, T_list[idx],conf.mode, conf.store_data, ) ) ) outputs=np.array(output) # cumulative lengths l = [np.cumsum([len(x) for x in t]) for t in traj] # params, gamma_var, _, _ = fetch_parameters_new() params, gamma_var, _, _ = parameters_linear_track(nr_states) if conf.mode == 'TD': w_stdps = np.array([[outputs[i][l_i] for l_i in l[i]] for i in range(len(outputs))]) w_stdps2 = np.mean(np.sum(w_stdps[:,:,:,:,:,:],axis=-2)/params['N_pre'],axis=-1) elif conf.mode == 'MC': w_stdps2 = outputs #np.array(outputs) elif conf.mode == 'MX': w_stdps2 = outputs if conf.store_data: pickle.dump(w_stdps2, open(newpath + '/Fig4_w_stdps2.pkl','wb')) T_mat = np.array( [[0,0.5,0,0.5], [0.5,0,0.5,0], [0,0.5,0,0.5], [0,0,0,1], ] ) M_theor = np.linalg.inv(np.eye(3)-gamma_var*T_mat[:-1,:-1]) plt.figure() root_mse_spiking,stderr = root_mse(w_stdps2, M_theor, 0) plt.plot(root_mse_spiking) plt.fill_between(range(len(stderr)), root_mse_spiking-stderr, root_mse_spiking+stderr,alpha=.3) if __name__ == '__main__': conf_figure4 = { 'experiment_folder': "simulations_figure4", # folder to store 'initial_trial_nr': 1, # initial trial number 'final_trial_nr': 3, # final trial number 'nr_epochs': 4, # number of epochs per trial 'cpu_percentage': 0.9, # percentage of cpus to use when running in parallel 'multiprocessing': False, # whether to run in parallel using multiprocessing 'store_data': False, # whether to store the experiment data 'mode': "MC", # mode, can be MC, TD or MX } run_simulation(conf_figure4)