From 180fbc6a6b818bf14b9d92f0b226eca4d76c7771 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@datajoint.com>
Date: Tue, 3 Dec 2024 09:38:27 -0600
Subject: [PATCH] feat: add `File` part table for external storage

---
 .../spike_sorting/si_spike_sorting.py         | 84 ++++++++++++++++++-
 1 file changed, 82 insertions(+), 2 deletions(-)

diff --git a/element_array_ephys/spike_sorting/si_spike_sorting.py b/element_array_ephys/spike_sorting/si_spike_sorting.py
index b4d17a27..13d48712 100644
--- a/element_array_ephys/spike_sorting/si_spike_sorting.py
+++ b/element_array_ephys/spike_sorting/si_spike_sorting.py
@@ -61,6 +61,14 @@ class PreProcessing(dj.Imported):
     execution_duration: float  # execution duration in hours
     """
 
+    class File(dj.Part):
+        definition = """
+        -> master
+        file_name: varchar(255)
+        ---
+        file: filepath@ephys-processed
+        """
+
     @property
     def key_source(self):
         return (
@@ -176,6 +184,14 @@ def make(self, key):
                 / 3600,
             }
         )
+        # Insert result files
+        self.File.insert(
+            [
+                {**key, "file_name": f.relative_to(recording_dir).as_posix(), "file": f}
+                for f in recording_dir.rglob("*")
+                if f.is_file()
+            ]
+        )
 
 
 @schema
@@ -189,6 +205,14 @@ class SIClustering(dj.Imported):
     execution_duration: float       # execution duration in hours
     """
 
+    class File(dj.Part):
+        definition = """
+        -> master
+        file_name: varchar(255)
+        ---
+        file: filepath@ephys-processed
+        """
+
     def make(self, key):
         execution_time = datetime.utcnow()
 
@@ -239,6 +263,18 @@ def _run_sorter():
                 / 3600,
             }
         )
+        # Insert result files
+        self.File.insert(
+            [
+                {
+                    **key,
+                    "file_name": f.relative_to(sorting_output_dir).as_posix(),
+                    "file": f,
+                }
+                for f in sorting_output_dir.rglob("*")
+                if f.is_file()
+            ]
+        )
 
 
 @schema
@@ -253,6 +289,14 @@ class PostProcessing(dj.Imported):
     do_si_export=0: bool       # whether to export to phy
     """
 
+    class File(dj.Part):
+        definition = """
+        -> master
+        file_name: varchar(255)
+        ---
+        file: filepath@ephys-processed
+        """
+
     def make(self, key):
         execution_time = datetime.utcnow()
 
@@ -290,7 +334,9 @@ def make(self, key):
         def _sorting_analyzer_compute():
             if not has_units:
                 log.info("No units found in sorting object. Skipping sorting analyzer.")
-                analyzer_output_dir.mkdir(parents=True, exist_ok=True)  # create empty directory anyway, for consistency
+                analyzer_output_dir.mkdir(
+                    parents=True, exist_ok=True
+                )  # create empty directory anyway, for consistency
                 return
 
             # Sorting Analyzer
@@ -316,7 +362,9 @@ def _sorting_analyzer_compute():
 
         _sorting_analyzer_compute()
 
-        do_si_export = postprocessing_params.get("export_to_phy", False) or postprocessing_params.get("export_report", False)
+        do_si_export = postprocessing_params.get(
+            "export_to_phy", False
+        ) or postprocessing_params.get("export_report", False)
 
         self.insert1(
             {
@@ -329,6 +377,17 @@ def _sorting_analyzer_compute():
                 "do_si_export": do_si_export and has_units,
             }
         )
+        self.File.insert(
+            [
+                {
+                    **key,
+                    "file_name": f.relative_to(analyzer_output_dir).as_posix(),
+                    "file": f,
+                }
+                for f in analyzer_output_dir.rglob("*")
+                if f.is_file()
+            ]
+        )
 
         # Once finished, insert this `key` into ephys.Clustering
         ephys.Clustering.insert1(
@@ -347,6 +406,14 @@ class SIExport(dj.Computed):
     execution_duration: float
     """
 
+    class File(dj.Part):
+        definition = """
+        -> master
+        file_name: varchar(255)
+        ---
+        file: filepath@ephys-processed
+        """
+
     @property
     def key_source(self):
         return PostProcessing & "do_si_export = 1"
@@ -409,3 +476,16 @@ def _export_report():
                 / 3600,
             }
         )
+        # Insert result files
+        for report_dirname in ("spikeinterface_report", "phy"):
+            self.File.insert(
+                [
+                    {
+                        **key,
+                        "file_name": f.relative_to(analyzer_output_dir).as_posix(),
+                        "file": f,
+                    }
+                    for f in (analyzer_output_dir / report_dirname).rglob("*")
+                    if f.is_file()
+                ]
+            )