[kushaangupta]

Fixes #4868

• Add max_fit_batch_memory="64Mi" for PCA fitting. The 64 MiB default is a proposed input-buffer budget, not a measured optimum or a bound on total process RAM.
• Read bounded groups of waveforms through slice_rows() before reshaping. This supports NumPy and Zarr storage and avoids first copying the whole unit.
• Include sparse padding when budgeting reads and exclude padded channels from the PCA rows.
• Also cap the augmented SVD matrix below the 32-bit LAPACK indexing limit, accounting for the previous components and mean-correction row.
• Preserve input order, current small-unit update boundaries, whitening, component count and the existing too-few-rows rule. Balance batch sizes to retain a short remainder; reject a budget that cannot support valid updates.

[JoeZiminski]

Thanks for this @kushaangupta, good catch that's funny that we are hitting the 32-bit indexing limit on these matrices 😂

This is a nice approach but I wonder if it could be more targetted to fix only the bug at hand, rather than extending to provide general memory handling. Fundamentally the main thing we want to avoid is passing too large matrices to partial fit and batch them instead, e.g. I think this could be achieved with something like the below helper:

import numpy as np
from sklearn.utils import gen_batches

def safe_partial_fit(pca, waveforms):
    n_samples, n_features = waveforms.shape

    # LAPACK indexing is limited to 2**31 - 1 elements. Allow for the extra
    # rows IncrementalPCA adds internally before computing the SVD.
    max_rows = np.iinfo(np.int32).max // n_features - (pca.n_components + 1)

    for batch in gen_batches(
        n_samples,
        max_rows,
        min_batch_size=pca.n_components,
    ):
        pca.partial_fit(waveforms[batch])

This is basically the approach IncrementalPCA is taking.

[alejoe91]

Thanks for this @kushaangupta, good catch that's funny that we are hitting the 32-bit indexing limit on these matrices 😂

This is a nice approach but I wonder if it could be more targetted to fix only the bug at hand, rather than extending to provide general memory handling. Fundamentally the main thing we want to avoid is passing too large matrices to partial fit and batch them instead, e.g. I think this could be achieved with something like the below helper:

import numpy as np
from sklearn.utils import gen_batches

def safe_partial_fit(pca, waveforms):
    n_samples, n_features = waveforms.shape

    # LAPACK indexing is limited to 2**31 - 1 elements. Allow for the extra
    # rows IncrementalPCA adds internally before computing the SVD.
    max_rows = np.iinfo(np.int32).max // n_features - (pca.n_components + 1)

    for batch in gen_batches(
        n_samples,
        max_rows,
        min_batch_size=pca.n_components,
    ):
        pca.partial_fit(waveforms[batch])

This is basically the approach IncrementalPCA is taking.

I agree with @JoeZiminski

@kushaangupta what do you think? Can you apply the suggested change?