nn_train_rutils_spikes.py

"""

"""

import os
import copy
import h5py
import yaml
import pickle

import numpy as np
import tensorflow as tf
from tensorflow.python import pywrap_tensorflow
from kinematics_decoding import set_kin_dimensions

from sklearn.metrics import explained_variance_score
from path_utils import MODELS_DIR

os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

gpu_options = tf.GPUOptions(allow_growth=True)

# gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.33) #per_process_gpu_memory_fraction=0.9)

def interp_single(model_act, mvt_duration):
    a_interp = np.interp(x = np.linspace(0, model_act.shape[0], mvt_duration),
                                 xp = np.linspace(0, model_act.shape[0], model_act.shape[0]),
                                 fp = model_act)
    return a_interp

def interpolate_neuron(pred_eval, trial_duration):
    """

    :param layer_act:
    :param ori_spike_data:
    :return:
    """
    from scipy import interpolate

    interp_act = []

    trial_interp_act = map(interp_single,list(pred_eval),[trial_duration]*pred_eval.shape[0])
    interp_act = np.asarray(list(trial_interp_act))

    return interp_act

def replace_nan_inf(data):
    # Replace nan and inf
    ind_nan = tf.where(tf.is_nan(data))
    ind_inf = tf.where(tf.is_inf(data))
    upd_nan = tf.tile([tf.constant(0,data.dtype)],[tf.shape(ind_nan)[0]])
    upd_inf = tf.tile([tf.constant(0,data.dtype)],[tf.shape(ind_inf)[0]])
    data = tf.scatter_nd_update(data, ind_inf, upd_inf)
    data = tf.scatter_nd_update(data, ind_nan, upd_nan)
    return data

def R_squared_masked_tf(y, y_pred, mask):
    '''
    R_squared computes the coefficient of determination.
    It is a measure of how well the observed outcomes are replicated by the model.
    '''
    n_time = tf.cast(tf.shape(y)[2], y.dtype)
    
    ## Compute time length for batch index
    N1 = tf.cast(tf.reduce_mean(tf.reduce_sum(mask,axis=2),axis=1), y.dtype)
    
    y_mask = y*mask
    y_pred_mask = y_pred*mask
    
    ### Compute the residual variance
    diff_array = tf.subtract(y_mask, y_pred_mask)
    mean_diff = tf.reduce_sum(diff_array,axis=2)/tf.expand_dims(N1,axis=1)
    diff_res_mask = (diff_array - tf.expand_dims(mean_diff,axis=2))*mask
    residual = tf.reduce_sum(tf.square(diff_res_mask), axis=2)/(tf.expand_dims(N1,axis=1))

    ### Compute the total variance
    mean_array = tf.reduce_sum(y_mask,axis=2)/tf.expand_dims(N1,axis=1)
    diff_total_mask = (y_mask - tf.expand_dims(mean_array,axis=2))*mask
    total = tf.reduce_sum(tf.square(diff_total_mask), axis=2)/(tf.expand_dims(N1,axis=1))
    
    ### Take only residual to minimize
    res = tf.div(residual, total)

    r2 = tf.subtract(1.0, res, name='explained_variance') 
    return residual, r2

def compute_mask(model, dataset):
    # stride = model.t_stride
    if (type(model).__name__ == 'ConvRModel') or (type(model).__name__ == 'ConvRModel_new'):
        stride = model.t_stride
        nlayers = model.nlayers
    else:
        stride = 1
        nlayers = 1

    nouttime = 400
    if isinstance(stride, (list)):
        flag_hom = False
    else:
        flag_hom = True

    for i in range(nlayers):
        if flag_hom:
            nouttime = int(np.ceil(nouttime/stride))
        else:
            nouttime = int(np.ceil(nouttime/stride[i]))

    if dataset.dataset_type == 'train':
        targets = dataset.train_targets
    else:
        targets = dataset.test_targets
    mask = np.zeros((targets.shape[0],nouttime,targets.shape[1]))


    for ind_batch in range(targets.shape[0]):
        start_idx = dataset.latents[ind_batch]['start_idx']
        end_idx = dataset.latents[ind_batch]['end_idx']
        for i in range(nlayers):
            if flag_hom:
                start_idx = int(np.ceil(start_idx/stride))
                end_idx = int(np.ceil(end_idx/stride))
            else:
                start_idx = int(np.ceil(start_idx/stride[i]))
                end_idx = int(np.ceil(end_idx/stride[i]))
        if start_idx == end_idx:    end_idx += 1
        if start_idx == mask.shape[1]:
            start_idx = start_idx - 1
            end_idx = end_idx - 1
        mask[ind_batch,start_idx:end_idx,:] = 1
    dataset.mask = mask.swapaxes(1,2)
    return

