Repository: SpikeInterface/spikeinterface
Pull request #3580 (2024-12-13 16:15:53)

Adding SharedMemoryTemplates
For template matching engines, it might be good to have some datastructures that would be shared in memory. For example the Templates (with the current PR), and maybe also the spatio-temporal overlaps (might be sorter specific, not sure we need a generic object for that)


Patch:
diff --git a/src/spikeinterface/core/template.py b/src/spikeinterface/core/template.py
index 3e3fcc738..5f59a18da 100644
--- a/src/spikeinterface/core/template.py
+++ b/src/spikeinterface/core/template.py
@@ -5,7 +5,10 @@ import json
 from dataclasses import dataclass, field, astuple, replace
 from probeinterface import Probe
 from pathlib import Path
+from multiprocessing.shared_memory import SharedMemory
+
 from .sparsity import ChannelSparsity
+from .core_tools import make_shared_array
 
 
 @dataclass
@@ -454,3 +457,88 @@ class Templates:
         assert self.probe is not None, "Templates.get_channel_locations() needs a probe to be set"
         channel_locations = self.probe.contact_positions
         return channel_locations
+
+
+class SharedMemoryTemplates(Templates):
+
+    def __init__(
+        self,
+        shm_name,
+        shape,
+        dtype,
+        sampling_frequency,
+        nbefore,
+        sparsity_mask,
+        channel_ids,
+        unit_ids,
+        probe,
+        is_scaled,
+        main_shm_owner=True,
+    ):
+
+        assert len(shape) == 3
+        assert shape[0] > 0, "SharedMemoryTemplates only supported with no empty templates"
+
+        self.shm = SharedMemory(shm_name, create=False)
+        templates_array = np.ndarray(shape=shape, dtype=dtype, buffer=self.shm.buf)
+
+        Templates.__init__(
+            self,
+            templates_array=templates_array,
+            sampling_frequency=sampling_frequency,
+            nbefore=nbefore,
+            sparsity_mask=sparsity_mask,
+            channel_ids=channel_ids,
+            unit_ids=unit_ids,
+            probe=probe,
+            is_scaled=is_scaled,
+        )
+
+        # self._serializability["memory"] = True
+        # self._serializability["json"] = False
+        # self._serializability["pickle"] = False
+
+        # this is very important for the shm.unlink()
+        # only the main instance need to call it
+        # all other instances that are loaded from dict are not the main owner
+        self.main_shm_owner = main_shm_owner
+
+        self._kwargs = dict(
+            shm_name=shm_name,
+            shape=shape,
+            sampling_frequency=sampling_frequency,
+            nbefore=self.nbefore,
+            sparsity_mask=self.sparsity_mask,
+            channel_ids=self.channel_ids,
+            unit_ids=self.unit_ids,
+            probe=self.probe,
+            is_scaled=self.is_scaled,
+            # this ensure that all dump/load will not be main shm owner
+            main_shm_owner=False,
+        )
+
+    def __del__(self):
+        self.shm.close()
+        if self.main_shm_owner:
+            self.shm.unlink()
+
+    @staticmethod
+    def from_templates(templates):
+        data = templates.get_dense_templates()
+        shm_templates, shm = make_shared_array(data.shape, data.dtype)
+        shm_templates[:] = data
+        shared_templates = SharedMemoryTemplates(
+            shm.name,
+            data.shape,
+            data.dtype,
+            templates.sampling_frequency,
+            templates.nbefore,
+            templates.sparsity_mask,
+            templates.channel_ids,
+            templates.unit_ids,
+            templates.probe,
+            templates.is_scaled,
+            main_shm_owner=True,
+        )
+        shm.close()
+        return shared_templates
diff --git a/src/spikeinterface/generation/__init__.py b/src/spikeinterface/generation/__init__.py
index 5bf42ecf0..4d744de0d 100644
--- a/src/spikeinterface/generation/__init__.py
+++ b/src/spikeinterface/generation/__init__.py
@@ -2,6 +2,7 @@ from .drift_tools import (
     move_dense_templates,
     interpolate_templates,
     DriftingTemplates,
+    SharedMemoryDriftingTemplates,
     InjectDriftingTemplatesRecording,
     make_linear_displacement,
 )
diff --git a/src/spikeinterface/generation/drift_tools.py b/src/spikeinterface/generation/drift_tools.py
index 9d2834035..ecfc84d25 100644
--- a/src/spikeinterface/generation/drift_tools.py
+++ b/src/spikeinterface/generation/drift_tools.py
@@ -7,6 +7,8 @@ import numpy as np
 from numpy.typing import ArrayLike
 from probeinterface import Probe
 from spikeinterface.core import BaseRecording, BaseRecordingSegment, BaseSorting, Templates
+from multiprocessing.shared_memory import SharedMemory
+from spikeinterface.core.core_tools import make_shared_array
 
 
 def interpolate_templates(templates_array, source_locations, dest_locations, interpolation_method="cubic"):
@@ -258,6 +260,85 @@ class DriftingTemplates(Templates):
         self.displacements = displacements
 
 
