decode_head_data.py

# coding: utf-8

# # Examples of all decoders (except Kalman Filter)
# 
# In this example notebook, we:
# 1. Import the necessary packages
# 2. Load a data file (spike trains and outputs we are predicting)
# 3. Preprocess the data for use in all decoders
# 4. Run all decoders and print the goodness of fit
# 5. Plot example decoded outputs
# 
# See "Examples_kf_decoder" for a Kalman filter example. <br>
# Because the Kalman filter utilizes different preprocessing, we don't include an example here. to keep this notebook more understandable

# ## 1. Import Packages
# 
# Below, we import both standard packages, and functions from the accompanying .py files

# In[2]:

#Import standard packages
import numpy as np
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import sys,os
from scipy import io
from scipy import stats,signal
import pickle
import h5py
#Import function to get the covariate matrix that includes spike history from previous bins
from preprocessing_funcs import get_spikes_with_history

#Import metrics
from metrics import get_R2
from metrics import get_rho

#Import decoder functions
from decoders import WienerCascadeDecoder
from decoders import WienerFilterDecoder
from decoders import DenseNNDecoder
from decoders import SimpleRNNDecoder
from decoders import GRUDecoder
from decoders import LSTMDecoder
#from decoders import XGBoostDecoder
from decoders import SVRDecoder

import h5py
from sklearn.externals import joblib
from sklearn.preprocessing import StandardScaler
from sklearn import linear_model
from sklearn.decomposition import PCA

import seaborn as sns
sns.set_style('white')

from scipy import stats,signal
#from transforms3d import euler

def filter(ephys,freq_range,filt_order = 4,filt_type='bandpass',fs=10.):
	
    # design Elliptic filter:

    [b,a] = signal.butter(filt_order,[freq/fs for freq in freq_range],btype=filt_type)
    
    filtered_trace = signal.filtfilt(b,a,ephys,axis=0)
    return filtered_trace

# def get_mean_lfp(neural_data):

# 	## get the 384 matrix down to 64 by averaging LFP power by band within each tetrode:
# 	x = np.arange(384)
# 	within_tet_band_starts = []
# 	for i in range(6):
# 	    within_tet_band_starts.append(x[i:24:6])
# 	within_tet_band_starts = np.asarray(within_tet_band_starts)
	
# 	tet_starts = x[0::24]

# 	for tet in range(16):
#     	for band in range(6):    
#         	#print x[within_tet_band_starts[band] + tet_starts[tet]]
#         	neural_data = 



