diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 225b37c9f0..6af397bcb5 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -1,4 +1,5 @@ import warnings +from contextlib import ExitStack import platform from pathlib import Path from tqdm.auto import tqdm @@ -10,7 +11,7 @@ import numpy as np from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension -from spikeinterface.core.core_tools import slice_rows, materialize_array +from spikeinterface.core.core_tools import slice_rows, materialize_array, convert_string_to_bytes from spikeinterface.core.job_tools import TimeSeriesChunkExecutor, _shared_job_kwargs_doc, fix_job_kwargs from spikeinterface.core.analyzer_extension_core import _inplace_sparse_realign_waveforms, _select_channels_sparse_data @@ -38,6 +39,13 @@ class ComputePrincipalComponents(AnalyzerExtension): If True, waveforms are pre-whitened dtype : dtype, default: "float32" Dtype of the pc scores + max_fit_batch_memory : str, default: "64Mi" + Maximum size of each waveform input buffer when fitting any PCA mode. + Uses memory strings such as "64Mi" or "100M". This is not a limit on total process memory: + reshaping, PCA, concurrent workers and storage can require additional buffers and workspace. + No additional waveform subsampling is performed; unit, spike and channel order are retained. + Units exceeding this limit are fitted in multiple incremental updates, which can change + the approximate PCA basis compared with a single update. {} Examples @@ -72,8 +80,11 @@ def _set_params( mode="by_channel_local", whiten=True, dtype="float32", + max_fit_batch_memory="64Mi", ): assert mode in _possible_modes, "Invalid mode!" + if convert_string_to_bytes(max_fit_batch_memory) <= 0: + raise ValueError("max_fit_batch_memory must be positive") # the sparsity in params is ONLY the injected sparsity and not the sorting_analyzer one params = dict( @@ -81,6 +92,7 @@ def _set_params( mode=mode, whiten=whiten, dtype=np.dtype(dtype), + max_fit_batch_memory=max_fit_batch_memory, ) return params @@ -444,103 +456,109 @@ def run_for_all_spikes(self, file_path=None, verbose=False, **job_kwargs): ) processor.run() + def _iter_fit_waveform_batches(self, progress_bar): + """Read bounded inputs and account for each mode's SVD dimensions.""" + p = self.params + analyzer = self.sorting_analyzer + waveforms = analyzer.get_extension("waveforms").data["waveforms"] + spikes = analyzer.get_extension("random_spikes").get_random_spikes() + sparsity = analyzer.sparsity + num_samples, max_num_channels = waveforms.shape[1:] + if max_num_channels == 0: + return + batch_memory = convert_string_to_bytes(p.get("max_fit_batch_memory", "64Mi")) + # slice_rows reads all packed channels, including any padding for sparse units. + waveforms_per_batch = batch_memory // (num_samples * max_num_channels * waveforms.dtype.itemsize) + units_loop = enumerate(analyzer.unit_ids) + if progress_bar: + units_loop = tqdm(units_loop, desc="Fitting PCA", total=len(analyzer.unit_ids)) + + for unit_ind, unit_id in units_loop: + spike_indices = np.flatnonzero(spikes["unit_index"] == unit_ind) + channel_inds = ( + np.arange(max_num_channels) if sparsity is None else sparsity.unit_id_to_channel_indices[unit_id] + ) + num_channels = channel_inds.size + if num_channels == 0: + continue + rows_per_waveform = num_channels if p["mode"] == "by_channel_global" else 1 + features_per_row = num_samples * num_channels if p["mode"] == "concatenated" else num_samples + if spike_indices.size * rows_per_waveform < p["n_components"]: + continue + # IPCA stacks the previous components and a mean-correction row before its SVD. + max_svd_rows = np.iinfo(np.int32).max // features_per_row - p["n_components"] - 1 + batch_size = min(waveforms_per_batch, max_svd_rows // rows_per_waveform) + min_batch_size = (p["n_components"] + rows_per_waveform - 1) // rows_per_waveform + if batch_size < min_batch_size: + raise ValueError("max_fit_batch_memory is too small to fit n_components waveform rows") + num_batches = (spike_indices.size + batch_size - 1) // batch_size + if spike_indices.size // num_batches < min_batch_size: + raise ValueError("Increase max_fit_batch_memory to avoid batches with fewer than n_components rows") + # Balance sizes so a short final batch is retained, with enough rows for PCA. + for batch_ind in range(num_batches): + start = batch_ind * spike_indices.size // num_batches + stop = (batch_ind + 1) * spike_indices.size // num_batches + wfs = slice_rows(waveforms, spike_indices[start:stop])[:, :, :num_channels] + yield wfs, channel_inds + del wfs + def _fit_by_channel_local(self, n_jobs, progress_bar, max_threads_per_worker, mp_context): from sklearn.decomposition import IncrementalPCA p = self.params - if mp_context is not None and platform.system() == "Windows": assert mp_context != "fork", "'fork' mp_context not supported on Windows!" elif mp_context == "fork" and platform.system() == "Darwin": warnings.warn('As of Python 3.8 "fork" is no longer considered safe on macOS') - unit_ids = self.sorting_analyzer.unit_ids channel_ids = self.sorting_analyzer.channel_ids - # there is one PCA per channel for independent fit per channel pca_models = [IncrementalPCA(n_components=p["n_components"], whiten=p["whiten"]) for _ in channel_ids] - - # fit - units_loop = enumerate(unit_ids) - if progress_bar: - units_loop = tqdm(units_loop, desc="Fitting PCA", total=len(unit_ids)) - - for unit_ind, unit_id in units_loop: - wfs, channel_inds, _ = self._get_sparse_waveforms(unit_id) - if len(wfs) < p["n_components"]: - continue - if n_jobs in (0, 1): - for wf_ind, chan_ind in enumerate(channel_inds): - pca = pca_models[chan_ind] - pca.partial_fit(wfs[:, :, wf_ind]) - else: - # create list of args to parallelize. For convenience, the max_threads_per_worker is passed - # as last argument - items = [ - (chan_ind, pca_models[chan_ind], wfs[:, :, wf_ind], max_threads_per_worker) - for wf_ind, chan_ind in enumerate(channel_inds) - ] - n_jobs = min(n_jobs, len(items)) - - with ProcessPoolExecutor( - max_workers=n_jobs, - mp_context=mp.get_context(mp_context), - ) as executor: + # Reuse workers across batches; each channel's updates remain sequential. + with ExitStack() as stack: + executor = None + if n_jobs not in (0, 1) and channel_ids.size: + executor = stack.enter_context( + ProcessPoolExecutor( + max_workers=min(n_jobs, channel_ids.size), mp_context=mp.get_context(mp_context) + ) + ) + for wfs, channel_inds in self._iter_fit_waveform_batches(progress_bar): + if executor is None: + for wf_ind, chan_ind in enumerate(channel_inds): + pca_models[chan_ind].partial_fit(wfs[:, :, wf_ind]) + else: + items = [ + (chan_ind, pca_models[chan_ind], wfs[:, :, wf_ind], max_threads_per_worker) + for wf_ind, chan_ind in enumerate(channel_inds) + ] results = executor.map(_partial_fit_one_channel, items) for chan_ind, pca_model_updated in results: pca_models[chan_ind] = pca_model_updated - + del items, results + del wfs return pca_models def _fit_by_channel_global(self, progress_bar): - p = self.params - unit_ids = self.sorting_analyzer.unit_ids - - # there is one unique PCA across channels from sklearn.decomposition import IncrementalPCA + p = self.params pca_model = IncrementalPCA(n_components=p["n_components"], whiten=p["whiten"]) - - # fit - units_loop = enumerate(unit_ids) - if progress_bar: - units_loop = tqdm(units_loop, desc="Fitting PCA", total=len(unit_ids)) - - # with 'by_channel_global' we can't parallelize over channels - for unit_ind, unit_id in units_loop: - wfs, _, _ = self._get_sparse_waveforms(unit_id) - shape = wfs.shape - if shape[0] * shape[2] < p["n_components"]: - continue - # avoid loop with reshape - wfs_concat = wfs.transpose(0, 2, 1).reshape(shape[0] * shape[2], shape[1]) + for wfs, _ in self._iter_fit_waveform_batches(progress_bar): + wfs_concat = wfs.transpose(0, 2, 1).reshape(-1, wfs.shape[1]) pca_model.partial_fit(wfs_concat) - + del wfs, wfs_concat return pca_model def _fit_concatenated(self, progress_bar): - - p = self.params - unit_ids = self.sorting_analyzer.unit_ids - - assert self.sorting_analyzer.sparsity is None, "For mode 'concatenated' waveforms need to be dense" - - # there is one unique PCA across channels from sklearn.decomposition import IncrementalPCA + assert self.sorting_analyzer.sparsity is None, "For mode 'concatenated' waveforms need to be dense" + p = self.params pca_model = IncrementalPCA(n_components=p["n_components"], whiten=p["whiten"]) - - # fit - units_loop = enumerate(unit_ids) - if progress_bar: - units_loop = tqdm(units_loop, desc="Fitting PCA", total=len(unit_ids)) - - for unit_ind, unit_id in units_loop: - wfs, _, _ = self._get_sparse_waveforms(unit_id) + for wfs, _ in self._iter_fit_waveform_batches(progress_bar): wfs_flat = wfs.reshape(wfs.shape[0], -1) - if len(wfs_flat) < p["n_components"]: - continue pca_model.partial_fit(wfs_flat) - + del wfs, wfs_flat return pca_model def _transform_waveforms(self, spikes, waveforms, pca_model, progress_bar): diff --git a/src/spikeinterface/postprocessing/tests/test_principal_component.py b/src/spikeinterface/postprocessing/tests/test_principal_component.py index 879f4ff055..160fd51894 100644 --- a/src/spikeinterface/postprocessing/tests/test_principal_component.py +++ b/src/spikeinterface/postprocessing/tests/test_principal_component.py @@ -14,6 +14,8 @@ class TestPrincipalComponentsExtension(AnalyzerExtensionCommonTestSuite): [ dict(mode="by_channel_local"), dict(mode="by_channel_global"), + dict(mode="by_channel_global", max_fit_batch_memory="32Ki"), + dict(mode="by_channel_local", max_fit_batch_memory="32Ki"), # mode concatenated cannot be tested here because it do not work with sparse=True ], ) @@ -246,3 +248,153 @@ def test_select_channels_sparse_pca(): if __name__ == "__main__": test = TestPrincipalComponentsExtension() test.test_get_projections(sparse=True) + + +def _make_pca_batch_analyzer(sparse): + from spikeinterface.core import ChannelSparsity, NumpyRecording, NumpySorting + + rng = np.random.default_rng(42) + recording = NumpyRecording( + [rng.normal(size=(400, 4)).astype("float32"), rng.normal(size=(200, 4)).astype("float32")], 1000.0 + ) + recording.set_dummy_probe_from_locations([[0, 0], [20, 0], [40, 0], [60, 0]]) + sorting = NumpySorting.from_unit_dict( + [ + {10: np.arange(20, 200, 10), 30: np.array([250]), 90: np.arange(210, 350, 20)}, + {10: np.arange(20, 100, 10), 30: np.array([], dtype="int64"), 90: np.array([], dtype="int64")}, + ], + 1000.0, + ) + sparsity = ( + ChannelSparsity( + np.array([[1, 1, 0, 1], [0, 1, 0, 0], [0, 0, 1, 0]], dtype=bool), sorting.unit_ids, recording.channel_ids + ) + if sparse + else None + ) + analyzer = create_sorting_analyzer( + sorting, recording, sparse=False, sparsity=sparsity, main_channel_indices=np.array([0, 1, 2]) + ) + analyzer.compute( + {"random_spikes": {"method": "all"}, "waveforms": {"ms_before": 3, "ms_after": 4}}, n_jobs=1, progress_bar=False + ) + return analyzer + + +@pytest.mark.parametrize( + "mode, sparse", + [ + ("by_channel_global", False), + ("by_channel_global", True), + ("by_channel_local", False), + ("by_channel_local", True), + ("concatenated", False), + ], +) +@pytest.mark.parametrize("storage", ["numpy", "zarr"]) +def test_pca_fit_batches(tmp_path, monkeypatch, mode, sparse, storage): + from sklearn.decomposition import IncrementalPCA + import spikeinterface.postprocessing.principal_component as pc_module + + analyzer = _make_pca_batch_analyzer(sparse) + waveforms_ext = analyzer.get_extension("waveforms") + waveforms = waveforms_ext.data["waveforms"] + if storage == "zarr": + import zarr + + stored = zarr.open_array( + str(tmp_path / "waveforms.zarr"), + mode="w", + shape=waveforms.shape, + dtype=waveforms.dtype, + chunks=(4, *waveforms.shape[1:]), + ) + stored[:] = waveforms + waveforms_ext.data["waveforms"] = stored + ext = ComputePrincipalComponents(analyzer) + ext.set_params(n_components=3, mode=mode) + keys = range(len(analyzer.channel_ids)) if mode == "by_channel_local" else [0] + original = {key: IncrementalPCA(n_components=3, whiten=True) for key in keys} + expected_inputs = {key: [] for key in keys} + for uid in analyzer.unit_ids: + wfs, channel_inds, _ = ext._get_sparse_waveforms(uid) + if mode == "by_channel_global": + inputs = {0: wfs.transpose(0, 2, 1).reshape(-1, wfs.shape[1])} + elif mode == "concatenated": + inputs = {0: wfs.reshape(len(wfs), -1)} + else: + inputs = {channel: wfs[:, :, position] for position, channel in enumerate(channel_inds)} + for key, data in inputs.items(): + if len(data) >= 3: + expected_inputs[key].append(data) + original[key].partial_fit(data) + + def fit(): + if mode == "by_channel_local": + models = ext._fit_by_channel_local(n_jobs=1, progress_bar=False, max_threads_per_worker=1, mp_context=None) + return dict(enumerate(models)) + if mode == "by_channel_global": + return {0: ext._fit_by_channel_global(progress_bar=False)} + return {0: ext._fit_concatenated(progress_bar=False)} + + # Units below the default budget keep exactly the previous fitting updates. + for key, model in fit().items(): + np.testing.assert_array_equal(model.components_, original[key].components_) + np.testing.assert_array_equal(model.singular_values_, original[key].singular_values_) + + batch_inputs, read_sizes = {}, [] + real_fit, real_slice = IncrementalPCA.partial_fit, pc_module.slice_rows + + def partial_fit(self, data, **kwargs): + batch_inputs.setdefault(id(self), []).append(data.copy()) + assert data.nbytes <= 512 + assert len(data) >= 3 + return real_fit(self, data, **kwargs) + + def slice_rows(array, indices): + selected = real_slice(array, indices) + read_sizes.append(selected.nbytes) + assert selected.nbytes <= 512, "A whole unit was read before batching" + return selected + + monkeypatch.setattr(IncrementalPCA, "partial_fit", partial_fit) + monkeypatch.setattr(pc_module, "slice_rows", slice_rows) + ext.params["max_fit_batch_memory"] = "0.5Ki" + fitted = fit() + assert any(len(batch_inputs[id(model)]) > len(expected_inputs[key]) for key, model in fitted.items()) + for key, model in fitted.items(): + expected = np.concatenate(expected_inputs[key]) + np.testing.assert_array_equal(np.concatenate(batch_inputs[id(model)]), expected) + assert model.n_samples_seen_ == len(expected) + np.testing.assert_allclose(model.mean_, np.mean(expected, axis=0), atol=1e-6) + assert read_sizes + + +def test_local_pca_fit_batches_match_between_serial_and_spawn(): + analyzer = _make_pca_batch_analyzer(sparse=True) + params = dict(n_components=3, mode="by_channel_local", max_fit_batch_memory="0.5Ki") + serial = analyzer.compute("principal_components", **params, n_jobs=1, progress_bar=False) + expected = serial.get_data().copy() + expected_models = [model.components_.copy() for model in serial.get_pca_model()] + parallel = analyzer.compute( + "principal_components", **params, n_jobs=2, mp_context="spawn", max_threads_per_worker=1, progress_bar=False + ) + np.testing.assert_allclose(parallel.get_data(), expected, atol=1e-6) + for model, components in zip(parallel.get_pca_model(), expected_models): + np.testing.assert_allclose(model.components_, components, atol=1e-6) + + +def test_pca_fit_batch_memory_validation(): + analyzer = _make_pca_batch_analyzer(sparse=True) + ext = ComputePrincipalComponents(analyzer) + for value in ("0Ki", "-1Ki"): + with pytest.raises(ValueError, match="max_fit_batch_memory must be positive"): + ext.set_params(mode="by_channel_global", max_fit_batch_memory=value) + ext.set_params(n_components=3, mode="by_channel_global", max_fit_batch_memory="0.001Ki") + with pytest.raises(ValueError, match="too small"): + ext._fit_by_channel_global(progress_bar=False) + # Seven one-channel waveforms cannot be partitioned into batches of at most + # three waveforms with at least three rows each; report an actionable error. + ext.set_params(n_components=3, mode="by_channel_global", max_fit_batch_memory="0.25Ki") + with pytest.raises(ValueError, match="Increase max_fit_batch_memory"): + ext._fit_by_channel_global(progress_bar=False)