class RDataset():
    """Defines a dataset object for regression, with simple routines to generate batches."""

    def __init__(self, path_to_data=None, path_to_latent=None, data=None, dataset_type='train', 
            key='spindle_info', fraction=None, target_key='endeffector_coords', n_out_time=400):
        """Set up the `Dataset` object.

        Arguments
        ---------
        path_to_data : str, absolute location of the dataset file.
        data: dict, optionally provide dataset directly as {'data': np.array, 'targets': np.array}
        dataset_type : {'train', 'test'} str, type of data that will be used along with the model.
        key : {'endeffector_coords', 'joint_coords', 'muscle_coords', 'spindle_firing'} str
        target_key: {'endeffector_coords', 'joint_coords'} str, specifies targets to regress

        """
        self.path_to_data = path_to_data
        self.path_to_latent = path_to_latent
        self.dataset_type = dataset_type
        self.key = key
        self.target_key = target_key
        self.train_data = self.train_targets = None
        self.val_data = self.val_targets = None
        self.test_data = self.test_targets = None
        self.n_out_time = n_out_time
        self.mask = None
        self.make_data(data)

        # For when I want to use only a fraction of the dataset to train!
        if fraction is not None:
            random_idx = np.random.permutation(self.train_data.shape[0])
            subset_num = int(fraction * random_idx.size)
            self.train_data = self.train_data[random_idx[:subset_num]]
            self.train_targets = self.train_targets[random_idx[:subset_num]]

    def make_data(self, mydata):
        """Load train/val or test splits into the `Dataset` instance.

        Returns
        -------
        if dataset_type == 'train' : loads train and val splits.
        if dataset_type == 'test' : loads the test split.

        """
        # Load dataset and compute mask
        if self.path_to_data is not None:
            datafile = h5py.File(self.path_to_data, 'r')
            datafile_latents = pickle.load(open(self.path_to_latent, 'rb'))

            if self.dataset_type == 'train':
                self.train_data = datafile[self.key]
                self.train_targets = datafile[self.target_key]
                self.train_data_mean = datafile['train_data_mean']

                self.latents = datafile_latents['all_latents_train']
            elif self.dataset_type == 'test':
                self.test_data = datafile[self.key]
                self.test_targets = datafile[self.target_key]

                self.latents = datafile_latents['all_latents_test']
        else: 
            data = mydata['data']
            labels = mydata['targets']

    def next_trainbatch(self, batch_size, step=0):
        """Returns a new batch of training data.

        Arguments
        ---------
        batch_size : int, size of training batch.
        step : int, step index in the epoch.

        Returns
        -------
        2-tuple of batch of training data and correspondig targets.

        """
        steps_per_epoch = self.train_data.shape[0] // batch_size
        if step == 0:
            exc = self.train_data.shape[0] % batch_size
            total_len = batch_size*steps_per_epoch
            poss_position = np.arange(0,total_len,batch_size)

            rand_idx = np.random.randint(1,total_len)
            poss_position[rand_idx:] = np.concatenate((poss_position[(rand_idx+1):] -batch_size + exc,np.array([self.train_data.shape[0]]) - batch_size),axis=0)
            self.shuffle_idx = np.random.permutation(poss_position)

        if step % steps_per_epoch == 0:
            rand_idx = np.random.randint(1,total_len)
            poss_position[rand_idx:] = np.concatenate((poss_position[(rand_idx+1):] -batch_size + exc,np.array([self.train_data.shape[0]]) - batch_size),axis=0)
            self.shuffle_idx = np.random.permutation(poss_position)
        mybatch_data = self.train_data[self.shuffle_idx[step]:self.shuffle_idx[step]+batch_size].astype('float32') 
        mybatch_targets = self.train_targets[self.shuffle_idx[step]:self.shuffle_idx[step]+batch_size] #batch_size*step:batch_size*(step+1)

        mybatch_mask = self.mask[self.shuffle_idx[step]:self.shuffle_idx[step]+batch_size]

        mybatch_targets = set_kin_dimensions(mybatch_targets, self.n_out_time)

        return (mybatch_data, mybatch_targets, mybatch_mask)

    def next_valbatch(self, batch_size, type='test', step=0):
        """Returns a new batch of validation or test data.

        Arguments
        ---------
        type : {'val', 'test'} str, type of data to return.

        """
        if type == 'train':
            mybatch_data = self.train_data[batch_size*step:batch_size*(step+1)]
            mybatch_targets = self.train_targets[batch_size*step:batch_size*(step+1)]
            mybatch_mask = self.mask[batch_size*step:batch_size*(step+1)]
        elif type == 'val':
            mybatch_data = self.val_data[batch_size*step:batch_size*(step+1)]
            mybatch_targets = self.val_targets[batch_size*step:batch_size*(step+1)]
            mybatch_mask = self.mask[batch_size*step:batch_size*(step+1)]
        elif type == 'test':
            mybatch_data = self.test_data[batch_size*step:batch_size*(step+1)]
            mybatch_targets = self.test_targets[batch_size*step:batch_size*(step+1)]
            mybatch_mask = self.mask[batch_size*step:batch_size*(step+1)]
            
        mybatch_targets = set_kin_dimensions(mybatch_targets, self.n_out_time) 

        return (mybatch_data, mybatch_targets, mybatch_mask)
    
    def set_outtime(self, outtime):
        self.n_out_time = outtime


