feat: VIVS - #32
Open
ori-kron-wis wants to merge 36 commits into
Open
Conversation
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds VIVSModule (BaseModuleClass) composing a fresh scvi.module.VAE (self.x_module) with an ImportanceScoreNet (self.xy_module). loss() switches on self._phase: "x" delegates to x_module's VAE loss, "xy" computes the CRT importance-score loss directly and leaves x_module untouched (frozen when pretrained, unused for grad in this phase). _get_inference_input/_get_generative_input/inference/generative all delegate to x_module, which is what lets VAEMixin.get_latent_representation work later without extra code. Deviations from brief needed to make the exact test code pass against the installed scvi-tools 1.5.0.post1: - test tensors dicts needed a "labels" key: VAE._get_generative_input unconditionally reads tensors[REGISTRY_KEYS.LABELS_KEY], with no .get() fallback, regardless of n_labels. - loss()'s xy-phase LossOutput(...) needed n_obs_minibatch=x.shape[0] explicit: reconstruction_loss=xy_out["loss"] is a 0-dim mean, and LossOutput.__post_init__ otherwise tries reconstruction_loss.shape[0] to infer it, which IndexErrors on a scalar. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Wires VIVSModule into the DestVI/SCVIVA-style model idiom: VAEMixin + UnsupervisedTrainingMixin + SpatialBaseModel, with setup_anndata registering X, batch, and the Y obsm field (n_Y drives n_responses). x_model plumbing is present but only the fresh-VAE constructor path is exercised here; pretrained reuse lands in a later task.
Trains the generative VAE over X first, freezes it, then trains the importance-score net for Y|X. Sequential order (not joint) is required for CRT validity: the knockoff sampler must not see Y. Also registers a dummy LABELS_KEY field in setup_anndata, needed because VAE._get_generative_input unconditionally reads it and the field was missing (unit tests masked this by hand-building tensors).
… delegation Adds test for Task 7 that validates VAEMixin.get_latent_representation works on VIVS instances purely through the delegation wiring from Task 4 (VIVSModule._get_inference_input and inference delegate to x_module). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
_sample_knockoffs previously re-ran the encoder (module.inference) on every Monte Carlo iteration, folding extra encoder-uncertainty variance into each null draw and diverging from the reference JAX implementation (vivs/_vivs.py:227-228,302), which samples z once per batch and only resamples the decoder's px noise across MC draws. Split _sample_knockoffs(x, batch_index) into _encode_for_knockoffs(x, batch_index) -> (z, library), run once per batch, and _sample_knockoffs(z, library, batch_index), called once per MC draw reusing the same z/library. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Implements _get_importance_vmap, vectorizing the per-gene knockoff substitution + statistic recomputation over the gene axis via torch.vmap(randomness="different"). Uses non-mutating torch.where gene substitution (verified in-place index_put_ breaks vmap once a dropout/BatchNorm submodule runs downstream). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ance Wrap the second (in-loop) `_get_importance_vmap` call site with the same try/except RuntimeError -> friendly OOM message used at the first call site, since it fires on the majority of MC samples for realistic n_mc_samples. Also pass device=x.device when building gene_ids in _get_importance_vmap so CUDA tensors don't hit a device-mismatch error that would otherwise be misreported as an OOM by the new guard. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Per-cell (unsummed) importance scores for a specific set of genes, substituted simultaneously with fresh decoder-noise draws per MC iteration (encode once per batch, resample px per iteration).
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Multi-resolution CRT: clusters genes via get_gene_groupings at each requested resolution (plus an always-appended full-resolution single-gene pass), then reruns the CRT with group-level knockoff substitution using a soft additive mask (x*(1-mask) + x_tilde*mask), reusing _encode_for_knockoffs/_sample_knockoffs/_crt_pvalue. Assembles per-gene pval/padj/cluster_assignment across resolutions into an xarray.Dataset. Fixes a brief bug found during TDD: torch.nn.functional.one_hot requires an int64 tensor, but hierarchy.fcluster's output and the appended np.arange(...).astype(np.int32) full-resolution grouping are int32 - cast to .long() before one_hot.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…rtance The .min(axis=1) cross-feature aggregation used to build genes_to_plot disagreed with the unchanged .all() reduction in is_cluster_detected, which still reduced over the feature dim too -- requiring near-unanimous significance across every response feature before a cluster counted as detected. Add a feature: int | str = 0 parameter and slice gene_results to a single feature (isel/sel) at the top of the function instead, matching how the original JAX reference implementation was actually used. Removes the .min(axis=1) hack entirely. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
docs/api/reference/ is gitignored and Sphinx-autosummary-generated for every other model; VIVS's .rst was force-added in the prior commit, making it the sole tracked file in that directory. Untrack it to match convention -- the file stays on disk, just no longer force-committed.
Adds docs/tutorials/VIVS_niche_gene_selection.ipynb, demonstrating VIVS paired with a pretrained scVIVA model (reused as VIVS's x_model knockoff sampler) to test gene dependence on scVIVA's niche_composition. Loads the same Xenium breast-cancer dataset as scVIVA_tutorial.ipynb (data-loading cell copied verbatim). Notebook is written but not executed (no GPU here); registered in docs/tutorials/index.md's toctree.
…n test obs_t was averaged across responses to (batch_n, 1) while tilde_t_mean kept the full (batch_n, n_responses_selected) shape, silently breaking the intended `tilde_t_mean - obs_t` comparison whenever more than one response is selected. Drop the mean reduction so obs_t matches tilde_t_mean's shape, matching the JAX reference semantics. Also add a save/load round-trip regression test for VIVS, which was planned in the design doc but never landed in the implementation plan. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e033878ced
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+283
to
+285
| obs_all_loss = self.module.xy_module(self.module.xy_input(x, batch_index), y)[ | ||
| "all_loss" | ||
| ] |
There was a problem hiding this comment.
When the model has been transferred to CUDA, the AnnData loader still yields CPU tensors and this direct xy_module call is not covered by auto_move_data, so the first linear layer fails with a CPU/CUDA device mismatch. The same direct-call pattern affects predict_t, get_cell_scores, and hierarchical importance; move each batch to the module device before invoking xy_module and create accumulators on that device.
Useful? React with 👍 / 👎.
| "`x_model` must already be trained (call `.train()` on it) before " | ||
| "being passed to VIVS." | ||
| ) | ||
| x_module = x_model.module |
There was a problem hiding this comment.
When a documented DestVI or ResolVI instance is supplied as x_model, assigning its module directly does not provide the VAE interface that the VIVS wrapper assumes. DestVI's MRDeconv inference requires augmented expression and its generative method requires ind_x, while ResolVI's Pyro module does not expose the delegated VAE inference/generative API, so training or the first knockoff request fails. Restrict accepted models to compatible VAE modules or add model-specific adapters.
Useful? React with 👍 / 👎.
| "xmax": xmax + 0.5, | ||
| "ymin": resolution_idx, | ||
| "ymax": resolution_idx + 1, | ||
| } |
There was a problem hiding this comment.
When no gene passes the threshold at base_resolution, plot_df is constructed from an empty list and consequently has no xmin, xmax, ymin, or ymax columns. The subsequent geom_rect mapping then fails instead of returning an empty plot, which is a normal outcome for a CRT result; initialize the expected columns or explicitly handle the no-discovery case.
Useful? React with 👍 / 👎.
…tion, grid search kwargs, empty-plot handling, and plotnine packaging
…ng — degenerate all-zero cells now normalize to 0 instead of NaN; real nonzero cells are unaffected
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
1 participant
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.
Description
Brief description of changes and motivation.
Type of Change
Testing
pytest tests/ -v)Checklist
ruff check src/ tests/)