+class SharedMemoryDriftingTemplates(DriftingTemplates):
+
+    def __init__(
+        self,
+        shm_name,
+        shape,
+        dtype,
+        templates_array_moved=None,
+        displacements=None,
+        main_shm_owner=True,
+        **static_kwargs,
+    ):
+
+        assert len(shape) == 4
+        assert shape[0] > 0, "SharedMemoryTemplates only supported with no empty templates"
+
+        self.shm = SharedMemory(shm_name, create=False)
+        templates_array_moved = np.ndarray(shape=shape, dtype=dtype, buffer=self.shm.buf)
+        self.static_kwargs = static_kwargs
+        DriftingTemplates.__init__(
+            self, templates_array_moved=templates_array_moved, displacements=displacements, **self.static_kwargs
+        )
+
+        # this is very important for the shm.unlink()
+        # only the main instance need to call it
+        # all other instances that are loaded from dict are not the main owner
+        self.main_shm_owner = main_shm_owner
+
+        self._kwargs = dict(
+            shm_name=shm_name,
+            shape=shape,
+            displacements=self.displacements,
+            static_kwargs=self.static_kwargs,
+            channel_ids=self.channel_ids,
+            # this ensure that all dump/load will not be main shm owner
+            main_shm_owner=False,
+        )
+
+    def __del__(self):
+        self.shm.close()
+        if self.main_shm_owner:
+            self.shm.unlink()
+
+    @staticmethod
+    def from_drifting_templates(drifting_templates):
+        assert (
+            drifting_templates.templates_array_moved is not None
+        ), "drifting_templates must have precomputed displacements"
+        data = drifting_templates.templates_array_moved
+        shm_templates, shm = make_shared_array(data.shape, data.dtype)
+        shm_templates[:] = data
+        static_kwargs = drifting_templates.to_dict()
+        init_kwargs = {
+            "templates_array": np.asarray(static_kwargs["templates_array"]),
+            "sparsity_mask": (
+                None if static_kwargs["sparsity_mask"] is None else np.asarray(static_kwargs["sparsity_mask"])
+            ),
+            "channel_ids": np.asarray(static_kwargs["channel_ids"]),
+            "unit_ids": np.asarray(static_kwargs["unit_ids"]),
+            "sampling_frequency": static_kwargs["sampling_frequency"],
+            "nbefore": static_kwargs["nbefore"],
+            "is_scaled": static_kwargs["is_scaled"],
+            "probe": (
+                static_kwargs["probe"] if static_kwargs["probe"] is None else Probe.from_dict(static_kwargs["probe"])
+            ),
+        }
+        shared_drifting_templates = SharedMemoryDriftingTemplates(
+            shm.name,
+            data.shape,
+            data.dtype,
+            shm_templates,
+            drifting_templates.displacements,
+            main_shm_owner=True,
+            **init_kwargs,
+        )
+        shm.close()
+        return shared_drifting_templates
+
+
 def make_linear_displacement(start, stop, num_step=10):
     """
     Generates 2D linear displacements between `start` and `stop` positions (included in returned displacements).


Test patch:
diff --git a/src/spikeinterface/core/tests/test_template_class.py b/src/spikeinterface/core/tests/test_template_class.py
index 4e0a0c856..b38d3cc2c 100644
--- a/src/spikeinterface/core/tests/test_template_class.py
+++ b/src/spikeinterface/core/tests/test_template_class.py
@@ -1,7 +1,7 @@
 import pytest
 import numpy as np
 import pickle
-from spikeinterface.core.template import Templates
+from spikeinterface.core.template import Templates, SharedMemoryTemplates
 from spikeinterface.core.sparsity import ChannelSparsity
 
 from probeinterface import generate_multi_columns_probe
@@ -171,6 +171,21 @@ def test_select_channels(template_type, is_scaled):
         assert np.array_equal(selected_template.sparsity_mask, template.sparsity_mask[:, selected_channel_ids_indices])
 
 
+@pytest.mark.parametrize("is_scaled", [True, False])
+@pytest.mark.parametrize("template_type", ["dense"])
+def test_shm_templates(template_type, is_scaled):
+    template = generate_test_template(template_type, is_scaled)
+    shm_templates = SharedMemoryTemplates.from_templates(template)
+
+    # Verify that the channel ids match
+    assert np.array_equal(shm_templates.channel_ids, template.channel_ids)
+    # Verify that the templates data matches
+    assert np.array_equal(shm_templates.templates_array, template.templates_array)
+
+    if template.sparsity_mask is not None:
+        assert np.array_equal(shm_templates.sparsity_mask, template.sparsity_mask)
+
+
 if __name__ == "__main__":
     # test_json_serialization("sparse")
     test_json_serialization("dense")
diff --git a/src/spikeinterface/generation/tests/test_drift_tools.py b/src/spikeinterface/generation/tests/test_drift_tools.py
index 5647b3393..466e9541a 100644
--- a/src/spikeinterface/generation/tests/test_drift_tools.py
+++ b/src/spikeinterface/generation/tests/test_drift_tools.py
@@ -6,6 +6,7 @@ from spikeinterface.generation import (
     move_dense_templates,
     make_linear_displacement,
     DriftingTemplates,
+    SharedMemoryDriftingTemplates,
     InjectDriftingTemplatesRecording,
 )
 from spikeinterface.core.generate import generate_templates, generate_sorting, NoiseGeneratorRecording
@@ -133,6 +134,17 @@ def test_DriftingTemplates():
     assert np.array_equal(drifting_templates_from_precomputed.displacements, drifting_templates.displacements)
 
 
+def test_SharedMemoryDriftingTemplates():
+    static_templates = make_some_templates()
+    drifting_templates = DriftingTemplates.from_static_templates(static_templates)
+    displacement = np.array([[5.0, 10.0]])
+    drifting_templates.precompute_displacements(displacement)
+    shm_drifting_templates = SharedMemoryDriftingTemplates.from_drifting_templates(drifting_templates)
+
+    assert np.array_equal(shm_drifting_templates.templates_array_moved, drifting_templates.templates_array_moved)
+    assert np.array_equal(shm_drifting_templates.displacements, drifting_templates.displacements)
+
+
 def test_InjectDriftingTemplatesRecording(create_cache_folder):
     cache_folder = create_cache_folder
     templates = make_some_templates()