class RTrainer:
    """Trains a `RModel` object with the given `Dataset` object."""

    def __init__(self, model=None, dataset=None, test_dataset=None, global_step=None):
        """Set up the `Trainer`.

        Arguments
        ---------
        model : an instance of `ConvRModel` or `RecurrentRModel` to be trained.
        dataset : an instance of `Dataset`, containing the train/val data splits.

        """
        self.model = model
        self.dataset = dataset
        self.test_dataset = test_dataset
        self.log_dir = model.model_path
        self.global_step = 0 if global_step == None else global_step
        self.session = None
        self.graph = None
        self.best_loss = 1e10
        self.validation_accuracy = 0

    def get_tensors_in_checkpoint_file(self, file_name,all_tensors=True,tensor_name=None):
        varlist=[]
        var_value =[]
        reader = pywrap_tensorflow.NewCheckpointReader(file_name)
        if all_tensors:
            var_to_shape_map = reader.get_variable_to_shape_map()
            for key in sorted(var_to_shape_map):
                varlist.append(key)
                var_value.append(reader.get_tensor(key))
        else:
            varlist.append(tensor_name)
            var_value.append(reader.get_tensor(tensor_name))
        return (varlist, var_value)
    
    def build_tensors_in_checkpoint_file(self, loaded_tensors):
        full_var_list = list()
        # Loop all loaded tensors
        for i, tensor_name in enumerate(loaded_tensors[0]):
            # Extract tensor
            if not 'Classifier' in tensor_name:
                try:
                    tensor_aux = self.graph.get_tensor_by_name(tensor_name+":0")
                    full_var_list.append(tensor_aux)
                except:
                    print('Not found: '+tensor_name)
        return full_var_list

    def build_graph(self, **kwargs):
        """Build training graph using the `Model`s predict function and setting up an optimizer."""
        
        _, ninputs, ntime, _ = self.dataset.train_data.shape
        _, ncoords, _ = self.dataset.train_targets.shape   #ncoords was 3

        if (type(self.model).__name__ == 'ConvRModel'):
            layers = self.model.nlayers
            stride = self.model.t_stride
            nouttime = 400
            for i in range(layers):
                nouttime = int(np.ceil(nouttime/stride))
            self.dataset.set_outtime(nouttime)
            self.test_dataset.set_outtime(nouttime)
        elif (type(self.model).__name__ == 'ConvRModel_new'):
            layers = self.model.nlayers
            stride = self.model.t_stride
            nouttime = 400
            for i in range(layers):
                nouttime = int(np.ceil(nouttime/stride[i]))
            self.dataset.set_outtime(nouttime)
            self.test_dataset.set_outtime(nouttime)
        else:
            nouttime = 400

        with tf.Graph().as_default() as self.graph:
            tf.set_random_seed(self.model.seed)
            # Placeholders
            self.learning_rate = tf.placeholder(tf.float32)
            self.X = tf.placeholder(tf.float32, shape=[self.batch_size, ninputs, ntime, 2], name="X")
            self.y = tf.placeholder(tf.float32, shape=[self.batch_size, nouttime, ncoords], name="y") # ALWAYS endeffectors for now.
            self.mask = tf.placeholder(tf.float32, shape=[self.batch_size, ncoords, nouttime], name="mask")

            # Set up optimizer, compute and apply gradients
            self.scores, _ = self.model.predict(self.X, is_training=True)

            ## Swap neuron and time for EV computation
            self.pred = tf.transpose(self.scores, [0, 2, 1])
            self.targ = tf.transpose(self.y, [0, 2, 1])

            self.res,self.r2 = R_squared_masked_tf(y=self.targ, y_pred=self.pred, mask=self.mask)

            rsquared_loss = tf.reduce_mean(self.res)
            self.train_op = tf.train.AdamOptimizer(self.learning_rate).minimize(rsquared_loss)
            
            # # Calculate metrics
            self.scores_eval, _ = self.model.predict(self.X, is_training=False)
            
            ## Swap neuron and time for EV computation
            self.pred_eval = tf.transpose(self.scores_eval, [0, 2, 1])
            self.targ_eval = tf.transpose(self.y, [0, 2, 1])

            res,self.r2_test = R_squared_masked_tf(y=self.targ_eval, y_pred=self.pred_eval, mask=self.mask)

            cond = tf.is_nan(self.r2_test)
            self.r2_masked_test = tf.where(cond, tf.zeros(tf.shape(self.r2_test)),self.r2_test)
            self.accuracy_test_op = tf.reduce_mean(tf.boolean_mask(self.r2_masked_test, tf.is_finite(self.r2_masked_test)), name="EV_test")

            cond = tf.is_nan(self.res)
            self.res = tf.where(cond, tf.zeros(tf.shape(self.res)),self.res)
            self.train_loss_op = tf.reduce_mean(self.res, name="train_loss")
            
            cond = tf.is_nan(self.r2)
            self.r2_masked = tf.where(cond, tf.zeros(tf.shape(self.r2)),self.r2)
            self.accuracy_op = tf.reduce_mean(tf.boolean_mask(self.r2_masked, tf.is_finite(self.r2_masked)), name="EV_train")

            tf.summary.scalar('Train_Loss', rsquared_loss)
            tf.summary.scalar('EV_train', self.accuracy_op)
            self.train_summary_op = tf.summary.merge_all()
            
            self.init = tf.global_variables_initializer()
            self.saver = tf.train.Saver()
            
            if not len(kwargs) == 0:
                varlist = self.get_tensors_in_checkpoint_file(file_name=kwargs['log_dir'],all_tensors=True,tensor_name=None)
                variables = self.build_tensors_in_checkpoint_file(varlist)
                
                self.loader = tf.train.Saver(variables)
            print('Built graph!')

    def load(self):
        self.saver.restore(self.session, os.path.join(self.log_dir, 'model.ckpt'))
        
    def save(self):
        self.saver.save(self.session, os.path.join(self.log_dir, 'model.ckpt'))

    def save_step(self,step):
        self.saver.save(self.session, os.path.join(self.log_dir, 'model_' + str(step) + '.ckpt'))
    
    def load_step(self, step, **kwargs):

        if len(kwargs) == 0:
            log_dir = self.log_dir
        else:
            log_dir = kwargs['log_dir']

        self.saver.restore(self.session, os.path.join(log_dir, 'model_' + str(step) + '.ckpt'))

    def make_model_name(self):
        if (type(self.model).__name__ == 'ConvRModel'):
            # Make model name
            if self.model.arch_type == 'spatial_temporal':
                kernels = ('-'.join(str(i) for i in self.model.n_skernels)) + '_' + ('-'.join(str(i) for i in self.model.n_tkernels))
            elif self.model.arch_type == 'temporal_spatial':
                kernels = ('-'.join(str(i) for i in self.model.n_tkernels)) + '_' + ('-'.join(str(i) for i in self.model.n_skernels))
            else:
                kernels = ('-'.join(str(i) for i in self.model.n_skernels))

            parts_name = [self.model.arch_type, str(self.model.nlayers), kernels,
                        ''.join(str(i) for i in [self.model.s_kernelsize, self.model.s_stride, self.model.t_kernelsize, self.model.t_stride])]

            # Create model directory
            name = '_'.join(parts_name)
        elif (type(self.model).__name__ == 'ConvRModel_new'):

            max_tstride = self.model.t_stride.count(2)**2
            max_sstride = self.model.s_stride.count(2)**2
            # Make model name
            if self.model.arch_type == 'spatial_temporal':
                kernels = ('-'.join(str(i) for i in self.model.n_skernels)) + '_' + ('-'.join(str(i) for i in self.model.n_tkernels))
            elif self.model.arch_type == 'temporal_spatial':
                kernels = ('-'.join(str(i) for i in self.model.n_tkernels)) + '_' + ('-'.join(str(i) for i in self.model.n_skernels))
            else:
                kernels = ('-'.join(str(i) for i in self.model.n_skernels))

            parts_name = [self.model.arch_type, str(self.model.nlayers), kernels,
                        ''.join(str(i) for i in [self.model.s_kernelsize, max_sstride, self.model.t_kernelsize, max_tstride])]

            # Create model directory
            name = '_'.join(parts_name)
        elif (type(self.model).__name__ == 'RecurrentRModel'):
            # Make model name
            units = ('-'.join(str(i) for i in self.model.nppfilters))
            parts_name = [self.model.rec_blocktype, str(self.model.npplayers), units, str(self.model.n_recunits)]

            # Create model directory
            name = '_'.join(parts_name)
            if self.model.seed is not None: name += '_' + str(self.model.seed)
        elif (type(self.model).__name__ == 'RecurrentRModel_new'):
            max_sstride = self.s_stride.count(2)**2
        
            # Make model name
            units = ('-'.join(str(i) for i in nppfilters))
            parts_name = [rec_blocktype, str(n_reclayers), str(npplayers), units, str(n_recunits),
                        ''.join(str(i) for i in [s_kernelsize, max_sstride])]

            # Create model directory
            name = '_'.join(parts_name)
            if self.model.seed is not None: name += '_' + str(self.model.seed)
        return name

    def train(self,
            num_epochs=10,
            learning_rate= 0.0005, #0.0005,
            batch_size=256,
            val_steps=200,
            early_stopping_epochs=10,
            retrain=False,
            retrain_same_init=False,
            old_exp_dir = None,
            normalize=False,
            verbose=True, 
            save_rand=False,
            window=5,
            latency=0):
        """Train the `Model` object.

        Arguments
        ---------
        num_epochs : int, number of epochs to train for.
        learning_rate : float, learning rate for Adam Optimizer.
        batch_size : int, size of batch to train on.
        val_steps : int, number of batches after which to perform validation.
        early_stopping_steps : int, number of steps for early stopping criterion.
        retrain : bool, train already existing model vs not.
        normalize : bool, whether to normalize training data or not.
        verbose : bool, print progress on screen.

        """
        steps_per_epoch = self.dataset.train_data.shape[0] // batch_size
        max_iter = num_epochs * steps_per_epoch
        early_stopping_steps = early_stopping_epochs * steps_per_epoch
        self.batch_size = batch_size

        ## Adjust mask
        compute_mask(self.model, self.dataset)
        compute_mask(self.model, self.test_dataset)

        if normalize:
            self.train_data_mean = float(self.dataset.train_data_mean)
            self.train_data_std = float(np.std(self.dataset.train_data))
        else:
            self.train_data_mean = self.train_data_std = 0
        train_params = {'train_mean': self.train_data_mean,
                        'train_std': self.train_data_std}
        test_params = {'test_accuracy': 0}

        if retrain_same_init:
            old_exp_dir = os.path.join(MODELS_DIR,'experiment_' + str(old_exp_dir))
            name = self.make_model_name()
            log_dir = os.path.join(old_exp_dir, name)
            log_dir = os.path.join(log_dir, 'model_0.ckpt')
            self.build_graph(log_dir = log_dir)
        else:
            self.build_graph()
            
        self.session = tf.Session(graph=self.graph, config=tf.ConfigProto(gpu_options=gpu_options))

        self.session.run(self.init)
        if retrain:
            self.load()
        
        if retrain_same_init:
            self.loader.restore(self.session, log_dir)


        if save_rand:
            self.save_step(0)
            self.model.is_training = False
            make_config_file(self.model, train_params, test_params, step = self.global_step) #, save_rand)

        # Create summaries
        self.train_summary = tf.summary.FileWriter(
            os.path.join(self.model.model_path, 'train'), graph=self.graph, flush_secs=30)
        self.val_summary = tf.summary.FileWriter(os.path.join(self.model.model_path, 'val'))

        # Define checkpoints
        try:
            if self.model.rec_blocktype == 'lstm':
                if batch_size == 512:
                    check1, check2, check3 = 5000, 10000, 15000
                else:
                    check1, check2, check3 = int(max_iter*0.1), int(max_iter*0.25), int(max_iter*0.35)
        except:
            if (self.model.arch_type == 'spatial_temporal') or (self.model.arch_type == 'temporal_spatial') or (self.model.arch_type == 'spatiotemporal'):
                if batch_size == 512:
                    check1, check2, check3 = 1000, 2500, 5000
                else:
                    check1, check2, check3 = int(max_iter*0.1), int(max_iter*0.25), int(max_iter*0.35)

        not_improved = 0
        end_training = 0

        for self.global_step in range(max_iter):
            
            # Training step
            batch_X, batch_y, batch_mask = self.dataset.next_trainbatch(
                batch_size, self.global_step % steps_per_epoch)
            feed_dict = {self.X: batch_X - self.train_data_mean,
                        self.y: batch_y, 
                        self.mask: batch_mask,
                        self.learning_rate: learning_rate}
            loss_train, acc_train, _ = self.session.run([self.train_loss_op, self.accuracy_op, self.train_op], feed_dict)

            # Validate/save periodically
            if self.global_step % val_steps == 0:
                # Summarize, print progress

                self.save_summary(feed_dict)
                if verbose:
                    print('Step : %4d, Validation EV : %.2f' % (self.global_step, acc_train))
                    print('best_loss:', self.best_loss, 'loss:', loss_train)

                if loss_train < self.best_loss:
                    self.best_loss = loss_train
                    self.validation_accuracy = acc_train
                    self.save()
                    train_params['train_loss'] = float(self.best_loss)
                    train_params['accuracy'] = float(acc_train)
                    not_improved = 0
                else:
                    not_improved += 1

                if not_improved >= early_stopping_steps:  #steps_per_epoch: #
                    learning_rate /= 4
                    print('lr:', learning_rate)
                    not_improved = 0
                    end_training += 1
                    self.load()

                if end_training == 2: #early_stopping_epochs:  #2
                    if self.global_step < 20*steps_per_epoch:
                        end_training = 1
                        not_improved = 0
                    else:
                        break

            if (self.global_step == check1) or (self.global_step == check2) or (self.global_step == check3):
                self.save_step(self.global_step)
                make_config_file(self.model, train_params, test_params, step = self.global_step)

        self.model.is_training = False

        ### Test the network
        self.save()
        test_accuracy, train_accuracy, result_dict = self.test_model()
        train_params['accuracy'] = float(train_accuracy)
        test_params = {'test_accuracy': float(test_accuracy)}

        make_config_file(self.model, train_params, test_params)
        self.session.close()
        return result_dict

    def save_summary(self, feed_dict):
        """Create and save summaries for training and validation."""
        train_summary = self.session.run(self.train_summary_op, feed_dict)
        self.train_summary.add_summary(train_summary, self.global_step)

    def eval(self):
        """Evaluate validation performance.
        
        Returns
        -------
        validation_loss : float, loss on the entire validation data
        validation_accuracy : float, accuracy on the validation data
        
        """
        num_iter = self.dataset.val_data.shape[0] // self.batch_size
        acc_val = np.zeros(num_iter)
        loss_val = np.zeros(num_iter)
        for i in range(num_iter):
            batch_X, batch_y = self.dataset.next_valbatch(self.batch_size, step=i)
            feed_dict = {self.X: batch_X - self.train_data_mean, self.y: batch_y} #res,r2
            loss_val[i], acc_val[i] = self.session.run([self.val_loss_op, self.accuracy_op], feed_dict)
        return loss_val.mean(), acc_val.mean()

    def test_model(self):
        """Evaluate test performance.
        
        Returns
        -------
        test_accuracy : float, accuracy on the test data
        
        """

        self.batch_size = 1
        result_dict = {}
        self.build_graph()
            
        self.session = tf.Session(graph=self.graph, config=tf.ConfigProto(gpu_options=gpu_options))

        self.session.run(self.init)

        num_iter_test = self.test_dataset.test_data.shape[0] // self.batch_size
        num_iter_train = self.dataset.train_data.shape[0] // self.batch_size
        self.load()
        
        ## TEST SET
        acc_test = []
        pred_y = []
        targ_y = []
        pred_y_notinterp = []
        targ_y_notinterp = []
        test_ev_list = []
        for i in range(num_iter_test):
            batch_X, batch_y, batch_mask = self.test_dataset.next_valbatch(self.batch_size, 'test', step=i)

            start_idx_targ = self.test_dataset.latents[i]['start_idx']
            end_idx_targ = self.test_dataset.latents[i]['end_idx']

            feed_dict = {self.X: batch_X - self.train_data_mean, 
                        self.y: batch_y,
                        self.mask: batch_mask}
            acc,pred_tmp,targ_tmp,r2_test = self.session.run([self.accuracy_test_op,self.pred_eval,self.targ_eval, self.r2_test], feed_dict)
            acc_test.append(acc)
            test_ev_list.append(r2_test)
            indexes = np.where(batch_mask[0,0,:] == 1)[0]
            start_idx = indexes[0]
            end_idx = indexes[-1]

            if start_idx == end_idx:    end_idx += 1

            pred_tmp_interp = interpolate_neuron(pred_tmp[0,:,start_idx:end_idx],end_idx_targ-start_idx_targ)

            pred_y.append(pred_tmp_interp)
            targ_y.append(self.test_dataset.test_targets[i,:,start_idx_targ:end_idx_targ])
            pred_y_notinterp.append(pred_tmp[0,:,start_idx:end_idx])
            targ_y_notinterp.append(targ_tmp[0,:,start_idx:end_idx])

        pred_y = np.hstack(pred_y).T
        targ_y = np.hstack(targ_y).T

        pred_y_notinterp = np.hstack(pred_y_notinterp).T
        targ_y_notinterp = np.hstack(targ_y_notinterp).T

        ## Compute EV score for concatenated test trials
        ev_score_test = explained_variance_score(targ_y, pred_y, multioutput='raw_values')
        print('EV scores test:', np.mean(ev_score_test))
        ev_score_test_notinterp = explained_variance_score(targ_y_notinterp, pred_y_notinterp, multioutput='raw_values')
        
        result_dict['ev_test'] = list(ev_score_test.astype(float))
        result_dict['ev_test_list'] = list(test_ev_list)
        result_dict['test_rates'] = list(targ_y.astype(float))
        result_dict['test_preds'] = list(pred_y.astype(float))
        result_dict['ev_test_notinterp'] = list(ev_score_test_notinterp.astype(float))
        result_dict['test_rates_notinterp'] = list(targ_y_notinterp.astype(float))
        result_dict['test_preds_notinterp'] = list(pred_y_notinterp.astype(float))

        ## TRAINING SET
        acc_train = []
        pred_y = []
        targ_y = []
        pred_y_notinterp = []
        targ_y_notinterp = []
        test_ev_list = []
        for i in range(num_iter_train):
            batch_X, batch_y, batch_mask = self.dataset.next_valbatch(self.batch_size, 'train', step=i)

            start_idx_targ = self.dataset.latents[i]['start_idx']
            end_idx_targ = self.dataset.latents[i]['end_idx']

            feed_dict = {self.X: batch_X - self.train_data_mean, 
                        self.y: batch_y,
                        self.mask: batch_mask}
            acc,pred_tmp,targ_tmp,r2_test = self.session.run([self.accuracy_test_op,self.pred_eval,self.targ_eval, self.r2_test], feed_dict)
            acc_train.append(acc)
            test_ev_list.append(r2_test)
            indexes = np.where(batch_mask[0,0,:] == 1)[0]
            start_idx = indexes[0]
            end_idx = indexes[-1]

            if start_idx == end_idx:    end_idx += 1

            pred_tmp_interp = interpolate_neuron(pred_tmp[0,:,start_idx:end_idx],end_idx_targ-start_idx_targ)
            pred_y.append(pred_tmp_interp)
            targ_y.append(self.dataset.train_targets[i,:,start_idx_targ:end_idx_targ])
            pred_y_notinterp.append(pred_tmp[0,:,start_idx:end_idx])
            targ_y_notinterp.append(targ_tmp[0,:,start_idx:end_idx])

        pred_y = np.hstack(pred_y).T
        targ_y = np.hstack(targ_y).T

        pred_y_notinterp = np.hstack(pred_y_notinterp).T
        targ_y_notinterp = np.hstack(targ_y_notinterp).T

        ## Compute EV score for concatenated test trials
        ev_score_test = explained_variance_score(targ_y, pred_y, multioutput='raw_values')
        print('EV scores train:', np.mean(ev_score_test))
        ev_score_test_notinterp = explained_variance_score(targ_y_notinterp, pred_y_notinterp, multioutput='raw_values')
        
        result_dict['ev_train'] = list(ev_score_test.astype(float))
        result_dict['ev_train_list'] = list(test_ev_list)

        result_dict['ev_train_notinterp'] = list(ev_score_test_notinterp.astype(float))

        # print('Validation Accuracy :',acc_val, 'loss:', loss_val)
        return np.mean(acc_test), np.mean(acc_train), result_dict


