Implement get_unit_spike_trains and performance improvements#4502
Open
alejoe91 wants to merge 24 commits intoSpikeInterface:mainfrom
Open
Implement get_unit_spike_trains and performance improvements#4502alejoe91 wants to merge 24 commits intoSpikeInterface:mainfrom
get_unit_spike_trains and performance improvements#4502alejoe91 wants to merge 24 commits intoSpikeInterface:mainfrom
Conversation
alejoe91
added
core
- Drop unused `return_times` parameter from get_unit_spike_trains_in_seconds - Clean up stale/truncated docstrings on get_unit_spike_train_in_seconds, get_unit_spike_trains, and get_unit_spike_trains_in_seconds - Fix UnitsSelectionSortingSegment.get_unit_spike_trains to re-key the returned dict with child unit ids (was returning parent-keyed dict, breaking whenever renamed_unit_ids differ from parent ids) - Fix test_get_unit_spike_trains: drop unused return_times kwarg, remove unused local variable, fix assertion.
The previous check `np.diff(self.ids_to_indices(self._renamed_unit_ids)).min() < 0`
was never `True`, because `ids_to_indices(self._renamed_unit_ids)` on a USS
always returns `[0, 1, ..., k-1]` (since `_main_ids == _renamed_unit_ids`), so the
diff was always positive and the lexsort branch was unreachable. Therefore the
cached spike vector was wrong whenever two units had co-temporal spikes and the
selection reordered them relative to the parent.
Replaced with a two-step check that attempt to avoid unneccessary lexsorts:
1. O(k) `_is_order_preserving_selection()` -- Checks if USS `._unit_ids` is
in the same relative order as in the parent. When True, the remapped vector
is guaranteed sorted (boolean filtering preserves order; the remap only
relabels unit_index values). This is the common case via `select_units()`
with a boolean mask.
2. O(n) `_is_spike_vector_sorted()` -- Checks if the remapped vector is still
sorted by (segment, sample, unit). Catches the case where the selection is
not order-preserving but no co-temporal (same exact sample) spikes exist.
Falls back to the original O(n log n) lexsort only when both checks fail.
`BaseSorting` builds the spike vector with a per-unit boolean scan over spike_clusters, which is (O(N*K)). If we already have the full flat spike time and spike cluster arrays, we can do a lot better by building the spike vector in one shot. (I think O(N log N) from the lexsort, which is also pessimistic, because the lexsort doesn't always need to happen. Under any circumstances I can dream of, K >> log N.) Since Phy/Kilosort segments already load the full flat arrays when the `PhyKilosortSorting` object is created, and keep them around as `._all_spikes` and `._all_clusters`, we can just use those! :) Also populates `_cached_spike_vector_segment_slices` directly, so that `BaseSorting`'s `_get_spike_vector_segment_slices()` lazy recomputation is skipped.
`BaseSortingSegment.get_unit_spike_trains()` loops over `get_unit_spike_train`, which is O(N*K) because each call is a boolean scan over _all_clusters/_all_spikes. If we know we are going to be getting all the trains, we can do it much faster. And if we can use numba, even faster still. In fact, even if we only want _some_ spike trains, it is still often faster to get all the trains and just discard the ones we don't need, than to get only the trains we need do unit-by-unit (because we only ever store or cache flat arrays of spike times/clusters). Note that **only the use_cache=False path is affected**; the use_cache=True triggers the computation of the spike vector, which I don't think can ever be the most efficient way to get spike trains.
…izations - Fixed test_compute_and_cache_spike_vector: was comparing an array to itself (to_spike_vector use_cache=False still returns the cached vector). Now explicitly calls the USS override and the BaseSorting implementation, and compares the two. - Added test_uss_get_unit_spike_trains_with_renamed_ids: also not a test of the optimization commits per se, but would have caught a mistake made along the way. Verifies get_unit_spike_trains returns child-keyed dicts (not parent-keyed). - Added test_spike_vector_sorted_after_reorder_with_cotemporal_spikes: verifies the USS spike vector is correctly sorted when the selection reverses unit order and co-temporal spikes exist. - Added test_phy_sorting_segment_get_unit_spike_trains: validates the new fast methods on PhySortingSegment. - Added test_phy_compute_and_cache_spike_vector: verifies the Phy override of _compute_and_cache_spike_vector matches BaseSorting implementation.
Contributor
|
@alejoe91 my changes PR'd to your fork whenever you're ready. The only thing I should point out that isn't in the commit messages: |
Additions to get_unit_spike_trains PR
for more information, see https://pre-commit.ci
Collaborator
|
I am curios on what prompted this? What profiling did you guys do? Any chance that we have a discussion here on the repo at least to know what where the performance benchmarks, reason and validation. |
Member
Author
|
@grahamfindlay is doing very long chronic recordings. He does all the processing and at a second iteration wants to load the phy sorting object, select some units, and get all the spike trains. Just caching the spike vector takes almost 4 minutes! Plus there were some additional lexsort that can be avoided and speed up computation. At least to give some context @h-mayorquin @grahamfindlay maybe you can add some more details on benchmarks and profiling? |
Contributor
|
Here are example timings for various operations using 1 example subject. This subject only has ~400 million spikes - I have some with many more. FWIW, you shouldn't need long chronic recordings to see tangible improvements from most of these changes. I must dig through notes but tested with 100M spikes and they were still clear gains. "Parent before" = The The two layers of
Comments:
There are rough edges with this PR I know about:
Another question I haven't resolved:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
use_cache(toget_unit_spike_train_in_seconds)to_reordered_spike_vectorselect_units@grahamfindlay
TODO
get_unit_spike_trainsforPhyKilosortSortingExtractor(maybe in follow up)