Add read_nwb_sorting_analyzer: build a SortingAnalyzer from an NWB Units table

[h-mayorquin]

read_nwb_sorting_analyzer builds a curatable SortingAnalyzer directly from an NWB (Neurodata Without Borders) Units table, populating each extension from what the file already stores rather than recomputing it. When the Units table has waveform_mean, the analyzer is built recordingless from those stored templates, plus the per-unit metrics and the electrodes region for sparsity; when the file has an accessible ElectricalSeries, it is used as the recording. It mirrors read_kilosort_as_analyzer in structure, injecting the templates, quality and template metrics, sparsity, and random_spikes extensions from the file's contents. This supersedes the earlier draft #4270.

The reads are deliberate: Units columns are classified from metadata and only the templates, the electrodes region, and the scalar metric/label columns are materialized, while the large per-spike ragged columns (spike times, amplitudes, depths) are never touched at build. The sorting is kept lazy so its spike times are read only on demand, which builds on #4662 (single bulk-read NWB spike vector) and the copy_sorting option (#4668); the recordingless case builds a lightweight placeholder recording from the electrode rel_x/rel_y geometry (via generate_ground_truth_recording, #4588) purely to carry probe geometry into the standard constructor, then drops it. Together these keep the build small and memory-light regardless of file size.

Because the reads are deliberate and the sorting stays lazy, the function works the same whether the file is local or streamed (stream_mode is passed through to the extractors). That makes the streamed case cheap: as a check, building from a real IBL (International Brain Laboratory) processed file on dandiset 000409, session 6713a4a7-faed-4df2-acab-ee4e63326f8d (898 units, 20.7M spikes), produced a curatable analyzer in about 14 s while transferring only ~22 MB, with the ~130 MB spike read deferred until a spike-based view needs it. This PR is the reader itself; a separate how-to PR will cover the streaming-from-DANDI workflow in depth. It depends on #4662 and the copy_sorting PR (#4668), which should merge first so this branch rebases down to just the reader.

from dandi.dandiapi import DandiAPIClient
from spikeinterface.extractors import read_nwb_sorting_analyzer

# IBL Brain Wide Map (dandiset 000409), one session's processed file (units + templates, no raw traces)
SESSION = "6713a4a7-faed-4df2-acab-ee4e63326f8d"
with DandiAPIClient() as client:
    dandiset = client.get_dandiset("000409", "draft")
    asset = next(a for a in dandiset.get_assets_by_glob(f"*{SESSION}*.nwb") if "desc-processed" in a.path)
    s3_url = asset.get_content_url(follow_redirects=1, strip_query=True)

# build the analyzer directly from the streamed Units table (no full download)
analyzer = read_nwb_sorting_analyzer(s3_url, stream_mode="remfile")

print(analyzer)
print("num_units:", analyzer.get_num_units())
print("templates:", analyzer.get_extension("templates").get_data().shape)
print("quality metrics:", list(analyzer.get_extension("quality_metrics").get_data().columns))

# open in spikeinterface-gui for curation (needs a display and `pip install spikeinterface-gui`)
# import spikeinterface_gui
# spikeinterface_gui.run_mainwindow(analyzer)

[alejoe91]

This is failing on https://dandiarchive.s3.amazonaws.com/blobs/c41/fae/c41fae67-e6e2-4dc9-adcb-6131d530b6cd since it has two probes and when you select groups it calls a select_units, which doesn't have the spike_times_data

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[4], line 1
----> 1 analyzer = se.read_nwb_sorting_analyzer(dandi_path_sorting, stream_mode="remfile", group_name="Probe00")

File ~/Documents/codes/spike_sorting/spikeinterface/spikeinterface/src/spikeinterface/extractors/nwbextractors.py:2108, in read_nwb_sorting_analyzer(file_path, t_start, sampling_frequency, electrical_series_path, unit_table_path, stream_mode, stream_cache_path, cache, storage_options, use_pynwb, group_name, compute_extra, compute_extra_params, extension_map, verbose)
   2106     analyzer_channel_ids = list(recording.get_channel_ids())
   2107 else:
