significant speed up of imputation - #3
Merged
Merged
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Pull request overview
This PR introduces a new Monte‑Carlo negative binomial normalization helper to accelerate (including on GPU) the computation of E[log1p(norm(X))] used by CellPin.impute(..., return_norm=True), and wires it into the imputation path with improved reproducibility controls.
Changes:
- Add
mc_log1p_normwith a threaded NumPy backend (CPU) and a Torch backend (GPU) in a newnb_samplingutility module. - Update
CellPin.imputeto use the new helper and addnb_seedfor reproducibleimputed_normgeneration. - Add targeted unit tests covering correctness vs a reference implementation, chunking behavior, and reproducibility.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/models/utils/test_nb_sampling.py | Adds unit tests for the new Monte‑Carlo NB normalization helper (reference matching, chunking, reproducibility). |
| tests/models/test_cellpin.py | Adds integration tests for impute(..., return_norm=True) and nb_seed reproducibility + area_key behavior. |
| src/cellpin/models/utils/nb_sampling.py | Implements chunked MC estimator with NumPy (threaded) and Torch (GPU) backends. |
| src/cellpin/models/utils/init.py | Exports mc_log1p_norm from cellpin.models.utils. |
| src/cellpin/models/cellpin_model.py | Switches imputed_norm computation to mc_log1p_norm and adds nb_seed parameter/docs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+144
to
+147
| if n_samples < 1: | ||
| raise ValueError(f"n_samples must be >= 1, got {n_samples}") | ||
|
|
||
| n_cells, n_genes = mu.shape |
Comment on lines
+99
to
+102
| acc = torch.zeros_like(mu_t) | ||
| for _ in range(n_samples): | ||
| lam = torch._standard_gamma(concentration, generator) / rate | ||
| draw = torch.poisson(lam, generator=generator) |
Comment on lines
+148
to
+151
| theta = np.asarray(theta, dtype=np.float64) | ||
| if theta.shape != (n_genes,): | ||
| raise ValueError(f"theta must have shape ({n_genes},), got {theta.shape}") | ||
| if scale is not None: |
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.
GPU 20X faster