Merge pull request #259 from katzlabbrandeis/214-improved-drift-plot1d-visualization

214 improved drift plot1d visualization

Author: abuzarmahmood

blech_autosort.sh

@@ -13,10 +13,6 @@ echo === Quality Assurance ===
 bash blech_run_QA.sh $DIR &&
 echo === Units Plot ===
 python blech_units_plot.py $DIR &&
-echo === Make PSTHs ===
-python blech_make_psth.py $DIR &&
-echo === Palatability Identity Setup ===
-python blech_palatability_identity_setup.py $DIR &&
-echo === Overlay PSTHs ===
-python blech_overlay_psth.py $DIR &&
+echo === Get unit characteristics ===
+python blech_units_characteristics.py $DIR && 
 echo === Done ===

blech_clust_post.sh

@@ -1,14 +1,11 @@
 DIR=$1
 BLECH_DIR=$HOME/Desktop/blech_clust
-echo Running Units Plot
-python $BLECH_DIR/blech_units_plot.py $DIR;
-echo Running Make Arrays
-python $BLECH_DIR/blech_make_arrays.py $DIR;
-echo Running Quality Assurance 
-bash $BLECH_DIR/blech_run_QA.sh $DIR;
-echo Running Make PSTHs
-python $BLECH_DIR/blech_make_psth.py $DIR;
-echo Running Palatability Identity Setup
-python $BLECH_DIR/blech_palatability_identity_setup.py $DIR;
-echo Running Overlay PSTH
-python $BLECH_DIR/blech_overlay_psth.py $DIR;
+echo === Make Arrays ===
+python $BLECH_DIR/blech_make_arrays.py $DIR &&
+echo === Quality Assurance === 
+bash $BLECH_DIR/blech_run_QA.sh $DIR &&
+echo === Units Plot ===
+python $BLECH_DIR/blech_units_plot.py $DIR &&
+echo === Get unit characteristics ===
+python $BLECH_DIR/blech_units_characteristics.py $DIR && 
+echo === Done ===

blech_post_process.py

@@ -135,7 +135,7 @@
 
 
 # Delete the raw node, if it exists in the hdf5 file, to cut down on file size
-if args.keep_raw == 'False':
+if args.keep_raw == False:
     repacked_bool = post_utils.delete_raw_recordings(hdf5_name)
 else:
     repacked_bool = False
@@ -415,6 +415,9 @@
         'autosort_outputs'
     )
 
+    # Create output directory if needed
+    if not os.path.exists(autosort_output_dir):
+        os.makedirs(autosort_output_dir)
 
     # Since this needs classifier output to run, check if it exists
     clf_list = glob('./spike_waveforms/electrode*/clf_prob.npy')

pipeline_testing/prefect_pipeline.py

@@ -352,9 +352,9 @@ def run_spike_test():
     select_clusters(data_dir)
     post_process(data_dir)
 
+    make_arrays(data_dir)
     quality_assurance(data_dir)
     units_plot(data_dir)
-    make_arrays(data_dir)
     units_characteristics(data_dir)
 
 @flow(log_prints=True)

utils/blech_post_process_utils.py

@@ -1449,9 +1449,6 @@ def auto_process_electrode(
             new_clust_names,
         )
 
-        # Create output directory if needed
-        if not os.path.exists(autosort_output_dir):
-            os.makedirs(autosort_output_dir)
 
         fig.savefig(
             os.path.join(

utils/ephys_data/ephys_data.py

@@ -517,9 +517,21 @@ def calc_firing_func(data):
     def get_firing_rates(self):
         """
         Converts spikes to firing rates
+
+        Requires:
+            - spikes
+            - firing_rate_params
+
+        Generates:
+            - firing_list : list of firing rates for each taste
+                - each element is a 3D array of shape (n_trials, n_neurons, n_timepoints)
+            - firing_array : 4D array of firing rates
+            - normalized_firing : 4D array of normalized firing rates
+            - all_firing_array : 3D array of all firing rates
+            - all_normalized_firing : 3D array of all normalized firing rates
         """
 
-        if self.spikes is None:
+        if 'spikes' not in dir(self):
             # raise Exception('Run method "get_spikes" first')
             print('No spikes found, getting spikes ...')
             self.get_spikes()

utils/qa_utils/drift_check.py

@@ -21,12 +21,15 @@
 import pingouin as pg
 import seaborn as sns
 import glob
+from sklearn.decomposition import PCA
+from umap import UMAP
 # Get script path
 script_path = os.path.realpath(__file__)
 script_dir_path = os.path.dirname(script_path)
 blech_path = os.path.dirname(os.path.dirname(script_dir_path))
 sys.path.append(blech_path)
 from utils.blech_utils import imp_metadata, pipeline_graph_check
+from utils.ephys_data import ephys_data
 
 def get_spike_trains(hf5_path):
     """
@@ -350,11 +353,49 @@ def array_to_df(array, dim_names):
         print('Post-stimulus limits: ' + str(stim_t) + ' to ' + str(stim_t+trial_duration) + ' ms', file=f)
         print('Trial Bin Count: ' + str(n_trial_bins), file=f)
         print('alpha: ' + str(alpha), file=f)
-        print('\n', file=f)
+        #print('\n', file=f)
         print(out_rows, file=f)
         print('\n', file=f)
         print('=== End Post-stimulus Drift Warning ===', file=f)
         print('\n', file=f)
 
+############################################################
+# Perform PCA on firing rates across trials
+############################################################
+dat = ephys_data.ephys_data(dir_name)
+dat.get_firing_rates()
+# each element is a 3D array of shape (n_trials, n_neurons, n_timepoints)
+firing_list = dat.firing_list
+# Normalize for each neuron
+n_neurons = firing_list[0].shape[1]
+norm_firing_list = []
+for i in range(len(firing_list)):
+    this_firing = firing_list[i]
+    norm_firing = np.zeros_like(this_firing)
+    for j in range(n_neurons):
+        norm_firing[:,j,:] = zscore(this_firing[:,j,:], axis=None)
+    norm_firing_list.append(norm_firing)
+
+# shape: (n_trials, n_neurons * n_timepoints)
+long_firing_list = [x.reshape(x.shape[0],-1) for x in norm_firing_list]
+
+# Perform PCA on long_firing_list
+pca_firing_list = [PCA(n_components=1, whiten=True).fit_transform(x) for x in long_firing_list]
+umap_firing_list = [UMAP(n_components=1).fit_transform(x) for x in long_firing_list]
+umap_zscore = [zscore(x, axis=None) for x in umap_firing_list]
+
+# Plot PCA and UMAP results
+fig, ax = plt.subplots(2, 1, figsize=(5, 5), sharex=True) 
+for i in range(len(pca_firing_list)):
+    ax[0].plot(pca_firing_list[i], alpha=0.7)
+    ax[1].plot(umap_zscore[i], alpha=0.7)
+    ax[0].set_title('PCA')
+    ax[1].set_title('UMAP')
+ax[-1].set_xlabel('Trial num')
+fig.suptitle('PCA and UMAP of Firing Rates \n' + basename)
+plt.tight_layout()
+plt.savefig(os.path.join(output_dir, 'pca_umap_firing_rates.png'))
+plt.close()
+
 # Write successful execution to log
 this_pipeline_check.write_to_log(script_path, 'completed')