-> 2108     analyzer_recording, analyzer_channel_ids = _make_placeholder_recording_from_electrodes(
   2109         sorting, electrodes_table, electrodes_indices, verbose=verbose
   2110     )
   2112 # Per-unit sparsity and channel map from the Units `electrodes` region. Each unit's `waveform_mean`
   2113 # is stored only on a subset of channels (near its peak); the region gives which channels those are.
   2114 # We use it both for the analyzer sparsity and to scatter the waveforms onto their true channel
   2115 # positions instead of stacking the sparse block densely.
   2116 sparsity, unit_local_channels = _make_sparsity_from_electrodes(
   2117     sorting, electrodes_table, electrodes_indices, analyzer_channel_ids
   2118 )

File ~/Documents/codes/spike_sorting/spikeinterface/spikeinterface/src/spikeinterface/extractors/nwbextractors.py:2200, in _make_placeholder_recording_from_electrodes(sorting, electrodes_table, electrodes_indices, verbose)
   2192 locations = np.array([electrodes_table_sliced["rel_x"][:], electrodes_table_sliced["rel_y"][:]]).T
   2194 # The recording length only needs to loosely bound the timeline for the recordingless GUI; nothing
   2195 # about curation depends on it being exact. Estimate it cheaply from the last stored spike time (one
   2196 # element, i.e. only the last spike_times chunk) rather than scanning the whole array for the true
   2197 # global maximum. spike_times is in seconds, so this is already a duration in seconds. It is only an
   2198 # approximation: because spike_times is concatenated per unit and not globally sorted, this is the
   2199 # last unit's last spike, not necessarily the latest spike overall.
-> 2200 last_spike_time = float(np.asarray(sorting._sorting_segments[0].spike_times_data[-1]))
   2201 duration = last_spike_time + 1.0
   2203 probe = Probe(si_units="um")

AttributeError: 'UnitsSelectionSortingSegment' object has no attribute 'spike_times_data'

[alejoe91]

This is failing on https://dandiarchive.s3.amazonaws.com/blobs/c41/fae/c41fae67-e6e2-4dc9-adcb-6131d530b6cd since it has two probes and when you select groups it calls a select_units, which doesn't have the spike_times_data

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[4], line 1
----> 1 analyzer = se.read_nwb_sorting_analyzer(dandi_path_sorting, stream_mode="remfile", group_name="Probe00")

File ~/Documents/codes/spike_sorting/spikeinterface/spikeinterface/src/spikeinterface/extractors/nwbextractors.py:2108, in read_nwb_sorting_analyzer(file_path, t_start, sampling_frequency, electrical_series_path, unit_table_path, stream_mode, stream_cache_path, cache, storage_options, use_pynwb, group_name, compute_extra, compute_extra_params, extension_map, verbose)
   2106     analyzer_channel_ids = list(recording.get_channel_ids())
   2107 else:
-> 2108     analyzer_recording, analyzer_channel_ids = _make_placeholder_recording_from_electrodes(
   2109         sorting, electrodes_table, electrodes_indices, verbose=verbose
   2110     )
   2112 # Per-unit sparsity and channel map from the Units `electrodes` region. Each unit's `waveform_mean`
   2113 # is stored only on a subset of channels (near its peak); the region gives which channels those are.
   2114 # We use it both for the analyzer sparsity and to scatter the waveforms onto their true channel
   2115 # positions instead of stacking the sparse block densely.
   2116 sparsity, unit_local_channels = _make_sparsity_from_electrodes(
   2117     sorting, electrodes_table, electrodes_indices, analyzer_channel_ids
   2118 )

