optimization_movie.py

import os, sys
import numpy as np
import scipy.linalg as spl
import scipy.stats as stt
import matplotlib.pyplot as pp
import pickle as pklupf.edu

os.system('clear')

obj_dir = './'
res_dir = 'res_movie/'
if not os.path.exists(res_dir):
	print 'create directory:', res_dir
	os.makedirs(res_dir)

print obj_dir, ';', res_dir

graph_format = 'png'

##################
# fMRI time series
ts_emp = np.load('rest_movie_ts.npy')

n_sub = 22 # number of subjects
n_cond = 2 # rest and movie conditions
n_session = 2 # 2 sessions per condition
N = 66 # number of ROIs
T = 300 # number of TRs of the recording

# time shifts for FC: 0, 1 and 2 TR
v_tau = np.arange(3,dtype=float)
n_tau = v_tau.size

# FC = spatiotemporal covariance of BOLD signals (average of the 2 sessions)
FC_emp = np.zeros([n_sub,n_cond,n_tau,N,N])
for i_sub in range(n_sub):
	for i_run in range(n_cond*n_session):
		ts_emp[i_sub,i_run,:,:] -= np.outer(ts_emp[i_sub,i_run,:,:].mean(1),np.ones(T)) # center the time series
for i_sub in range(n_sub):
	for i_cond in range(n_cond):
		for i_session in range(n_session):
			for i_tau in range(n_tau):
				FC_emp[i_sub,i_cond,i_tau,:,:] += np.tensordot(ts_emp[i_sub,i_cond*2+i_session,:,0:T-n_tau+1],ts_emp[i_sub,i_cond*2+i_session,:,i_tau:T-n_tau+1+i_tau],axes=(1,1)) / float((T-n_tau)*n_session)

FC_emp *= 0.5/FC_emp[:,0,0,:,:].mean()
print 'max FC value (most of the distribution should be between 0 and 1):', FC_emp.mean()

# time constant for BOLD autocovariances
slopes = np.zeros([n_sub,n_cond,N])
for i_sub in range(n_sub):
	for i_cond in range(n_cond):
		for i in range(N):
			ac_tmp = np.maximum(FC_emp[i_sub,i_cond,:,i,i],1e-10) # autocovariance for time shifts in v_tau; with lower bound to avoid negative values (cf. log)
			slopes[i_sub,i_cond,i] = np.polyfit(v_tau,np.log(ac_tmp),1)[0] # slope of autocovariance for ROI i

tau_x = -1./slopes.mean(2) # inverse of negative slope of autocovariance


#################
# structural data
SC_anat = np.load(obj_dir+'SC_anat.npy')

lim_SC = 0. # limit DTI value to determine SC (only connections with larger values are tuned)

# mask for existing connections for EC
mask_EC = np.zeros([N,N],dtype=bool) # EC weights to tune
mask_EC[SC_anat>lim_SC] = True
for i in range(N):
	mask_EC[i,i] = False # no self connection
	mask_EC[i,N-1-i] = False # additional interhemispheric connections
print 'EC density:', mask_EC.sum()/float(N*(N-1))

# diagonal mask for input noise matrix (here, no input cross-correlation)
mask_Sigma = np.eye(N,dtype=bool)


##############
# optimization

# optimzation rates (to avoid explosion of activity, Sigma is tuned quicker)
epsilon_EC = 0.0005
epsilon_Sigma = 0.05

min_val_EC = 0. # minimal value for tuned EC elements
max_val_EC = 1. # maximal value for tuned EC elements
min_val_Sigma = 0. # minimal value for tuned Sigma elements


i_tau = 1 # time shift for optimization (in TR; can be 1 or 2)
tau = v_tau[i_tau]
print 'opt with time shift', tau, 'TR'

EC_mod = np.zeros([n_sub,n_cond,N,N])
Sigma_mod = np.zeros([n_sub,n_cond,N,N])
FC0_mod = np.zeros([n_sub,n_cond,N,N])
FCtau_mod = np.zeros([n_sub,n_cond,N,N])

