Fix bugs in testing auto-sorting

Author: abuzarmahmood

blech_post_process.py

@@ -65,10 +65,13 @@
         description = 'Spike extraction and sorting script')
 parser.add_argument('dir_name',
                     help = 'Directory containing data files')
-parser.add_argument('--show-plot', '-p', 
-        help = 'Show waveforms while iterating (True/False)', default = 'True')
 parser.add_argument('--sort-file', '-f', help = 'CSV with sorted units',
                     default = None)
+parser.add_argument('--show-plot', 
+        help = 'Show waveforms while iterating', 
+                    action = 'store_true')
+parser.add_argument('--keep-raw', help = 'Keep raw data in hdf5 file',
+                    action = 'store_true')
 args = parser.parse_args()
 
 ############################################################
@@ -132,7 +135,11 @@
 
 
 # Delete the raw node, if it exists in the hdf5 file, to cut down on file size
-repacked_bool = post_utils.delete_raw_recordings(hdf5_name)
+if args.keep_raw == 'False':
+    repacked_bool = post_utils.delete_raw_recordings(hdf5_name)
+else:
+    repacked_bool = False
+    print('=== Keeping raw data in hdf5 file ===')
 
 # Open the hdf5 file
 if repacked_bool:
@@ -193,12 +200,12 @@
         energy,
         amplitudes,
         predictions,
-    ) = post_utils.load_data_from_disk(electrode_num, num_clusters)
+    ) = post_utils.load_data_from_disk(dir_name, electrode_num, num_clusters)
 
     # Re-show images of neurons so dumb people like Abu can make sure they
     # picked the right ones
     #if ast.literal_eval(args.show_plot):
-    if args.show_plot == 'True':
+    if args.show_plot: 
         post_utils.gen_select_cluster_plot(electrode_num, num_clusters, clusters)
 
     ############################################################

blech_units_plot.py

@@ -10,7 +10,6 @@
 
 # Import 3rd part code
 from utils import blech_waveforms_datashader
-from utils import memory_monitor as mm
 from utils.blech_utils import imp_metadata, pipeline_graph_check
 from utils.blech_process_utils import gen_isi_hist
 

pipeline_testing/prefect_pipeline.py

@@ -214,17 +214,17 @@ def select_clusters(data_dir):
     raise_error_if_error(process,stderr,stdout)
 
 @task(log_prints=True)
-def post_process(data_dir, use_file = True):
+def post_process(data_dir, use_file = True, keep_raw = False):
     script_name = 'blech_post_process.py'
-    dir_pos = data_dir
-    plot_flag = '-p ' + 'False'
     sorted_units_path = glob(os.path.join(data_dir, '*sorted_units.csv'))[0]
     if use_file:
         file_flag = '-f' + sorted_units_path
+        run_list = ["python", script_name, data_dir, file_flag]
     else:
-        file_flag = ''
-    process = Popen(["python", script_name, dir_pos, plot_flag, file_flag],
-                               stdout = PIPE, stderr = PIPE)
+        run_list = ["python", script_name, data_dir]
+    if keep_raw:
+        run_list.append('--keep-raw')
+    process = Popen(run_list, stdout = PIPE, stderr = PIPE)
     stdout, stderr = process.communicate()
     raise_error_if_error(process,stderr,stdout)
 
@@ -342,7 +342,8 @@ def run_spike_test():
     change_waveform_classifier(use_classifier=1)
     change_auto_params(use_auto=1)
     run_jetstream_bash(data_dir)
-    post_process(data_dir, use_file = False)
+    # Keep raw in the first pass so jetstream step can be rerun
+    post_process(data_dir, use_file = False, keep_raw = True)
 
     # Run with classifier disabled and manual sorting 
     change_waveform_classifier(use_classifier=0)