From 4ca51710ddf3f4516a364e90f1de82a59ded2f3d Mon Sep 17 00:00:00 2001 From: Kushaan Gupta Date: Wed, 14 Oct 2026 11:25:16 -0400 Subject: [PATCH 1/2] Bound global PCA fitting batches to avoid LAPACK index overflow Add a configurable 64 MiB waveform input budget for by_channel_global fitting, read batches before reshaping, and cap augmented SVD dimensions below the 32-bit LAPACK limit. Preserve input ordering, sparse channel handling and existing update boundaries for small units. Document numerical implications and cover NumPy/Zarr reads, complete input coverage and extension round trips. --- .../postprocessing/principal_component.py | 49 ++++++-- .../tests/test_principal_component.py | 111 ++++++++++++++++++ 2 files changed, 153 insertions(+), 7 deletions(-) diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 225b37c9f0..5e7ef232bf 100644 --- a/src/spikeinterface/postprocessing/principal_component.py +++ b/src/spikeinterface/postprocessing/principal_component.py @@ -10,7 +10,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 +38,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 in "by_channel_global" mode. + Uses memory strings such as "64Mi" or "100M". This is not a limit on total process memory: + reshaping, PCA and the storage backend 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. Other modes ignore this limit. {} Examples @@ -72,8 +79,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 +91,7 @@ def _set_params( mode=mode, whiten=whiten, dtype=np.dtype(dtype), + max_fit_batch_memory=max_fit_batch_memory, ) return params @@ -505,15 +516,39 @@ def _fit_by_channel_global(self, progress_bar): if progress_bar: units_loop = tqdm(units_loop, desc="Fitting PCA", total=len(unit_ids)) + waveforms = self.sorting_analyzer.get_extension("waveforms").data["waveforms"] + spikes = self.sorting_analyzer.get_extension("random_spikes").get_random_spikes() + sparsity = self.sorting_analyzer.sparsity + num_samples, max_num_channels = waveforms.shape[1:] + if max_num_channels == 0: + return pca_model + 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) + # IPCA stacks the previous components and a mean-correction row before its SVD. + max_svd_rows = np.iinfo(np.int32).max // num_samples - p["n_components"] - 1 + # 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"]: + spike_indices = np.flatnonzero(spikes["unit_index"] == unit_ind) + num_channels = max_num_channels if sparsity is None else sparsity.unit_id_to_channel_indices[unit_id].size + if spike_indices.size * num_channels < p["n_components"]: continue - # avoid loop with reshape - wfs_concat = wfs.transpose(0, 2, 1).reshape(shape[0] * shape[2], shape[1]) - pca_model.partial_fit(wfs_concat) + batch_size = min(waveforms_per_batch, max_svd_rows // num_channels) + min_batch_size = (p["n_components"] + num_channels - 1) // num_channels + 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] + wfs_concat = wfs.transpose(0, 2, 1).reshape(-1, num_samples) + pca_model.partial_fit(wfs_concat) + del wfs, wfs_concat return pca_model diff --git a/src/spikeinterface/postprocessing/tests/test_principal_component.py b/src/spikeinterface/postprocessing/tests/test_principal_component.py index 879f4ff055..02e171a203 100644 --- a/src/spikeinterface/postprocessing/tests/test_principal_component.py +++ b/src/spikeinterface/postprocessing/tests/test_principal_component.py @@ -14,6 +14,7 @@ class TestPrincipalComponentsExtension(AnalyzerExtensionCommonTestSuite): [ dict(mode="by_channel_local"), dict(mode="by_channel_global"), + dict(mode="by_channel_global", max_fit_batch_memory="32Ki"), # mode concatenated cannot be tested here because it do not work with sparse=True ], ) @@ -246,3 +247,113 @@ def test_select_channels_sparse_pca(): if __name__ == "__main__": test = TestPrincipalComponentsExtension() test.test_get_projections(sparse=True) + + +def _make_global_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("sparse", [False, True]) +@pytest.mark.parametrize("storage", ["numpy", "zarr"]) +def test_global_pca_fit_batches(tmp_path, monkeypatch, sparse, storage): + from sklearn.decomposition import IncrementalPCA + import spikeinterface.postprocessing.principal_component as pc_module + + analyzer = _make_global_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="by_channel_global") + expected_inputs = [] + original = IncrementalPCA(n_components=3, whiten=True) + for uid in analyzer.unit_ids: + wfs, _, _ = ext._get_sparse_waveforms(uid) + flat = wfs.transpose(0, 2, 1).reshape(-1, wfs.shape[1]) + if len(flat) >= 3: + expected_inputs.append(flat) + original.partial_fit(flat) + # Units below the default budget keep exactly the previous fitting updates. + unchanged = ext._fit_by_channel_global(progress_bar=False) + np.testing.assert_array_equal(unchanged.components_, original.components_) + np.testing.assert_array_equal(unchanged.singular_values_, original.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.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 = ext._fit_by_channel_global(progress_bar=False) + expected = np.concatenate(expected_inputs) + np.testing.assert_array_equal(np.concatenate(batch_inputs), expected) + assert len(batch_inputs) > len(expected_inputs) + assert fitted.n_samples_seen_ == len(expected) + np.testing.assert_allclose(fitted.mean_, np.mean(expected, axis=0), atol=1e-6) + assert read_sizes + + +def test_global_pca_fit_batch_memory_validation(): + analyzer = _make_global_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) From f3481cc9ca26b23da28371d670bc61face382e62 Mon Sep 17 00:00:00 2001 From: Kushaan Gupta Date: Wed, 14 Oct 2026 11:51:57 -0400 Subject: [PATCH 2/2] Bound local and concatenated PCA updates as well as global fitting Share waveform batching with mode-specific PCA row and feature counts. Preserve each channel model update order and reuse local workers across bounded batches. Cover dense and sparse NumPy/Zarr fitting across all modes, and compare serial with spawn-worker local results. --- .../postprocessing/principal_component.py | 167 ++++++++---------- .../tests/test_principal_component.py | 91 +++++++--- 2 files changed, 141 insertions(+), 117 deletions(-) diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py index 5e7ef232bf..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 @@ -39,12 +40,12 @@ class ComputePrincipalComponents(AnalyzerExtension): 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 in "by_channel_global" mode. + 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 and the storage backend can require additional buffers and workspace. + 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. Other modes ignore this limit. + the approximate PCA basis compared with a single update. {} Examples @@ -455,87 +456,39 @@ def run_for_all_spikes(self, file_path=None, verbose=False, **job_kwargs): ) processor.run() - def _fit_by_channel_local(self, n_jobs, progress_bar, max_threads_per_worker, mp_context): - from sklearn.decomposition import IncrementalPCA - + def _iter_fit_waveform_batches(self, progress_bar): + """Read bounded inputs and account for each mode's SVD dimensions.""" 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: - results = executor.map(_partial_fit_one_channel, items) - for chan_ind, pca_model_updated in results: - pca_models[chan_ind] = pca_model_updated - - 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 - - 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)) - - waveforms = self.sorting_analyzer.get_extension("waveforms").data["waveforms"] - spikes = self.sorting_analyzer.get_extension("random_spikes").get_random_spikes() - sparsity = self.sorting_analyzer.sparsity + 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 pca_model + 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) - # IPCA stacks the previous components and a mean-correction row before its SVD. - max_svd_rows = np.iinfo(np.int32).max // num_samples - p["n_components"] - 1 + units_loop = enumerate(analyzer.unit_ids) + if progress_bar: + units_loop = tqdm(units_loop, desc="Fitting PCA", total=len(analyzer.unit_ids)) - # with 'by_channel_global' we can't parallelize over channels for unit_ind, unit_id in units_loop: spike_indices = np.flatnonzero(spikes["unit_index"] == unit_ind) - num_channels = max_num_channels if sparsity is None else sparsity.unit_id_to_channel_indices[unit_id].size - if spike_indices.size * num_channels < p["n_components"]: + 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 - batch_size = min(waveforms_per_batch, max_svd_rows // num_channels) - min_batch_size = (p["n_components"] + num_channels - 1) // num_channels + 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 @@ -546,36 +499,66 @@ def _fit_by_channel_global(self, progress_bar): 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] - wfs_concat = wfs.transpose(0, 2, 1).reshape(-1, num_samples) - pca_model.partial_fit(wfs_concat) - del wfs, wfs_concat + yield wfs, channel_inds + del wfs - return pca_model - - def _fit_concatenated(self, progress_bar): + def _fit_by_channel_local(self, n_jobs, progress_bar, max_threads_per_worker, mp_context): + from sklearn.decomposition import IncrementalPCA p = self.params - unit_ids = self.sorting_analyzer.unit_ids + 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') - assert self.sorting_analyzer.sparsity is None, "For mode 'concatenated' waveforms need to be dense" + channel_ids = self.sorting_analyzer.channel_ids + pca_models = [IncrementalPCA(n_components=p["n_components"], whiten=p["whiten"]) for _ in channel_ids] + # 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 - # there is one unique PCA across channels + def _fit_by_channel_global(self, progress_bar): from sklearn.decomposition import IncrementalPCA + p = self.params pca_model = IncrementalPCA(n_components=p["n_components"], whiten=p["whiten"]) + 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 - # fit - units_loop = enumerate(unit_ids) - if progress_bar: - units_loop = tqdm(units_loop, desc="Fitting PCA", total=len(unit_ids)) + def _fit_concatenated(self, progress_bar): + from sklearn.decomposition import IncrementalPCA - for unit_ind, unit_id in units_loop: - wfs, _, _ = self._get_sparse_waveforms(unit_id) + 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"]) + 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 02e171a203..160fd51894 100644 --- a/src/spikeinterface/postprocessing/tests/test_principal_component.py +++ b/src/spikeinterface/postprocessing/tests/test_principal_component.py @@ -15,6 +15,7 @@ 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 ], ) @@ -249,7 +250,7 @@ def test_select_channels_sparse_pca(): test.test_get_projections(sparse=True) -def _make_global_pca_batch_analyzer(sparse): +def _make_pca_batch_analyzer(sparse): from spikeinterface.core import ChannelSparsity, NumpyRecording, NumpySorting rng = np.random.default_rng(42) @@ -280,13 +281,22 @@ def _make_global_pca_batch_analyzer(sparse): return analyzer -@pytest.mark.parametrize("sparse", [False, True]) +@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_global_pca_fit_batches(tmp_path, monkeypatch, sparse, storage): +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_global_pca_batch_analyzer(sparse) + analyzer = _make_pca_batch_analyzer(sparse) waveforms_ext = analyzer.get_extension("waveforms") waveforms = waveforms_ext.data["waveforms"] if storage == "zarr": @@ -302,25 +312,41 @@ def test_global_pca_fit_batches(tmp_path, monkeypatch, sparse, storage): stored[:] = waveforms waveforms_ext.data["waveforms"] = stored ext = ComputePrincipalComponents(analyzer) - ext.set_params(n_components=3, mode="by_channel_global") - expected_inputs = [] - original = IncrementalPCA(n_components=3, whiten=True) + 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, _, _ = ext._get_sparse_waveforms(uid) - flat = wfs.transpose(0, 2, 1).reshape(-1, wfs.shape[1]) - if len(flat) >= 3: - expected_inputs.append(flat) - original.partial_fit(flat) + 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. - unchanged = ext._fit_by_channel_global(progress_bar=False) - np.testing.assert_array_equal(unchanged.components_, original.components_) - np.testing.assert_array_equal(unchanged.singular_values_, original.singular_values_) + 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 = [], [] + batch_inputs, read_sizes = {}, [] real_fit, real_slice = IncrementalPCA.partial_fit, pc_module.slice_rows def partial_fit(self, data, **kwargs): - batch_inputs.append(data.copy()) + batch_inputs.setdefault(id(self), []).append(data.copy()) assert data.nbytes <= 512 assert len(data) >= 3 return real_fit(self, data, **kwargs) @@ -334,17 +360,32 @@ def slice_rows(array, indices): monkeypatch.setattr(IncrementalPCA, "partial_fit", partial_fit) monkeypatch.setattr(pc_module, "slice_rows", slice_rows) ext.params["max_fit_batch_memory"] = "0.5Ki" - fitted = ext._fit_by_channel_global(progress_bar=False) - expected = np.concatenate(expected_inputs) - np.testing.assert_array_equal(np.concatenate(batch_inputs), expected) - assert len(batch_inputs) > len(expected_inputs) - assert fitted.n_samples_seen_ == len(expected) - np.testing.assert_allclose(fitted.mean_, np.mean(expected, axis=0), atol=1e-6) + 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_global_pca_fit_batch_memory_validation(): - analyzer = _make_global_pca_batch_analyzer(sparse=True) +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"):