Add read_nwb_sorting_analyzer: build a SortingAnalyzer from an NWB Units table - #4645
h-mayorquin wants to merge 46 commits into
Conversation
…load_analyzer_from_nwb
…erface into load_analyzer_from_nwb
# Conflicts: # src/spikeinterface/extractors/nwbextractors.py
for more information, see https://pre-commit.ci
…eberto' into load_analyzer_nwb_heberto
h-mayorquin
changed the title
|
This is failing on |
…ikeinterface into load_analyzer_nwb_heberto
Fixed in last couple of commits |
# Conflicts: # src/spikeinterface/extractors/nwbextractors.py
|
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:
Two behavior changes to flag: with a |
|
NeurodataWithoutBorders/pynwb#2162 It was on the code (as a comment but not here) |
|
This is impressive. |
| print("Could not load recording, proceeding without it") | ||
| recording = None | ||
|
|
||
| t_start_tmp = 0 if t_start is None else t_start |
There was a problem hiding this comment.
| 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?
There was a problem hiding this comment.
should it raise an error up-front if you try?
| if extension_map is not None: | ||
| resolved_extension_map.update(extension_map) | ||
| # try to read recording object to get the analyzer | ||
| try: |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Same I was having some errors detailed below, removing the try/accept lead to somewhat confusing error like: ValueError: Multiple ElectricalSeries found in the file. Please specify the 'electrical_series_path' argument:Available options are: [].
debug code
from dandi.dandiapi import DandiAPIClient
from spikeinterface.extractors import read_nwb_sorting_analyzer
DANDISET = "000253"
VERSION = "0.240923.1441"
with DandiAPIClient() as client:
dandiset = client.get_dandiset(DANDISET, VERSION)
asset = next(a for a in dandiset.get_assets() if a.path.endswith(".nwb"))
s3_url = asset.get_content_url(follow_redirects=1, strip_query=True)
import remfile
import h5py
analyzer = read_nwb_sorting_analyzer(
s3_url,
stream_mode="remfile",
sampling_frequency=30_000.0,
# group_name="18005110031 1-281"
)
|
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. |
|
Thanks for trying this, @alejoe91 I am gonna do another pass on Thursday maybe it can make it for release : ) |
|
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? |
JoeZiminski
left a comment
JoeZiminski
left a comment
There was a problem hiding this comment.
Hey @h-mayorquin this is very cool, is very nice to be able to load a NWB file directly into an analyzer! + was interesting read, some general comments below.
One thing I was unsure about was what exactly the recording is used for. AFAICT It's loaded so it can be attached to the analyzer, and also for the channel_ids. We also need it for t_start but this is through a parallel mechanism (the NWBSortingExtractor) which I found slightly confusing (to allow a single source of truth, can we use the loaded recording (if successful) also for the t_start?).
Even if the recording can be loaded, is it a good idea to expose it through the analyzer? Is the intention that the user can do things like compute("waveforms") using the dummy randomly selected spikes, and/or recompute randomly selected spikes, waveforms, templates etc? If not I would unset the recording from the analzer to make clearer to the user this isn't a true analyzer but instead a handy wrapper for the precomputed data.
On this point more generally (outside the scope of this PR, but for discussion), now with this PR and the kilosort extractor, there are a few cases where an analyzer can be loaded from external sources to produce something close to an analyzer but which doesn't have the same guarantees. I wonder if it would be nice to have this as a distinct construct, like PrecomputedAnalyzer or ImmutableAnalzyer to make it very clear that this thing is not your usual analyzer. For example get_probe() when no recording will return the probe generated from the unit table, but in general a user would expect the full probe geometry as used in the recording. Similarly the extensions e.g. "templates" are set directly on the extension data object without going through a formal API (as it doesn't exist). This is okay but might be brittle going forward in case the analyzer shape changes. This is all outside of scope but just a note that if the pattern of loading external things into analysers in this way continues, we should probably formalise it.
I tried testing on some other random datasets, below are some issues (using this code below):
from dandi.dandiapi import DandiAPIClient
from spikeinterface.extractors import read_nwb_sorting_analyzer
DANDISET = "000469"
VERSION = "0.240123.1806"
with DandiAPIClient() as client:
dandiset = client.get_dandiset(DANDISET, VERSION)
asset = next(a for a in dandiset.get_assets() if a.path.endswith(".nwb"))
s3_url = asset.get_content_url(follow_redirects=1, strip_query=True)
analyzer = read_nwb_sorting_analyzer(
s3_url,
stream_mode="remfile",
# sampling_frequency=30_000.0,
# group_name="18005110031 1-281"
)
Often I had to manually pass the sampling_frequency, I'm not sure why it was failing to detect this on the recordings I tried. For
DANDISET = "000363"
VERSION = "0.231012.2129"
I got
# File "....spikeinterface\src\spikeinterface\extractors\nwbextractors.py", line 2373, in _make_templates
# dense_templates[unit_index][:, positions] = waveform_mean[unit_index][:, :k]
it didn't seem to like that these are 1D waveform means, but i think this is supported by NWB spec and so would make sense to handle.
Similarly for
DANDISET = "000469"
VERSION = "0.240123.1806"
it didn't like that the electrode values were scalar not list, but I think this is supported in NWB spec (I think related issue to above)
Finally with
DANDISET = "000253"
VERSION = "0.240923.1441"
The electrodes couldn't be found and it failed. In this case (any maybe other unsupported cases) it can fail with a message indicating the information it was expecting, but could not find.
| return 0 | ||
| return int(np.max(spikes_in_segment["sample_index"])) | ||
|
|
||
| def get_last_spike_time(self, segment_index: int | None = None) -> float: |
There was a problem hiding this comment.
Is this / the version in NwbSortingSegment used anywhere?
| print("Could not load recording, proceeding without it") | ||
| recording = None | ||
|
|
||
| t_start_tmp = 0 if t_start is None else t_start |
There was a problem hiding this comment.
should it raise an error up-front if you try?
| print("Could not load recording, proceeding without it") | ||
| recording = None | ||
|
|
||
| t_start_tmp = 0 if t_start is None else t_start |
There was a problem hiding this comment.
At this stage if the recording was successfully loaded do we want to read t_start directly from it and set electrical_series_path to None below? (rather than perform it twice)
| load_unit_properties=False, # columns are read deliberately below, into their extensions | ||
| ) | ||
|
|
||
| sorting = sorting_tmp |
There was a problem hiding this comment.
Makes sense to do this but maybe an explanatory comment, could sorting_tmp be called instead sorting_orig or something?
| ) | ||
|
|
||
| sorting = sorting_tmp | ||
| # Recordingless case: leave t_start at 0 (set when the sorting was constructed). NWB spike times are |
There was a problem hiding this comment.
recordingless case (+ when t_start is not passed)?
| positions = [ | ||
| position_of_channel_id[electrode_row_to_channel_id[int(electrode_row)]] | ||
| for electrode_row in region | ||
| if electrode_row_to_channel_id[int(electrode_row)] in position_of_channel_id |
There was a problem hiding this comment.
Is it not the case that if this is not true, something has gone wrong? e.g. a channel id exists in the units table that does not exist on the full probe?
| # `read_nwb_sorting_analyzer`); `None` disables an extension. This is only for real analyzer extensions; | ||
| # sparsity, sorting properties, and recording metadata are handled separately by the reader. | ||
| # | ||
| # The "typed_container" source is a hook: the typed reader currently lives in ndx-spikesorting, so for |
There was a problem hiding this comment.
I initially found this a bit confusing, with the hooks ignored and and unit_locations is computed by default through a different mechanism compute_extra: List[str] | None = ["unit_locations"] but reading this you might anticipate it be read from UnitsMetrics. For clarity could these hooks be removed and implemented together when ndx-spikesorting is supported
| rescale_templates_to_uV: bool = True, | ||
| verbose: bool = False, | ||
| ) -> SortingAnalyzer: | ||
| # extension_map overrides (per extension) merge over DEFAULT_EXTENSION_MAP; see its docstring. |
There was a problem hiding this comment.
Could expand the docstring here, in particular to explan what the recording is used for and when (I think channel_ids, and t_start) and what you can / cannot do with the produced analyzer (e.g. can I compute("waveforms") directly from the recording, given the dummy random_spikes="all", remake templates etc?
| if extension_map is not None: | ||
| resolved_extension_map.update(extension_map) | ||
| # try to read recording object to get the analyzer | ||
| try: |
There was a problem hiding this comment.
Same I was having some errors detailed below, removing the try/accept lead to somewhat confusing error like: ValueError: Multiple ElectricalSeries found in the file. Please specify the 'electrical_series_path' argument:Available options are: [].
debug code
from dandi.dandiapi import DandiAPIClient
from spikeinterface.extractors import read_nwb_sorting_analyzer
DANDISET = "000253"
VERSION = "0.240923.1441"
with DandiAPIClient() as client:
dandiset = client.get_dandiset(DANDISET, VERSION)
asset = next(a for a in dandiset.get_assets() if a.path.endswith(".nwb"))
s3_url = asset.get_content_url(follow_redirects=1, strip_query=True)
import remfile
import h5py
analyzer = read_nwb_sorting_analyzer(
s3_url,
stream_mode="remfile",
sampling_frequency=30_000.0,
# group_name="18005110031 1-281"
)
| check_recordings_equal(recording_backend, recording_pynwb) | ||
|
|
||
|
|
||
| def _make_units_nwb(path, n_units=6, n_ch=8, n_samp=30, with_std=False): |
There was a problem hiding this comment.
Not sure how difficult it would be in the SI testing infrastructure but it would be nice to also test this on some very small real world recordings, across pynwb and the other path, and check against manually verified ground truth (e.g. label and quality metric values, template values).
6 participants
read_nwb_sorting_analyzerbuilds a curatableSortingAnalyzerdirectly 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 haswaveform_mean, the analyzer is built recordingless from those stored templates, plus the per-unit metrics and theelectrodesregion for sparsity; when the file has an accessible ElectricalSeries, it is used as the recording. It mirrorsread_kilosort_as_analyzerin 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
electrodesregion, 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 thecopy_sortingoption (#4668); the recordingless case builds a lightweight placeholder recording from the electroderel_x/rel_ygeometry (viagenerate_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_modeis 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, session6713a4a7-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 thecopy_sortingPR (#4668), which should merge first so this branch rebases down to just the reader.