Visitar URL original
Merge pull request #1955 from h-mayorquin/blackrock_lazy_waveform_rea… · NeuralEnsemble/python-neo@57f8d78 · GitHub
Skip to content

Commit 57f8d78

Browse files
authored
Merge pull request #1955 from h-mayorquin/blackrock_lazy_waveform_reading
BlackrockRawIO: read `.nev` spike waveforms lazily
2 parents 479d276 + b527229 commit 57f8d78

1 file changed

Lines changed: 41 additions & 12 deletions

File tree

‎neo/rawio/blackrockrawio.py‎

Lines changed: 41 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1028,23 +1028,33 @@ def _get_spike_raw_waveforms(self, block_index, seg_index, unit_index, t_start,
10281028
)
10291029
unit_spikes = all_spikes[mask]
10301030

1031+
timestamp = unit_spikes["timestamp"]
1032+
sl = self._get_timestamp_slice(timestamp, seg_index, t_start, t_stop)
1033+
unit_spikes = unit_spikes[sl]
1034+
1035+
# Read the waveforms of only these spikes, from the positions of their packets in the file
1036+
packet_size_bytes = int(self._nev_basic_header["bytes_in_data_packets"])
1037+
spike_packets = self._create_mmap_view(
1038+
fid=self._get_nev_file_handle(),
1039+
dtype=NEV_PACKET_DATA_TYPES_BY_SPEC[self._nev_spec]["Spikes"](packet_size_bytes),
1040+
offset=int(self._nev_basic_header["bytes_in_headers"]),
1041+
num_samples=self._nev_packet_count,
1042+
)
1043+
raw_waveforms = spike_packets["waveform"][unit_spikes["position"]]
1044+
10311045
wf_dtype = self._nev_params("waveform_dtypes")[channel_id]
10321046
wf_size = self._nev_params("waveform_size")[channel_id]
10331047
wf_byte_size = np.dtype(wf_dtype).itemsize * wf_size
10341048

10351049
dt1 = [
1036-
("extra", "S{}".format(unit_spikes["waveform"].dtype.itemsize - wf_byte_size)),
1050+
("extra", "S{}".format(raw_waveforms.dtype.itemsize - wf_byte_size)),
10371051
("ch_waveform", "S{}".format(wf_byte_size)),
10381052
]
10391053

1040-
waveforms = unit_spikes["waveform"].view(dt1)["ch_waveform"].flatten().view(wf_dtype)
1054+
waveforms = raw_waveforms.view(dt1)["ch_waveform"].flatten().view(wf_dtype)
10411055

10421056
waveforms = waveforms.reshape(int(unit_spikes.size), 1, int(wf_size))
10431057

1044-
timestamp = unit_spikes["timestamp"]
1045-
sl = self._get_timestamp_slice(timestamp, seg_index, t_start, t_stop)
1046-
waveforms = waveforms[sl]
1047-
10481058
return waveforms
10491059

10501060
def _event_count(self, block_index, seg_index, event_channel_index):
@@ -1996,15 +2006,34 @@ def _read_nev_data(self, spec, filename):
19962006

19972007
event_segment_ids = self._get_event_segment_ids(raw_data, masks, spec)
19982008

1999-
# Extract data for each packet type using view-then-mask pattern
2000-
# Strategy: reinterpret entire raw_data array with each packet type's structure, then filter
2001-
# All NEV data packets are fixed-width, so temporarily viewing "Spikes" data as "Comments"
2002-
# structure is safe - we immediately filter to keep only packets that actually match.
2003-
# This avoids creating copies of large data arrays during the parsing process.
2009+
# All NEV data packets are fixed-width, so raw_data is viewed with the structure of each
2010+
# packet type and the mask copies the packets of that type.
2011+
# Spike packets are most of the file, so for them only the fields used to select spikes are
2012+
# copied, with the position of each packet. _get_spike_raw_waveforms reads the waveforms
2013+
# from those positions.
2014+
self._nev_packet_count = n_packets
20042015
data = {}
20052016
for data_type in packet_identifiers:
20062017
mask = masks[data_type]
2007-
data[data_type] = (raw_data.view(types[data_type])[mask], event_segment_ids[mask])
2018+
packets = raw_data.view(types[data_type])
2019+
if data_type == "Spikes":
2020+
positions = np.flatnonzero(mask)
2021+
position_dtype = "uint32" if n_packets <= 2**32 else "uint64"
2022+
spikes = np.empty(
2023+
positions.size,
2024+
dtype=[
2025+
("timestamp", ts_format),
2026+
("packet_id", "uint16"),
2027+
("unit_class_nb", "uint8"),
2028+
("position", position_dtype),
2029+
],
2030+
)
2031+
for field in ("timestamp", "packet_id", "unit_class_nb"):
2032+
spikes[field] = packets[field][positions]
2033+
spikes["position"] = positions
2034+
data[data_type] = (spikes, event_segment_ids[mask])
2035+
else:
2036+
data[data_type] = (packets[mask], event_segment_ids[mask])
20082037

20092038
return data
20102039

0 commit comments

Comments
 (0)