def load_data(head_file,neural_data_file):

	head_data = h5py.File(head_file,'r')
	
	#jerk_power = head_data['jerk_power'][:]

	dx = head_data['dx'][:]
	dy = head_data['dy'][:]
	dz = head_data['dz'][:]
	oz = head_data['oz'][:]
	oy = head_data['oy'][:]
	ox = head_data['ox'][:]
	ay = head_data['ay'][:]
	ax = head_data['ax'][:]
	az = head_data['az'][:]

	xyz = np.sqrt(ax**2 + ay**2 + az**2)

	# theta = np.rad2deg(np.arctan(ax/ay))

	# dx_neg = np.empty(dx.shape)
	# dx_pos = np.empty(dx.shape)
	
	# dx_neg[np.where(dx < 0)[0]] = dx[np.where(dx < 0)[0]]

	# dx_pos[np.where(dx > 0)[0]] = dx[np.where(dx > 0)[0]]


	# left = dx_neg**2
	# right = dx_pos**2

	head_data.close()
	#xy_acc = head_data['xy_acc']
	#theta = head_data['theta']
	#time = head_data['time']


	


	#y = np.vstack([xyz,oz,dx,dy,dz,ax,ay,az,ox,oy]).T
	#y_name = ['xyz','oz','dx','dy','dz','ax','ay','az','ox','oy']

	y = np.vstack([xyz]).T
	y_name = ['xyz']



	
	#y = np.unwrap(np.unwrap(np.deg2rad(y)))

	#y = np.vstack([ox,oy,dx,dy,ax,ay,az]).T
	#y_name = ['ox','oy','dx','dy','ax','ay','az']


	


	neural_data_file = h5py.File(neural_data_file,'r') 

	### determine if it's spikes or LFPs:
	print(neural_data_file.keys())

	if neural_data_file.keys()[0].find('spikes') == 1:
		print('Loading Spikes')
		neural_data = neural_data_file['sorted_spikes'][:]
		
	elif neural_data_file.keys()[0].find('lfp_power') == 1:
		print('Loading LFPs')
		neural_data = neural_data_file['lfp_power'][:]
	else:
		print('Loading something?')
		key = neural_data_file.keys()[0]
		neural_data = neural_data_file[key][:]
		

	neural_data_file.close()

	### make sure neural data is in right shape (samples x channels):
	if neural_data.shape[0] < neural_data.shape[1]:
		neural_data = neural_data.T


	## limit the size of the neural data?
	
	#q = []

	#for i in range(1000000): 
	#	q.append(euler.euler2quat(np.deg2rad(ox[i]),np.deg2rad(oz[i]),np.deg2rad(oy[i]))) ### yaw,pitch,roll

	#y = np.asarray(q)  #np.vstack([q]).T
	#y_name = ['quaternion1','quaternion2','quaternion3','quaternion4']

	#dx = np.gradient(filter(signal.medfilt(np.unwrap(np.deg2rad(ox)),[21]),[1],filt_type='lowpass',fs=100.))
	#dy = np.gradient(filter(signal.medfilt(oy,[21]),[1],filt_type='lowpass',fs=100.))
	#dz = np.gradient(filter(signal.medfilt(oz,[21]),[1],filt_type='lowpass',fs=100.))

	#y = np.vstack([left,right]).T
	#y_name = ['left','right']

	# y = np.vstack([jerk_power]).T

	# y_name = ['jerk_power']

	cuttoff = int(1e6)
	if neural_data.shape[0] > cuttoff:
		print('Truncating  data to %d points' % cuttoff)
		neural_data = neural_data[0:cuttoff,:]
		y = y[0:cuttoff,:]
	

	#spikes_file = h5py.File('all_sorted_spikes.hdf5','r') 

	#spikes = spikes_file['sorted_spikes'][:]

	#spikes_file.close()	

	print('Shape of head data = ', y.shape)
	print('Shape of neural_data = ', neural_data.shape)

	# for i in range(len(y_name)):
	# #	y[:,i] = signal.medfilt(y[:,i],[21])
	# 	y[:,i] = filter(y[:,i],[1.],filt_type='lowpass',fs=100.)



	#idx = 1000 #int(y.shape[0]/2)
	#print 'max idx = ', idx
	#return y[0:idx,:], neural_data[0:idx,:],y_name
	return y, neural_data,y_name
	
	
def sample_dx_uniformly(derivative,num_points=10000):
    ################### sample the dx distribution evenly: ####################
    derivative = np.squeeze(derivative)
    bins = 10000
    hist,edges = np.histogram(derivative,bins=bins,normed=True)
    
    bins_where_values_from = np.searchsorted(edges,derivative)
    
    bin_weights = 1/(hist/sum(hist))
    
    inv_weights = bin_weights[bins_where_values_from-1]
    
    dx_idx = np.arange(0,len(derivative),1)

    sampled_dx_idx = np.random.choice(dx_idx,size=num_points,replace=False,p =inv_weights/sum(inv_weights)  )

    sampled_dx = np.random.choice(derivative,size=num_points,replace=False,p =inv_weights/sum(inv_weights)  )


    f,axarr = plt.subplots(2,dpi=600,sharex=True)

    axarr[0].hist(derivative,bins=200)
    axarr[0].set_ylabel('d_yaw \n original')

    axarr[1].hist(sampled_dx,bins=200)
    axarr[1].set_ylabel('d_yaw \n resampled')

    sns.despine(left=True,bottom=True)

    f.savefig('resampled_original_histograms.pdf')

    return sampled_dx_idx	

