diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py index 2f228432b0..e698302ee1 100644 --- a/src/spikeinterface/core/recording_tools.py +++ b/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)