def evaluate_model(model, dataset, batch_size=200):
    """Evaluation routine for trained models.

    Arguments
    ---------
    model : the `Conv`, `Affine` or `Recurrent` model to be evaluated. The test data is 
        assumed to be defined within the model.dataset object.
    dataset : the `Dataset` object on which the model is to be evaluated.

    Returns
    -------
    accuracy : float, Classification accuracy of the model on the given dataset.

    """
    # Data handling
    nsamples, ninputs, ntime, _ = dataset.test_data.shape
    _, ncoords, _ = dataset.test_targets.shape
    
    if type(model).__name__ == 'ConvRModel':
            layers = model.nlayers
            stride = model.t_stride
            nouttime = 400
            for i in range(layers):
                nouttime = int(np.ceil(nouttime/stride))
            dataset.set_outtime(nouttime)
    elif (type(model).__name__ == 'ConvRModel_new'):
        layers = model.nlayers
        stride = model.t_stride
        nouttime = 400
        for i in range(layers):
            nouttime = int(np.ceil(nouttime/stride[i]))
        dataset.set_outtime(nouttime)
    else:
        nouttime = 400
    
    num_steps = nsamples // batch_size

    # Retrieve training mean, if data was normalized
    path_to_config_file = os.path.join(model.model_path, 'config.yaml')
    with open(path_to_config_file, 'r') as myfile:
        model_config = yaml.load(myfile)
    train_mean = model_config['train_mean']

    mygraph = tf.Graph()
    with mygraph.as_default():
        # Declare placeholders for input data and labels
        X = tf.placeholder(tf.float32, shape=[batch_size, ninputs, ntime, 2], name="X")
        y = tf.placeholder(tf.float32, shape=[batch_size, nouttime, ncoords], name="y")
        mask = tf.placeholder(tf.float32, shape=[batch_size, ncoords, nouttime], name="y")

        # Compute scores and accuracy
        # # Calculate metrics
        scores_eval, _ = model.predict(X, is_training=False)
        
        ## Swap neuron and time for EV computation
        pred = tf.transpose(scores_eval, [0, 2, 1])
        targ = tf.transpose(y, [0, 2, 1])

        res,r2 = R_squared_masked_tf(y=targ, y_pred=pred, mask=dataset.mask)
        accuracy = tf.reduce_mean(tf.boolean_mask(r2, tf.is_finite(r2)), name="EV_val")

        # Test the `model`!
        restorer = tf.train.Saver()
        myconfig = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True, gpu_options=gpu_options)
        with tf.Session(config=myconfig) as sess:
            ckpt_filepath = os.path.join(model.model_path, 'model.ckpt')
            restorer.restore(sess, ckpt_filepath)
            
            test_accuracy = []
            for step in range(num_steps):         
                batch_x, batch_y, batch_mask = dataset.next_valbatch(batch_size, 'test', step)
                acc = sess.run([accuracy], feed_dict={X: batch_x - train_mean, y: batch_y, mask:batch_mask})
                test_accuracy.append(acc)

    return np.mean(test_accuracy)