def preprocess(jerk,neural_data):
	# ## 3. Preprocess Data

	# ### 3A. User Inputs
	# The user can define what time period to use spikes from (with respect to the output).

	# In[25]:

	bins_before=10 #How many bins of neural data prior to the output are used for decoding
	bins_current=1 #Whether to use concurrent time bin of neural data
	bins_after=10 #How many bins of neural data after the output are used for decoding


	# ### 3B. Format Covariates

	# #### Format Input Covariates

	# In[26]:

	# Format for recurrent neural networks (SimpleRNN, GRU, LSTM)
	# Function to get the covariate matrix that includes spike history from previous bins
	X=get_spikes_with_history(neural_data,bins_before,bins_after,bins_current)

	# Format for Wiener Filter, Wiener Cascade, XGBoost, and Dense Neural Network
	#Put in "flat" format, so each "neuron / time" is a single feature
	X_flat=X.reshape(X.shape[0],(X.shape[1]*X.shape[2]))



	# #### Format Output Covariates

	# In[79]:

	#Set decoding output
	#y=jerk_power
	y=jerk

	# for i in [0]: # range(y.shape[1])
	# 	resamp_idx = sample_dx_uniformly(y[:,i],num_points=10000)
	# 	resamp_idx = np.sort(resamp_idx)

	# 	y = y[resamp_idx,:]
	# 	X_flat = X_flat[resamp_idx,:]

	# print '###################### getting non-zero values ######################'
	# non_zeros = np.where(abs(y) >= 0.5 )[0]
	# print y.shape
	# y = y[non_zeros,:]
	# print y.shape
	# print X_flat.shape
	# X_flat = X_flat[non_zeros,:]

	# ### 3C. Split into training / testing / validation sets
	# Note that hyperparameters should be determined using a separate validation set. 
	# Then, the goodness of fit should be be tested on a testing set (separate from the training and validation sets).

	# #### User Options

	# In[32]:

	#Set what part of data should be part of the training/testing/validation sets
	training_range=[0, 0.5]
	testing_range=[0.7, 0.85]
	valid_range=[0.5, 1]


	# #### Split Data

	# In[81]:

	num_examples=X_flat.shape[0]

	#Note that each range has a buffer of"bins_before" bins at the beginning, and "bins_after" bins at the end
	#This makes it so that the different sets don't include overlapping neural data
	training_set=np.arange(np.int(np.round(training_range[0]*num_examples))+bins_before,np.int(np.round(training_range[1]*num_examples))-bins_after)
	testing_set=np.arange(np.int(np.round(testing_range[0]*num_examples))+bins_before,np.int(np.round(testing_range[1]*num_examples))-bins_after)
	valid_set=np.arange(np.int(np.round(valid_range[0]*num_examples))+bins_before,np.int(np.round(valid_range[1]*num_examples))-bins_after)

	#Get training data
	X_train=X[training_set,:,:]
	X_flat_train=X_flat[training_set,:]


	y_train=y[training_set,:]

	#Get testing data
	X_test=X[testing_set,:,:]
	X_flat_test=X_flat[testing_set,:]

	y_test=y[testing_set,:]

	#Get validation data
	X_valid=X[valid_set,:,:]
	X_flat_valid=X_flat[valid_set,:]


	y_valid=y[valid_set,:]


	# ### 3D. Process Covariates
	# We normalize (z_score) the inputs and zero-center the outputs.
	# Parameters for z-scoring (mean/std.) should be determined on the training set only, and then these z-scoring parameters are also used on the testing and validation sets.

	# In[ ]:




	# In[82]:

	#Z-score "X" inputs. 
	# X_train_mean=np.nanmean(X_train,axis=0)
	# X_train_std=np.nanstd(X_train,axis=0)
	# X_train=(X_train-X_train_mean)/X_train_std
	# X_test=(X_test-X_train_mean)/X_train_std
	# X_valid=(X_valid-X_train_mean)/X_train_std


	#Z-score "X_flat" inputs. 
	X_flat_train_mean=np.nanmean(X_flat_train,axis=0)
	X_flat_train_std=np.nanstd(X_flat_train,axis=0)
	X_flat_train=(X_flat_train-X_flat_train_mean)/X_flat_train_std
	X_flat_test=(X_flat_test-X_flat_train_mean)/X_flat_train_std
	X_flat_valid=(X_flat_valid-X_flat_train_mean)/X_flat_train_std

	#Z-score  outputs
	y_train_mean=np.nanmean(y_train,axis=0)

	y_train_std=np.nanstd(y_train,axis=0)


	#### 
	y_train=(y_train-y_train_mean) #/y_train_std
	y_test=(y_test-y_train_mean) #/y_train_std
	y_valid=(y_valid-y_train_mean) #/y_train_std

	


	return X_flat_train,X_flat_valid,X_train,X_valid,y_train,y_valid,y_train_mean,y_train_std

# ## 4. Run Decoders

