dPCA_nans_removed.py

import numpy as np

n_stim                = 3
n_responses           = 2
trial_size            = 10
n_trials_per_stimtype = 50
n_neurons             = 4
single_trial_data = np.zeros([n_trials_per_stimtype, n_neurons, n_stim, n_responses, trial_size])
single_trial_data.fill(np.NaN)

for stim in [0, 1, 2]:
    for trial in range(n_trials_per_stimtype):
        decision = np.random.randint(2)
        single_trial_data[trial, :, stim, decision, :] = np.random.uniform(size=[n_neurons, trial_size])

trial_average_data = np.nanmean(single_trial_data,0)
trial_average_data -= np.mean(trial_average_data.reshape((n_neurons,-1)),1)[:,None,None,None]

# Replace NaNs with the mean across all trials, stimuli, and responses
single_trial_noNaN = np.zeros_like(single_trial_data)

for neuron in range(n_neurons):
    for trial in range(n_trials_per_stimtype):
        for stimulus in range(n_stim):
            for response in range(n_responses):

                # if there is missing data, replace with the average trial of that neuron
                # across all trials, stimuli, and responses
                if np.isnan(single_trial_data[trial, neuron, stimulus, response, :]).any():
                    print(trial, neuron, stimulus, response)
                    single_trial_noNaN[trial, neuron, stimulus, response, :] =\
                        np.nanmean(single_trial_data[:, neuron, :, :, :], axis=(0, 1, 2))

                # otherwise, insert the existing data
                else:
                    single_trial_noNaN[trial, neuron, stimulus, response, :] =\
                        single_trial_data[trial, neuron, stimulus, response, :]



from dPCA import dPCA
dpca = dPCA.dPCA(labels='srt',regularizer='auto', n_components=1)
dpca.protect = ['t']
Z = dpca.fit_transform(trial_average_data, trialX=single_trial_noNaN)