Merge pull request #3669 from alejoe91/revert-sharedmem-templates

Revert PR #3580

Author: samuelgarcia

src/spikeinterface/core/template.py

@@ -5,10 +5,8 @@
 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
@@ -457,88 +455,3 @@ def get_channel_locations(self) -> np.ndarray:
         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

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, SharedMemoryTemplates
+from spikeinterface.core.template import Templates
 from spikeinterface.core.sparsity import ChannelSparsity
 
 from probeinterface import generate_multi_columns_probe
@@ -171,21 +171,6 @@ 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")

src/spikeinterface/generation/__init__.py

@@ -2,7 +2,6 @@
     move_dense_templates,
     interpolate_templates,
     DriftingTemplates,
-    SharedMemoryDriftingTemplates,
     InjectDriftingTemplatesRecording,
     make_linear_displacement,
 )

src/spikeinterface/generation/drift_tools.py

@@ -7,8 +7,6 @@
 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"):
@@ -260,85 +258,6 @@ def precompute_displacements(self, displacements, **interpolation_kwargs):
         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).

src/spikeinterface/generation/tests/test_drift_tools.py

@@ -6,7 +6,6 @@
     move_dense_templates,
     make_linear_displacement,
     DriftingTemplates,
-    SharedMemoryDriftingTemplates,
     InjectDriftingTemplatesRecording,
 )
 from spikeinterface.core.generate import generate_templates, generate_sorting, NoiseGeneratorRecording
@@ -134,25 +133,13 @@ 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()
     probe = templates.probe
 
     # drifting templates
     drifting_templates = DriftingTemplates.from_static_templates(templates)
-    channel_locations = probe.contact_positions
 
     num_units = templates.unit_ids.size
     sampling_frequency = templates.sampling_frequency