def BayesianRidge_model(X_train,X_valid,y_train,y_test,y_name, y_train_mean,y_train_std):
 
	model_name = 'BayesianRidge'
	print('head items to fit are: ', y_name)
		# In[ ]:
	for head_item in range(len(y_name)):

		y_train_item = y_train[:,head_item]
		#y_train_item = np.reshape(y_train_item,[y_train.shape[0],1])

		y_test_item = y_test[:,head_item]
		#y_test_item = np.reshape(y_test_item,[y_test_item.shape[0],1])
		print('********************************** Fitting %s on %s Data **********************************' % (model_name,y_name[head_item]))
		#Declare model
		model = linear_model.BayesianRidge(compute_score=True)

		

		#Fit model
		model.fit(X_train,y_train_item)

		#Get predictions
		y_valid_predicted=model.predict(X_valid)


		training_prediction=model.predict(X_train)

		R2s_training=get_R2(y_train_item,training_prediction)
		print('R2 on training set = ', R2s_training)

		#Get metric of fit
		R2s=get_R2(y_test_item,y_valid_predicted)
		print('R2s:', R2s)
		print('saving prediction ...')
		np.savez(y_name[head_item] + '_%s_ypredicted.npz' % model_name,y_test=y_test_item,y_prediction=y_valid_predicted,
			y_train_=y_train_item,training_prediction=training_prediction,
			y_train_mean=y_train_mean[head_item],y_train_std=y_train_std[head_item])
		#print 'saving model ...'
		joblib.dump(model, y_name[head_item] + '_%s.pkl' % model_name) 
		print('plotting results...')
		plot_results(y_test_item,y_valid_predicted,y_name[head_item],R2s,model_name=model_name)

	return model



def ridgeCV_model(X_train,X_valid,y_train,y_test,y_name, y_train_mean,y_train_std):
 
	print('head items to fit are: ', y_name)
		# In[ ]:
	for head_item in range(len(y_name)):

		y_train_item = y_train[:,head_item]
		y_train_item = np.reshape(y_train_item,[y_train.shape[0],1])

		y_test_item = y_test[:,head_item]
		y_test_item = np.reshape(y_test_item,[y_test_item.shape[0],1])
		print('********************************** Fitting RidgeCV on %s Data **********************************' % y_name[head_item])
		#Declare model
		model = linear_model.RidgeCV(alphas=[0.1, 1.0, 10.0],normalize=True,fit_intercept=True)

		

		#Fit model
		model.fit(X_train,y_train_item)

		#Get predictions
		y_valid_predicted=model.predict(X_valid)


		training_prediction=model.predict(X_train)

		R2s_training=get_R2(y_train_item,training_prediction)
		print('R2 on training set = ', R2s_training)

		#Get metric of fit
		R2s=get_R2(y_test_item,y_valid_predicted)
		print('R2s:', R2s)
		print('saving prediction ...')
		np.savez(y_name[head_item] + '_RidgeCV_ypredicted.npz',y_test=y_test_item,y_prediction=y_valid_predicted,
			y_train_=y_train_item,training_prediction=training_prediction,
			y_train_mean=y_train_mean[head_item],y_train_std=y_train_std[head_item])
		#print 'saving model ...'
		joblib.dump(model, y_name[head_item] + '_Ridge.pkl') 
		print('plotting results...')
		plot_results(y_test_item,y_valid_predicted,y_name[head_item],R2s,model_name='RidgeCV')

	return model



def Wiener(X_flat_train,X_flat_valid,y_train,y_valid):
	#Declare model
	model_wf=WienerFilterDecoder()

	#Fit model
	model_wf.fit(X_flat_train,y_train)

	#Get predictions
	y_valid_predicted_wf=model_wf.predict(X_flat_valid)

	#Get metric of fit
	R2s_wf=get_R2(y_valid,y_valid_predicted_wf)
	print('R2s:', R2s_wf)

	#plot_results(y_valid,y_valid_predicted_wf)

	return model_wf

