import numpy as np
import pylab as plt
from utils_learning_rules import fetch_parameters_new
from utils_environment import T_random_walk_linear_track
from utils_nn import run_td_lambda_new_continuousTime, run_MC
import pickle
import multiprocessing
from multiprocessing import Pool
from main_compareMLmodel12_random_walk_parall import root_mse
import datetime, os, string
def parall_fun(traj):
traj,nr,ini,path,T_lists = traj
print(nr)
params, gamma, eta, lambda_var = fetch_parameters_new()
params['trajectories'] = traj
params['offline'] = False
params['T_lists'] = T_lists
del params['T']
# print(gamma)
M_stdp = run_td_lambda_new_continuousTime(**params)#,ini='random')
# M_TD = theor_TD_lambda(traj, gamma, lambda_var, params['no_states'], eta)
with open(path+f'/TD_traj_{ini+nr}.pkl','wb') as f:
pickle.dump(M_stdp,f)
return M_stdp, np.zeros_like(M_stdp) #M_TD
def parall_fun_MC(traj):
traj,nr,ini,path,T_lists = traj
print(nr)
params, gamma_var, eta, _ = fetch_parameters_new()
temporal_same = 2#.1#params['tau_plus'] * np.log(params['A_plus'])
temporal_between = -np.log(gamma_var)*params['tau_plus']
params['temporal_same'] = temporal_same
params['temporal_between'] = temporal_between#15#temporal_between
# print('temporal_same: ' + str(temporal_same))
# print('temporal_between: ' + str(temporal_between))
params['eta_stdp'] = eta/(params['A_plus']*np.exp(-temporal_same/params['tau_plus']))
eta_replay = params['eta_stdp'] *params['A_plus']*np.exp(-temporal_same/params['tau_plus'])
gamma_replay = np.exp(-temporal_between/params['tau_plus'])
# print('Eta replay: '+str(eta_replay))
# print('Gamma replay: '+str(gamma_replay))
params['trajectories'] = traj
_,M_stdp = run_MC(**params,spike_noise_prob=0.2)#,ini='random')
# M_theor, theor_TD_lambda(traj, gamma_replay, 1, 4, eta_replay, M_TD_lambda=[])
with open(path+f'/MC_traj_{ini+nr}.pkl','wb') as f:
pickle.dump(M_stdp,f)
# with open(path+f'/MCtheor_traj_{ini+nr}.pkl','wb') as f:
# pickle.dump(M_theor,f)
return M_stdp, np.zeros_like(M_stdp)
if __name__ == '__main__':
# comment_script = input("Comment: ")
plt.close('all')
newpath = str(datetime.datetime.now())
newpath= newpath.translate(str.maketrans(string.punctuation, '_'*len(string.punctuation)))
newpath=newpath.replace(" ", "_")
#store comment about experiment
# g = open(newpath+'/comment.txt','w')
# g.write(comment_script)
# g.close()
# #store name experiment in file
# f = open("list_experiments", "a")
# f.write(newpath+ "\n")
# f.close()
#store copy of the script
# filename = __file__
# filename = filename.split("/")1
# filename = filename[-1]
# shutil.copy(filename, newpath)
# traj = pickle.load(open('1000trajectories.pkl', 'rb'))
ini_run = 0
fin_run = 10
assert fin_run>ini_run, 'should have more than 1 run!'
tot_runs = fin_run-ini_run
tot_epochs = 50
# traj = [[t[:] for t in x[:tot_epochs]] for x in traj[ini_run:fin_run]]
#traj = [t[:] for t in traj[:1][:tot_epochs]]*tot_runs
# traj_linear_track = np.repeat([np.repeat(np.arange(0, 1).reshape(1,-1), tot_epochs, axis=0)], tot_runs, axis=0)
traj_linear_track = np.repeat([np.repeat([[0, 1, 2, 3]], tot_epochs, axis=0)], tot_runs, axis=0)
traj = traj_linear_track
T_list = [[100,100,100,100]]*tot_epochs
mul_proc = True
mode = 'MC'
if mode == 'TD':
fun_to_run = parall_fun
newpath = 'TD_runs/'+newpath
elif mode == 'MC':
fun_to_run = parall_fun_MC
newpath = 'MC_runs/'+newpath
### make directory
os.makedirs(newpath)
if mul_proc:
cc = int(.75*multiprocessing.cpu_count())
pool = Pool(cc)
output = pool.map(fun_to_run, zip(traj,range(len(traj)),[ini_run]*len(traj),
[newpath]*len(traj), [np.array(T_list)]*len(traj) ) )
output=np.array(output)
else:
output = fun_to_run((traj[0],0,ini_run,
newpath, np.array(T_list) ))
output=np.array([output])
outputs = output[:,0]
# pickle.dump(output, open(newpath + '/output_stdp.pkl','wb'))
#pickle.dump(output_td, open(newpath + '/output_td.pkl','wb'))
l = [np.cumsum([len(x) for x in t]) for t in traj]
params, gamma_var, _, _ = fetch_parameters_new()
if 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 mode == 'MC':
w_stdps2 = np.array(outputs)
# pickle.dump(w_stdps, open(newpath + '/w_stdps.pkl','wb'))
#
pickle.dump(w_stdps2, open(newpath + '/Fig3_w_stdps2.pkl','wb'))
init = np.tile(np.expand_dims(np.expand_dims(np.eye(4),0),0),[tot_runs,1,1,1])
w_stdps2 = np.concatenate([init,w_stdps2],axis=1)
T_mat = T_random_walk_linear_track(params['no_states']-1)
M_theor = np.linalg.inv(np.eye(params['no_states']-1)-gamma_var*T_mat[:-1, :-1])
#
fig = plt.figure()
for w1 in range(params['no_states']-1):
for w2 in range(params['no_states']-1):
if 3 not in [w1,w2]:
mn = w_stdps2[:, :, w1, w2].mean(axis=0)
# std = np.std(w_stdps2[:, :, w1, w2],axis=0)
plt.plot(mn)
# plt.fill_between(range(len(std)), mn-std, mn+std,alpha=.3)
# plt.plot([0,len(w_stdps2[0,:,0,0])],[M_theor[w1, w2],M_theor[w1, w2]])
# fig.savefig(newpath + '/weights.png')
fig2 = plt.figure()
root_mse_TD_spiking,std = root_mse(w_stdps2[:,:,:-1,:-1], M_theor, 0)
plt.plot(root_mse_TD_spiking)
plt.fill_between(range(len(std)), root_mse_TD_spiking-std, root_mse_TD_spiking+std,alpha=.3)
# fig2.savefig(newpath + '/MSE_TD_spiking.png')
# pickle.dump(root_mse_TD_spiking, open('rmse_TD_id.pkl','wb'))
fs = 16
from matplotlib import cm
from matplotlib.colors import ListedColormap
top = cm.get_cmap('Greys', 128)
bottom = cm.get_cmap('magma', 128*2)
newcolors = np.vstack((top(np.linspace(0, 1, 128)),
bottom(np.linspace(0, 1, 128*2))[:128]))
newcmp = ListedColormap(newcolors, name='OrangeBlue')
fig3 = plt.figure()
w_stdp_matrix = plt.imshow(w_stdps2[:,-2,:,:].mean(axis=0),cmap=newcmp,vmin=0,vmax=2)
cbar = plt.colorbar()
cbar.ax.tick_params(labelsize=fs)
plt.yticks(ticks=[0,1,2,3],labels=[1,2,3,4],fontsize=fs)
plt.xticks(ticks=[0,1,2,3],labels=[1,2,3,4],fontsize=fs)
plt.xlabel('Future state (CA1)',fontsize=fs)
plt.ylabel('Current state (CA3)',fontsize=fs)
# fig3.savefig('fig3_.eps',bbox_inches='tight',format='eps')
# import seaborn as sn
k = np.shape(w_stdps2)[-1]
from scipy import interpolate
for i in range(k):
data = w_stdps2[:,-1,:,i].mean(axis=0)
plt.figure()
plt.bar(x = list(range(1,k+1)), height=data,color='b',alpha=.5)#,kde=True,bins=4)
x = list(range(0,k+2))
y = np.concatenate([[0],data,[0]])
f2 = interpolate.interp1d(x, y, kind='quadratic')
xnew = np.linspace(0,k+1,51)
plt.fill_between(xnew[:10*(i+2)], 0, f2(xnew)[:10*(i+2)])
# data = [[float(idx)]*int(100*val) for idx,val in enumerate(w_stdps2[:,-1,:,i].mean(axis=0))]
# sn.kdeplot(x=np.concatenate(data),fill=True,bw_adjust=2.5)
