scbutterfly: replace dense TF-IDF and cut redundant copies - #71
Merged
VladimirShitov merged 1 commit intoSep 23, 2026
Merged
Conversation
scButterfly's TFIDF tiles the per-cell and per-peak totals into dense (n_peaks, n_cells) float64 matrices and divides densely, so it costs O(n_peaks * n_cells) regardless of how sparse the input is -- 37 GiB per array at NeurIPS2021 Multiome scale, with three or four alive at once. That is what drove the ~265 GB peak and the OOM on the 2022 Multiome datasets. Patch it with a sparse O(nnz) equivalent, and trim the wrapper's own avoidable allocations: - sparse TF-IDF monkeypatch: same arithmetic in the same order, on the CSR .data array. Bit-identical, since scButterfly binarizes first and the totals are exact integers in float32. - _placeholder_block gathers rows on the CSR instead of densifying the whole training block first (>12 GiB of ATAC at Multiome scale). - _to_counts_X builds a minimal AnnData instead of a full .copy(), which was duplicating the normalized layer and obsm['gene_activity'] only for sc.concat to discard them. - skip the no-op var reindex whose .copy() duplicated the matrix. - scatter_to_target copies in column blocks rather than one numpy call per feature (up to ~116k of them for ATAC). - suppress test_model's unconditional PCA + neighbour graph over both predicted matrices; neither is read back. - log torch version, CUDA build and device at startup, so a silent CPU fallback shows up in the run log instead of only as a slow run. - request highsharedmem, matching the other torch methods. Verified bit-identical to the previous implementation: TF-IDF in float32 and float64 over binarized and raw counts, every touched helper against the real test h5ads, and build_butterfly end to end in both directions. Both component tests pass, and the predict test loads weights trained by the previous code, so the reconstructed architecture -- and therefore the preprocessing -- is confirmed unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
benjaminfreyuu
marked this pull request as draft
September 22, 2026 22:09
benjaminfreyuu
marked this pull request as ready for review
September 22, 2026 22:09
Contributor
Author
|
@VladimirShitov could you take a look at the PR? |
VladimirShitov
approved these changes
Sep 23, 2026
VladimirShitov
left a comment
Contributor
There was a problem hiding this comment.
Overall looks good to me. Funny solution with turning off PCA and neighbors computation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
scButterfly's TFIDF tiles the per-cell and per-peak totals into dense (n_peaks, n_cells) float64 matrices and divides densely, so it costs O(n_peaks * n_cells) regardless of how sparse the input is -- 37 GiB per array at NeurIPS2021 Multiome scale, with three or four alive at once. That is what likely drove the ~265 GB peak and the OOM on the 2022 Multiome datasets.
Patch it with a sparse O(nnz) equivalent, and trim the wrapper's own avoidable allocations:
Verified bit-identical to the previous implementation: TF-IDF in float32 and float64 over binarized and raw counts, every touched helper against the real test h5ads, and build_butterfly end to end in both directions. Both component tests pass, and the predict test loads weights trained by the previous code, so the reconstructed architecture -- and therefore the preprocessing -- is confirmed unchanged.
Open:
Local testing showed CUDA: Not available; but local machine does not have CUDA gpu support, so proper testing on gpu needed.
Describe your changes
Checklist before requesting a review
[x ] I have performed a self-review of my code
Check the correct box. Does this PR contain:
Proposed changes are described in the CHANGELOG.md
[x ] CI Tests succeed and look good!