Merge pull request #2513 from DradeAW/astype_rounding

Added `round` option to `recording.astype`

src/spikeinterface/core/baserecording.py

@@ -647,10 +647,10 @@ def binary_compatible_with(
         # good job you pass all crucible
         return True
 
-    def astype(self, dtype):
+    def astype(self, dtype, round: bool | None = None):
         from ..preprocessing.astype import astype
 
-        return astype(self, dtype=dtype)
+        return astype(self, dtype=dtype, round=round)
 
 
 class BaseRecordingSegment(BaseSegment):

src/spikeinterface/preprocessing/astype.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import numpy as np
+
 from ..core.core_tools import define_function_from_class
 from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
 from .filter import fix_dtype
@@ -11,6 +13,9 @@ class AstypeRecording(BasePreprocessor):
     Converts a recording to another dtype on the fly.
 
     For recording with an unsigned dtype, please use the `unsigned_to_signed` preprocessing function.
+
+    If `round` is True, will round the values to the nearest integer.
+    If `round` is None, will round in the case of float to integer conversion.
     """
 
     name = "astype"
@@ -19,20 +24,26 @@ def __init__(
         self,
         recording,
         dtype=None,
+        round: bool | None = None,
     ):
         dtype_ = fix_dtype(recording, dtype)
         BasePreprocessor.__init__(self, recording, dtype=dtype_)
 
+        if round is None:
+            round = np.issubdtype(dtype, np.integer)
+
         for parent_segment in recording._recording_segments:
             rec_segment = AstypeRecordingSegment(
                 parent_segment,
                 dtype,
+                round,
             )
             self.add_recording_segment(rec_segment)
 
         self._kwargs = dict(
             recording=recording,
             dtype=dtype_.str,
+            round=round,
         )
 
 
@@ -41,14 +52,18 @@ def __init__(
         self,
         parent_recording_segment,
         dtype,
+        round: bool,
     ):
         BasePreprocessorSegment.__init__(self, parent_recording_segment)
         self.dtype = dtype
+        self.round = round
 
     def get_traces(self, start_frame, end_frame, channel_indices):
         if channel_indices is None:
             channel_indices = slice(None)
         traces = self.parent_recording_segment.get_traces(start_frame, end_frame, channel_indices)
+        if self.round:
+            np.round(traces, out=traces)
         return traces.astype(self.dtype, copy=False)
 
 

src/spikeinterface/preprocessing/tests/test_astype.py

@@ -15,7 +15,9 @@ def test_astype():
     traces = (rng.randn(10000, 4) * 100).astype("float32")
     rec_float32 = NumpyRecording(traces, sampling_frequency=30000)
     traces_int16 = traces.astype("int16")
-    np.testing.assert_array_equal(traces_int16, astype(rec_float32, "int16").get_traces())
+    np.testing.assert_array_equal(traces_int16, astype(rec_float32, "int16", round=False).get_traces())
+    traces_int16_rounded = traces.round().astype("int16")
+    np.testing.assert_array_equal(traces_int16_rounded, astype(rec_float32, "int16").get_traces())
     traces_float64 = traces.astype("float64")
     np.testing.assert_array_equal(traces_float64, astype(rec_float32, "float64").get_traces())