def WienerCascade(X_train,X_valid,y_train,y_test,y_name, y_train_mean,y_train_std):
	# ### 4B. Wiener Cascade (Linear Nonlinear Model)

	#print 'head items to fit are: ', y_name
		# In[ ]:
	for head_item in range(len(y_name)):

		y_train_item = y_train[:,head_item]
		y_train_item = np.reshape(y_train_item,[y_train.shape[0],1])

		y_test_item = y_test[:,head_item]
		y_test_item = np.reshape(y_test_item,[y_test_item.shape[0],1])
		#print '********************************** Fitting WienerCascade on %s Data **********************************' % y_name[head_item]
		#Declare model
		model = WienerCascadeDecoder(degree=3)

		

		#Fit model
		model.fit(X_train,y_train_item)

		#Get predictions
		y_valid_predicted=model.predict(X_valid)


		training_prediction=model.predict(X_train)

		R2s_training=get_R2(y_train_item,training_prediction)
		print('R2 on training set = ', R2s_training)

		#Get metric of fit
		R2s=get_R2(y_test_item,y_valid_predicted)
		print('R2s:', R2s)
		#print 'saving prediction ...'
		np.savez(y_name[head_item] + '_WienerCascade_ypredicted.npz',y_test=y_test_item,y_prediction=y_valid_predicted,
			y_train_=y_train_item,training_prediction=training_prediction,
			y_train_mean=y_train_mean[head_item],y_train_std=y_train_std[head_item])
		#print 'saving model ...'
		joblib.dump(model, y_name[head_item] + '_WienerCascade.pkl') 
		#print 'plotting results...'
		plot_results(y_test_item,y_valid_predicted,y_name[head_item],R2s,model_name='WienerCascade')

	return model

def XGBoost():
	# ### 4C. XGBoost (Extreme Gradient Boosting)

	# In[ ]:

	#Declare model
	model_xgb=XGBoostDecoder(max_depth=3,num_round=200,eta=0.3,gpu=-1) 

	#Fit model
	model_xgb.fit(X_flat_train, y_train)

	#Get predictions
	y_valid_predicted_xgb=model_xgb.predict(X_flat_valid)

	#Get metric of fit
	R2s_xgb=get_R2(y_valid,y_valid_predicted_xgb)
	print('R2s:', R2s_xgb)

def SVR(X_flat_train,X_flat_valid,y_train,y_valid,y_name):
	# ### 4D. SVR (Support Vector Regression)

	# In[40]:

	#The SVR works much better when the y values are normalized, so we first z-score the y values
	#They have previously been zero-centered, so we will just divide by the stdev (of the training set)
	y_train_std=np.nanstd(y_train,axis=0)
	y_zscore_train=y_train/y_train_std
	#y_zscore_test=y_test/y_train_std
	y_zscore_valid=y_valid/y_train_std

	#Declare model
	model_svr=SVRDecoder(C=.1, max_iter=10000,gamma=1e-5)
	
	#Fit model
	

	for head_item in range(len(y_name)):
		### fit one at a time and save/plot the results 
		#print '########### Fitting SVR on %s data ###########' % y_name[head_item]

		y_zscore_train_item = y_zscore_train[:,head_item]
		y_zscore_train_item = np.reshape(y_zscore_train_item,[y_zscore_train.shape[0],1])
		#print 'shape of y_zscore_train_item = ', y_zscore_train_item.shape

		y_zscore_valid_item = y_zscore_valid[:,head_item]
		y_zscore_valid_item = np.reshape(y_zscore_valid_item,[y_zscore_valid_item.shape[0],1])

		model_svr.fit(X_flat_train,y_zscore_train_item)

		#Get predictions
		y_zscore_valid_predicted_svr=model_svr.predict(X_flat_valid)

		#Get metric of fit
		R2s_svr=get_R2(y_zscore_valid_item,y_zscore_valid_predicted_svr)
		print(y_name[head_item], 'R2:', R2s_svr)

		np.savez(y_name[head_item] + '_svr_ypredicted.npz',y_zscore_valid=y_zscore_valid_item,y_zscore_valid_predicted_svr=y_zscore_valid_predicted_svr)


		plot_results(y_zscore_valid_item,y_zscore_valid_predicted_svr,y_name[head_item],R2s_svr)