# Auxiliary Functions

def train_val_split(data, labels):
    num_train = int(0.9*data.shape[0])
    train_data, train_labels = data[:num_train], labels[:num_train]
    val_data, val_labels = data[num_train:], labels[num_train:]

    return (train_data, train_labels, val_data, val_labels)


def make_config_file(model, train_params, test_params, **kwargs):
    """Make a configuration file for the given model, created after training.

    Given a `ConvModel`, `AffineModel` or `RecurrentModel` instance, generates a 
    yaml file to save the configuration of the model.

    """
    mydict = copy.copy(model.__dict__)
    # Convert to python native types for better readability
    for (key, value) in mydict.items():
        if isinstance(value, np.generic):
            mydict[key] = float(value)
        elif isinstance(value, list) or isinstance(value, np.ndarray):
            mydict[key] = [int(item) for item in value]

    # Save yaml file in the model's path
    if len(kwargs) == 0:
        path_to_yaml_file = os.path.join(model.model_path, 'config.yaml')
    else:
        path_to_yaml_file = os.path.join(model.model_path, 'config_' + str(kwargs['step']) + '.yaml')

    # Save yaml file in the model's path
    # path_to_yaml_file = os.path.join(model.model_path, 'config.yaml')
    with open(path_to_yaml_file, 'w') as myfile:
        yaml.dump(mydict, myfile, default_flow_style=False)
        yaml.dump(train_params, myfile, default_flow_style=False)
        yaml.dump(test_params, myfile, default_flow_style=False)

    return