File ~/Documents/codes/spike_sorting/spikeinterface/spikeinterface/src/spikeinterface/extractors/nwbextractors.py:2200, in _make_placeholder_recording_from_electrodes(sorting, electrodes_table, electrodes_indices, verbose)
   2192 locations = np.array([electrodes_table_sliced["rel_x"][:], electrodes_table_sliced["rel_y"][:]]).T
   2194 # The recording length only needs to loosely bound the timeline for the recordingless GUI; nothing
   2195 # about curation depends on it being exact. Estimate it cheaply from the last stored spike time (one
   2196 # element, i.e. only the last spike_times chunk) rather than scanning the whole array for the true
   2197 # global maximum. spike_times is in seconds, so this is already a duration in seconds. It is only an
   2198 # approximation: because spike_times is concatenated per unit and not globally sorted, this is the
   2199 # last unit's last spike, not necessarily the latest spike overall.
-> 2200 last_spike_time = float(np.asarray(sorting._sorting_segments[0].spike_times_data[-1]))
   2201 duration = last_spike_time + 1.0
   2203 probe = Probe(si_units="um")

AttributeError: 'UnitsSelectionSortingSegment' object has no attribute 'spike_times_data'

Fixed in last couple of commits

[h-mayorquin]

Thanks for fixing the multi-group bug and the random-spikes handling. I built on that and fixed a few more things on the branch:

• Multi-group root cause (bytes vs str): the group_name="Probe00" crash was a byte-string comparison. NWB stores group_name and channel_name as HDF5 byte strings, so the group mask matched zero units. I now decode them to str at the source in _create_df_from_nwb_table, which also fixes the object-dtype channel-id rejection.
• Build laziness, duration: the get_last_spike_time override sat on NwbSortingSegment instead of NwbSortingExtractor, so it never resolved and the placeholder duration fell through to to_spike_vector, reading the full ~130 MB spike array at build. Replaced it with a cheap spike_times[-1] * 2 bound, since the placeholder only needs to loosely bound the timeline.
• Build laziness, count: the group path called count_total_num_spikes() on the UnitsSelectionSorting wrapper, which materializes the parent's whole spike vector (~482 MB). Replaced it with a count read straight from spike_times_index. Together these keep the build lazy: single-probe ~150 MB down to 22 MB, two-probe 482 MB down to 47 MB.
• Init latency: cut the extractor init from ~591 network round-trips to ~35 by not walking the whole file to locate the Units table and the ElectricalSeries. Canonical /units and /acquisition checks, plus defaulting t_start = 0.
• waveform_unit: dropped the plumbing since the schema fixes that attribute to volts and pynwb never writes the intended value (your issue Update readme to point to the latest release docs #2162). Infer the volts-to-microvolts scale from the template magnitude instead, on by default.

Two behavior changes to flag: with a /units table present we no longer raise on multiple Units tables, and an unnamed ElectricalSeries no longer raises since t_start now defaults to 0.

[h-mayorquin]

NeurodataWithoutBorders/pynwb#2162

It was on the code (as a comment but not here)

[samuelgarcia]

This is impressive.
OK for me.

[chrishalcrow]

Suggested change

	t_start_tmp = 0 if t_start is None else t_start
	if electrical_series_path is None:
	t_start_tmp = 0 if t_start is None else t_start
	else:
	t_start_tmp = None

You can't pass an electical_series_path and a t_start. Something like this?

[chrishalcrow]

This try/except tricked me for a bit. I wasn't passing an electrical_series_path so this recording load was failing silently, then there were confusing downstream bugs. Not sure the solution.

[chrishalcrow]

I was trying this out on non-IBL datasets and I think this code makes quite a lot of assumptions about how the data is structured. E.g. all datasets I tried do not store an electrode index in their units table. I've not managed to load anything non-IBL - I failed to load any Allen datasets, but maybe @alejoe91 knows a good one to try.

So I think the current function name is over promising, and we should change it to be more IBL focused. Or maybe shelve this for 0.105 and work to try and make it more general.

[h-mayorquin]

Thanks for trying this, @alejoe91 I am gonna do another pass on Thursday maybe it can make it for release : )

[JoeZiminski]

Thanks @h-mayorquin will check this out today / tomorrow! I am curious on Chris' message and the NWB provenance. Does this load NWB files when they strictly adhere to the NWB format as described here? And the issue is that in many cases data is stored in NWB files but if not according to the exact schema?