def DNN(X_train,X_valid,y_train,y_test,y_name):
	# ### 4E. Dense Neural Network

	
	print('head items to fit are: ', y_name)
	# In[ ]:
	for head_item in range(len(y_name)):

		y_train_item = y_train[:,head_item]
		y_train_item = np.reshape(y_train_item,[y_train.shape[0],1])

		y_test_item = y_test[:,head_item]
		y_test_item = np.reshape(y_test_item,[y_test_item.shape[0],1])
		print('********************************** Fitting DNN on %s Data **********************************' % y_name[head_item])
		#Declare model
		model_dnn=DenseNNDecoder(units=[128,64,32],num_epochs=15)

		#Fit model
		model_dnn.fit(X_train,y_train_item)

		#Get predictions
		y_valid_predicted=model_dnn.predict(X_valid)

		#Get metric of fit
		R2s=get_R2(y_test_item,y_valid_predicted)
		print('R2s:', R2s)
		#print 'saving prediction ...'
		np.savez(y_name[head_item] + '_DNN_ypredicted.npz',y_test=y_test_item,y_prediction=y_valid_predicted)
		#print 'saving model ...'
		#joblib.dump(model_dnn, y_name[head_item] + '_LSTM.pkl') 
		print('plotting results...')
		plot_results(y_test_item,y_valid_predicted,y_name[head_item],R2s,model_name='DNN')

	return model_dnn

def RNN(X_train,y_train,X_valid,y_valid,y_name):
	model_name = 'RNN'
	# ### 4F. Simple RNN
	#print '############################# RUNNING RNN #############################'
	# In[ ]:

	#Declare model
	model_rnn=SimpleRNNDecoder(units=400,dropout=0,num_epochs=100)

	for head_item in range(len(y_name)):
		### fit one at a time and save/plot the results 
		#print '########### Fitting RNN on %s data ###########' % y_name[head_item]

		y_train_item = y_train[:,head_item]
		y_train_item = np.reshape(y_train_item,[y_train.shape[0],1])
		#print 'shape of y_train_item = ', y_train_item.shape

		y_valid_item = y_valid[:,head_item]
		y_valid_item = np.reshape(y_valid_item,[y_valid_item.shape[0],1])

		model_rnn.fit(X_train,y_train_item)

		#Get predictions
		y_valid_predicted_rnn=model_rnn.predict(X_valid)

		#Get metric of fit
		R2s_rnn=get_R2(y_valid_item,y_valid_predicted_rnn)
		print(y_name[head_item], 'R2:', R2s_rnn)

		np.savez(y_name[head_item] + '_rnn_ypredicted.npz',y_valid=y_valid_item,y_valid_predicted_rnn=y_valid_predicted_rnn)


		plot_results(y_valid_item,y_valid_predicted_rnn,y_name[head_item],R2s_rnn,model_name)


def GRU():
	# ### 4G. GRU (Gated Recurrent Unit)

	# In[ ]:

	#Declare model
	model_gru=GRUDecoder(units=400,dropout=0,num_epochs=5)

	#Fit model
	model_gru.fit(X_train,y_train)

	#Get predictions
	y_valid_predicted_gru=model_gru.predict(X_valid)

	#Get metric of fit
	R2s_gru=get_R2(y_valid,y_valid_predicted_gru)
	print('R2s:', R2s_gru)

def run_LSTM(X_train,X_valid,y_train,y_test,y_name, y_train_mean,y_train_std):
	# ### 4H. LSTM (Long Short Term Memory)
	#print 'head items to fit are: ', y_name
	# In[ ]:
	for head_item in range(len(y_name)):

		y_train_item = y_train[:,head_item]
		y_train_item = np.reshape(y_train_item,[y_train.shape[0],1])

		y_test_item = y_test[:,head_item]
		y_test_item = np.reshape(y_test_item,[y_test_item.shape[0],1])
		#print '********************************** Fitting Deep Net on %s Data **********************************' % y_name[head_item]
		#Declare model
		model_lstm=LSTMDecoder(dropout=0.25,num_epochs=5)

		model_lstm.get_means(y_train_mean,y_train_std) ### for un-zscoring during loss calculation ??? 

		#Fit model
		model_lstm.fit(X_train,y_train_item)

		#Get predictions
		y_valid_predicted_lstm=model_lstm.predict(X_valid)


		training_prediction=model_lstm.predict(X_train)

		R2s_training=get_R2(y_train_item,training_prediction)
		#print 'R2 on training set = ', R2s_training

		#Get metric of fit
		R2s_lstm=get_R2(y_test_item,y_valid_predicted_lstm)
		print('R2s:', R2s_lstm)
		#print 'saving prediction ...'




		np.savez(y_name[head_item] + '_LSTM_ypredicted.npz',y_test=y_test_item,y_prediction=y_valid_predicted_lstm,
			y_train_=y_train_item,training_prediction=training_prediction,
			y_train_mean=y_train_mean[head_item],y_train_std=y_train_std[head_item])
		#print 'saving model ...'
		#joblib.dump(model_lstm, y_name[head_item] + '_LSTM.pkl') 
		#print 'plotting results...'
		plot_results(y_test_item,y_valid_predicted_lstm,y_name[head_item],R2s_lstm,model_name='LSTM')

	return model_lstm


