Merge pull request #41 from SpikeInterface/msc_dumpable

Unified file creation/deletion

Author: alejoe91

spikecomparison/groundtruthstudy.py

@@ -157,13 +157,11 @@ def aggregate_dataframes(self, copy_into_folder=True, **karg_thresh):
         perfs = self.aggregate_performance_by_units()
 
         dataframes['perf_by_units'] = perfs.reset_index()
-        # dataframes['perf_pooled_with_average'] = perfs.reset_index().groupby(['rec_name', 'sorter_name']).mean().reset_index()
         dataframes['count_units'] = self.aggregate_count_units(**karg_thresh).reset_index()
 
         if copy_into_folder:
             tables_folder = self.study_folder / 'tables'
-            if not os.path.exists(tables_folder):
-                os.makedirs(str(tables_folder))
+            tables_folder.mkdir(parents=True, exist_ok=True)
 
             for name, df in dataframes.items():
                 df.to_csv(str(tables_folder / (name + '.csv')), sep='\t', index=False)
@@ -190,11 +188,11 @@ def get_units_snr(self, rec_name=None, **snr_kargs):
         rec_name = self._check_rec_name(rec_name)
 
         metrics_folder = self.study_folder / 'metrics'
-        if not (os.path.exists(metrics_folder)):
-            os.makedirs(str(metrics_folder))
+        metrics_folder.mkdir(parents=True, exist_ok=True)
+
         filename = metrics_folder / ('SNR ' + rec_name + '.txt')
 
-        if os.path.exists(filename):
+        if filename.is_file():
             snr = pd.read_csv(filename, sep='\t', index_col=None)
             snr = snr.set_index('gt_unit_id')
         else:

spikecomparison/multisortingcomparison.py

@@ -79,8 +79,7 @@ def compute_subgraphs(self):
 
     def dump(self, save_folder):
         save_folder = Path(save_folder)
-        if not save_folder.is_dir():
-            os.makedirs(str(save_folder))
+        save_folder.mkdir(parents=True, exist_ok=True)
         filename = str(save_folder / 'multicomparison.gpickle')
         nx.write_gpickle(self.graph, filename)
         kwargs = {'delta_time': self.delta_time, 'sampling_frequency': self.sampling_frequency,

spikecomparison/studytools.py

@@ -24,13 +24,11 @@
 
 from spikesorters.sorterlist import sorter_dict
 
-# TODO change this when sorters will be remove from toolkit
 from spikesorters import run_sorters, iter_output_folders, iter_sorting_output
 from .comparisontools import _perf_keys
 from .groundtruthcomparison import compare_sorter_to_ground_truth
 
 
-
 def setup_comparison_study(study_folder, gt_dict):
     """
     Based on a dict of (recording, sorting) create the study folder.
@@ -46,13 +44,13 @@ def setup_comparison_study(study_folder, gt_dict):
     """
 
     study_folder = Path(study_folder)
-    assert not os.path.exists(study_folder), 'study_folder already exists'
+    assert not study_folder.is_dir(), "'study_folder' already exists. Please remove it"
 
-    os.makedirs(str(study_folder))
-    os.makedirs(str(study_folder / 'raw_files'))
-    os.makedirs(str(study_folder / 'ground_truth'))
-    os.makedirs(str(study_folder / 'sortings'))
-    os.makedirs(str(study_folder / 'sortings/run_log'))
+    study_folder.mkdir(parents=True, exist_ok=True)
+    (study_folder / 'raw_files').mkdir(parents=True, exist_ok=True)
+    (study_folder / 'ground_truth').mkdir(parents=True, exist_ok=True)
+    (study_folder / 'sortings').mkdir(parents=True, exist_ok=True)
+    (study_folder / 'sortings' / 'run_log').mkdir(parents=True, exist_ok=True)
 
     for rec_name, (recording, sorting_gt) in gt_dict.items():
         # write recording as binary format + json + prb
@@ -123,7 +121,6 @@ def get_one_recording(study_folder, rec_name):
         info = json.load(f)
     rec = se.BinDatRecordingExtractor(raw_filename, info['sample_rate'], info['num_chan'],
                                       info['dtype'], time_axis=info['time_axis'])
-    # rec = rec.load_probe_file(prb_filename)
     load_probe_file_inplace(rec, prb_filename)
 
     return rec
@@ -250,8 +247,7 @@ def copy_sortings_to_npz(study_folder):
     sorter_folders = study_folder / 'sorter_folders'
     sorting_folders = study_folder / 'sortings'
 
-    if not os.path.exists(sorting_folders):
-        os.makedirs(str(sorting_folders))
+    sorting_folders.mkdir(parents=True, exist_ok=True)
 
     for rec_name, sorter_name, output_folder in iter_output_folders(sorter_folders):
         SorterClass = sorter_dict[sorter_name]
@@ -261,13 +257,12 @@ def copy_sortings_to_npz(study_folder):
             sorting = SorterClass.get_result_from_folder(output_folder)
             se.NpzSortingExtractor.write_sorting(sorting, npz_filename)
         except:
-            if os.path.exists(npz_filename):
-                os.remove(npz_filename)
-        if os.path.exists(output_folder / 'spikeinterface_log.json'):
+            if npz_filename.is_file():
+                npz_filename.unlink()
+        if (output_folder / 'spikeinterface_log.json').is_file():
             shutil.copyfile(output_folder / 'spikeinterface_log.json', sorting_folders / 'run_log' / (fname + '.json'))
 
 
-
 def iter_computed_names(study_folder):
     sorting_folder = Path(study_folder) / 'sortings'
     for filename in os.listdir(sorting_folder):
@@ -299,8 +294,7 @@ def collect_run_times(study_folder):
     log_folder = sorting_folders / 'run_log'
     tables_folder = study_folder / 'tables'
 
-    if not os.path.exists(tables_folder):
-        os.makedirs(str(tables_folder))
+    tables_folder.mkdir(parents=True, exist_ok=True)
 
     run_times = []
     for filename in os.listdir(log_folder):
@@ -441,8 +435,6 @@ def aggregate_performances_table(study_folder, exhaustive_gt=False, **karg_thres
     return dataframes
 
 
-
-
 def load_probe_file_inplace(recording, probe_file):
     '''
     This is a locally modified version of spikeextractor.extraction_tools.load_probe_file.