Merge pull request #2796 from h-mayorquin/improve_write_to_binary

Only assign memmap within boundaries for `write_binary`

src/spikeinterface/core/recording_tools.py

@@ -152,33 +152,40 @@ def _write_binary_chunk(segment_index, start_frame, end_frame, worker_ctx):
     cast_unsigned = worker_ctx["cast_unsigned"]
     file = worker_ctx["file_dict"][segment_index]
 
-    # Open the memmap
-    # What we need is the file_path
     num_channels = recording.get_num_channels()
-    num_frames = recording.get_num_frames(segment_index=segment_index)
-    shape = (num_frames, num_channels)
     dtype_size_bytes = np.dtype(dtype).itemsize
-    data_size_bytes = dtype_size_bytes * num_frames * num_channels
 
-    # Offset (The offset needs to be multiple of the page size)
-    # The mmap offset is associated to be as big as possible but still a multiple of the page size
-    # The array offset takes care of the reminder
-    mmap_offset, array_offset = divmod(byte_offset, mmap.ALLOCATIONGRANULARITY)
-    mmmap_length = data_size_bytes + array_offset
-    memmap_obj = mmap.mmap(file.fileno(), length=mmmap_length, access=mmap.ACCESS_WRITE, offset=mmap_offset)
+    # Calculate byte offsets for the start and end frames relative to the entire recording
+    start_byte = byte_offset + start_frame * num_channels * dtype_size_bytes
+    end_byte = byte_offset + end_frame * num_channels * dtype_size_bytes
 
-    array = np.ndarray.__new__(np.ndarray, shape=shape, dtype=dtype, buffer=memmap_obj, order="C", offset=array_offset)
-    # apply function
+    # The mmap offset must be a multiple of mmap.ALLOCATIONGRANULARITY
+    memmap_offset, start_offset = divmod(start_byte, mmap.ALLOCATIONGRANULARITY)
+    memmap_offset *= mmap.ALLOCATIONGRANULARITY
+
+    # This maps in bytes the region of the memmap that corresponds to the chunk
+    length = (end_byte - start_byte) + start_offset
+    memmap_obj = mmap.mmap(file.fileno(), length=length, access=mmap.ACCESS_WRITE, offset=memmap_offset)
+
+    # To use numpy semantics we use the array interface of the memmap object
+    num_frames = end_frame - start_frame
+    shape = (num_frames, num_channels)
+    memmap_array = np.ndarray(shape=shape, dtype=dtype, buffer=memmap_obj, offset=start_offset)
+
+    # Extract the traces and store them in the memmap array
     traces = recording.get_traces(
         start_frame=start_frame, end_frame=end_frame, segment_index=segment_index, cast_unsigned=cast_unsigned
     )
+
     if traces.dtype != dtype:
         traces = traces.astype(dtype, copy=False)
-    array[start_frame:end_frame, :] = traces
 
-    # Close the memmap
+    memmap_array[...] = traces
+
     memmap_obj.flush()
 
+    memmap_obj.close()
+
 
 write_binary_recording.__doc__ = write_binary_recording.__doc__.format(_shared_job_kwargs_doc)