def plot_results(y_valid,y_valid_predicted,y_name,R2s,params='_',model_name='SVR'):
    print('y_valid shape = ',y_valid.shape)
    print('y_valid_predicted shape = ', y_valid_predicted.shape)
    print(stats.pearsonr(y_valid,y_valid_predicted))
    f, axarr = plt.subplots(2,dpi=600)
    axarr[0].set_title(model_name +' Model of %s. R^2 = %f. r = %f ' % (y_name,R2s,stats.pearsonr(y_valid,y_valid_predicted)[0] ))


    axarr[0].plot(y_valid,linewidth=0.1)
    axarr[0].set_ylabel('Head Data')

    axarr[0].plot(y_valid_predicted,linewidth=0.1,color='red')

    axarr[1].set_title(params)
    axarr[1].scatter(y_valid,y_valid_predicted,alpha=0.05,marker='o')
    #axarr[1].set_title('R2 = ' + str(R2s))
    axarr[1].set_xlabel('Actual')
    axarr[1].set_ylabel('Predicted')
    axarr[1].axis('equal')

    sns.despine(left=True,bottom=True)


    save_folder = './plots/' + model_name + '/'
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)

    f.savefig(save_folder + model_name + '_%s.pdf' % y_name)




# In[ ]:
if __name__ == "__main__":

	model_type = sys.argv[1] ## wiener or lstm

	head_file = sys.argv[2]
	neural_data_file = sys.argv[3]

	head_data,neural_data,y_name = load_data(head_file,neural_data_file)

	#X_flat_train,X_flat_valid,X_train,X_valid,y_train,y_valid, y_train_mean,y_train_std = preprocess(head_data,neural_data)

	
	if model_type == 'lstm':
		data_model = run_LSTM(X_train,X_valid,y_train,y_valid,y_name, y_train_mean,y_train_std)
	elif model_type == 'wiener':
		data_model = Wiener(X_flat_train,X_flat_valid,y_train,y_valid)
	elif model_type == 'svr':
		data_model = SVR(X_flat_train,X_flat_valid,y_train,y_valid,y_name)
	elif model_type == 'rnn':
		RNN(X_train,y_train,X_valid,y_valid,y_name,y_name)
	elif model_type == 'dnn':
		data_model = DNN(X_flat_train,X_flat_valid,y_train,y_valid,y_name)
	

	elif model_type == 'ridge':
		X_flat_train,X_flat_valid,X_train,X_valid,y_train,y_valid, y_train_mean,y_train_std = preprocess(head_data,neural_data)
		data_model = ridgeCV_model(X_flat_train,X_flat_valid,y_train,y_valid,y_name, y_train_mean,y_train_std)
	
	elif model_type == 'ridge_band_time_search':

		band_names = ['delta','theta','alpha','beta','low_gamma','high_gamma']

		## neural data = samples x channels
		for i in range(6):

			X_flat_train,X_flat_valid,X_train,X_valid,y_train,y_valid, y_train_mean,y_train_std = preprocess(head_data,neural_data[:,i::6])

			

			tmp_name = [name + '_' + band_names[i] for name in y_name]

			data_model = ridgeCV_model(X_flat_train,X_flat_valid,y_train,y_valid,tmp_name, y_train_mean,y_train_std)
	




	elif model_type == 'WienerCascade':
		data_model = WienerCascade(X_flat_train,X_flat_valid,y_train,y_valid,y_name, y_train_mean,y_train_std)
	elif model_type == 'BayesianRidge':
		data_model = BayesianRidge_model(X_flat_train,X_flat_valid,y_train,y_valid,y_name, y_train_mean,y_train_std)
	

	#with open('model_' + model_type + '_rawjerk','wb') as f:
	#	pickle.dump(data_model,f)