Skip to content

High-performance scaling optimizations for large-scale datasets - #308

Merged
Yenaled merged 1 commit into
pachterlab:developmentfrom
svdrecbd:optimize-scaling-and-parity
Mar 8, 2026
Merged

High-performance scaling optimizations for large-scale datasets#308
Yenaled merged 1 commit into
pachterlab:developmentfrom
svdrecbd:optimize-scaling-and-parity

Conversation

@svdrecbd

Copy link
Copy Markdown
Contributor

This PR introduces critical algorithmic and I/O optimizations to the kb_python core processing paths. By replacing sequential Python loops with vectorized operations and implementing memory-safe streaming, these changes enable the tool to scale gracefully to multi-million cell datasets.

Optimizations:

  • collapse_anndata (Vectorized): Replaced manual sparse column loops with a single sparse matrix-matrix multiplication (X @ S).
    • Performance: Achieved a ~2900x speedup on a 100k sample and reduced full-scale (6.7M cells) processing time from ~1 hour to 3.4 seconds.
  • do_sum_matrices (1-Pass Merge): Switched to a single-pass streaming merge with a temporary body file.
    • Performance: Improved runtime by 1.72x while maintaining O(1) memory safety.
    • Integrity: Uses strict integer arithmetic when processing count matrices to ensure bit-exact precision (addressing float rounding risks).
  • generate_kite_fasta (KITE Ref Gen): Optimized collision detection from $O(N^2)$ to $O(N)$ using dictionary-based lookups.
    • Performance: Reduced runtime for 10,000 features from ~80s to ~3.5s.
  • validate_mtx: Replaced full-file loading with header-only inspection using scipy.io.mminfo.
    • Performance: Reduced validation time by ~1800x, turning a linear file read into a constant-time check.

Dataset Used for Validation:
All performance metrics and parity checks were conducted using the following real-world dataset:

Scientific Parity & Compliance:

  • Numerical Parity: Confirmed exact scientific parity ((X_old != X_new).nnz == 0) across all optimized paths using the 6.7M cell dataset.
  • MatrixMarket Compliance: Body lines are explicitly formatted as integers when multimapping is disabled, ensuring strict compliance with the .mtx format.
  • Regression Suite: All existing tests pass. Added new regression tests for do_sum_matrices covering integer formatting and complex sparse overlapping merges.

@svdrecbd

svdrecbd commented Feb 6, 2026

Copy link
Copy Markdown
Contributor Author

bumping this, if reviewers need additional context or information please feel free to ask!

@Yenaled

Yenaled commented Feb 6, 2026

Copy link
Copy Markdown
Collaborator

Don't worry -- I will merge in the next release!

@svdrecbd

svdrecbd commented Feb 6, 2026

Copy link
Copy Markdown
Contributor Author

very cool, thanks for the clarification!

@Yenaled
Yenaled changed the base branch from master to development March 8, 2026 07:37
@Yenaled
Yenaled merged commit 373e649 into pachterlab:development Mar 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants