[gviejo]

Sorted-array storage for TsGroup (issue #662)

Summary

This PR changes the internal storage of TsGroup. On main, a TsGroup is a dict with one Ts/Tsd for each unit. On this branch, it holds the spikes of all the units in one array, sorted by time:

Array	Content
_times	the timestamps of all the units, sorted (stable for equal times)
_clusters	the key of the unit of each timestamp
_data	the values, if the members are Tsd/TsdFrame/TsdTensor
_is_tsd, _columns	for each unit: whether it holds values, and the shared columns

So a group operation runs once on the merged array, not once for each unit. The PR also adds a lazy TsGroup for data on disk, and it moves the process functions to the new storage.

The branch has 22 commits after dev, including a merge of dev. It changes 29 files: about 7,460 lines added and 1,740 removed.

Main changes

1. TsGroup (core/ts_group.py)

• The constructor merges the units, sorts them once and restricts them once to the time support.
• New cached indices:
  • _cluster_positions: the position of the unit of each spike
  • _ragged_index: the spikes of each unit, in time order, for fast access to one unit
• count, restrict, value_from, to_tsd, get, time_diff and the other methods work on the merged arrays, with numba kernels.
• tsg[k] builds the unit object on each call, from _ragged_index.
• group.data is now a read-only view (_TsGroupDictView).
• The rate metadata is now read-only.
• The docstrings of the public methods and of the class are rewritten.

2. New public entry points

• TsGroup.from_ragged_arrays(...) and TsGroup.from_sorted_arrays(...) build a group from flat arrays.
• If the arrays are on disk (h5py/zarr, checked with utils.is_lazy_array) and lazy=True, they return a lazy group. In all other cases, they return a regular TsGroup.

3. Lazy TsGroup (core/lazy_ts_group.py, private)

• The class _LazyTsGroup has two sources: _RaggedArraySource (NWB style) and _SortedArraySource (for example zarr).
• It reads only the spikes of the selected units, or of the selected time windows.
• A method of the lazy group reads the spikes it needs into a regular TsGroup, then calls the same method on it. The process functions use group._load_in_memory(ep).
• Lazy groups hold no values for now. The documentation will come later.

4. NWB (io/interface_nwb.py)

• nwb['units'] now returns a lazy group that reads the NWB ragged arrays.
• NWBFile.close() now works with pynwb objects.

5. The matching kernel of value_from (core/_jitted_functions.py)

• jitvaluefrom is now the only matching kernel. It does one merge scan for each epoch, and both value_from and compute_tuning_curves call it.
• The tie rules do not change: _valuefrom_match reproduces the rules of the old kernel, and the tests check this against a brute-force oracle.
• The binary search and its threshold are removed for now.

6. Process functions

• Correlograms:
  • Auto, cross (one group or two groups) and event correlograms all use one parallel kernel, _cross_correlograms. It counts with integer times, in units of the time precision.
  • Thus the bins have no float error: a lag exactly on a bin edge always goes into the bin on its right.
  • The lags are calculated from the integer bins, so they are exact.
  • A binsize below the time precision raises a ValueError.
  • The float kernel _cross_correlogram is removed.
• Tuning curves:
  • One pass over the merged spikes: jitvaluefrom matches the spikes, then _jitcount_spikes counts them for each unit and bin.
  • The feature is binned only once.
  • New docstrings with equations and type hints, for compute_tuning_curves, compute_response_per_epoch and compute_mutual_information.
• Perievent, randomize and warping: these functions now use the merged arrays and new numba kernels, and they read lazy groups.

7. Other fixes

• A __getattr__ recursion in the metadata mixin, when the object is not fully built.
• Errors in the TsGroup constructor.

Performance (1000 units × 2,000 spikes, branch against main)

Faster	Speedup
get(a, b)	130×
crosscorrelogram (100 units) / autocorrelogram (100 units)	51–62× / 16×
perievent / warp_tensor	56× / 644×
compute_tuning_curves (1D, 2D, 100 epochs)	40–43×
restrict / value_from	16× / 16×
nwb['units']	11×
count / to_tsd / time_diff	6× / 9× / 4×

Slower	Ratio	Cause
TsGroup(dict)	0.4× (187 ms)	the stable sort of all the spikes (123 ms)
tsg[5 units]	0.1× (8 ms)	_take scans all the spikes with np.isin
tsg[k], loop over units	0.0× (0.015 ms each)	one new object for each call
first count of a lazy NWB group	0.2×	it reads all the spikes and builds the group
jitter_timestamps / resample_timestamps	0.5× / 0.8×	probably the construction sort (not checked)
value_from with very sparse inputs	up to 7×	the merge scan reads the full target

Tests

New test file: test_lazy_ts_group.py. Updated test files: test_ts_group.py, test_correlograms.py, test_jitted.py, test_metadata.py, test_npz_file.py, test_nwb.py, test_perievent.py, test_randomize.py, test_restrict_routing.py, test_tuning_curves.py, test_value_from_kernel.py and test_warping.py. The last full run passed: 6047 passed, 280 skipped.

[codecov]

Codecov Report

❌ Patch coverage is 82.37650% with 132 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pynapple/process/correlograms.py	62.50%	30 Missing and 3 partials ⚠️
pynapple/process/perievent.py	59.42%	28 Missing ⚠️
pynapple/core/lazy_ts_group.py	92.89%	16 Missing and 9 partials ⚠️
pynapple/process/warping.py	37.14%	22 Missing ⚠️
pynapple/process/randomize.py	85.54%	6 Missing and 6 partials ⚠️
pynapple/process/tuning_curves.py	79.66%	9 Missing and 3 partials ⚠️

Files with missing lines	Coverage Δ
pynapple/core/_core_functions.py	91.54% <100.00%> (+1.89%)	⬆️
pynapple/core/base_class.py	98.06% <100.00%> (-0.03%)	⬇️
pynapple/core/metadata_class.py	93.35% <100.00%> (+0.27%)	⬆️
pynapple/core/time_series.py	93.34% <100.00%> (-0.36%)	⬇️
pynapple/core/ts_group.py	91.15% <ø> (-1.86%)	⬇️
pynapple/io/interface_nwb.py	84.34% <100.00%> (+0.56%)	⬆️
pynapple/process/randomize.py	85.00% <85.54%> (+0.38%)	⬆️
pynapple/process/tuning_curves.py	94.00% <79.66%> (-3.02%)	⬇️
pynapple/process/warping.py	71.95% <37.14%> (-28.05%)	⬇️
pynapple/core/lazy_ts_group.py	92.89% <92.89%> (ø)
... and 2 more

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.