Repository: SpikeInterface/spikeinterface
Pull request #485 (2022-01-31 09:53:25)

Improve portability of WaveformExtractor
Fixes https://github.com/SpikeInterface/spikeinterface/issues/473


Patch:
diff --git a/spikeinterface/core/base.py b/spikeinterface/core/base.py
index 576f87216..f9ca3afd3 100644
--- a/spikeinterface/core/base.py
+++ b/spikeinterface/core/base.py
@@ -4,6 +4,7 @@ from copy import deepcopy
 import weakref
 import json
 import pickle
+import os
 import random
 import string
 import warnings
@@ -649,12 +650,14 @@ def _make_paths_relative(d, relative):
             if "path" in k:
                 # paths can be str or list of str
                 if isinstance(d[k], str):
-                    d[k] = str(Path(d[k]).relative_to(relative))
+                    # we use os.path.relpath here to allow for relative paths with respect to shared root
+                    d[k] = os.path.relpath(str(d[k]), start=str(relative.absolute()))
                 else:
                     assert isinstance(d[k], list), "Paths can be strings or lists in kwargs"
                     relative_paths = []
                     for path in d[k]:
-                        relative_paths.append(str(Path(path).relative_to(relative)))
+                        # we use os.path.relpath here to allow for relative paths with respect to shared root
+                        relative_paths.append(os.path.relpath(str(path), start=str(relative.absolute())))
                     d[k] = relative_paths
         return d
 
diff --git a/spikeinterface/core/waveform_extractor.py b/spikeinterface/core/waveform_extractor.py
index 8fbcaf157..97ec6e95d 100644
--- a/spikeinterface/core/waveform_extractor.py
+++ b/spikeinterface/core/waveform_extractor.py
@@ -86,9 +86,11 @@ class WaveformExtractor:
     @classmethod
     def load_from_folder(cls, folder):
         folder = Path(folder)
-        assert folder.is_dir(), f'This folder do not exists {folder}'
-        recording = load_extractor(folder / 'recording.json')
-        sorting = load_extractor(folder / 'sorting.json')
+        assert folder.is_dir(), f'This folder does not exists {folder}'
+        recording = load_extractor(folder / 'recording.json',
+                                   base_folder=folder)
+        sorting = load_extractor(folder / 'sorting.json',
+                                 base_folder=folder)
         we = cls(recording, sorting, folder)
 
         for mode in _possible_template_modes:
@@ -100,7 +102,8 @@ class WaveformExtractor:
         return we
 
     @classmethod
-    def create(cls, recording, sorting, folder, remove_if_exists=False):
+    def create(cls, recording, sorting, folder, remove_if_exists=False,
+               use_relative_path=False):
         folder = Path(folder)
         if folder.is_dir():
             if remove_if_exists:
@@ -108,11 +111,16 @@ class WaveformExtractor:
             else:
                 raise FileExistsError('Folder already exists')
         folder.mkdir(parents=True)
+        
+        if use_relative_path:
+            relative_to = folder
+        else:
+            relative_to = None
 
         if recording.is_dumpable:
-            recording.dump(folder / 'recording.json', relative_to=None)
+            recording.dump(folder / 'recording.json', relative_to=relative_to)
         if sorting.is_dumpable:
-            sorting.dump(folder / 'sorting.json', relative_to=None)
+            sorting.dump(folder / 'sorting.json', relative_to=relative_to)
 
         return cls(recording, sorting, folder)
 
@@ -811,6 +819,7 @@ def extract_waveforms(recording, sorting, folder,
                       overwrite=False,
                       return_scaled=True,
                       dtype=None,
+                      use_relative_path=False,
                       **job_kwargs):
     """
     Extracts waveform on paired Recording-Sorting objects.
@@ -843,7 +852,11 @@ def extract_waveforms(recording, sorting, folder,
         If True and recording has gain_to_uV/offset_to_uV properties, waveforms are converted to uV.
     dtype: dtype or None
         Dtype of the output waveforms. If None, the recording dtype is maintained.
-
+    use_relative_path: bool
+        If True, the recording and sorting paths are relative to the waveforms folder. 
+        This allows portability of the waveform folder provided that the relative paths are the same, 
+        but forces all the data files to be in the same drive.
+        Default is False.
 
     {}
 
@@ -860,7 +873,7 @@ def extract_waveforms(recording, sorting, folder,
     if load_if_exists and folder.is_dir():
         we = WaveformExtractor.load_from_folder(folder)
     else:
-        we = WaveformExtractor.create(recording, sorting, folder)
+        we = WaveformExtractor.create(recording, sorting, folder, use_relative_path=use_relative_path)
         we.set_params(ms_before=ms_before, ms_after=ms_after, max_spikes_per_unit=max_spikes_per_unit, dtype=dtype,
                       return_scaled=return_scaled)
         we.run_extract_waveforms(**job_kwargs)


Test patch:
diff --git a/spikeinterface/core/tests/test_waveform_extractor.py b/spikeinterface/core/tests/test_waveform_extractor.py
index 025237fc1..0fa8b3f08 100644
--- a/spikeinterface/core/tests/test_waveform_extractor.py
+++ b/spikeinterface/core/tests/test_waveform_extractor.py
@@ -182,6 +182,57 @@ def test_sparsity():
         assert temp_diff.shape[-1] == len(sparsity_diff[unit])
 
 
+def test_portability():
+    durations = [30, 40]
+    sampling_frequency = 30000.
+    
+    folder_to_move = Path("original_folder")
+    if folder_to_move.is_dir():
+        shutil.rmtree(folder_to_move)
+    folder_to_move.mkdir()
+    folder_moved = Path("moved_folder")
+    if folder_moved.is_dir():
+        shutil.rmtree(folder_moved)
+    # folder_moved.mkdir()
+        
+    # 2 segments
+    num_channels = 2
+    recording = generate_recording(num_channels=num_channels, durations=durations, 
+                                   sampling_frequency=sampling_frequency)
+    recording.annotate(is_filtered=True)
+    folder_rec = folder_to_move / "rec"
+    recording = recording.save(folder=folder_rec)
+    num_units = 15
+    sorting = generate_sorting(num_units=num_units, sampling_frequency=sampling_frequency, durations=durations)
+    folder_sort = folder_to_move / "sort"
+    sorting = sorting.save(folder=folder_sort)
+
+    wf_folder = folder_to_move / "waveform_extractor"
+    if wf_folder.is_dir():
+        shutil.rmtree(wf_folder)
+
+    # save with relative paths
+    we = extract_waveforms(recording, sorting, wf_folder, use_relative_path=True)
+
+    # move all to a separate folder
+    shutil.copytree(folder_to_move, folder_moved)
+    wf_folder_moved = folder_moved / "waveform_extractor"
+    we_loaded = extract_waveforms(recording, sorting, wf_folder_moved, load_if_exists=True)
+    
+    assert we_loaded.recording is not None
+    assert we_loaded.sorting is not None
+
+    assert np.allclose(
+        we.recording.get_channel_ids(), we_loaded.recording.get_channel_ids())
+    assert np.allclose(
+        we.sorting.get_unit_ids(), we_loaded.sorting.get_unit_ids())
+
+
+    for unit in we.sorting.get_unit_ids():
+        wf = we.get_waveforms(unit_id=unit)
+        wf_loaded = we_loaded.get_waveforms(unit_id=unit)
+
+        assert np.allclose(wf, wf_loaded)
 
 
 if __name__ == '__main__':
@@ -189,3 +240,4 @@ if __name__ == '__main__':
     test_WaveformExtractor()
     test_extract_waveforms()
     test_sparsity()
+    test_portability()