Fit PCA in batches to overcome LAPACK memory overflow

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,15 +80,19 @@ 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(
             n_components=n_components,
             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):

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)