for i_sub in range(n_sub):
	for i_cond in range(n_cond):
		print
		print 'sub', i_sub, '; cond', i_cond

		# initial EC
		EC = np.zeros([N,N]) # initial connectivity
		Sigma = np.eye(N)  # initial noise

		# record best fit (matrix distance between model and empirical FC)
		best_dist = 1e10

		# objective FC matrices (empirical)
		FC0_obj = FC_emp[i_sub,i_cond,0,:,:]
		FCtau_obj = FC_emp[i_sub,i_cond,i_tau,:,:]

		stop_opt = False
		i_opt = 0
		while not stop_opt:

			# calculate Jacobian of dynamical system
			J = -np.eye(N)/tau_x[i_sub,:].mean() + EC

			# calculate FC0 and FCtau for model
			FC0 = spl.solve_lyapunov(J,-Sigma)
			FCtau = np.dot(FC0,spl.expm(J.T*tau))

			# matrices of model error
			Delta_FC0 = FC0_obj-FC0
			Delta_FCtau = FCtau_obj-FCtau

			# calculate error between model and empirical data for FC0 and FC_tau (matrix distance)
			dist_FC_tmp = 0.5*(np.sqrt((Delta_FC0**2).sum()/(FC0_obj**2).sum())+np.sqrt((Delta_FCtau**2).sum()/(FCtau_obj**2).sum()))

			# calculate Pearson correlation between model and empirical data for FC0 and FC_tau
			Pearson_FC_tmp = 0.5*(stt.pearsonr(FC0.reshape(-1),FC0_obj.reshape(-1))[0]+stt.pearsonr(FCtau.reshape(-1),FCtau_obj.reshape(-1))[0])

			# record best model parameters
			if dist_FC_tmp<best_dist:
				best_dist = dist_FC_tmp
				best_Pearson = Pearson_FC_tmp
				i_best = i_opt
				EC_mod_tmp = np.array(EC)
				Sigma_mod_tmp = np.array(Sigma)
				FC0_mod_tmp = np.array(FC0)
				FCtau_mod_tmp = np.array(FCtau)
			else:
				stop_opt = i_opt>100

			# Jacobian update
			Delta_J = np.dot(np.linalg.pinv(FC0),Delta_FC0+np.dot(Delta_FCtau,spl.expm(-J.T*tau))).T/tau

			# update EC (recurrent connectivity)
			EC[mask_EC] += epsilon_EC * Delta_J[mask_EC]
			EC[mask_EC] = np.clip(EC[mask_EC],min_val_EC,max_val_EC)

			# update Sigma (input variances)
			Delta_Sigma = -np.dot(J,Delta_FC0)-np.dot(Delta_FC0,J.T)
			Sigma[mask_Sigma] += epsilon_Sigma * Delta_Sigma[mask_Sigma]
			Sigma[mask_Sigma] = np.maximum(Sigma[mask_Sigma],min_val_Sigma)

			# check for stop
			if not stop_opt:
				if (i_opt)%50==0:
					print 'opt step:', i_opt
					print 'dist FC:', dist_FC_tmp, '; Pearson FC:', Pearson_FC_tmp
				i_opt += 1
			else:
				print 'stop at step', i_opt, 'with best FC dist:', best_dist, '; best FC Pearson:', best_Pearson

		EC_mod[i_sub,i_cond,:,:] = EC_mod_tmp
		Sigma_mod[i_sub,i_cond,:,:] = Sigma_mod_tmp
		FC0_mod[i_sub,i_cond,:,:] = FC0_mod_tmp
		FCtau_mod[i_sub,i_cond,:,:] = FCtau_mod_tmp


# save results
np.save(res_dir+'FC_emp.npy',FC_emp) # empirical spatiotemporal FC
np.save(res_dir+'mask_EC.npy',mask_EC) # mask of optimized connections
np.save(res_dir+'mask_Sigma.npy',mask_Sigma) # mask of optimized Sigma elements

np.save(res_dir+'EC_mod.npy',EC_mod) # estimated EC matrices
np.save(res_dir+'Sigma_mod.npy',Sigma_mod) # estimated Sigma matrices
np.save(res_dir+'FC0_mod.npy',FC0_mod) # model FC0
np.save(res_dir+'FCtau_mod.npy',FCtau_mod) # model FCtau (tau = 1 or 2 TR)

# various compiled results (if needed)
if False:
	np.save(res_dir+'mean_FC0_rest.npy',FC_emp[:,0,0,:,:].mean(0))
	np.save(res_dir+'mean_FC1_rest.npy',FC_emp[:,0,1,:,:].mean(0))
	np.save(res_dir+'mean_EC_rest.npy',EC_mod[:,0,:,:].mean(0))
	np.save(res_dir+'mean_Sigma_rest.npy',Sigma_mod[:,0,:,:].mean(0))
	np.save(res_dir+'mean_FC0_movie.npy',FC_emp[:,1,0,:,:].mean(0))
	np.save(res_dir+'mean_FC1_movie.npy',FC_emp[:,1,1,:,:].mean(0))
	np.save(res_dir+'mean_EC_movie.npy',EC_mod[:,1,:,:].mean(0))
	np.save(res_dir+'mean_Sigma_movie.npy',Sigma_mod[:,1,:,:].mean(0))