Merge pull request #195 from ttngu207/dev_spikeinterface_v101

feat: separate `export` (phy and report) into a separate table

Author: ttngu207

element_array_ephys/spike_sorting/si_spike_sorting.py

@@ -239,6 +239,7 @@ class PostProcessing(dj.Imported):
     ---
     execution_time: datetime   # datetime of the start of this step
     execution_duration: float  # execution duration in hours
+    do_si_export=0: bool       # whether to export to phy
     """
 
     def make(self, key):
@@ -295,22 +296,6 @@ def _sorting_analyzer_compute():
 
             sorting_analyzer.compute(extensions_to_compute, **job_kwargs)
 
-            # Save to phy format
-            if postprocessing_params.get("export_to_phy", False):
-                si.exporters.export_to_phy(
-                    sorting_analyzer=sorting_analyzer,
-                    output_folder=analyzer_output_dir / "phy",
-                    use_relative_path=True,
-                    **job_kwargs,
-                )
-            # Generate spike interface report
-            if postprocessing_params.get("export_report", True):
-                si.exporters.export_report(
-                    sorting_analyzer=sorting_analyzer,
-                    output_folder=analyzer_output_dir / "spikeinterface_report",
-                    **job_kwargs,
-                )
-
         _sorting_analyzer_compute()
 
         self.insert1(
@@ -321,10 +306,87 @@ def _sorting_analyzer_compute():
                     datetime.utcnow() - execution_time
                 ).total_seconds()
                 / 3600,
+                "do_si_export": postprocessing_params.get("export_to_phy", False)
+                or postprocessing_params.get("export_report", False),
             }
         )
 
         # Once finished, insert this `key` into ephys.Clustering
         ephys.Clustering.insert1(
             {**key, "clustering_time": datetime.utcnow()}, allow_direct_insert=True
         )
+
+
+@schema
+class SIExport(dj.Computed):
+    """A SpikeInterface export report and to Phy"""
+
+    definition = """
+    -> PostProcessing
+    ---
+    execution_time: datetime
+    execution_duration: float
+    """
+
+    @property
+    def key_source(self):
+        return PostProcessing & "do_si_export = 1"
+
+    def make(self, key):
+        execution_time = datetime.utcnow()
+
+        clustering_method, output_dir, params = (
+            ephys.ClusteringTask * ephys.ClusteringParamSet & key
+        ).fetch1("clustering_method", "clustering_output_dir", "params")
+        output_dir = find_full_path(ephys.get_ephys_root_data_dir(), output_dir)
+        sorter_name = clustering_method.replace(".", "_")
+
+        postprocessing_params = params["SI_POSTPROCESSING_PARAMS"]
+
+        job_kwargs = postprocessing_params.get(
+            "job_kwargs", {"n_jobs": -1, "chunk_duration": "1s"}
+        )
+
+        analyzer_output_dir = output_dir / sorter_name / "sorting_analyzer"
+        sorting_analyzer = si.load_sorting_analyzer(folder=analyzer_output_dir)
+
+        @memoized_result(
+            uniqueness_dict=postprocessing_params,
+            output_directory=analyzer_output_dir / "phy",
+        )
+        def _export_to_phy():
+            # Save to phy format
+            si.exporters.export_to_phy(
+                sorting_analyzer=sorting_analyzer,
+                output_folder=analyzer_output_dir / "phy",
+                use_relative_path=True,
+                **job_kwargs,
+            )
+
+        @memoized_result(
+            uniqueness_dict=postprocessing_params,
+            output_directory=analyzer_output_dir / "spikeinterface_report",
+        )
+        def _export_report():
+            # Generate spike interface report
+            si.exporters.export_report(
+                sorting_analyzer=sorting_analyzer,
+                output_folder=analyzer_output_dir / "spikeinterface_report",
+                **job_kwargs,
+            )
+
+        if postprocessing_params.get("export_report", False):
+            _export_report()
+        if postprocessing_params.get("export_to_phy", False):
+            _export_to_phy()
+
+        self.insert1(
+            {
+                **key,
+                "execution_time": execution_time,
+                "execution_duration": (
+                    datetime.utcnow() - execution_time
+                ).total_seconds()
+                / 3600,
+            }
+        )