Replace int() truncation with shared ms_to_samples() (#4484)

Author: Antorhythms

src/spikeinterface/benchmark/tests/common_benchmark_testing.py

@@ -16,6 +16,7 @@
     estimate_templates,
     Templates,
     create_sorting_analyzer,
+    ms_to_samples,
 )
 from spikeinterface.generation import generate_drifting_recording
 
@@ -54,8 +55,8 @@ def make_dataset(job_kwargs={}):
 def compute_gt_templates(recording, gt_sorting, ms_before=2.0, ms_after=3.0, return_in_uV=False, **job_kwargs):
     spikes = gt_sorting.to_spike_vector()  # [spike_indices]
     fs = recording.sampling_frequency
-    nbefore = int(ms_before * fs / 1000)
-    nafter = int(ms_after * fs / 1000)
+    nbefore = ms_to_samples(ms_before, fs)
+    nafter = ms_to_samples(ms_after, fs)
     templates_array = estimate_templates(
         recording,
         spikes,

src/spikeinterface/core/__init__.py

@@ -90,6 +90,7 @@
     read_python,
     write_python,
     normal_pdf,
+    ms_to_samples,
 )
 from .job_tools import (
     get_best_job_kwargs,

src/spikeinterface/core/analyzer_extension_core.py

@@ -21,6 +21,7 @@
 from .template import Templates
 from .sorting_tools import random_spikes_selection, select_sorting_periods_mask, spike_vector_to_indices
 from .job_tools import fix_job_kwargs, split_job_kwargs
+from .core_tools import ms_to_samples
 
 
 class ComputeRandomSpikes(AnalyzerExtension):
@@ -170,11 +171,11 @@ class ComputeWaveforms(AnalyzerExtension):
 
     @property
     def nbefore(self):
-        return int(self.params["ms_before"] * self.sorting_analyzer.sampling_frequency / 1000.0)
+        return ms_to_samples(self.params["ms_before"], self.sorting_analyzer.sampling_frequency)
 
     @property
     def nafter(self):
-        return int(self.params["ms_after"] * self.sorting_analyzer.sampling_frequency / 1000.0)
+        return ms_to_samples(self.params["ms_after"], self.sorting_analyzer.sampling_frequency)
 
     def _run(self, verbose=False, **job_kwargs):
         self.data.clear()
@@ -540,12 +541,12 @@ def _compute_and_append_from_waveforms(self, operators):
 
     @property
     def nbefore(self):
-        nbefore = int(self.params["ms_before"] * self.sorting_analyzer.sampling_frequency / 1000.0)
+        nbefore = ms_to_samples(self.params["ms_before"], self.sorting_analyzer.sampling_frequency)
         return nbefore
 
     @property
     def nafter(self):
-        nafter = int(self.params["ms_after"] * self.sorting_analyzer.sampling_frequency / 1000.0)
+        nafter = ms_to_samples(self.params["ms_after"], self.sorting_analyzer.sampling_frequency)
         return nafter
 
     def _select_extension_data(self, unit_ids):

src/spikeinterface/core/core_tools.py

@@ -757,3 +757,8 @@ def is_path_remote(path: str | Path) -> bool:
         Whether the path is a remote path.
     """
     return "s3://" in str(path) or "gcs://" in str(path)
+
+
+def ms_to_samples(ms: float, sampling_frequency: float) -> int:
+    """Convert a duration in milliseconds to the nearest number of samples."""
+    return round(ms * sampling_frequency / 1000.0)

src/spikeinterface/core/generate.py

@@ -12,7 +12,7 @@
 
 from spikeinterface.core import BaseRecording, BaseRecordingSegment, BaseSorting
 from .snippets_tools import snippets_from_sorting
-from .core_tools import define_function_from_class
+from .core_tools import define_function_from_class, ms_to_samples
 
 
 def _ensure_seed(seed):
@@ -1598,8 +1598,8 @@ def generate_single_fake_waveform(
     assert ms_after > depolarization_ms + repolarization_ms
     assert ms_before > depolarization_ms
 
-    nbefore = int(sampling_frequency * ms_before / 1000.0)
-    nafter = int(sampling_frequency * ms_after / 1000.0)
+    nbefore = ms_to_samples(ms_before, sampling_frequency)
+    nafter = ms_to_samples(ms_after, sampling_frequency)
     width = nbefore + nafter
     wf = np.zeros(width, dtype=dtype)
 
@@ -1776,8 +1776,8 @@ def generate_templates(
 
     num_units = units_locations.shape[0]
     num_channels = channel_locations.shape[0]
-    nbefore = int(sampling_frequency * ms_before / 1000.0)
-    nafter = int(sampling_frequency * ms_after / 1000.0)
+    nbefore = ms_to_samples(ms_before, sampling_frequency)
+    nafter = ms_to_samples(ms_after, sampling_frequency)
     width = nbefore + nafter
 
     if upsample_factor is not None:
@@ -2451,8 +2451,8 @@ def generate_ground_truth_recording(
             upsample_factor = templates.shape[3]
             upsample_vector = rng.integers(0, upsample_factor, size=num_spikes)
 
-    nbefore = int(ms_before * sampling_frequency / 1000.0)
-    nafter = int(ms_after * sampling_frequency / 1000.0)
+    nbefore = ms_to_samples(ms_before, sampling_frequency)
+    nafter = ms_to_samples(ms_after, sampling_frequency)
     assert (nbefore + nafter) == templates.shape[1]
 
     # construct recording

src/spikeinterface/core/node_pipeline.py

@@ -14,6 +14,7 @@
 from spikeinterface.core import BaseRecording, get_chunk_with_margin
 from spikeinterface.core.job_tools import ChunkRecordingExecutor, fix_job_kwargs, _shared_job_kwargs_doc
 from spikeinterface.core import get_channel_distances
+from spikeinterface.core.core_tools import ms_to_samples
 
 
 class PipelineNode:
@@ -314,8 +315,8 @@ def __init__(
         PipelineNode.__init__(self, recording=recording, parents=parents, return_output=return_output)
         self.ms_before = ms_before
         self.ms_after = ms_after
-        self.nbefore = int(ms_before * recording.get_sampling_frequency() / 1000.0)
-        self.nafter = int(ms_after * recording.get_sampling_frequency() / 1000.0)
+        self.nbefore = ms_to_samples(ms_before, recording.get_sampling_frequency())
+        self.nafter = ms_to_samples(ms_after, recording.get_sampling_frequency())
         self.neighbours_mask = None
 
 

src/spikeinterface/core/sparsity.py

@@ -7,6 +7,7 @@
 from .sorting_tools import random_spikes_selection
 from .job_tools import _shared_job_kwargs_doc
 from .waveform_tools import estimate_templates_with_accumulator
+from .core_tools import ms_to_samples
 
 _sparsity_doc = """
     method : str
@@ -784,8 +785,8 @@ def estimate_sparsity(
         probe = recording.create_dummy_probe_from_locations(chan_locs)
 
     if method != "by_property":
-        nbefore = int(ms_before * recording.sampling_frequency / 1000.0)
-        nafter = int(ms_after * recording.sampling_frequency / 1000.0)
+        nbefore = ms_to_samples(ms_before, recording.sampling_frequency)
+        nafter = ms_to_samples(ms_after, recording.sampling_frequency)
 
         num_samples = [recording.get_num_samples(seg_index) for seg_index in range(recording.get_num_segments())]
         random_spikes_indices = random_spikes_selection(

src/spikeinterface/core/tests/test_loading.py

@@ -5,6 +5,7 @@
     generate_ground_truth_recording,
     create_sorting_analyzer,
     load,
+    ms_to_samples,
     SortingAnalyzer,
     Templates,
     aggregate_channels,
@@ -71,7 +72,7 @@ def generate_templates_object():
     templates = Templates(
         templates_array=templates_arr,
         sampling_frequency=sampling_frequency,
-        nbefore=int(ms_before * sampling_frequency / 1000),
+        nbefore=ms_to_samples(ms_before, sampling_frequency),
         probe=probe,
     )
     return templates

src/spikeinterface/core/tests/test_waveform_tools.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from spikeinterface.core import generate_recording, generate_sorting, generate_ground_truth_recording
+from spikeinterface.core import generate_recording, generate_sorting, generate_ground_truth_recording, ms_to_samples
 from spikeinterface.core.waveform_tools import (
     extract_waveforms_to_buffers,
     extract_waveforms_to_single_buffer,
@@ -56,8 +56,8 @@ def test_waveform_tools(create_cache_folder):
     recording, sorting = get_dataset()
     sampling_frequency = recording.sampling_frequency
 
-    nbefore = int(3.0 * sampling_frequency / 1000.0)
-    nafter = int(4.0 * sampling_frequency / 1000.0)
+    nbefore = ms_to_samples(3.0, sampling_frequency)
+    nafter = ms_to_samples(4.0, sampling_frequency)
 
     dtype = recording.get_dtype()
     # return_in_uV = False
@@ -164,8 +164,8 @@ def test_estimate_templates_with_accumulator():
     ms_before = 1.0
     ms_after = 1.5
 
-    nbefore = int(ms_before * recording.sampling_frequency / 1000.0)
-    nafter = int(ms_after * recording.sampling_frequency / 1000.0)
+    nbefore = ms_to_samples(ms_before, recording.sampling_frequency)
+    nafter = ms_to_samples(ms_after, recording.sampling_frequency)
 
     spikes = sorting.to_spike_vector()
     # take one spikes every 10
@@ -218,8 +218,8 @@ def test_estimate_templates():
     ms_before = 1.0
     ms_after = 1.5
 
-    nbefore = int(ms_before * recording.sampling_frequency / 1000.0)
-    nafter = int(ms_after * recording.sampling_frequency / 1000.0)
+    nbefore = ms_to_samples(ms_before, recording.sampling_frequency)
+    nafter = ms_to_samples(ms_after, recording.sampling_frequency)
 
     spikes = sorting.to_spike_vector()
     # take one spikes every 10

src/spikeinterface/core/waveforms_extractor_backwards_compatibility.py

@@ -21,6 +21,7 @@
 from .sparsity import ChannelSparsity
 from .sortinganalyzer import SortingAnalyzer, load_sorting_analyzer
 from .loading import load
+from .core_tools import ms_to_samples
 from .analyzer_extension_core import ComputeRandomSpikes, ComputeWaveforms, ComputeTemplates
 
 _backwards_compatibility_msg = """####
@@ -162,12 +163,12 @@ def unit_ids(self) -> np.ndarray:
     @property
     def nbefore(self) -> int:
         ms_before = self.sorting_analyzer.get_extension("waveforms").params["ms_before"]
-        return int(ms_before * self.sampling_frequency / 1000.0)
+        return ms_to_samples(ms_before, self.sampling_frequency)
 
     @property
     def nafter(self) -> int:
         ms_after = self.sorting_analyzer.get_extension("waveforms").params["ms_after"]
-        return int(ms_after * self.sampling_frequency / 1000.0)
+        return ms_to_samples(ms_after, self.sampling_frequency)
 
     @property
     def nsamples(self) -> int:
@@ -522,8 +523,8 @@ def _read_old_waveforms_extractor_binary(folder, sorting):
         else:
             max_num_channel = np.max(np.sum(sparsity.mask, axis=1))
 
-        nbefore = int(params["ms_before"] * sorting.sampling_frequency / 1000.0)
-        nafter = int(params["ms_after"] * sorting.sampling_frequency / 1000.0)
+        nbefore = ms_to_samples(params["ms_before"], sorting.sampling_frequency)
+        nafter = ms_to_samples(params["ms_after"], sorting.sampling_frequency)
 
         waveforms = np.zeros((num_spikes, nbefore + nafter, max_num_channel), dtype=params["dtype"])
         # then read waveforms per units