Rename DriftCorrection to MotionDrift + add decorators + refactor from_spikeinterface()

Author: FrancescoNegri

spec/ndx-motion-drift.extensions.yaml

@@ -42,8 +42,9 @@ groups:
     target_type: ElectricalSeries
     doc: Link to the ElectricalSeries that is the source of the motion drift 
       estimation.
-- neurodata_type_def: DriftCorrection
+- neurodata_type_def: MotionDrift
   neurodata_type_inc: NWBDataInterface
+  default_name: MotionDrift
   doc: Displacement data from one or more channels. This can be used to store 
     the output of motion drift estimation and correction algorithms.
   groups:

src/pynwb/ndx_motion_drift/__init__.py

@@ -1,5 +1,6 @@
 from importlib.resources import files
-from pynwb import load_namespaces, get_class
+from pynwb import load_namespaces, get_class, register_class
+from pynwb.core import MultiContainerInterface
 
 # Get path to the namespace.yaml file with the expected location when installed not in editable mode
 __location_of_this_file = files(__name__)
@@ -13,11 +14,23 @@
 load_namespaces(str(__spec_path))
 
 DisplacementSeries = get_class("DisplacementSeries", "ndx-motion-drift")
-DriftCorrection = get_class("DriftCorrection", "ndx-motion-drift")
+@register_class("MotionDrift", "ndx-motion-drift")
+class MotionDrift(MultiContainerInterface):
+    """
+    Displacement data from one or more channels. This can be used to store the output of motion drift estimation and correction algorithms.
+    """
+
+    __clsconf__ = [{
+        'attr': 'displacement_series',
+        'type': DisplacementSeries,
+        'add': 'add_displacement_series',
+        'get': 'get_displacement_series',
+        'create': 'create_displacement_series',
+    }]
 
 __all__ = [
     "DisplacementSeries",
-    "DriftCorrection",
+    "MotionDrift",
 ]
 
 # Remove these functions/modules from the package

src/pynwb/ndx_motion_drift/io.py

@@ -1,23 +1,29 @@
 import json
 import numpy as np
 from spikeinterface.core import BaseRecording, Motion
-from pynwb import get_class
+from spikeinterface.extractors.nwbextractors import _retrieve_electrical_series_pynwb
+from pynwb import get_class, NWBFile
+from pynwb.ecephys import ElectricalSeries
 
 
-def from_spikeinterface(recording: BaseRecording, motion: Motion, **kwargs):
-    assert isinstance(recording, BaseRecording), f"The input must be a SpikeInterface BaseRecording object, not {type(recording)}."
+def from_spikeinterface(
+        motion: Motion, nwbfile: NWBFile,
+        source_electrical_series: ElectricalSeries | str, electrodes,
+        **kwargs
+    ):
+    assert isinstance(motion, Motion), f"The input must be a SpikeInterface Motion object, not {type(motion)}."
+    assert isinstance(nwbfile, NWBFile), f"The input must be a PyNWB NWBFile object, not {type(nwbfile)}."
+    assert isinstance(source_electrical_series, (ElectricalSeries, str)), f"The source_electrical_series must be either a PyNWB ElectricalSeries object or the path of an ElectricalSeries in the NWBFile, not {type(source_electrical_series)}."
 
-    recording_tree = [recording]
-    while recording_tree[-1].get_parent() is not None:
-        recording_tree.append(recording_tree[-1].get_parent())
-    has_electrical_series = [hasattr(r, "electrical_series") for r in recording_tree]
-    assert any(has_electrical_series), "The recording must have an ElectricalSeries to link the DisplacementSeries to."
+    if isinstance(source_electrical_series, str):
+        source_electrical_series = _retrieve_electrical_series_pynwb(nwbfile, source_electrical_series)
 
-    assert isinstance(motion, Motion), f"The input must be a SpikeInterface Motion object, not {type(motion)}."
+    if isinstance(electrodes, (list, np.ndarray)):
+        electrodes = nwbfile.create_electrode_table_region(
+            description="Electrodes for which the motion drift was estimated.",
+            region=np.where(np.isin(nwbfile.electrodes["channel_name"].data[:], electrodes))[0].tolist()
+        )
 
-    recording_nwb = recording_tree[has_electrical_series.index(True)]
-    nwbfile = recording_nwb._nwbfile
-    
     DisplacementSeries = get_class("DisplacementSeries", "ndx-motion-drift")
 
     displacement_series = DisplacementSeries(
@@ -32,11 +38,8 @@ def from_spikeinterface(recording: BaseRecording, motion: Motion, **kwargs):
         conversion=kwargs.get("conversion", 1.0),
         offset=kwargs.get("offset", 0.0),
         resolution=kwargs.get("resolution", -1.0),
-        source_electricalseries=recording_nwb.electrical_series,
-        electrodes=nwbfile.create_electrode_table_region(
-            description="Electrodes for which the motion drift was estimated.",
-            region=np.where(np.isin(nwbfile.electrodes["channel_name"].data[:], recording.channel_ids))[0].tolist()
-        ),
+        source_electricalseries=source_electrical_series,
+        electrodes=electrodes,
     )
 
     return displacement_series

src/spec/create_extension_spec.py

@@ -72,9 +72,10 @@ def main():
         ],
     )
 
-    drift_correction = NWBGroupSpec(
-        neurodata_type_def="DriftCorrection",
+    motion_drift = NWBGroupSpec(
+        neurodata_type_def="MotionDrift",
         neurodata_type_inc="NWBDataInterface",
+        default_name="MotionDrift",
         doc="Displacement data from one or more channels. This can be used to store the output of motion drift estimation and correction algorithms.",
         groups=[
             NWBGroupSpec(
@@ -85,7 +86,7 @@ def main():
         ]
     )
 
-    new_data_types = [displacement_series, drift_correction]
+    new_data_types = [displacement_series, motion_drift]
 
     # export the spec to yaml files in the root spec folder
     output_dir = str((Path(__file__).parent.parent.parent / "spec").absolute())