diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index dac93d8..defa075 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -83,6 +83,8 @@ jobs: run: uvx hatch run ${{ matrix.env.name }}:coverage xml - name: Upload coverage uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} # Check that all tests defined above pass. This makes it easy to set a single "required" test in branch # protection instead of having to update it frequently. See https://github.com/re-actors/alls-green#why. diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 18937f1..9366fbd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -11,7 +11,7 @@ repos: hooks: - id: prettier - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.11.13 + rev: v0.12.2 hooks: - id: ruff types_or: [python, pyi, jupyter] diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f22ff0..320b168 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +## [0.1.9] - 2025-07-01 + +- Add DRVI-APnoEXP baseline +- Imorove documnetation for all classes and functions in repository + ## [0.1.8] - 2025-06-22 - Discretize latent dimension values in MI for benchmarking due to [this bug](https://github.com/scikit-learn/scikit-learn/issues/30772). diff --git a/README.md b/README.md index 15b3f8b..d7e1bb8 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,7 @@ Please refer to the [documentation][link-docs]. In particular, the We recommend running DRVI on a recent Linux distribution. DRVI is actively tested on the latest LTS version of Ubuntu (currently 24.04 LTS). -[//]: # "TODO: remove ubuntu version later" + For optimal performance, we highly recommend using a GPU with CUDA capabilities. While CPU-based systems are supported, GPU-powered systems are strongly recommended for optimal performance. @@ -49,7 +49,7 @@ Python installed, we recommend installing [Mambaforge](https://github.com/conda- There are several options to install drvi: -[//]: # "TODO: remove install time!" + 1. Install the latest release of `drvi-py` from [PyPI][link-pypi], which should take around two minutes: @@ -74,8 +74,8 @@ See the [changelog][changelog]. ## Contact -[//]: # "TODO: make clear where to ask questions:" -[//]: # "For questions and help requests, you can reach out in the [scverse discourse][scverse-discourse]." + + If you found a bug, please use the [issue tracker][issue-tracker]. diff --git a/docs/api/metrics.md b/docs/api/metrics.md index fe9563c..3880540 100644 --- a/docs/api/metrics.md +++ b/docs/api/metrics.md @@ -4,6 +4,7 @@ We have implemented a user-friendly class for evaluation of disentanglement with ```{eval-rst} .. module:: drvi.utils.metrics + :no-index: .. currentmodule:: drvi.utils .. autosummary:: @@ -17,6 +18,7 @@ The following functions represent the similarity functions used in benchmarking: ```{eval-rst} .. module:: drvi.utils.metrics + :no-index: .. currentmodule:: drvi.utils .. autosummary:: @@ -32,6 +34,7 @@ The following functions represent the aggregation functions used in benchmarking ```{eval-rst} .. module:: drvi.utils.metrics + :no-index: .. currentmodule:: drvi.utils .. autosummary:: diff --git a/docs/api/model.md b/docs/api/model.md index 61d0333..055638a 100644 --- a/docs/api/model.md +++ b/docs/api/model.md @@ -1,7 +1,23 @@ # Model +The DRVI (Disentangled Representation Variational Inference) model is a model designed for single-cell omics data analysis. It provides disentangled latent representations that separate individual biological processes, enabling better interpretation and downstream analysis. + +## Overview + +DRVI extends the standard variational autoencoder architecture with specialized decoder architecture. The model learns disentangled representations and separates different sources of variation in the data, such as: + +- **Biological factors**: Cell types, developmental processes, perturbation responses, signaling pathways +- **Technical factors**: Background expressions, technical stress responses + +## Core Components + +### DRVI model + +This is the main model class that can be used to define, train, and evaluate the model on an anndata. `DRVI` passes any extra argument to `DRVIModule` in initialization. Accordingly, we suggest to check its documentation (below) for additional configurations. + ```{eval-rst} .. module:: drvi.model + :no-index: .. currentmodule:: drvi .. autosummary:: @@ -9,5 +25,57 @@ :toctree: generated model.DRVI +``` + +### DRVIModule + +This is the pytorch neural network module and contains DRVI logic. + +```{eval-rst} +.. module:: drvi.model +.. currentmodule:: drvi + +.. autosummary:: + :nosignatures: + :toctree: generated + model.DRVIModule ``` + +## Usage Example + +```python +import anndata as ad +from drvi.model import DRVI + +# Load your data +adata = ad.read_h5ad("your_data.h5ad") + +# Setup anndata +DRVI.setup_anndata( + adata, + layer="counts", + categorical_covariate_keys=["batch"], + is_count_data=True, +) + +# Initialize the model +model = DRVI( + adata, + categorical_covariates=["batch"], + n_latent=64, + encoder_dims=[128, 128], + decoder_dims=[128, 128], +) + +# Train the model +model.train( + max_epochs=400, + early_stopping=False, +) + +# Get disentangled representations +latent = model.get_latent_representation() + +# Please check tutorials for more details and downstream steps +``` diff --git a/docs/api/plotting.md b/docs/api/plotting.md index 07249af..820d3eb 100644 --- a/docs/api/plotting.md +++ b/docs/api/plotting.md @@ -1,7 +1,22 @@ # Plotting +The DRVI plotting module provides visualization tools for analyzing latent representations and interpretability. These functions help researchers understand the disentangled representations learned by the DRVI model and their biological implications. + +## Overview + +The plotting module is organized into several categories: + +- **Latent Visualization**: Functions for exploring and visualizing latent dimensions +- **Interpretability Analysis**: Tools for understanding how latent dimensions affect gene expression +- **Utility Functions**: Additional utility functions. + +## Latent Dimension Analysis + +The core functions are: + ```{eval-rst} .. module:: drvi.utils.plotting + :no-index: .. currentmodule:: drvi.utils .. autosummary:: @@ -11,9 +26,142 @@ plotting.plot_latent_dimension_stats plotting.plot_latent_dims_in_umap plotting.plot_latent_dims_in_heatmap - plotting.make_heatmap_groups - plotting.differential_vars_heatmap +``` + +### plot_latent_dimension_stats + +Analyzes and visualizes statistics of latent dimensions to understand their properties and importance. + +- Plots multiple statistics (reconstruction effect, max value, mean, std) across dimension ranking +- Distinguishes between vanished and non-vanished dimensions + +**Use Cases:** + +- Identify which latent dimensions are most important for reconstruction +- Understand the distribution of activation values across dimensions +- Detect vanished dimensions that contribute little to the model + +### plot_latent_dims_in_umap + +Visualizes latent dimensions as continuous variables on UMAP embeddings to understand their spatial distribution. +For each latent dimension, one UMAP plot wil be generated. + +**Use Cases:** + +- Understand how latent dimensions relate to cell clustering +- Identify spatial patterns in latent dimension activation + +### plot_latent_dims_in_heatmap + +Creates heatmap visualizations of latent dimensions across different cell groups or conditions. + +- Groups cells by categorical variables (e.g., cell types, conditions) +- Supports balanced sampling for better visualization +- Configurable ordering and filtering of dimensions + +**Use Cases:** + +- Compare latent dimension activation across cell types +- Identify condition-specific latent patterns + +## Interpretability and Differential Effects + +The core functions are: + +```{eval-rst} +.. module:: drvi.utils.plotting + :no-index: +.. currentmodule:: drvi.utils + +.. autosummary:: + :nosignatures: + :toctree: generated + plotting.show_top_differential_vars - plotting.show_differential_vars_scatter_plot plotting.plot_relevant_genes_on_umap + plotting.show_differential_vars_scatter_plot + plotting.differential_vars_heatmap +``` + +### show_top_differential_vars + +Displays bar plots of the top relevant expressed genes for each latent dimension. + +- Shows top N genes with highest score per dimension +- Support for gene symbol mapping + +**Use Cases:** + +- Identify the most important genes for each biological process +- Compare gene effects across different latent dimensions + +### plot_relevant_genes_on_umap + +Visualizes the expression of top relevant genes for each dimension on UMAP embeddings. + +- Shows genes most affected by selected latent dimensions +- Automatic title generation and layout + +**Use Cases:** + +- Understand expression patterns of key genes +- Validate biological interpretation of latent dimensions + +### show_differential_vars_scatter_plot + +Creates scatter plots, allowing users to understand the scoring function under the hood. + +- Compares two main effect values (min_possible and max_possible) +- Colors genes by the combined effect + +**Use Cases:** + +- Visualize max_possible and min_possible effects +- Understand the scoring function + +### differential_vars_heatmap + +Generates comprehensive heatmaps showing how genes respond to latent dimension traversals. + +- Shows stepwise effects across all latent dimensions and genes +- Groups genes by their maximum effect dimension + +**Use Cases:** + +- Observe the sparsity of the identified modules + +## Utility Functions + +The core functions are: + +```{eval-rst} +.. module:: drvi.utils.plotting + :no-index: +.. currentmodule:: drvi.utils + +.. autosummary:: + :nosignatures: + :toctree: generated + + plotting.make_balanced_subsample + plotting.cmap ``` + +### make_balanced_subsample + +Creates balanced subsamples of AnnData objects with respect to a categorical variable. + +- Equal sampling from each category +- Configurable minimum sample size per category + +**Use Cases:** + +- Create balanced samples for heatmap visualization + +## Custom Colormaps + +The module provides specialized colormaps designed for biological data visualization: + +- **cmap.saturated_red_blue_cmap**: Enhanced red-blue diverging colormap for differential effects +- **cmap.saturated_just_sky_cmap**: Sky-blue colormap for positive-only effects +- **cmap.saturated_sky_cmap**: Sky-blue colormap on the positive side and gray colormap for the negative side diff --git a/docs/api/tools.md b/docs/api/tools.md index 6a3cbf5..fc61144 100644 --- a/docs/api/tools.md +++ b/docs/api/tools.md @@ -1,7 +1,19 @@ # Tools +The DRVI tools module provides utilities for analyzing and interpreting latent representations. + +## Overview + +The tools module is organized into two main categories: + +- **Latent Dimension Analysis**: Functions for analyzing and characterizing latent dimensions +- **Interpretability Tools**: Functions for understanding how latent dimensions affect gene expression + +## Latent Dimension Analysis + ```{eval-rst} .. module:: drvi.utils.tools + :no-index: .. currentmodule:: drvi.utils .. autosummary:: @@ -9,8 +21,74 @@ :toctree: generated tools.set_latent_dimension_stats +``` + +### set_latent_dimension_stats + +Analyzes and characterizes latent dimensions by computing various statistics. + +- Calculates basic statistics and reconstruction effect for each dimension +- Identifies vanished dimensions that contribute little to the model +- Provides ranking and ordering of dimensions by importance +- Essential for understanding + +**Use Cases:** + +- Identify and filter out non-informative dimensions +- Rank dimensions for downstream analysis and visualization + +## Interpretability Tools + +```{eval-rst} +.. module:: drvi.utils.tools + :no-index: +.. currentmodule:: drvi.utils + +.. autosummary:: + :nosignatures: + :toctree: generated + tools.traverse_latent tools.calculate_differential_vars tools.get_split_effects tools.iterate_on_top_differential_vars ``` + +### traverse_latent + +Performs systematic traversals through latent dimensions to understand their effects on gene expression. + +- Systematically varies each latent dimension while keeping others fixed +- Generates synthetic data points across the latent space +- Enables analysis of how individual dimensions affect gene expression +- Should be used with the next function + +**Use Cases:** + +- Understand how each latent dimension affects gene expression + +### calculate_differential_vars + +Identifies genes that are differentially affected by latent dimension changes (traverses). + +- Computes various differential effect metrics (max_possible, min_possible, combined_score) +- Identifies genes most relevant to each latent dimension + +**Use Cases:** + +- Identify genes related to specific biological processes +- Quantify the strength of gene-latent dimension relationships +- Generate gene lists for downstream biological analysis + +### get_split_effects + +This function is simply the combination of `traverse_latent` and `calculate_differential_vars`. + +### iterate_on_top_differential_vars + +Iterative over top relevant genes. + +**Use Cases:** + +- Can be used to construct a for loop over top relevant dimensions. +- It can be used along other tools for biological interpretations of latent dimensions diff --git a/docs/conf.py b/docs/conf.py index bbce5b0..4b3e7bf 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -37,7 +37,7 @@ html_context = { "display_github": True, # Integrate GitHub "github_user": "theislab", - "github_repo": "https://github.com/theislab/drvi", + "github_repo": "drvi", "github_version": "main", "conf_py_path": "/docs/", } @@ -51,6 +51,7 @@ "sphinx_copybutton", "sphinx.ext.autodoc", "sphinx.ext.intersphinx", + "sphinx.ext.viewcode", "sphinx.ext.autosummary", "sphinx.ext.napoleon", "sphinxcontrib.bibtex", @@ -65,7 +66,7 @@ autosummary_generate = True autodoc_member_order = "groupwise" default_role = "literal" -napoleon_google_docstring = False +napoleon_google_docstring = True napoleon_numpy_docstring = True napoleon_include_init_with_doc = False napoleon_use_rtype = True # having a separate entry generally helps readability @@ -100,7 +101,7 @@ "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None), "scipy": ("https://docs.scipy.org/doc/scipy/", None), # Deep learning - "torch": ("https://docs.pytorch.org/docs/main", None), + "torch": ("https://docs.pytorch.org/docs/stable", None), "lightning": ("https://lightning.ai/docs/pytorch/stable/", None), # Special "anndata": ("https://anndata.readthedocs.io/en/stable/", None), @@ -130,6 +131,7 @@ "use_repository_button": True, "path_to_docs": "docs/", "navigation_with_keys": False, + "use_source_button": True, } pygments_style = "default" @@ -138,4 +140,33 @@ # If building the documentation fails because of a missing link that is outside your control, # you can add an exception to this list. # ("py:class", "igraph.Graph"), + # PyTorch references that are not properly resolved + ("py:class", "Module"), + ("py:class", "torch.nn.modules.Module"), + ("py:class", "torch.nn.Parameter"), + ("py:class", "Tensor"), + ("py:class", "optional"), + ("py:class", "torch.utils.hooks.RemovableHandle"), + ("py:class", "Dropout"), + ("py:class", "BatchNorm"), + ("py:class", "Parameter"), + ("py:attr", "state_dict"), + ("py:attr", "strict"), + ("py:attr", "assign"), + ("py:attr", "persistent"), + ("py:attr", "grad_input"), + ("py:attr", "grad_output"), + ("py:attr", "requires_grad"), + ("py:attr", "dst_type"), + ("py:attr", "dtype"), + ("py:attr", "device"), + ("py:attr", "non_blocking"), + ("py:func", "add_module"), + ("py:func", "register_module_forward_hook"), + ("py:func", "register_module_forward_pre_hook"), + ("py:func", "register_module_full_backward_hook"), + ("py:func", "register_module_full_backward_pre_hook"), + # DRVI internal references + ("py:class", "drvi.scvi_tools_based.merlin_data._data.MerlinData"), + ("py:class", "drvi.nn_modules.layer.factory.LayerFactory"), ] diff --git a/docs/contributing.md b/docs/contributing.md index af96564..476089e 100644 --- a/docs/contributing.md +++ b/docs/contributing.md @@ -21,7 +21,7 @@ pip install -e ".[dev,test,doc]" This package uses [pre-commit][] to enforce consistent code-styles. On every commit, pre-commit checks will either automatically fix issues with the code, or raise an error message. -See [pre-commit checks](template_usage.md#pre-commit-checks) for a full list of checks enabled for this repository. +See the [pre-commit documentation][pre-commit] for a full list of checks enabled for this repository. To enable pre-commit locally, simply run diff --git a/docs/references.bib b/docs/references.bib index 7387b1b..f7a665d 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -20,3 +20,16 @@ @ARTICLE{Moinfar2024-cx year = 2024, language = "en" } + +@ARTICLE{Lotfollahi21, + title = "Mapping single-cell data to reference atlases by transfer learning", + author = "Lotfollahi, Mohammad and Naghipourfar, Mohsen and Luecken, Malte D and Khajavi, Matin and Büttner, Maren and Wagenstetter, Marco and Avsec, Žiga and Gayoso, Adam and Yosef, Nir and Interlandi, Marta and Rybakov, Sergei and Misharin, Alexander V and Theis, Fabian J", + journal = "Nat. Biotechnol.", + publisher = "Springer Science and Business Media LLC", + volume = 40, + number = 1, + pages = "121--130", + month = jan, + year = 2022, + language = "en" +} diff --git a/docs/tutorials/external b/docs/tutorials/external index 8b42475..f0be942 160000 --- a/docs/tutorials/external +++ b/docs/tutorials/external @@ -1 +1 @@ -Subproject commit 8b42475c80a32c660f586f45fadaff1255ac164c +Subproject commit f0be942c3c821f2dfd90b6e1e0407eea13c7eb82 diff --git a/pyproject.toml b/pyproject.toml index dbf4bdb..fb946fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ requires = ["hatchling"] [project] name = "drvi-py" -version = "0.1.8" +version = "0.1.9" description = "Unsupervised Deep Disentangled Representation of Single-Cell Omics" readme = "README.md" license = { file = "LICENSE" } diff --git a/src/drvi/nn_modules/embedding.py b/src/drvi/nn_modules/embedding.py index 430959c..f742d04 100644 --- a/src/drvi/nn_modules/embedding.py +++ b/src/drvi/nn_modules/embedding.py @@ -1,5 +1,7 @@ import logging from collections import OrderedDict +from collections.abc import Callable +from typing import Any import numpy as np import pandas as pd @@ -9,14 +11,70 @@ class FreezableEmbedding(nn.Embedding): - def __init__(self, num_embeddings, embedding_dim, n_freeze_x=0, n_freeze_y=0, **kwargs): + """Embedding layer with partial freezing capability. + + This embedding layer allows freezing specific regions of the embedding + matrix, which is useful for transfer learning and fine-tuning scenarios. + + Parameters + ---------- + num_embeddings + Number of embeddings (vocabulary size). + embedding_dim + Dimension of each embedding vector. + n_freeze_x + Number of embedding indices to freeze (from the beginning). + n_freeze_y + Number of embedding dimensions to freeze (from the beginning). + **kwargs + Additional arguments passed to nn.Embedding. + + Notes + ----- + The freezing mechanism works by masking gradients during backpropagation. + Frozen regions have their gradients set to zero, preventing parameter updates. + + Examples + -------- + >>> import torch + >>> # Create embedding with frozen first 5 indices and first 10 dimensions + >>> emb = FreezableEmbedding(100, 50, n_freeze_x=5, n_freeze_y=10) + >>> # Test forward pass + >>> indices = torch.randint(0, 100, (10,)) + >>> output = emb(indices) + >>> print(output.shape) # torch.Size([10, 50]) + >>> # Test backward pass (frozen regions won't be updated) + >>> loss = output.sum() + >>> loss.backward() + >>> print(emb.weight.grad[0, 0]) # Frozen region + >>> print(emb.weight.grad[10, 20]) # Non-frozen region + """ + + def __init__( + self, num_embeddings: int, embedding_dim: int, n_freeze_x: int = 0, n_freeze_y: int = 0, **kwargs: Any + ) -> None: self._freeze_hook = None self.n_freeze_x = None self.n_freeze_y = None super().__init__(num_embeddings, embedding_dim, **kwargs) self.freeze(n_freeze_x, n_freeze_y) - def freeze(self, n_freeze_x=0, n_freeze_y=0): + def freeze(self, n_freeze_x: int = 0, n_freeze_y: int = 0) -> None: + """Set the freezing configuration for the embedding. + + Parameters + ---------- + n_freeze_x + Number of embedding indices to freeze (from the beginning). + n_freeze_y + Number of embedding dimensions to freeze (from the beginning). + + Notes + ----- + This method registers a backward hook that masks gradients for + the specified frozen regions. The hook is automatically removed + and re-registered when this method is called. + """ self.n_freeze_x = n_freeze_x self.n_freeze_y = n_freeze_y @@ -26,8 +84,26 @@ def freeze(self, n_freeze_x=0, n_freeze_y=0): if self.n_freeze_x > 0 and self.n_freeze_y > 0: self._freeze_hook = self.weight.register_hook(self.partial_freeze_backward_hook) - def partial_freeze_backward_hook(self, grad): + def partial_freeze_backward_hook(self, grad: torch.Tensor) -> torch.Tensor: + """Backward hook to mask gradients for frozen regions. + + Parameters + ---------- + grad + Gradient tensor for the embedding weights. + + Returns + ------- + torch.Tensor + Masked gradient tensor. + + Notes + ----- + This hook creates a mask that zeros out gradients for the frozen + regions while preserving gradients for the non-frozen regions. + """ with torch.no_grad(): + assert self.n_freeze_x is not None and self.n_freeze_y is not None mask = F.pad( torch.zeros(self.n_freeze_x, self.n_freeze_y, device=grad.device), (0, self.embedding_dim - self.n_freeze_y, 0, self.num_embeddings - self.n_freeze_x), @@ -35,7 +111,14 @@ def partial_freeze_backward_hook(self, grad): ) return grad * mask - def __repr__(self): + def __repr__(self) -> str: + """String representation of the embedding layer. + + Returns + ------- + str + String representation showing embedding dimensions and freeze status. + """ if self._freeze_hook is None: return f"Emb({self.num_embeddings}, {self.embedding_dim})" else: @@ -43,14 +126,48 @@ def __repr__(self): class MultiEmbedding(nn.Module): + """Multi-embedding layer that combines multiple embedding tables. + + This module combines multiple embedding layers with different vocabulary + sizes and embedding dimensions into a single layer. + + Parameters + ---------- + n_embedding_list + List of vocabulary sizes for each embedding table. + embedding_dim_list + List of embedding dimensions for each embedding table. + init_method + Initialization method for embedding weights. + normalization + Normalization method to apply to concatenated embeddings. + **kwargs + Additional arguments passed to FreezableEmbedding. + + Notes + ----- + The embeddings are concatenated along the last dimension. If normalization + is specified, it's applied to the concatenated result. + + Examples + -------- + >>> import torch + >>> # Create multi-embedding with different vocab sizes and dimensions + >>> multi_emb = MultiEmbedding(n_embedding_list=[100, 50, 200], embedding_dim_list=[32, 16, 64], normalization="l2") + >>> # Test forward pass + >>> indices = torch.randint(0, 100, (10, 3)) # 3 indices per sample + >>> output = multi_emb(indices) + >>> print(output.shape) # torch.Size([10, 112]) # 32+16+64=112 + """ + def __init__( self, n_embedding_list: list[int], embedding_dim_list: list[int], - init_method="xavier_uniform", - normalization=None, - **kwargs, - ): + init_method: str | Callable = "xavier_uniform", + normalization: str | None = None, + **kwargs: Any, + ) -> None: super().__init__() assert len(n_embedding_list) == len(embedding_dim_list) @@ -64,7 +181,26 @@ def __init__( self.normalization = normalization self.reset_parameters(init_method) - def reset_parameters(self, init_method): + def reset_parameters(self, init_method: str | Callable) -> None: + """Reset parameters of all embedding layers. + + Parameters + ---------- + init_method + Initialization method to apply to each embedding layer. + + Notes + ----- + Supported initialization methods: + - "xavier_uniform": Xavier uniform initialization + - "xavier_normal": Xavier normal initialization + - "uniform": Uniform initialization in [-1, 1] + - "normal": Normal initialization + - "zero": Zero initialization + - "one": One initialization + - None: No initialization + - callable: Custom initialization function + """ # TODO: research the correct distribution considering emb_side and layer_sizes for emb in self.emb_list: if init_method is None: @@ -87,18 +223,71 @@ def reset_parameters(self, init_method): raise NotImplementedError() def forward(self, index_list: torch.Tensor) -> torch.Tensor: + """Forward pass through the multi-embedding layer. + + Parameters + ---------- + index_list + Tensor of indices with shape (..., n_embeddings). + + Returns + ------- + torch.Tensor + Concatenated embeddings with shape (..., total_embedding_dim). + + Notes + ----- + The input should have the same number of indices as there are + embedding tables. Each index is used to look up embeddings from + the corresponding table, and the results are concatenated. + """ assert index_list.shape[-1] == len(self.emb_list) emb = torch.concat([emb(index_list[..., i]) for i, emb in enumerate(self.emb_list)], dim=-1) if self.normalization is None: return emb elif self.normalization == "l2": return F.normalize(emb, p=2, dim=1) + else: + raise NotImplementedError() @classmethod - def from_pretrained(cls, feature_embedding_instance): + def from_pretrained(cls, feature_embedding_instance: Any) -> "MultiEmbedding": + """Create MultiEmbedding from a pretrained feature embedding. + + Parameters + ---------- + feature_embedding_instance + Pretrained feature embedding instance. + + Returns + ------- + MultiEmbedding + New MultiEmbedding instance with pretrained weights. + + Raises + ------ + NotImplementedError + This method is not yet implemented. + """ raise NotImplementedError() - def load_weights_from_trained_module(self, other, freeze_old=False): + def load_weights_from_trained_module(self, other: "MultiEmbedding", freeze_old: bool = False) -> None: + """Load weights from another MultiEmbedding module. + + Parameters + ---------- + other + Source MultiEmbedding module to load weights from. + freeze_old + Whether to freeze the loaded weights. + + Notes + ----- + This method transfers weights from a source module, handling cases + where the target module has larger vocabulary sizes or embedding + dimensions. The transferred weights are placed in the top-left + corner of each embedding matrix. + """ assert len(self.emb_list) >= len(other.emb_list) if len(self.emb_list) > len(other.emb_list): logging.warning(f"Extending feature embedding {other} to {self} with more feature categories.") @@ -123,19 +312,52 @@ def load_weights_from_trained_module(self, other, freeze_old=False): if freeze_old: self_emb.freeze(other_emb.num_embeddings, other_emb.embedding_dim) - def freeze_top_embs(self, n_freeze_list): + def freeze_top_embs(self, n_freeze_list: list[int]) -> None: + """Freeze the top embeddings of each embedding table. + + Parameters + ---------- + n_freeze_list + Number of top embeddings to freeze for each table. + + Notes + ----- + This method freezes the first n embeddings from each embedding + table, which is useful for transfer learning scenarios. + """ for emb, n_freeze in zip(self.emb_list, n_freeze_list, strict=False): emb.freeze(n_freeze, emb.embedding_dim) @property - def num_embeddings(self): + def num_embeddings(self) -> list[int]: + """Get list of vocabulary sizes. + + Returns + ------- + list[int] + List of vocabulary sizes for each embedding table. + """ return [emb.num_embeddings for emb in self.emb_list] @property - def embedding_dim(self): + def embedding_dim(self) -> int: + """Get total embedding dimension. + + Returns + ------- + int + Sum of all embedding dimensions. + """ return sum(emb.embedding_dim for emb in self.emb_list) - def __repr__(self): + def __repr__(self) -> str: + """String representation of the multi-embedding layer. + + Returns + ------- + str + String representation showing all embedding tables and normalization. + """ repr_text = "cat(" + ", ".join(repr(emb) for emb in self.emb_list) + ")" if self.normalization: repr_text = f"{self.normalization}({repr_text})" @@ -143,7 +365,41 @@ def __repr__(self): class FeatureEmbedding(nn.Module): - def __init__(self, vocab_list: list[list[str]], embedding_dims: list[int], **kwargs): + """Feature embedding layer for categorical features. + + This module creates embeddings for categorical features based on + vocabulary lists. It handles the mapping from categorical values + to embedding indices automatically. + + Parameters + ---------- + vocab_list + List of vocabulary lists, one for each categorical feature. + embedding_dims + List of embedding dimensions for each feature. + **kwargs + Additional arguments passed to MultiEmbedding. + + Notes + ----- + This module automatically creates vocabulary mappings and handles + the conversion from categorical values to indices. It also provides + caching for efficient repeated lookups. + + Examples + -------- + >>> import numpy as np + >>> # Create feature embedding for categorical features + >>> vocab_list = [["A", "B", "C"], ["X", "Y"], ["1", "2", "3", "4"]] + >>> embedding_dims = [16, 8, 32] + >>> feature_emb = FeatureEmbedding(vocab_list, embedding_dims) + >>> # Test with categorical data + >>> sentences = np.array([["A", "X", "1"], ["B", "Y", "2"]]) + >>> output = feature_emb(sentences) + >>> print(output.shape) # torch.Size([2, 56]) # 16+8+32=56 + """ + + def __init__(self, vocab_list: list[list[str]], embedding_dims: list[int], **kwargs: Any) -> None: super().__init__() assert len(vocab_list) == len(embedding_dims) self.device_container = nn.Parameter(torch.tensor([])) @@ -152,15 +408,34 @@ def __init__(self, vocab_list: list[list[str]], embedding_dims: list[int], **kwa n_vocab_list = [len(vocab) for vocab in self.vocab_list] self.multi_emb = self.define_embeddings(n_vocab_list, embedding_dims, **kwargs) - self.__index_cache = {} - self.__vocab_map_list_cache = None + self.__index_cache: dict[str, torch.Tensor] = {} + self.__vocab_map_list_cache: list[dict[str, int]] | None = None + + def reset_cache(self) -> None: + """Reset the internal caches. - def reset_cache(self): + Notes + ----- + This method clears both the index cache and vocabulary mapping cache. + It should be called when the vocabulary changes. + """ self.__index_cache = {} self.__vocab_map_list_cache = None @property - def vocab_map_list(self): + def vocab_map_list(self) -> list[dict[str, int]]: + """Get list of vocabulary mappings. + + Returns + ------- + list[dict] + List of dictionaries mapping categorical values to indices. + + Notes + ----- + This property is cached for efficiency. The cache is automatically + invalidated when reset_cache() is called. + """ if self.__vocab_map_list_cache is None: n_vocab_list = [len(vocab) for vocab in self.vocab_list] self.__vocab_map_list_cache = [ @@ -170,13 +445,53 @@ def vocab_map_list(self): return self.__vocab_map_list_cache @staticmethod - def define_embeddings(n_embedding_list, embedding_dims, **kwargs): + def define_embeddings(n_embedding_list: list[int], embedding_dims: list[int], **kwargs: Any) -> MultiEmbedding: + """Define the underlying embedding layers. + + Parameters + ---------- + n_embedding_list + List of vocabulary sizes. + embedding_dims + List of embedding dimensions. + **kwargs + Additional arguments for MultiEmbedding. + + Returns + ------- + MultiEmbedding + Multi-embedding layer. + """ return MultiEmbedding(n_embedding_list, embedding_dims, **kwargs) - def reset_parameters(self, init_method): + def reset_parameters(self, init_method: str | Callable) -> None: + """Reset parameters of the embedding layers. + + Parameters + ---------- + init_method : str or callable + Initialization method to apply. + """ self.multi_emb.reset_parameters(init_method) - def _get_index_from_sentences(self, index_sentences: np.ndarray): + def _get_index_from_sentences(self, index_sentences: np.ndarray) -> torch.Tensor: + """Convert categorical sentences to embedding indices. + + Parameters + ---------- + index_sentences + Array of categorical values with shape (..., n_features). + + Returns + ------- + torch.Tensor + Tensor of embedding indices with shape (..., n_features). + + Notes + ----- + This method uses the vocabulary mappings to convert categorical + values to their corresponding embedding indices. + """ assert index_sentences.shape[-1] == len(self.vocab_map_list) mapping_list = map(lambda mapping: np.vectorize(lambda key: mapping[key]), self.vocab_map_list) @@ -185,7 +500,26 @@ def _get_index_from_sentences(self, index_sentences: np.ndarray): ) return indices.to(self.device_container.device) - def forward(self, index_sentences: np.ndarray, index_cache_key=None): + def forward(self, index_sentences: np.ndarray, index_cache_key: str | None = None) -> torch.Tensor: + """Forward pass through the feature embedding layer. + + Parameters + ---------- + index_sentences + Array of categorical values with shape (..., n_features). + index_cache_key + Cache key for storing computed indices. + + Returns + ------- + torch.Tensor + Concatenated embeddings with shape (..., total_embedding_dim). + + Notes + ----- + If a cache key is provided, the computed indices are stored for + future use, improving efficiency for repeated lookups. + """ if index_cache_key is not None and index_cache_key in self.__index_cache: return self.multi_emb(self.__index_cache[index_cache_key]) indices = self._get_index_from_sentences(index_sentences) @@ -193,21 +527,76 @@ def forward(self, index_sentences: np.ndarray, index_cache_key=None): self.__index_cache[index_cache_key] = indices return self.multi_emb(indices) - def get_extra_state(self): + def get_extra_state(self) -> dict[str, Any]: + """Get extra state for serialization. + + Returns + ------- + dict + Extra state information including vocabulary lists. + """ return {"vocab_list": self.vocab_list} - def set_extra_state(self, state): + def set_extra_state(self, state: dict[str, Any]) -> None: + """Set extra state from serialization. + + Parameters + ---------- + state + Extra state information. + """ self.vocab_list = state["vocab_list"] self.__vocab_map_list_cache = None @classmethod - def from_numpy_array(cls, sentences_array: np.ndarray, embedding_dims: np.ndarray | list, **kwargs): + def from_numpy_array( + cls, sentences_array: np.ndarray, embedding_dims: list[int], **kwargs: Any + ) -> "FeatureEmbedding": + """Create FeatureEmbedding from a numpy array. + + Parameters + ---------- + sentences_array + Array of categorical sentences with shape (n_samples, n_features). + embedding_dims + Embedding dimensions for each feature. + **kwargs + Additional arguments for FeatureEmbedding. + + Returns + ------- + FeatureEmbedding + New FeatureEmbedding instance. + + Notes + ----- + This method automatically extracts unique values for each feature + to create the vocabulary lists. + """ word_list = sentences_array.transpose().tolist() vocab_list = [list(OrderedDict.fromkeys(words)) for words in word_list] return cls(vocab_list, embedding_dims, **kwargs) @classmethod - def from_pandas_dataframe(cls, sentences_df: pd.DataFrame, embedding_dims: pd.DataFrame | list, **kwargs): + def from_pandas_dataframe( + cls, sentences_df: pd.DataFrame, embedding_dims: list[int], **kwargs: Any + ) -> "FeatureEmbedding": + """Create FeatureEmbedding from a pandas DataFrame. + + Parameters + ---------- + sentences_df + DataFrame of categorical sentences. + embedding_dims + Embedding dimensions for each feature. + **kwargs + Additional arguments for FeatureEmbedding. + + Returns + ------- + FeatureEmbedding + New FeatureEmbedding instance. + """ if isinstance(embedding_dims, pd.DataFrame): assert sentences_df.columns == embedding_dims.columns assert len(embedding_dims) == 1 @@ -215,10 +604,41 @@ def from_pandas_dataframe(cls, sentences_df: pd.DataFrame, embedding_dims: pd.Da return cls.from_numpy_array(sentences_df.values, embedding_dims, **kwargs) @classmethod - def from_pretrained(cls, feature_embedding_instance): + def from_pretrained(cls, feature_embedding_instance: "FeatureEmbedding") -> "FeatureEmbedding": + """Create FeatureEmbedding from a pretrained instance. + + Parameters + ---------- + feature_embedding_instance + Pretrained FeatureEmbedding instance. + + Returns + ------- + FeatureEmbedding + New FeatureEmbedding instance with pretrained weights. + + Raises + ------ + NotImplementedError + This method is not yet implemented. + """ raise NotImplementedError() - def load_weights_from_trained_module(self, other, freeze_old=False): + def load_weights_from_trained_module(self, other: "FeatureEmbedding", freeze_old: bool = False) -> None: + """Load weights from another FeatureEmbedding module. + + Parameters + ---------- + other + Source FeatureEmbedding module to load weights from. + freeze_old + Whether to freeze the loaded weights. + + Notes + ----- + This method transfers weights from the source module's multi_emb + to this module's multi_emb. + """ assert isinstance(other, self.__class__) assert len(self.vocab_list) >= len(other.vocab_list) @@ -230,8 +650,22 @@ def load_weights_from_trained_module(self, other, freeze_old=False): self.multi_emb.load_weights_from_trained_module(other.multi_emb, freeze_old=freeze_old) @property - def embedding_dim(self): + def embedding_dim(self) -> int: + """Get total embedding dimension. + + Returns + ------- + int + Sum of all embedding dimensions. + """ return self.multi_emb.embedding_dim - def __repr__(self): + def __repr__(self) -> str: + """String representation of the feature embedding layer. + + Returns + ------- + str + String representation showing vocabulary sizes and embedding dimensions. + """ return f"str2index -> {repr(self.multi_emb)}" diff --git a/src/drvi/nn_modules/encodig.py b/src/drvi/nn_modules/encodig.py index c435ee0..99ef19e 100644 --- a/src/drvi/nn_modules/encodig.py +++ b/src/drvi/nn_modules/encodig.py @@ -1,4 +1,5 @@ from collections import OrderedDict +from typing import Any import numpy as np import pandas as pd @@ -10,12 +11,62 @@ class MultiOneHotEncoding(nn.Module): - def __init__(self, n_embedding_list: list[int], **kwargs): + """Multi-one-hot encoding layer for categorical features. + + This module creates one-hot encodings for multiple categorical features + and concatenates them. It's similar to MultiEmbedding but uses one-hot + encoding instead of learned embeddings. + + Parameters + ---------- + n_embedding_list + List of vocabulary sizes for each categorical feature. + **kwargs + Additional arguments (unused, kept for compatibility). + + Notes + ----- + This module creates one-hot encodings for each categorical feature and + concatenates them along the last dimension. The total output dimension + is the sum of all vocabulary sizes. + + Examples + -------- + >>> import torch + >>> # Create multi-one-hot encoding + >>> encoding = MultiOneHotEncoding([3, 2, 4]) + >>> # Test with indices + >>> indices = torch.tensor([[0, 1, 2], [2, 0, 1]]) + >>> output = encoding(indices) + >>> print(output.shape) # torch.Size([2, 9]) # 3+2+4=9 + >>> # Check one-hot encoding + >>> print(output[0]) # [1,0,0, 0,1, 0,0,1,0] for indices [0,1,2] + """ + + def __init__(self, n_embedding_list: list[int], **kwargs: Any) -> None: super().__init__() self.n_embedding_list = n_embedding_list self.device_container = nn.Parameter(torch.tensor([])) def forward(self, index_list: torch.Tensor) -> torch.Tensor: + """Forward pass through the multi-one-hot encoding layer. + + Parameters + ---------- + index_list + Tensor of indices with shape (..., n_features). + + Returns + ------- + torch.Tensor + Concatenated one-hot encodings with shape (..., total_vocab_size). + + Notes + ----- + Each index is converted to a one-hot encoding, and all encodings + are concatenated along the last dimension. The output dimension + is the sum of all vocabulary sizes. + """ assert index_list.shape[-1] == len(self.n_embedding_list) result = torch.concat( [ @@ -26,34 +77,150 @@ def forward(self, index_list: torch.Tensor) -> torch.Tensor: ) return result.to(self.device_container.device) - def get_extra_state(self): + def get_extra_state(self) -> dict[str, Any]: + """Get extra state for serialization. + + Returns + ------- + dict + Extra state information including vocabulary sizes. + """ return {"n_embedding_list": self.n_embedding_list} - def set_extra_state(self, state): + def set_extra_state(self, state: dict[str, Any]) -> None: + """Set extra state from serialization. + + Parameters + ---------- + state + Extra state information. + """ self.n_embedding_list = state["n_embedding_list"] @property - def embedding_dim(self): + def embedding_dim(self) -> int: + """Get total embedding dimension. + + Returns + ------- + int + Sum of all vocabulary sizes. + + Notes + ----- + For one-hot encoding, the embedding dimension equals the vocabulary size + since each category is represented by a single dimension. + """ return sum(self.n_embedding_list) class FeatureOneHotEncoding(FeatureEmbedding): - def __init__(self, vocab_list: list[list[str]], **kwargs): + """Feature one-hot encoding layer for categorical features. + + This module creates one-hot encodings for categorical features based on + vocabulary lists. It's similar to FeatureEmbedding but uses one-hot + encoding instead of learned embeddings. + + Parameters + ---------- + vocab_list + List of vocabulary lists, one for each categorical feature. + **kwargs + Additional arguments passed to FeatureEmbedding. + + Notes + ----- + This module automatically creates vocabulary mappings and handles + the conversion from categorical values to one-hot encodings. It + inherits most functionality from FeatureEmbedding but uses one-hot + encoding instead of learned embeddings. + + Examples + -------- + >>> import numpy as np + >>> # Create feature one-hot encoding + >>> vocab_list = [["A", "B", "C"], ["X", "Y"], ["1", "2", "3", "4"]] + >>> encoding = FeatureOneHotEncoding(vocab_list) + >>> # Test with categorical data + >>> sentences = np.array([["A", "X", "1"], ["B", "Y", "2"]]) + >>> output = encoding(sentences) + >>> print(output.shape) # torch.Size([2, 9]) # 3+2+4=9 + >>> # Check one-hot encoding for first sample + >>> print(output[0]) # One-hot encoding for ["A", "X", "1"] + """ + + def __init__(self, vocab_list: list[list[str]], **kwargs: Any) -> None: n_vocab_list = [len(vocab) for vocab in vocab_list] super().__init__(vocab_list, n_vocab_list, **kwargs) @staticmethod - def define_embeddings(n_embedding_list, embedding_dims, **kwargs): + def define_embeddings(n_embedding_list: list[int], embedding_dims: list[int], **kwargs: Any) -> MultiOneHotEncoding: + """Define the underlying one-hot encoding layers. + + Parameters + ---------- + n_embedding_list + List of vocabulary sizes. + embedding_dims + List of embedding dimensions (must equal vocabulary sizes). + **kwargs + Additional arguments for MultiOneHotEncoding. + + Returns + ------- + MultiOneHotEncoding + Multi-one-hot encoding layer. + + Notes + ----- + For one-hot encoding, the embedding dimension must equal the + vocabulary size since each category is represented by a single + dimension in the one-hot vector. + """ for n_key, dim in zip(n_embedding_list, embedding_dims, strict=False): assert n_key == dim return MultiOneHotEncoding(n_embedding_list, **kwargs) @classmethod - def from_numpy_array(cls, sentences_array: np.ndarray, **kwargs): + def from_numpy_array(cls, sentences_array: np.ndarray, **kwargs: Any) -> "FeatureOneHotEncoding": + """Create FeatureOneHotEncoding from a numpy array. + + Parameters + ---------- + sentences_array + Array of categorical sentences with shape (n_samples, n_features). + **kwargs + Additional arguments for FeatureOneHotEncoding. + + Returns + ------- + FeatureOneHotEncoding + New FeatureOneHotEncoding instance. + + Notes + ----- + This method automatically extracts unique values for each feature + to create the vocabulary lists. The embedding dimensions are set + to match the vocabulary sizes for one-hot encoding. + """ word_list = sentences_array.transpose().tolist() vocab_list = [list(OrderedDict.fromkeys(words)) for words in word_list] return cls(vocab_list, **kwargs) @classmethod - def from_pandas_dataframe(cls, sentences_df: pd.DataFrame, **kwargs): - return cls.from_numpy_array(sentences_df.values) + def from_pandas_dataframe(cls, sentences_df: pd.DataFrame, **kwargs: Any) -> "FeatureOneHotEncoding": + """Create FeatureOneHotEncoding from a pandas DataFrame. + + Parameters + ---------- + sentences_df : pandas.DataFrame + DataFrame of categorical sentences. + **kwargs + Additional arguments for FeatureOneHotEncoding. + + Returns + ------- + FeatureOneHotEncoding + New FeatureOneHotEncoding instance. + """ + return cls.from_numpy_array(sentences_df.values, **kwargs) diff --git a/src/drvi/nn_modules/feature_interface.py b/src/drvi/nn_modules/feature_interface.py index 9f77a58..6ba538b 100644 --- a/src/drvi/nn_modules/feature_interface.py +++ b/src/drvi/nn_modules/feature_interface.py @@ -1,5 +1,7 @@ import re from collections import namedtuple +from collections.abc import Generator +from typing import Any import anndata as ad import numpy as np @@ -9,7 +11,57 @@ class FeatureInfoList: - def __init__(self, feature_info_str_list: list[str], axis="var", total_dim=None, default_dim=None): + """A list of feature information for managing feature dimensions and metadata. + + This class parses and manages feature information strings that specify + feature names, dimensions, and keywords. It's used to organize features + in multi-modal datasets where different features may have different + dimensionalities. + + Parameters + ---------- + feature_info_str_list + List of feature information strings in the format "name[@dim][!keyword1!keyword2...]". + axis + Whether features are stored in AnnData.var or AnnData.obs. + total_dim + Total dimensionality to distribute among features with unspecified dimensions. + default_dim + Default dimension to assign to features with unspecified dimensions. + + Notes + ----- + Feature information strings follow the pattern: + - "name": Feature name only + - "name@dim": Feature name with specified dimension + - "name!keyword": Feature name with keyword + - "name@dim!keyword1!keyword2": Feature name with dimension and keywords + + Examples + -------- + >>> # Create feature info list with explicit dimensions + >>> feature_list = FeatureInfoList(["gene@1000", "protein@50", "metadata@10"]) + >>> print(feature_list.names) # ['gene', 'protein', 'metadata'] + >>> print(feature_list.dims) # [1000, 50, 10] + >>> # Create feature info list with total dimension + >>> feature_list = FeatureInfoList( + ... [ + ... "gene@1000", + ... "protein", # dimension will be inferred + ... "metadata", + ... ], + ... total_dim=1100, + ... ) + >>> print(feature_list.dims) # [1000, 50, 50] + """ + + def __init__( + self, + feature_info_str_list: list[str], + axis: str = "var", + total_dim: int | None = None, + default_dim: int | None = None, + ) -> None: assert axis in ["var", "obs"] assert total_dim is None or default_dim is None self.feature_info_list = list(self.parse(feature_info_str_list)) @@ -23,22 +75,63 @@ def __init__(self, feature_info_str_list: list[str], axis="var", total_dim=None, self._fill_with_default_dim(default_dim) @staticmethod - def parse(feature_info_list): + def parse(feature_info_list: list[str]) -> Generator[FeatureInfo, int | None, list[tuple[str, ...]] | None]: + """Parse feature information strings into FeatureInfo objects. + + Parameters + ---------- + feature_info_list + List of feature information strings to parse. + + Yields + ------ + FeatureInfo + Parsed feature information objects. + + Notes + ----- + The parsing uses a regex pattern to extract: + - name: Feature name (required) + - dim: Dimension (optional, converted to int) + - keywords: List of keywords (optional, separated by !) + """ pattern = re.compile(r"\A(?P\w+)(@(?P\d*))?(?P(!\w+)+)?\Z") for feature_info in feature_info_list: match = pattern.match(feature_info) + if match is None: + raise ValueError(f"Invalid feature info format: {feature_info}") name, dim, kw = match.group("name"), match.group("dim"), match.group("kw") dim = None if dim is None else int(dim) kw = () if kw is None else tuple(kw[1:].split("!")) yield FeatureInfo(name, dim, kw) - def _fill_with_default_dim(self, default_dim): + def _fill_with_default_dim(self, default_dim: int) -> None: + """Fill unspecified dimensions with a default value. + + Parameters + ---------- + default_dim + Default dimension to assign to features with unspecified dimensions. + """ for i in range(len(self.feature_info_list)): if self.feature_info_list[i].dim is None: self.feature_info_list[i] = self.feature_info_list[i]._replace(dim=default_dim) - def _fill_with_total_dim(self, total_dim): + def _fill_with_total_dim(self, total_dim: int) -> None: + """Distribute total dimension among features with unspecified dimensions. + + Parameters + ---------- + total_dim + Total dimensionality to distribute. + + Notes + ----- + This method evenly distributes the remaining dimension among + features that don't have specified dimensions. The distribution + must result in integer dimensions. + """ n_none_dims = sum([(fi.dim is None) + 0.0 for fi in self.feature_info_list]) if n_none_dims > 0: remaining_dims = total_dim - sum([fi.dim for fi in self.feature_info_list if fi.dim is not None]) @@ -48,18 +141,58 @@ def _fill_with_total_dim(self, total_dim): assert sum([fi.dim for fi in self.feature_info_list]) == total_dim @property - def names(self): + def names(self) -> list[str]: + """Get list of feature names. + + Returns + ------- + list[str] + List of feature names. + """ return [fi.name for fi in self.feature_info_list] @property - def dims(self): + def dims(self) -> list[int]: + """Get list of feature dimensions. + + Returns + ------- + list[int] + List of feature dimensions. + """ return [fi.dim for fi in self.feature_info_list] @property - def keywords_list(self): + def keywords_list(self) -> list[tuple[str, ...]]: + """Get list of feature keywords. + + Returns + ------- + list[tuple] + List of keyword tuples for each feature. + """ return [fi.keywords for fi in self.feature_info_list] - def get_possible_values_array(self, data: dict[str, ad.AnnData]): + def get_possible_values_array(self, data: dict[str, ad.AnnData] | ad.AnnData) -> np.ndarray: + """Get unique combinations of feature values across datasets. + + Parameters + ---------- + data + Dictionary of AnnData objects or single AnnData object. + + Returns + ------- + numpy.ndarray + Array of unique feature value combinations. + + Notes + ----- + This method extracts unique combinations of feature values from + the specified axis (obs or var) across all provided datasets. + It's useful for understanding the range of possible values for + categorical features. + """ keys = self.names if keys is None or len(keys) == 0: return np.array([]) @@ -76,11 +209,32 @@ def get_possible_values_array(self, data: dict[str, ad.AnnData]): possible_vals.append(df[list(keys)].drop_duplicates(subset=keys)) return pd.concat(possible_vals).drop_duplicates(subset=keys).to_numpy() - def __len__(self): + def __len__(self) -> int: + """Get the number of features. + + Returns + ------- + int + Number of features in the list. + """ return len(self.feature_info_list) - def __iter__(self): + def __iter__(self) -> Any: + """Iterate over feature information objects. + + Yields + ------ + FeatureInfo + Feature information objects. + """ return self.feature_info_list.__iter__() - def __repr__(self): + def __repr__(self) -> str: + """String representation of the feature info list. + + Returns + ------- + str + String representation showing the feature information. + """ return repr(self.feature_info_list) diff --git a/src/drvi/nn_modules/freezable.py b/src/drvi/nn_modules/freezable.py index c1a3304..38a7179 100644 --- a/src/drvi/nn_modules/freezable.py +++ b/src/drvi/nn_modules/freezable.py @@ -1,16 +1,110 @@ +from typing import Any + from torch import nn def freezable(base_norm_class): + """Decorator to create freezable versions of normalization layers. + + This decorator creates a wrapper class around a base normalization class + that adds freezing capability. When frozen, the normalization layer + operates in evaluation mode regardless of the model's training state. + + Parameters + ---------- + base_norm_class + The base normalization class to make freezable. + + Returns + ------- + type + A new class that inherits from base_norm_class with freezing capability. + + Notes + ----- + The freezable wrapper adds: + - A `_freeze` attribute to track freeze status + - A `freeze()` method to control freezing + - Modified `forward()` method that respects freeze status + + When frozen, the normalization layer: + - Temporarily switches to evaluation mode + - Uses running statistics instead of batch statistics + - Prevents parameter updates + + Examples + -------- + >>> import torch + >>> from torch import nn + >>> # Create a freezable batch norm + >>> FreezableBN = freezable(nn.BatchNorm1d) + >>> bn = FreezableBN(10) + >>> # Test normal operation + >>> x = torch.randn(5, 10) + >>> output1 = bn(x) + >>> # Freeze the layer + >>> bn.freeze(True) + >>> output2 = bn(x) + >>> # The outputs will be different due to different normalization behavior + """ + class FreezableNormClass(base_norm_class): - def __init__(self, *args, **kwargs): + """Freezable wrapper for normalization layers. + + This class adds freezing capability to normalization layers by + temporarily switching to evaluation mode when frozen. + + Parameters + ---------- + *args + Arguments passed to the base normalization class. + **kwargs + Keyword arguments passed to the base normalization class. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: self._freeze = False super().__init__(*args, **kwargs) - def freeze(self, freeze_status=True): + def freeze(self, freeze_status: bool = True) -> None: + """Set the freeze status of the normalization layer. + + Parameters + ---------- + freeze_status + Whether to freeze the normalization layer. + + Notes + ----- + When frozen, the layer will operate in evaluation mode during + forward passes, using running statistics instead of batch statistics. + This is useful for inference or when you want to prevent the + normalization parameters from being updated. + """ self._freeze = freeze_status - def forward(self, *args, **kwargs): + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Forward pass with freeze-aware behavior. + + Parameters + ---------- + *args + Arguments passed to the base normalization forward method. + **kwargs + Keyword arguments passed to the base normalization forward method. + + Returns + ------- + torch.Tensor + Output of the normalization layer. + + Notes + ----- + If the layer is frozen, it temporarily switches to evaluation mode + for the forward pass, then restores the original training state. + This ensures that frozen layers use running statistics regardless + of the model's overall training state. + """ training_status = self.training if self._freeze: self.train(False) @@ -21,11 +115,5 @@ def forward(self, *args, **kwargs): return FreezableNormClass -@freezable -class FreezableBatchNorm1d(nn.BatchNorm1d): - pass - - -@freezable -class FreezableLayerNorm(nn.LayerNorm): - pass +FreezableBatchNorm1d = freezable(nn.BatchNorm1d) +FreezableLayerNorm = freezable(nn.LayerNorm) diff --git a/src/drvi/nn_modules/layer/factory.py b/src/drvi/nn_modules/layer/factory.py index 04d517a..9201cf8 100644 --- a/src/drvi/nn_modules/layer/factory.py +++ b/src/drvi/nn_modules/layer/factory.py @@ -1,3 +1,5 @@ +from typing import Any, Literal + from torch import nn from drvi.nn_modules.layer.linear_layer import StackedLinearLayer @@ -5,19 +7,136 @@ class LayerFactory: - def __init__(self, intermediate_arch="SAME", residual_preferred=False): + """Abstract base class for creating neural network layers. + + This class provides a factory pattern for creating different types of + neural network layers. It supports both normal layers and stacked layers + with configurable architectures and residual connections. + + Parameters + ---------- + intermediate_arch + Architecture type for intermediate layers: + - "SAME": Use the same architecture as defined in subclasses + - "FC": Use fully connected layers + residual_preferred + Whether to wrap layers in residual connections when input and output + dimensions match. + + Notes + ----- + This is an abstract base class. Subclasses must implement: + - `_get_normal_layer`: Creates normal layers + - `_get_stacked_layer`: Creates stacked layers + + The factory pattern allows for easy switching between different layer + architectures while maintaining a consistent interface. + """ + + def __init__(self, intermediate_arch: Literal["SAME", "FC"] = "SAME", residual_preferred: bool = False) -> None: assert intermediate_arch in ["SAME", "FC"] self.intermediate_arch = intermediate_arch self.residual_preferred = residual_preferred - def _get_normal_layer(self, d_in, d_out, bias=True, **kwargs): + def _get_normal_layer(self, d_in: int, d_out: int, bias: bool = True, **kwargs: Any) -> nn.Module: + """Create a normal layer (to be implemented by subclasses). + + Parameters + ---------- + d_in + Input dimension. + d_out + Output dimension. + bias + Whether to include bias term. + **kwargs + Additional keyword arguments. + + Returns + ------- + nn.Module + The created layer. + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses. + """ raise NotImplementedError() - def _get_stacked_layer(self, d_channel, d_in, d_out, bias=True, **kwargs): + def _get_stacked_layer(self, d_channel: int, d_in: int, d_out: int, bias: bool = True, **kwargs: Any) -> nn.Module: + """Create a stacked layer (to be implemented by subclasses). + + Parameters + ---------- + d_channel + Number of channels/splits. + d_in + Input dimension. + d_out + Output dimension. + bias + Whether to include bias term. + **kwargs + Additional keyword arguments. + + Returns + ------- + nn.Module + The created stacked layer. + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses. + """ raise NotImplementedError() - def get_normal_layer(self, d_in, d_out, bias=True, intermediate_layer=None, **kwargs): + def get_normal_layer( + self, d_in: int, d_out: int, bias: bool = True, intermediate_layer: bool | None = None, **kwargs: Any + ) -> nn.Module: + """Create a normal layer with optional residual connection. + + Parameters + ---------- + d_in + Input dimension. + d_out + Output dimension. + bias + Whether to include bias term. + intermediate_layer + Whether this is an intermediate layer. If None, defaults to True. + **kwargs + Additional keyword arguments passed to layer creation. + + Returns + ------- + nn.Module + The created layer, optionally wrapped in a residual connection. + + Notes + ----- + The layer creation logic depends on the `intermediate_arch` setting: + - If `intermediate_layer=True` and `intermediate_arch="FC"`: Creates a linear layer + - Otherwise: Uses the subclass-specific `_get_normal_layer` method + + If `residual_preferred=True` and `d_in == d_out`, the layer is wrapped + in a `SimpleResidual` connection. + + Examples + -------- + >>> factory = FCLayerFactory() + >>> # Create intermediate layer with FC architecture + >>> factory.intermediate_arch = "FC" + >>> layer = factory.get_normal_layer(64, 128, intermediate_layer=True) + >>> print(type(layer)) # + >>> # Create with residual connection + >>> factory.residual_preferred = True + >>> layer = factory.get_normal_layer(64, 64) + >>> print(type(layer)) # + """ if intermediate_layer is None: intermediate_layer = True if intermediate_layer and self.intermediate_arch == "FC": @@ -31,7 +150,60 @@ def get_normal_layer(self, d_in, d_out, bias=True, intermediate_layer=None, **kw layer = SimpleResidual(layer) return layer - def get_stacked_layer(self, d_channel, d_in, d_out, bias=True, intermediate_layer=None, **kwargs): + def get_stacked_layer( + self, + d_channel: int, + d_in: int, + d_out: int, + bias: bool = True, + intermediate_layer: bool | None = None, + **kwargs: Any, + ) -> nn.Module: + """Create a stacked layer with optional residual connection. + + Parameters + ---------- + d_channel + Number of channels/splits for the stacked layer. + d_in + Input dimension. + d_out + Output dimension. + bias + Whether to include bias term. + intermediate_layer + Whether this is an intermediate layer. If None, defaults to True. + **kwargs + Additional keyword arguments passed to layer creation. + + Returns + ------- + nn.Module + The created stacked layer, optionally wrapped in a residual connection. + + Notes + ----- + The layer creation logic depends on the `intermediate_arch` setting: + - If `intermediate_layer=True` and `intermediate_arch="FC"`: Creates a `StackedLinearLayer` + - Otherwise: Uses the subclass-specific `_get_stacked_layer` method + + If `residual_preferred=True` and `d_in == d_out`, the layer is wrapped + in a `SimpleResidual` connection. + + Stacked layers are useful for processing multiple splits of the input + in parallel, which is common in disentanglement models. + + Examples + -------- + >>> factory = FCLayerFactory() + >>> # Create stacked layer with 4 splits + >>> layer = factory.get_stacked_layer(4, 64, 128) + >>> print(type(layer)) # + >>> # Create with residual connection + >>> factory.residual_preferred = True + >>> layer = factory.get_stacked_layer(4, 64, 64) + >>> print(type(layer)) # + """ if intermediate_layer is None: intermediate_layer = True if intermediate_layer and self.intermediate_arch == "FC": @@ -47,14 +219,118 @@ def get_stacked_layer(self, d_channel, d_in, d_out, bias=True, intermediate_laye class FCLayerFactory(LayerFactory): - def __init__(self, intermediate_arch="SAME", residual_preferred=False): + """Factory for creating fully connected neural network layers. + + This factory creates standard fully connected (linear) layers and stacked + linear layers. It inherits from LayerFactory and implements the abstract + methods to provide concrete layer creation functionality. + + Parameters + ---------- + intermediate_arch + Architecture type for intermediate layers: + - "SAME": Use fully connected layers (same as "FC") + - "FC": Use fully connected layers + residual_preferred + Whether to wrap layers in residual connections when input and output + dimensions match. + + Notes + ----- + This factory creates: + - Normal layers: `nn.Linear` layers + - Stacked layers: `StackedLinearLayer` for processing multiple splits + + The "SAME" and "FC" architectures are equivalent for this factory since + both create fully connected layers. + + Examples + -------- + >>> # Create factory with default settings + >>> factory = FCLayerFactory() + >>> # Create a simple linear layer + >>> layer = factory.get_normal_layer(64, 128) + >>> print(type(layer)) # + >>> # Create a stacked layer + >>> stacked_layer = factory.get_stacked_layer(4, 64, 128) + >>> print(type(stacked_layer)) # + """ + + def __init__(self, intermediate_arch: Literal["SAME", "FC"] = "SAME", residual_preferred: bool = False) -> None: super().__init__(intermediate_arch=intermediate_arch, residual_preferred=residual_preferred) - def _get_normal_layer(self, d_in, d_out, bias=True, **kwargs): + def _get_normal_layer(self, d_in: int, d_out: int, bias: bool = True, **kwargs: Any) -> nn.Linear: + """Create a fully connected layer. + + Parameters + ---------- + d_in + Input dimension. + d_out + Output dimension. + bias + Whether to include bias term. + **kwargs + Additional keyword arguments (ignored for linear layers). + + Returns + ------- + nn.Linear + A fully connected linear layer. + + Examples + -------- + >>> factory = FCLayerFactory() + >>> layer = factory._get_normal_layer(64, 128) + >>> print(layer.weight.shape) # torch.Size([128, 64]) + >>> print(layer.bias.shape) # torch.Size([128]) + """ return nn.Linear(d_in, d_out, bias=bias) - def _get_stacked_layer(self, d_channel, d_in, d_out, bias=True, **kwargs): + def _get_stacked_layer( + self, d_channel: int, d_in: int, d_out: int, bias: bool = True, **kwargs: Any + ) -> StackedLinearLayer: + """Create a stacked linear layer. + + Parameters + ---------- + d_channel + Number of channels/splits for the stacked layer. + d_in + Input dimension. + d_out + Output dimension. + bias + Whether to include bias term. + **kwargs + Additional keyword arguments (ignored for stacked linear layers). + + Returns + ------- + StackedLinearLayer + A stacked linear layer that processes multiple splits in parallel. + + Notes + ----- + The `StackedLinearLayer` applies the same linear transformation to + multiple splits of the input, which is useful for disentanglement + models where different splits may represent different factors of variation. + + Examples + -------- + >>> factory = FCLayerFactory() + >>> layer = factory._get_stacked_layer(4, 64, 128) + >>> print(layer.weight.shape) # torch.Size([4, 128, 64]) + >>> print(layer.bias.shape) # torch.Size([4, 128]) + """ return StackedLinearLayer(d_channel, d_in, d_out, bias=bias) - def __str__(self): + def __str__(self) -> str: + """String representation of the factory. + + Returns + ------- + str + A string describing the factory configuration. + """ return f"FCLayerFactory(residual_preferred={self.residual_preferred})" diff --git a/src/drvi/nn_modules/layer/linear_layer.py b/src/drvi/nn_modules/layer/linear_layer.py index 4847194..c1df86b 100644 --- a/src/drvi/nn_modules/layer/linear_layer.py +++ b/src/drvi/nn_modules/layer/linear_layer.py @@ -1,18 +1,74 @@ import math +from typing import Any import torch from torch import nn class StackedLinearLayer(nn.Module): + """A stacked linear layer that applies multiple linear transformations in parallel. + + This layer applies a linear transformation to multiple channels/splits + of the input. It's particularly useful in additive decoders where + different splits should be calculated in parallel. + + Parameters + ---------- + n_channels + Number of channels/splits to process in parallel. + in_features + Number of input features per channel. + out_features + Number of output features per channel. + bias + Whether to include bias terms for each channel. + device + Device to place the layer on. + dtype + Data type for the layer parameters. + + Notes + ----- + The layer maintains separate weight and bias parameters for each channel: + - Weight shape: (n_channels, in_features, out_features) + - Bias shape: (n_channels, out_features) if bias=True, None otherwise + + The forward pass applies the transformation to each channel independently: + output[b, c, o] = sum_i(input[b, c, i] * weight[c, i, o]) + bias[c, o] + + This is equivalent to applying n_channels separate linear layers in parallel, + which is more efficient than using separate nn.Linear layers. + + Examples + -------- + >>> import torch + >>> # Create a stacked linear layer with 4 channels + >>> layer = StackedLinearLayer(n_channels=4, in_features=64, out_features=128) + >>> # Input shape: (batch_size, n_channels, in_features) + >>> x = torch.randn(32, 4, 64) + >>> # Forward pass + >>> output = layer(x) + >>> print(output.shape) # torch.Size([32, 4, 128]) + >>> # Each channel has its own parameters + >>> print(layer.weight.shape) # torch.Size([4, 64, 128]) + >>> print(layer.bias.shape) # torch.Size([4, 128]) + """ + __constants__ = ["n_channels", "in_features", "out_features"] n_channels: int in_features: int out_features: int weight: torch.Tensor + bias: torch.Tensor | None def __init__( - self, n_channels: int, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None + self, + n_channels: int, + in_features: int, + out_features: int, + bias: bool = True, + device: Any = None, + dtype: Any = None, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -27,28 +83,121 @@ def __init__( self.reset_parameters() def reset_parameters(self) -> None: + """Reset the layer parameters to their initial values. + + This method reinitializes both weights and biases using the same + initialization strategy as the default nn.Linear layer. + + Notes + ----- + The initialization follows PyTorch's default linear layer initialization: + - Weights: Uniform distribution in [-1/sqrt(in_features), 1/sqrt(in_features)] + - Biases: Uniform distribution in [-1/sqrt(in_features), 1/sqrt(in_features)] + + This ensures that the variance of the output is approximately preserved + across the layer. + """ self._init_weight() - if self.bias is not None: - self._init_bias() + self._init_bias() def _init_weight(self) -> None: + """Initialize the weight parameters. + + Notes + ----- + Uses the same initialization as default nn.Linear: + Uniform distribution in [-1/sqrt(in_features), 1/sqrt(in_features)] + + This initialization helps maintain the variance of activations + across the network, which is important for training stability. + """ # Same as default nn.Linear (https://github.com/pytorch/pytorch/issues/57109) fan_in = self.in_features bound = 1 / math.sqrt(fan_in) nn.init.uniform_(self.weight, -bound, bound) def _init_bias(self) -> None: - fan_in = self.in_features - bound = 1 / math.sqrt(fan_in) - nn.init.uniform_(self.bias, -bound, bound) + """Initialize the bias parameters. + + Notes + ----- + Uses the same initialization as default nn.Linear: + Uniform distribution in [-1/sqrt(in_features), 1/sqrt(in_features)] + + The bias initialization is independent of the weight initialization + and helps ensure that the layer can learn appropriate offsets. + """ + if self.bias is not None: + fan_in = self.in_features + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) def forward(self, input: torch.Tensor) -> torch.Tensor: + r"""Forward pass through the stacked linear layer. + + Parameters + ---------- + input + Input tensor with shape (batch_size, n_channels, in_features). + + Returns + ------- + torch.Tensor + Output tensor with shape (batch_size, n_channels, out_features). + + Notes + ----- + The forward pass applies the linear transformation to each channel: + + .. math:: + \text{output}[b, c, o] = \\sum_{i} \text{input}[b, c, i] \\cdot \text{weight}[c, i, o] + \text{bias}[c, o] + + where: + - b: batch index + - c: channel index + - i: input feature index + - o: output feature index + + The computation is performed efficiently using torch.einsum, which + is equivalent to applying n_channels separate linear transformations + in parallel. + + Examples + -------- + >>> import torch + >>> # Create layer + >>> layer = StackedLinearLayer(n_channels=3, in_features=10, out_features=5) + >>> # Input: batch_size=2, n_channels=3, in_features=10 + >>> x = torch.randn(2, 3, 10) + >>> # Forward pass + >>> output = layer(x) + >>> print(output.shape) # torch.Size([2, 3, 5]) + """ mm = torch.einsum("bci,cio->bco", input, self.weight) if self.bias is not None: mm = mm + self.bias # They will broadcast well return mm def extra_repr(self) -> str: + """String representation for printing the layer. + + Returns + ------- + str + A string containing the layer's configuration parameters. + + Notes + ----- + This method is used by PyTorch's __repr__ method to provide + a detailed string representation of the layer, which is useful + for debugging and understanding the layer's configuration. + + Examples + -------- + >>> layer = StackedLinearLayer(n_channels=4, in_features=64, out_features=128, bias=True) + >>> print(layer) + StackedLinearLayer(in_features=64, out_features=128, n_channels=4, bias=True) + """ return ( f"in_features={self.in_features}, out_features={self.out_features}, " f"n_channels={self.n_channels}, bias={self.bias is not None}" diff --git a/src/drvi/nn_modules/layer/structures.py b/src/drvi/nn_modules/layer/structures.py index 560b3e7..9d45e3c 100644 --- a/src/drvi/nn_modules/layer/structures.py +++ b/src/drvi/nn_modules/layer/structures.py @@ -1,10 +1,74 @@ +import torch from torch import nn class SimpleResidual(nn.Module): - def __init__(self, layer): + """Simple residual connection that adds the input to the layer output. + + This module implements a basic residual connection by adding the input + directly to the output of the wrapped layer. It's useful for creating + skip connections that help with gradient flow in deep networks. + + Parameters + ---------- + layer : nn.Module + The layer to be wrapped with a residual connection. + + Notes + ----- + The forward pass computes: output = input + layer(input) + + This is the simplest form of residual connection, where the input is + added directly to the layer's output. For this to work properly, the + layer should preserve the input shape (i.e., input.shape == layer(input).shape). + + Examples + -------- + >>> import torch + >>> from torch import nn + >>> # Create a simple linear layer + >>> linear_layer = nn.Linear(64, 64) + >>> # Wrap it with a residual connection + >>> residual_layer = SimpleResidual(linear_layer) + """ + + def __init__(self, layer: nn.Module): + """Initialize the residual connection. + + Parameters + ---------- + layer + The layer to be wrapped with a residual connection. + + Notes + ----- + The layer should preserve the input shape for the residual connection + to work properly. This means layer(input).shape should equal input.shape. + """ super().__init__() self.layer = layer - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + r"""Forward pass through the residual connection. + + Parameters + ---------- + x + Input tensor. The shape should be compatible with the wrapped layer. + + Returns + ------- + torch.Tensor + Output tensor with the same shape as the input. + + Notes + ----- + The forward pass computes the residual connection: + + .. math:: + \text{output} = \text{input} + \text{layer}(\text{input}) + + This is equivalent to adding a skip connection from the input to the + output of the wrapped layer. + """ return x + self.layer(x) diff --git a/src/drvi/nn_modules/noise_model.py b/src/drvi/nn_modules/noise_model.py index 27690ff..acb355c 100644 --- a/src/drvi/nn_modules/noise_model.py +++ b/src/drvi/nn_modules/noise_model.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Any, Literal import torch from scvi.distributions import NegativeBinomial @@ -7,34 +7,160 @@ class NoiseModel: - def __init__(self): + """Abstract base class for noise models in variational autoencoders. + + This class defines the interface for different noise models that can be + used in the decoder of variational autoencoders. Each noise model specifies + how the output distribution should be parameterized and how the data should + be transformed. + + Notes + ----- + Subclasses must implement: + - `parameters`: Property defining the required parameters + - `dist`: Method creating the output distribution + + The noise model is responsible for: + 1. Defining what parameters the decoder should output + 2. Transforming the decoder outputs into distribution parameters + 3. Creating the appropriate distribution for the data + """ + + def __init__(self) -> None: pass @property - def parameters(self): + def parameters(self) -> dict[str, str]: + """Get the parameter specification for this noise model. + + Returns + ------- + dict + Dictionary mapping parameter names to their specifications. + Common specifications include: + - "no_transformation": Parameter is used directly + - "per_feature": Parameter is learned per feature + - "fixed=value": Parameter is fixed to a specific value + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses. + """ raise NotImplementedError() @property - def main_param(self): + def main_param(self) -> str: + """Get the main parameter name for this noise model. + + Returns + ------- + str + Name of the main parameter (usually "mean"). + """ return "mean" - def initial_transformation(self, x, x_mask=1.0): + def initial_transformation( + self, x: torch.Tensor, x_mask: torch.Tensor | None = None + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Apply initial transformation to input data. + + Parameters + ---------- + x + Input data tensor. + x_mask + Mask for the input data. + + Returns + ------- + tuple + (transformed_x, aux_info) where aux_info contains additional + information needed for the distribution. + """ x = x aux_info = {} return x, aux_info - def dist(self, aux_info, parameters, lib_y): + def dist(self, aux_info: dict[str, Any], parameters: dict[str, torch.Tensor], lib_y: torch.Tensor) -> Distribution: + """Create the output distribution. + + Parameters + ---------- + aux_info + Additional information from initial_transformation. + parameters + Parameters from the decoder network. + lib_y + Library size tensor. + + Returns + ------- + Distribution + The output distribution. + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses. + """ # lib_y is so strange here. I should think about it. raise NotImplementedError() -def calculate_library_size(x, x_mask=1.0): +def calculate_library_size(x: torch.Tensor, x_mask: torch.Tensor | None = None) -> torch.Tensor: + """Calculate the library size for each sample. + + Parameters + ---------- + x + Input data tensor. + x_mask + Mask for the input data. + + Returns + ------- + torch.Tensor + Library size for each sample. + + Notes + ----- + The library size is the sum of all features for each sample, + optionally weighted by the mask. + """ return (x_mask * x).sum(dim=-1) if x_mask is not None else x.sum(dim=-1) def library_size_normalization( - x, lib_size, library_normalization: Literal["none", "x_lib", "x_loglib", "div_lib_x_loglib", "x_loglib_all"] -): + x: torch.Tensor, + lib_size: torch.Tensor, + library_normalization: Literal["none", "x_lib", "x_loglib", "div_lib_x_loglib", "x_loglib_all"], +) -> torch.Tensor: + """Normalize data by library size. + + Parameters + ---------- + x + Input data tensor. + lib_size + Library size tensor. + library_normalization + Normalization method to apply. + + Returns + ------- + torch.Tensor + Normalized data tensor. + + Notes + ----- + Different normalization methods: + - "none": No normalization + - "x_lib": Divide by library size and scale by 1e4 + - "x_loglib": No normalization (same as "none") + - "div_lib_x_loglib": Divide by library size and scale by 1e4 + - "x_loglib_all": Divide by log of library size and scale by 1e1 + """ # TODO: remove any library_normalization but 'none', 'x_lib' if library_normalization in ["none", "x_loglib"]: x = x @@ -48,11 +174,35 @@ def library_size_normalization( def library_size_correction( - x, - lib_size, + x: torch.Tensor, + lib_size: torch.Tensor, library_normalization: Literal["none", "x_lib", "x_loglib", "div_lib_x_loglib", "x_loglib_all"], - log_space=False, -): + log_space: bool = False, +) -> torch.Tensor: + """Apply library size correction to data. + + Parameters + ---------- + x : torch.Tensor + Input data tensor. + lib_size : torch.Tensor + Library size tensor. + library_normalization : {"none", "x_lib", "x_loglib", "div_lib_x_loglib", "x_loglib_all"} + Normalization method that was applied. + log_space : bool, default=False + Whether the data is in log space. + + Returns + ------- + torch.Tensor + Corrected data tensor. + + Notes + ----- + This function reverses the library size normalization to recover + the original scale of the data. The correction depends on whether + the data is in log space or not. + """ # TODO: remove any library_normalization but 'none', 'x_lib' if library_normalization in ["none"]: x = x @@ -72,6 +222,23 @@ def library_size_correction( class NormalNoiseModel(NoiseModel): + """Normal (Gaussian) noise model for continuous data. + + This noise model assumes the data follows a normal distribution. + It can handle both fixed and learnable variance parameters. + + Parameters + ---------- + model_var + Variance modeling strategy: + - "fixed": Use fixed variance of 1e-2 + - "fixed=value": Use fixed variance of specified value + - "dynamic": Learn variance per sample + - "feature": Learn variance per feature + eps + Small constant added to variance for numerical stability. + """ + def __init__(self, model_var="fixed", eps=1e-8): super().__init__() self.model_var = model_var @@ -79,6 +246,13 @@ def __init__(self, model_var="fixed", eps=1e-8): @property def parameters(self): + """Get parameter specification for normal noise model. + + Returns + ------- + dict + Parameter specification with mean and variance parameters. + """ if self.model_var == "fixed": var_desc = "fixed=1e-2" elif self.model_var.startswith("fixed="): @@ -95,11 +269,41 @@ def parameters(self): } def initial_transformation(self, x, x_mask=1.0): + """Apply initial transformation for normal noise model. + + Parameters + ---------- + x + Input data tensor. + x_mask + Mask for the input data. + + Returns + ------- + tuple + (transformed_x, aux_info) where aux_info is empty for normal model. + """ x = x aux_info = {} return x, aux_info def dist(self, aux_info, parameters, lib_y): + """Create normal distribution. + + Parameters + ---------- + aux_info + Additional information (unused for normal model). + parameters + Dictionary containing 'mean' and 'var' parameters. + lib_y + Library size tensor (unused for normal model). + + Returns + ------- + Normal + Normal distribution with specified mean and variance. + """ var = parameters["var"] if self.model_var: var = torch.nan_to_num(torch.exp(var), posinf=100, neginf=0) + self.eps @@ -107,6 +311,33 @@ def dist(self, aux_info, parameters, lib_y): class PoissonNoiseModel(NoiseModel): + """Poisson noise model for count data. + + This noise model assumes the data follows a Poisson distribution, + which is appropriate for count data like gene expression counts. + + Parameters + ---------- + mean_transformation + Transformation to apply to the mean parameter: + - "exp": Exponential transformation + - "softmax": Softmax transformation + library_normalization + Library size normalization method. + + Examples + -------- + >>> import torch + >>> # Create Poisson noise model + >>> model = PoissonNoiseModel(mean_transformation="exp") + >>> print(model.parameters) + {'mean': 'no_transformation'} + >>> # Test with sample data + >>> x = torch.randint(0, 100, (10, 5)) + >>> transformed_x, aux_info = model.initial_transformation(x) + >>> print(transformed_x.shape) # torch.Size([10, 5]) + """ + def __init__(self, mean_transformation="exp", library_normalization="x_lib"): super().__init__() self.mean_transformation = mean_transformation @@ -114,11 +345,32 @@ def __init__(self, mean_transformation="exp", library_normalization="x_lib"): @property def parameters(self): + """Get parameter specification for Poisson noise model. + + Returns + ------- + dict + Parameter specification with mean parameter. + """ return { "mean": "no_transformation", } def initial_transformation(self, x, x_mask=1.0): + """Apply initial transformation for Poisson noise model. + + Parameters + ---------- + x : torch.Tensor + Input data tensor. + x_mask : torch.Tensor, default=1.0 + Mask for the input data. + + Returns + ------- + tuple + (transformed_x, aux_info) where aux_info contains library size. + """ x = x aux_info = {} aux_info["x_library_size"] = calculate_library_size(x, x_mask) @@ -127,6 +379,22 @@ def initial_transformation(self, x, x_mask=1.0): return x, aux_info def dist(self, aux_info, parameters, lib_y): + """Create Poisson distribution. + + Parameters + ---------- + aux_info : dict + Additional information from initial_transformation. + parameters : dict + Dictionary containing 'mean' parameter. + lib_y : torch.Tensor + Library size tensor. + + Returns + ------- + Poisson + Poisson distribution with transformed mean. + """ mean = parameters["mean"] library_size = lib_y @@ -142,15 +410,37 @@ def dist(self, aux_info, parameters, lib_y): class NegativeBinomialNoiseModel(NoiseModel): + """Negative binomial noise model for overdispersed count data. + + This noise model assumes the data follows a negative binomial distribution, + which is appropriate for gene expression data that exhibits overdispersion. + + Parameters + ---------- + dispersion : {"feature"}, default="feature" + Dispersion parameter modeling strategy. + mean_transformation : {"exp", "softmax", "softplus", "none"}, default="exp" + Transformation to apply to the mean parameter. + library_normalization : {"none", "x_lib", "x_loglib", "div_lib_x_loglib", "x_loglib_all"}, default="x_lib" + Library size normalization method. + """ + def __init__(self, dispersion="feature", mean_transformation="exp", library_normalization="x_lib"): super().__init__() - assert mean_transformation in ["exp", "softmax"] + assert mean_transformation in ["exp", "softmax", "softplus", "none"] self.dispersion = dispersion self.mean_transformation = mean_transformation self.library_normalization = library_normalization @property def parameters(self): + """Get parameter specification for negative binomial noise model. + + Returns + ------- + dict + Parameter specification with mean and dispersion parameters. + """ params = { "mean": "no_transformation", } @@ -161,6 +451,20 @@ def parameters(self): return params def initial_transformation(self, x, x_mask=1.0): + """Apply initial transformation for negative binomial noise model. + + Parameters + ---------- + x : torch.Tensor + Input data tensor. + x_mask : torch.Tensor, default=1.0 + Mask for the input data. + + Returns + ------- + tuple + (transformed_x, aux_info) where aux_info contains library size. + """ x = x aux_info = {} aux_info["x_library_size"] = calculate_library_size(x, x_mask) @@ -169,6 +473,22 @@ def initial_transformation(self, x, x_mask=1.0): return x, aux_info def dist(self, aux_info, parameters, lib_y): + """Create negative binomial distribution. + + Parameters + ---------- + aux_info : dict + Additional information from initial_transformation. + parameters : dict + Dictionary containing 'mean' and 'r' parameters. + lib_y : torch.Tensor + Library size tensor. + + Returns + ------- + NegativeBinomial + Negative binomial distribution with transformed parameters. + """ mean = parameters["mean"] r = 1.0 + parameters["r"] library_size = lib_y @@ -179,6 +499,13 @@ def dist(self, aux_info, parameters, lib_y): elif self.mean_transformation == "softmax": trans_mean = torch.softmax(mean, dim=-1) trans_mean = library_size.unsqueeze(-1) * trans_mean + # `softplus` and `none` for ablation. Useless in practice. + elif self.mean_transformation == "softplus": + trans_mean = F.softplus(mean) + trans_mean = library_size_correction(trans_mean, library_size, self.library_normalization, log_space=False) + elif self.mean_transformation == "none": + trans_mean = mean + trans_mean = library_size_correction(trans_mean, library_size, self.library_normalization, log_space=False) else: raise NotImplementedError() trans_r = torch.exp(r) @@ -186,6 +513,29 @@ def dist(self, aux_info, parameters, lib_y): class LogNegativeBinomial(Distribution): + """Log-space negative binomial distribution. + + This distribution represents a negative binomial distribution + parameterized in log space for numerical stability. + + Parameters + ---------- + log_m : torch.Tensor + Log of the mean parameter. + log_r : torch.Tensor + Log of the dispersion parameter. + eps : float, default=1e-8 + Small constant for numerical stability. + validate_args : bool, default=False + Whether to validate distribution arguments. + + Notes + ----- + This distribution is parameterized in log space to avoid numerical + issues with very small or large values. The actual mean and dispersion + are computed as exp(log_m) and exp(log_r) respectively. + """ + def __init__(self, log_m, log_r, eps: float = 1e-8, validate_args=False) -> None: self.log_m = log_m self.log_r = log_r @@ -194,21 +544,80 @@ def __init__(self, log_m, log_r, eps: float = 1e-8, validate_args=False) -> None @property def mean(self): + """Get the mean of the distribution. + + Returns + ------- + torch.Tensor + Mean of the distribution. + """ return torch.exp(self.log_m) @property def theta(self): + """Get the dispersion parameter. + + Returns + ------- + torch.Tensor + Dispersion parameter of the distribution. + """ return torch.exp(self.log_r) @property def variance(self): - return self.mean + (self.mean**2) / self.theta + """Get the variance of the distribution. + + Returns + ------- + torch.Tensor + Variance of the distribution. + """ + m = self.mean + r = self.theta + return m + m * m / r def sample(self, sample_shape): - raise NotImplementedError() + """Generate samples from the distribution. + + Parameters + ---------- + sample_shape : torch.Size + Shape of the samples to generate. + + Returns + ------- + torch.Tensor + Samples from the distribution. + """ + raise NotImplementedError("Sampling not implemented for LogNegativeBinomial") @staticmethod def negative_binomial_log_ver(k, m_log, r_log, eps=1e-8): + """Compute log probability for negative binomial in log space. + + Parameters + ---------- + k : torch.Tensor + Observed values. + m_log : torch.Tensor + Log of mean parameter. + r_log : torch.Tensor + Log of dispersion parameter. + eps : float, default=1e-8 + Small constant for numerical stability. + + Returns + ------- + torch.Tensor + Log probability of the observations. + + Notes + ----- + This function computes the log probability of negative binomial + observations when the parameters are given in log space. It uses + the log-gamma function for numerical stability. + """ # r :D r = torch.exp(r_log) @@ -226,6 +635,27 @@ def log_prob(self, value: torch.Tensor) -> torch.Tensor: class LogNegativeBinomialNoiseModel(NoiseModel): + """Log-space negative binomial noise model. + + This noise model uses a negative binomial distribution parameterized + in log space for better numerical stability. + + Parameters + ---------- + dispersion : {"feature"}, default="feature" + Dispersion parameter modeling strategy. + mean_transformation : {"none"}, default="none" + Transformation to apply to the mean parameter. + library_normalization : {"none", "x_lib", "x_loglib", "div_lib_x_loglib", "x_loglib_all"}, default="x_lib" + Library size normalization method. + + Notes + ----- + This model is similar to NegativeBinomialNoiseModel but uses log-space + parameterization for better numerical stability, especially when dealing + with very small or large values. + """ + def __init__(self, dispersion="feature", mean_transformation="none", library_normalization="x_lib"): super().__init__() self.dispersion = dispersion @@ -234,6 +664,13 @@ def __init__(self, dispersion="feature", mean_transformation="none", library_nor @property def parameters(self): + """Get parameter specification for log negative binomial noise model. + + Returns + ------- + dict + Parameter specification with mean and dispersion parameters. + """ params = { "mean": "no_transformation", } @@ -244,6 +681,20 @@ def parameters(self): return params def initial_transformation(self, x, x_mask=1.0): + """Apply initial transformation for log negative binomial noise model. + + Parameters + ---------- + x : torch.Tensor + Input data tensor. + x_mask : torch.Tensor, default=1.0 + Mask for the input data. + + Returns + ------- + tuple + (transformed_x, aux_info) where aux_info contains library size. + """ x = x aux_info = {} aux_info["x_library_size"] = calculate_library_size(x, x_mask) @@ -252,6 +703,22 @@ def initial_transformation(self, x, x_mask=1.0): return x, aux_info def dist(self, aux_info, parameters, lib_y): + """Create log-space negative binomial distribution. + + Parameters + ---------- + aux_info : dict + Additional information from initial_transformation. + parameters : dict + Dictionary containing 'mean' and 'r' parameters. + lib_y : torch.Tensor + Library size tensor. + + Returns + ------- + LogNegativeBinomial + Log-space negative binomial distribution. + """ mean = parameters["mean"] r = 1.0 + parameters["r"] library_size = lib_y diff --git a/src/drvi/nn_modules/prior.py b/src/drvi/nn_modules/prior.py index 2d24560..d0155f5 100644 --- a/src/drvi/nn_modules/prior.py +++ b/src/drvi/nn_modules/prior.py @@ -1,3 +1,6 @@ +from collections.abc import Callable +from typing import Any + import torch from torch import nn from torch.distributions import Normal, kl_divergence @@ -6,36 +9,154 @@ class Prior(nn.Module): - def __init__(self): + """Abstract base class for prior distributions in variational autoencoders. + + This class defines the interface for different prior distributions that can be + used in variational autoencoders. The prior distribution is used to regularize + the latent space and compute the KL divergence with the posterior. + + Notes + ----- + Subclasses must implement: + - `kl`: Method to compute KL divergence with the posterior distribution + """ + + def __init__(self) -> None: super().__init__() - def kl(self, qz): + def kl(self, qz: Normal) -> torch.Tensor: + """Compute KL divergence between posterior and prior. + + Parameters + ---------- + qz + The posterior distribution (usually Normal). + + Returns + ------- + torch.Tensor + KL divergence between posterior and prior. + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses. + """ raise NotImplementedError() class StandardPrior(Prior): - def __init__(self): + """Standard normal prior distribution. + + This is the most commonly used prior in variational autoencoders. + It assumes a standard normal distribution N(0, I) for the latent variables. + + Examples + -------- + >>> import torch + >>> from torch.distributions import Normal + >>> # Create standard prior + >>> prior = StandardPrior() + >>> # Create a posterior distribution + >>> qz = Normal(torch.randn(10, 5), torch.ones(10, 5)) + >>> # Compute KL divergence + >>> kl = prior.kl(qz) + >>> print(kl.shape) # torch.Size([10, 5]) + """ + + def __init__(self) -> None: super().__init__() - def kl(self, qz): + def kl(self, qz: Normal) -> torch.Tensor: + """Compute KL divergence with standard normal prior. + + Parameters + ---------- + qz + The posterior distribution (must be Normal). + + Returns + ------- + torch.Tensor + KL divergence between qz and N(0, I). + + Notes + ----- + For Normal distributions, the KL divergence has a closed-form solution: + KL(q||p) = 0.5 * (μ² + σ² - log(σ²) - 1) + where μ and σ are the mean and standard deviation of q. + """ # 1 x N assert isinstance(qz, Normal) return kl_divergence(qz, Normal(torch.zeros_like(qz.mean), torch.ones_like(qz.mean))) class VampPrior(Prior): + """VampPrior (Variational Mixture of Posteriors) prior distribution. + + This prior uses a mixture of posteriors as the prior distribution, + which can capture more complex structure in the data than a standard + normal prior. + + This is adapted from https://github.com/jmtomczak/intro_dgm/main/vaes/vae_priors_example.ipynb + + Parameters + ---------- + n_components + Number of mixture components. + encoder + Encoder network to compute posterior parameters for pseudo-inputs. + model_input + Dictionary containing input data for pseudo-inputs. + trainable_keys + Keys in model_input that should be trainable parameters. + fixed_keys + Keys in model_input that should be fixed parameters. + input_type + Type of input format expected by the encoder. + preparation_function + Function to prepare inputs for the encoder. + + Notes + ----- + VampPrior learns a set of pseudo-inputs and uses the posterior + distributions of these pseudo-inputs as mixture components for + the prior. This allows the prior to adapt to the data distribution. + + The prior is computed as: + p(z) = Σ_k w_k * q(z|u_k) + where u_k are the pseudo-inputs and w_k are mixing weights. + + Examples + -------- + >>> import torch + >>> from torch import nn + >>> # Create a simple encoder + >>> class SimpleEncoder(nn.Module): + ... def __init__(self): + ... super().__init__() + ... self.fc = nn.Linear(10, 4) + ... + ... def forward(self, x): + ... return self.fc(x), torch.ones_like(self.fc(x)) + >>> # Create model input + >>> model_input = {"x": torch.randn(5, 10)} + >>> # Create VampPrior + >>> prior = VampPrior(n_components=5, encoder=SimpleEncoder(), model_input=model_input, trainable_keys=("x",)) + """ + # Adapted from https://github.com/jmtomczak/intro_dgm/main/vaes/vae_priors_example.ipynb # K - components, I - inputs, L - latent, N - samples def __init__( self, - n_components, - encoder, - model_input, - trainable_keys=("x",), - fixed_keys=(), - input_type="scvi", - preparation_function=None, - ): + n_components: int, + encoder: nn.Module, + model_input: dict[str, Any], + trainable_keys: tuple[str, ...] = ("x",), + fixed_keys: tuple[str, ...] = (), + input_type: str = "scvi", + preparation_function: Callable | None = None, + ) -> None: super().__init__() self.encoder = encoder @@ -59,7 +180,19 @@ def __init__( # mixing weights self.w = torch.nn.Parameter(torch.zeros(n_components, 1, 1)) # K x 1 x 1 - def get_params(self): + def get_params(self) -> tuple[torch.Tensor, torch.Tensor]: + """Get the parameters of the mixture components. + + Returns + ------- + tuple + (means, variances) of the mixture components. + + Notes + ----- + This method computes the posterior parameters for each pseudo-input + by passing them through the encoder network. + """ # u->encoder->mean, var original_mode = self.encoder.training self.encoder.train(False) @@ -67,6 +200,8 @@ def get_params(self): z = self.encoder({**self.pi_aux_data, **self.pi_tensor_data}) output = z["qz_mean"], z["qz_var"] elif self.input_type == "scvi": + if self.preparation_function is None: + raise ValueError("preparation_function must be provided for scvi input type") x, args, kwargs = self.preparation_function({**self.pi_aux_data, **self.pi_tensor_data}) q_m, q_v, latent = self.encoder(x, *args, **kwargs) output = q_m, q_v @@ -76,7 +211,25 @@ def get_params(self): self.encoder.train(original_mode) return output # (K x L), (K x L) - def log_prob(self, z): + def log_prob(self, z: torch.Tensor) -> torch.Tensor: + """Compute log probability of latent variables under the prior. + + Parameters + ---------- + z + Latent variables with shape (N, L). + + Returns + ------- + torch.Tensor + Log probability with shape (N, L). + + Notes + ----- + The log probability is computed as: + log p(z) = log Σ_k w_k * q(z|u_k) + where the sum is computed using logsumexp for numerical stability. + """ # Mixture of gaussian computed on K x N x L z = z.unsqueeze(0) # 1 x N x L @@ -94,32 +247,103 @@ def log_prob(self, z): return log_prob # N x L - def kl(self, qz): + def kl(self, qz: Normal) -> torch.Tensor: + """Compute KL divergence using Monte Carlo estimation. + + Parameters + ---------- + qz + The posterior distribution. + + Returns + ------- + torch.Tensor + KL divergence between posterior and VampPrior. + + Notes + ----- + The KL divergence is estimated using Monte Carlo sampling: + KL(q||p) = E_q[log q(z) - log p(z)] + where z ~ q(z) is sampled from the posterior. + """ assert isinstance(qz, Normal) z = qz.rsample() return qz.log_prob(z) - self.log_prob(z) - def get_extra_state(self): + def get_extra_state(self) -> dict[str, Any]: + """Get extra state for serialization. + + Returns + ------- + dict + Extra state information. + """ return { "pi_aux_data": self.pi_aux_data, "input_type": self.input_type, } - def set_extra_state(self, state): + def set_extra_state(self, state: dict[str, Any]) -> None: + """Set extra state from serialization. + + Parameters + ---------- + state + Extra state information. + """ self.pi_aux_data = state["pi_aux_data"] self.input_type = state["input_type"] class GaussianMixtureModelPrior(Prior): + """Gaussian Mixture Model prior distribution. + + This prior uses a mixture of Gaussian distributions with learnable + parameters. Unlike VampPrior, the mixture components are not tied + to the encoder network. + + Parameters + ---------- + n_components + Number of mixture components. + n_latent + Dimensionality of the latent space. + data + Initial means and variances as (means, variances). + trainable_priors + Whether the prior parameters should be trainable. + + Notes + ----- + This prior learns the means, variances, and mixing weights of a + Gaussian mixture model directly. It's simpler than VampPrior but + may not capture data-specific structure as well. + + The prior is computed as: + p(z) = Σ_k w_k * N(z|μ_k, σ_k²) + where μ_k, σ_k², and w_k are learnable parameters. + + Examples + -------- + >>> import torch + >>> # Create GMM prior + >>> prior = GaussianMixtureModelPrior(n_components=3, n_latent=5) + >>> # Create a posterior distribution + >>> qz = torch.distributions.Normal(torch.randn(10, 5), torch.ones(10, 5)) + >>> # Compute KL divergence + >>> kl = prior.kl(qz) + >>> print(kl.shape) # torch.Size([10, 5]) + """ + # Based on VampPrior class def __init__( self, - n_components, - n_latent, - data=None, - trainable_priors=True, - ): + n_components: int, + n_latent: int, + data: tuple[torch.Tensor, torch.Tensor] | None = None, + trainable_priors: bool = True, + ) -> None: # Do we need python 2 compatibility? super().__init__() @@ -135,7 +359,25 @@ def __init__( # mixing weights self.w = torch.nn.Parameter(torch.zeros(self.p_m.shape[0], 1, 1)) # K x 1 x 1 - def log_prob(self, z): + def log_prob(self, z: torch.Tensor) -> torch.Tensor: + """Compute log probability of latent variables under the prior. + + Parameters + ---------- + z + Latent variables with shape (N, L). + + Returns + ------- + torch.Tensor + Log probability with shape (N, L). + + Notes + ----- + The log probability is computed as: + log p(z) = log Σ_k w_k * N(z|μ_k, σ_k²) + where the sum is computed using logsumexp for numerical stability. + """ # Mixture of gaussian computed on K x N x L z = z.unsqueeze(0) # 1 x N x L @@ -151,7 +393,25 @@ def log_prob(self, z): return log_prob # N x L - def kl(self, qz): + def kl(self, qz: Normal) -> torch.Tensor: + """Compute KL divergence using Monte Carlo estimation. + + Parameters + ---------- + qz + The posterior distribution. + + Returns + ------- + torch.Tensor + KL divergence between posterior and GMM prior. + + Notes + ----- + The KL divergence is estimated using Monte Carlo sampling: + KL(q||p) = E_q[log q(z) - log p(z)] + where z ~ q(z) is sampled from the posterior. + """ assert isinstance(qz, Normal) z = qz.rsample() return qz.log_prob(z) - self.log_prob(z) diff --git a/src/drvi/scvi_tools_based/model/_drvi.py b/src/drvi/scvi_tools_based/model/_drvi.py index 1f232eb..4e6b60f 100644 --- a/src/drvi/scvi_tools_based/model/_drvi.py +++ b/src/drvi/scvi_tools_based/model/_drvi.py @@ -1,6 +1,6 @@ import logging from collections.abc import Sequence -from typing import Literal +from typing import Any, Literal import numpy as np from anndata import AnnData @@ -29,13 +29,12 @@ class DRVI(VAEMixin, DRVIArchesMixin, UnsupervisedTrainingMixin, BaseModelClass, GenerativeMixin): - """ - DRVI model based on scvi-tools skelethon + """DRVI model based on scvi-tools framework for disentangled representation learning. Parameters ---------- adata - AnnData object that has been registered via :meth:`~drvi.model.DRVI.setup_anndata`. + AnnData object or MerlinData object that has been registered via :meth:`~drvi.model.DRVI.setup_anndata`. n_latent Dimensionality of the latent space. encoder_dims @@ -43,19 +42,21 @@ class DRVI(VAEMixin, DRVIArchesMixin, UnsupervisedTrainingMixin, BaseModelClass, decoder_dims Number of nodes in hidden layers of the decoder. prior - Prior model. defaults to normal. + Prior model type. prior_init_obs - When initializing priors from data, these observations are used. + When using "gmm_x" or "vamp_x" priors, these observations are used to initialize the prior parameters. + Number of observations must match the x value in the prior name. categorical_covariates - Categorical Covariates as a list of texts. You can specify emb dimension by appending @dim to each cpvariate. + List of categorical covariates to condition on. Each covariate can specify its embedding dimension + by appending @dim (e.g. "batch@32"). Default embedding dimension is 10. **model_kwargs - Keyword args for :class:`~drvi.model.DRVIModule` + Additional keyword arguments passed to :class:`~drvi.model.DRVIModule`. Examples -------- >>> adata = anndata.read_h5ad(path_to_anndata) - >>> drvi.DRVI.setup_anndata(adata, categorical_covariate_keys=["batch"]) - >>> vae = drvi.DRVI(adata) + >>> drvi.model.DRVI.setup_anndata(adata, categorical_covariate_keys=["batch"]) + >>> vae = drvi.model.DRVI(adata) >>> vae.train() >>> adata.obsm["latent"] = vae.get_latent_representation() """ @@ -70,7 +71,7 @@ def __init__( prior_init_obs: np.ndarray | None = None, categorical_covariates: list[str] = (), **model_kwargs, - ): + ) -> None: super().__init__(adata) # TODO: Remove later. Currently used to detect autoreload problems sooner. @@ -81,7 +82,7 @@ def __init__( else: raise ValueError( "Only AnnData and MerlinData are supported. " - "If you have passed an instalce of MerlinData and still get this error, " + "If you have passed an instance of MerlinData and still get this error, " "make sure merlin is installed as a dependency." ) @@ -181,7 +182,31 @@ def setup_merlin_data( categorical_covariate_keys: list[str] | None = None, continuous_covariate_keys: list[str] | None = None, **kwargs, - ): + ) -> None: + """Setup MerlinData for use with DRVI. + + Parameters + ---------- + merlin_data + MerlinData object to register. + labels_key + Key in `merlin_data` for labels. + layer + key in `merlin_data` to use as input. + is_count_data + Whether the data is count data. + categorical_covariate_keys + List of categorical covariate keys in `merlin_data`. + continuous_covariate_keys + List of continuous covariate keys in `merlin_data`. + **kwargs + Additional keyword arguments passed to the MerlinDataManager registration. + + Returns + ------- + None + This method sets up the class for use with MerlinData and does not return anything. + """ setup_method_args = cls._get_setup_method_args(**locals()) setup_method_args["drvi_version"] = drvi.__version__ @@ -197,13 +222,13 @@ def setup_merlin_data( def _make_data_loader( self, - adata, + adata: AnnData | MerlinData, indices: Sequence[int] | None = None, batch_size: int | None = None, shuffle: bool = False, - data_loader_class=None, + data_loader_class: type | None = None, **data_loader_kwargs, - ): + ) -> Any: """Create a AnnDataLoader object for data iteration. Parameters @@ -213,13 +238,18 @@ def _make_data_loader( indices Indices of cells in adata to use. If `None`, all cells are used. batch_size - Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`. + Minibatch size for data loading into model. shuffle - Whether observations are shuffled each iteration though + Whether observations are shuffled each iteration though. data_loader_class - Class to use for data loader + Class to use for data loader. data_loader_kwargs - Kwargs to the class-specific data loader class + Kwargs to the class-specific data loader class. + + Returns + ------- + Any + Data loader object for iteration. """ if isinstance(adata, AnnData): return super()._make_data_loader( @@ -244,6 +274,6 @@ def _make_data_loader( else: raise ValueError( "Only AnnData and MerlinData are supported. " - "If you have passed an instalce of MerlinData and still get this error, " + "If you have passed an instance of MerlinData and still get this error, " "make sure merlin is installed as a dependency." ) diff --git a/src/drvi/scvi_tools_based/model/base/_archesmixin.py b/src/drvi/scvi_tools_based/model/base/_archesmixin.py index 445ae8b..4e8d5c8 100644 --- a/src/drvi/scvi_tools_based/model/base/_archesmixin.py +++ b/src/drvi/scvi_tools_based/model/base/_archesmixin.py @@ -1,4 +1,5 @@ import logging +from collections.abc import Sequence import torch from anndata import AnnData @@ -211,18 +212,18 @@ def load_query_data( def _set_params_online_update( - module, - reloaded_tensor_keys, - unfrozen, - previous_n_cats_per_cov, - n_cats_per_cov, - freeze_dropout, - freeze_shared_emb, - freeze_encoder, - freeze_decoder, - freeze_batchnorm_encoder, - freeze_batchnorm_decoder, -): + module: nn.Module, + reloaded_tensor_keys: list[str], + unfrozen: bool, + previous_n_cats_per_cov: Sequence[int] | None, + n_cats_per_cov: Sequence[int] | None, + freeze_dropout: bool, + freeze_shared_emb: bool, + freeze_encoder: bool, + freeze_decoder: bool, + freeze_batchnorm_encoder: bool, + freeze_batchnorm_decoder: bool, +) -> None: """Freeze parts of network for scArches.""" # do nothing if unfrozen if unfrozen: @@ -234,7 +235,7 @@ def _set_params_online_update( assert tuple(module.shared_covariate_emb.num_embeddings) == tuple(n_cats_per_cov) module.shared_covariate_emb.freeze_top_embs(previous_n_cats_per_cov) - def requires_grad(key): + def requires_grad(key: str) -> bool: if not freeze_decoder and "decoder" in key: return True if not freeze_encoder and "z_encoder" in key: @@ -273,4 +274,4 @@ def requires_grad(key): if freeze_batchnorm: print(f"Freezing normalization layer {key}") mod.freeze(freeze_batchnorm) - mod.momentum = 0 + mod.momentum = 0.0 diff --git a/src/drvi/scvi_tools_based/model/base/_generative_mixin.py b/src/drvi/scvi_tools_based/model/base/_generative_mixin.py index 70c0c92..ebda5eb 100644 --- a/src/drvi/scvi_tools_based/model/base/_generative_mixin.py +++ b/src/drvi/scvi_tools_based/model/base/_generative_mixin.py @@ -1,5 +1,6 @@ import logging from collections.abc import Callable, Sequence +from typing import Any import numpy as np import scvi @@ -12,7 +13,25 @@ class GenerativeMixin: - """Mixins to interpret generative part of the method.""" + """Mixin class for generative model interpretation and analysis. + + This mixin provides methods for analyzing the generative part of variational + autoencoder models. It includes utilities for decoding latent representations, + iterating over model outputs, and analyzing the effects of different model + components on the reconstruction. + + Notes + ----- + This mixin is designed to work with scVI-based models that have a generative + component. It provides high-level interfaces for: + - Decoding latent samples to reconstruct data + - Iterating over model outputs with custom functions + - Analyzing reconstruction effects of different model splits + - Computing maximum effects across the data distribution + + All methods operate in inference mode and handle batching automatically + to manage memory usage with large datasets. + """ @torch.inference_mode() def iterate_on_decoded_latent_samples( @@ -23,32 +42,69 @@ def iterate_on_decoded_latent_samples( lib: np.ndarray | None = None, cat_values: np.ndarray | None = None, cont_values: np.ndarray | None = None, - batch_size=scvi.settings.batch_size, + batch_size: int = scvi.settings.batch_size, map_cat_values: bool = False, ) -> np.ndarray: - r"""Iterate over decoder outputs and aggregate the results. + """Iterate over decoder outputs and aggregate the results. + + This method processes latent samples through the generative model in batches, + applies a custom function to each batch output, and aggregates the results. Parameters ---------- z - Latent samples. + Latent samples with shape (n_samples, n_latent). step_func Function to apply to the decoder output at each step. - generative_outputs and a store variable are given to the function + Should accept (generative_outputs, store) as arguments. aggregation_func - Function to aggregate the step results. + Function to aggregate the step results from the store. + Should accept the store list and return the final result. lib - Library size array. + Library size array with shape (n_samples,). + If None, defaults to 1e4 for all samples. cat_values - Categorical covariates (required if model has categorical key). + Categorical covariates with shape (n_samples, n_cat_covs). + Required if model has categorical covariates. cont_values - Continuous covariates. + Continuous covariates with shape (n_samples, n_cont_covs). batch_size - Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`. + Minibatch size for data loading into model. map_cat_values - Whether to map categorical covariates in cat_values to integers based on anndata manager pipeline + Whether to map categorical covariates to integers based on + the AnnData manager pipeline. + + Returns + ------- + np.ndarray + Aggregated results from processing all latent samples. + + Notes + ----- + This method operates in inference mode and processes data in batches + to manage memory usage. The step_func receives the generative outputs + and a store variable for accumulating results. + + If map_cat_values is True, categorical values are automatically mapped + to integers using the model's category mappings. + + Examples + -------- + >>> import numpy as np + >>> # Define custom step function to extract means + >>> def extract_means(gen_output, store): + ... store.append(gen_output["params"]["mean"].detach().cpu()) + >>> # Define aggregation function to concatenate results + >>> def concatenate_results(store): + ... return torch.cat(store, dim=0).numpy() + >>> # Process latent samples + >>> z = np.random.randn(50, 32) # assuming 32 latent dimensions + >>> result = model.iterate_on_decoded_latent_samples( + ... z=z, step_func=extract_means, aggregation_func=concatenate_results + ... ) + >>> print(result.shape) # (50, n_genes) """ - store = [] + store: list[Any] = [] self.module.eval() if cat_values is not None and map_cat_values: @@ -99,30 +155,64 @@ def decode_latent_samples( lib: np.ndarray | None = None, cat_values: np.ndarray | None = None, cont_values: np.ndarray | None = None, - batch_size=scvi.settings.batch_size, + batch_size: int = scvi.settings.batch_size, map_cat_values: bool = False, ) -> np.ndarray: - r"""Return the distribution produces by the decoder for the given latent samples. + r"""Return the distribution produced by the decoder for the given latent samples. - This is typically considered as :math:`p(x \mid z)`. + This method computes :math:`p(x \mid z)`, the reconstruction distribution + for given latent samples. It returns the mean of the reconstruction + distribution for each sample. Parameters ---------- z - Latent samples. + Latent samples with shape (n_samples, n_latent). lib - Library size array. + Library size array with shape (n_samples,). + If None, defaults to 1e4 for all samples. cat_values - Categorical covariates (required if model has categorical key). + Categorical covariates with shape (n_samples, n_cat_covs). + Required if model has categorical covariates. cont_values - Continuous covariates. + Continuous covariates with shape (n_samples, n_cont_covs). batch_size - Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`. - map_cat_valuse - Whether to map categorical covariates in cat_values to integers based on anndata manager pipeline + Minibatch size for data loading into model. + map_cat_values + Whether to map categorical covariates to integers based on + the AnnData manager pipeline. + + Returns + ------- + np.ndarray + Reconstructed means with shape (n_samples, n_genes). + + Notes + ----- + This method is equivalent to computing the expected value of the + reconstruction distribution :math:`E[p(x \mid z)]`. It's useful for: + - Generating synthetic data from latent samples + - Analyzing model reconstructions + - Visualizing the generative capabilities of the model + + Examples + -------- + >>> import numpy as np + >>> # Generate random latent samples + >>> z = np.random.randn(100, 32) # assuming 32 latent dimensions + >>> # Decode to get reconstructed means + >>> reconstructed = model.decode_latent_samples(z) + >>> print(reconstructed.shape) # (100, n_genes) + >>> # With categorical covariates + >>> cat_covs = np.array([0, 1, 0, 1] * 25) # batch labels + >>> reconstructed = model.decode_latent_samples(z, cat_values=cat_covs) """ - step_func = lambda gen_output, store: store.append(gen_output["params"]["mean"].detach().cpu()) - aggregation_func = lambda store: torch.cat(store, dim=0).numpy(force=True) + + def step_func(gen_output: dict[str, Any], store: list[Any]) -> None: + store.append(gen_output["params"]["mean"].detach().cpu()) + + def aggregation_func(store: list[Any]) -> np.ndarray: + return torch.cat(store, dim=0).numpy(force=True) return self.iterate_on_decoded_latent_samples( z=z, @@ -145,29 +235,67 @@ def iterate_on_ae_output( batch_size: int | None = None, deterministic: bool = False, ) -> np.ndarray: - r"""Iterate over autoencoder outputs and aggregate the results. + """Iterate over autoencoder outputs and aggregate the results. + + This method processes data through the full autoencoder (encoder + decoder) + and applies custom functions to analyze the outputs. Parameters ---------- adata AnnData object with equivalent structure to initial AnnData. - If `None`, defaults to the AnnData object used to initialize the model. + If None, defaults to the AnnData object used to initialize the model. step_func Function to apply to the autoencoder output at each step. - inference_outputs, generative_outputs, losses, and a store variable are given to the function + Should accept (inference_outputs, generative_outputs, losses, store) + as arguments. aggregation_func - Function to aggregate the step results. + Function to aggregate the step results from the store. + Should accept the store list and return the final result. indices - Indices of cells in adata to use. If `None`, all cells are used. + Indices of cells in adata to use. If None, all cells are used. batch_size - Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`. + Minibatch size for data loading into model. + Defaults to scvi.settings.batch_size. deterministic - Makes model fully deterministic (e.g. no sampling in the bottleneck). + Makes model fully deterministic (e.g., no sampling in the bottleneck). + + Returns + ------- + np.ndarray + Aggregated results from processing all data through the autoencoder. + + Notes + ----- + This method processes data through both the encoder and decoder components + of the model. The step_func receives: + - inference_outputs: Outputs from the encoder (latent variables, etc.) + - generative_outputs: Outputs from the decoder (reconstruction parameters) + - losses: Training losses for the batch + - store: List for accumulating results + + When deterministic=True, the model operates without stochastic sampling, + which is useful for reproducible analysis. + + Examples + -------- + >>> import anndata as ad + >>> # Define function to extract latent means + >>> def extract_latent_means(inference_outputs, generative_outputs, losses, store): + ... store.append(inference_outputs["qz_m"].detach().cpu()) + >>> # Define aggregation function + >>> def concatenate_latents(store): + ... return torch.cat(store, dim=0).numpy() + >>> # Process data through autoencoder + >>> latents = model.iterate_on_ae_output( + ... adata=adata, step_func=extract_latent_means, aggregation_func=concatenate_latents, deterministic=True + ... ) + >>> print(latents.shape) # (n_cells, n_latent) """ adata = self._validate_anndata(adata) data_loader = self._make_data_loader(adata=adata, indices=indices, batch_size=batch_size) - store = [] + store: list[Any] = [] try: if deterministic: self.module.fully_deterministic = True @@ -190,26 +318,62 @@ def get_reconstruction_effect_of_each_split( add_to_counts: float = 1.0, aggregate_over_cells: bool = True, deterministic: bool = True, - **kwargs, + **kwargs: Any, ) -> np.ndarray: - r"""Return the effect of each split on the reconstructed per sample. + """Return the effect of each split on the reconstructed expression per sample. + + This method analyzes how different model splits contribute to the + reconstruction of gene expression values. Parameters ---------- adata AnnData object with equivalent structure to initial AnnData. - If `None`, defaults to the AnnData object used to initialize the model. + If None, defaults to the AnnData object used to initialize the model. add_to_counts Value to add to the counts before computing the logarithm. + Used for numerical stability in log-space calculations. aggregate_over_cells Whether to aggregate the effect over cells and return a value per dimension. + If False, returns per-cell effects. deterministic - Makes model fully deterministic (e.g. no sampling in the bottleneck). - kwargs + Makes model fully deterministic (e.g., no sampling in the bottleneck). + **kwargs Additional keyword arguments for the `iterate_on_ae_output` method. + + Returns + ------- + np.ndarray + Effect of each split on reconstruction. + Shape depends on aggregate_over_cells: + - If True: (n_splits,) - aggregated effects per split + - If False: (n_cells, n_splits) - per-cell effects + + Notes + ----- + This method computes the contribution of each model split to the + reconstruction. The calculation depends on the model's split_aggregation: + + - "logsumexp": Uses log-space softmax aggregation + - "sum": Uses absolute value summation + + The effect is computed by analyzing how each split contributes to + the final reconstruction parameters. + + Examples + -------- + >>> # Get aggregated effects across all cells + >>> effects = model.get_reconstruction_effect_of_each_split() + >>> print(effects.shape) # (n_splits,) + >>> print("Effect of each split:", effects) + >>> # Get per-cell effects + >>> cell_effects = model.get_reconstruction_effect_of_each_split(aggregate_over_cells=False) + >>> print(cell_effects.shape) # (n_cells, n_splits) """ - def calculate_effect(inference_outputs, generative_outputs, losses, store): + def calculate_effect( + inference_outputs: dict[str, Any], generative_outputs: dict[str, Any], losses: Any, store: list[Any] + ) -> None: if self.module.split_aggregation == "logsumexp": log_mean_params = generative_outputs["original_params"]["mean"] # n_samples x n_splits x n_genes log_mean_params = F.pad( @@ -224,9 +388,9 @@ def calculate_effect(inference_outputs, generative_outputs, losses, store): ) # n_samples x n_splits else: raise NotImplementedError("Only logsumexp and sum aggregations are supported for now.") - return store.append(effect_share.detach().cpu()) + store.append(effect_share.detach().cpu()) - def aggregate_effects(store): + def aggregate_effects(store: list[Any]) -> np.ndarray: return torch.cat(store, dim=0).numpy(force=True) output = self.iterate_on_ae_output( @@ -242,29 +406,29 @@ def aggregate_effects(store): return output - # @torch.inference_mode() def get_max_effect_of_splits_within_distribution( self, adata: AnnData | None = None, add_to_counts: float = 1.0, deterministic: bool = True, - **kwargs, - ): - r""" - Return the max effect of each split on the reconstructed expression params for all genes. + **kwargs: Any, + ) -> np.ndarray: + """Return the maximum effect of each split on the reconstructed expression params for all genes. - These values are empirical and inexact for de reasoning of dimensions. + This method computes the maximum contribution of each split across all + samples in the dataset, providing a global view of split importance. Parameters ---------- adata AnnData object with equivalent structure to initial AnnData. - If `None`, defaults to the AnnData object used to initialize the model. + If None, defaults to the AnnData object used to initialize the model. add_to_counts Value to add to the counts before computing the logarithm. + Used for numerical stability in log-space calculations. deterministic - Makes model fully deterministic (e.g. no sampling in the bottleneck). - kwargs + Makes model fully deterministic (e.g., no sampling in the bottleneck). + **kwargs Additional keyword arguments for the `iterate_on_ae_output` method. Returns @@ -273,11 +437,25 @@ def get_max_effect_of_splits_within_distribution( Max effect of each split on the reconstructed expression params for all genes - ------- - Example usage - ------- + Maximum effect of each split on the reconstructed expression params. + Shape: (n_splits, n_genes) + + Notes + ----- + This function is experimental. Please use interpretability pipeline or DE instead. + + The calculation depends on the model's split_aggregation: + - "logsumexp": Uses log-space softmax aggregation + - "sum": Uses absolute value summation - import utils.plotting._interpretability >>> effects = model.get_max_effect_of_splits_within_distribution(add_to_counts=0.1) + Examples + -------- + >>> import utils.plotting._interpretability + >>> + >>> # Get empirical maximum effects + >>> max_effects = model.get_max_effect_of_splits_within_distribution(add_to_counts=0.1) + >>> print(max_effects.shape) # (n_splits, n_genes) + >>> >>> effect_data = ( ... pd.DataFrame( ... effects, @@ -292,7 +470,9 @@ def get_max_effect_of_splits_within_distribution( >>> utils.plotting._interpretability._umap_of_relevant_genes(adata, embed, plot_info, dim_subset=["DR 1"]) """ - def calculate_effect(inference_outputs, generative_outputs, losses, store): + def calculate_effect( + inference_outputs: dict[str, Any], generative_outputs: dict[str, Any], losses: Any, store: list[Any] + ) -> None: if self.module.split_aggregation == "logsumexp": log_mean_params = generative_outputs["original_params"]["mean"] # n_samples x n_splits x n_genes log_mean_params = F.pad( @@ -309,7 +489,7 @@ def calculate_effect(inference_outputs, generative_outputs, losses, store): else: store[0] = np.maximum(store[0], effect_share) - def aggregate_effects(store): + def aggregate_effects(store: list[Any]) -> np.ndarray: return store[0] output = self.iterate_on_ae_output( diff --git a/src/drvi/scvi_tools_based/module/_drvi.py b/src/drvi/scvi_tools_based/module/_drvi.py index bbba0ba..ab6d545 100644 --- a/src/drvi/scvi_tools_based/module/_drvi.py +++ b/src/drvi/scvi_tools_based/module/_drvi.py @@ -1,5 +1,5 @@ from collections.abc import Callable, Iterable, Sequence -from typing import Literal +from typing import Any, Literal import numpy as np import torch @@ -23,74 +23,79 @@ class DRVIModule(BaseModuleClass): - """ - Skeleton Variational auto-encoder model. - - Here we implement a basic version of scVI's underlying VAE :cite:p:`Lopez18`. - This implementation is for instructional purposes only. + """DRVI (Disentangled Representation Variational Inference) pytorch module. Parameters ---------- n_input - Number of input genes + Number of input genes. n_latent - Dimensionality of the latent space + Dimensionality of the latent space. n_split_latent - Number of splits in the latent space. -1 means split all dimensions (n_split_latent=n_latent). + Number of splits in the latent space. -1 means split all dimensions + (n_split_latent=n_latent). split_aggregation - How to aggregate splits in the last layer of the decoder + How to aggregate splits in the last layer of the decoder. split_method - How to make splits. - "split" for splitting the latent space, - "power" for transforming the latent space to n_split vectors of size n_latent - "split_map" for splitting the latent space then map each to latent space using unique transformations + How to make splits: + - "split" : Split the latent space + - "power" : Transform the latent space to n_split vectors of size n_latent + - "split_map" : Split the latent space then map each to latent space using unique transformations decoder_reuse_weights - Were to reuse the weights of the decoder layers when using splitting - Possible values are 'everywhere', 'last', 'intermediate', 'nowhere'. Defaults to "everywhere"/ + Where to reuse the weights of the decoder layers when using splitting. encoder_dims Number of nodes in hidden layers of the encoder. decoder_dims Number of nodes in hidden layers of the decoder. n_cats_per_cov - Number of categories for each categorical covariate + Number of categories for each categorical covariate. n_continuous_cov - Number of continuous covariates + Number of continuous covariates. encode_covariates - Whether to concatenate covariates to expression in encoder + Whether to concatenate covariates to expression in encoder. deeply_inject_covariates - Whether to concatenate covariates into output of hidden layers in encoder/decoder. This option - only applies when `n_layers` >= 1. The covariates are concatenated to the input of subsequent hidden layers. + Whether to concatenate covariates into output of hidden layers in encoder/decoder. + This option only applies when `n_layers` >= 1. The covariates are concatenated + to the input of subsequent hidden layers. categorical_covariate_dims - Emb dim of covariate keys if applicable + Embedding dimension of covariate keys if applicable. covariate_modeling_strategy - The strategy model takes to remove covariates + The strategy model takes to remove covariates. use_batch_norm Whether to use batch norm in layers. affine_batch_norm Whether to use affine batch norm in layers. use_layer_norm Whether to use layer norm in layers. + fill_in_the_blanks_ratio + Ratio for fill-in-the-blanks training. input_dropout_rate - Dropout rate to apply to the input + Dropout rate to apply to the input. encoder_dropout_rate - Dropout rate to apply to each of the encoder hidden layers + Dropout rate to apply to each of the encoder hidden layers. decoder_dropout_rate - Dropout rate to apply to each of the decoder hidden layers + Dropout rate to apply to each of the decoder hidden layers. gene_likelihood - gene likelihood model + Gene likelihood model. Options include: + - "normal", "normal_v", "normal_sv" : Normal distributions + - "poisson", "poisson_orig" : Poisson distributions + - "nb", "nb_sv", "nb_libnorm", "nb_loglib_rec", "nb_libnorm_loglib_rec", "nb_loglibnorm_all", "nb_orig", "nb_softmax", "nb_softplus", "nb_none", "nb_orig_libnorm" : Negative binomial distributions + - "pnb", "pnb_sv", "pnb_softmax" : Log negative binomial distributions prior - Prior model. defaults to normal. + Prior model. prior_init_dataloader Dataloader constructed to initialize the prior (or maintain in vamp). var_activation - The activation function to ensure positivity of the variatinal distribution. Defaults to "exp". + The activation function to ensure positivity of the variational distribution. + Options include "exp", "pow2", "2sig" or a custom callable. mean_activation - The activation function at the end of mean encoder. Defaults to "identity". - Possible values are "identity", "relu", "leaky_relu", "leaky_relu_{slope}", "elu", "elu_{min_vaule}". + The activation function at the end of mean encoder. + Options include "identity", "relu", "leaky_relu", "leaky_relu_{slope}", + "elu", "elu_{min_value}" or a custom callable. encoder_layer_factory - A layer Factory instance for build encoder layers + A layer Factory instance for building encoder layers. decoder_layer_factory - A layer Factory instance for build decoder layers + A layer Factory instance for building decoder layers. extra_encoder_kwargs Extra keyword arguments passed into encoder. extra_decoder_kwargs @@ -141,6 +146,8 @@ def __init__( "nb_loglibnorm_all", "nb_orig", "nb_softmax", + "nb_softplus", + "nb_none", "nb_orig_libnorm", "pnb", "pnb_sv", @@ -148,13 +155,13 @@ def __init__( ] = "pnb_softmax", prior: Literal["normal", "gmm_x", "vamp_x"] = "normal", prior_init_dataloader: DataLoader | None = None, - var_activation: Callable | Literal["exp", "pow2", "2sig"] | Callable = "exp", + var_activation: Callable | Literal["exp", "pow2", "2sig"] = "exp", mean_activation: Callable | str = "identity", - encoder_layer_factory: LayerFactory = None, - decoder_layer_factory: LayerFactory = None, - extra_encoder_kwargs: dict | None = None, - extra_decoder_kwargs: dict | None = None, - ): + encoder_layer_factory: LayerFactory | None = None, + decoder_layer_factory: LayerFactory | None = None, + extra_encoder_kwargs: dict[str, Any] | None = None, + extra_decoder_kwargs: dict[str, Any] | None = None, + ) -> None: super().__init__() self.n_latent = n_latent if n_split_latent == -1: @@ -236,7 +243,24 @@ def __init__( self.prior = self._construct_prior(prior, prior_init_dataloader) self.fully_deterministic = False - def _construct_gene_likelihood_module(self, gene_likelihood): + def _construct_gene_likelihood_module(self, gene_likelihood: str) -> Any: + """Construct the gene likelihood module based on the specified type. + + Parameters + ---------- + gene_likelihood + Type of gene likelihood model to construct. + + Returns + ------- + object + Constructed gene likelihood module. + + Raises + ------ + NotImplementedError + If the gene likelihood type is not supported. + """ if gene_likelihood == "normal": return NormalNoiseModel(model_var="fixed=1") elif gene_likelihood == "normal_v": @@ -261,6 +285,14 @@ def _construct_gene_likelihood_module(self, gene_likelihood): return NegativeBinomialNoiseModel( dispersion="feature", mean_transformation="softmax", library_normalization="none" ) + elif gene_likelihood in ["nb_softplus"]: + return NegativeBinomialNoiseModel( + dispersion="feature", mean_transformation="softplus", library_normalization="none" + ) + elif gene_likelihood in ["nb_none"]: + return NegativeBinomialNoiseModel( + dispersion="feature", mean_transformation="none", library_normalization="none" + ) elif gene_likelihood == "nb_orig_libnorm": return NegativeBinomialNoiseModel( dispersion="feature", mean_transformation="softmax", library_normalization="x_lib" @@ -274,7 +306,28 @@ def _construct_gene_likelihood_module(self, gene_likelihood): else: raise NotImplementedError() - def _construct_prior(self, prior, prior_init_dataloader=None): + def _construct_prior(self, prior: str, prior_init_dataloader: DataLoader | None = None) -> Any: + """Construct the prior model based on the specified type. + + Parameters + ---------- + prior + Type of prior model to construct. + prior_init_dataloader + Dataloader for initializing the prior (required for some prior types). + + Returns + ------- + object + Constructed prior model. + + Raises + ------ + ValueError + If VaMP prior is specified without a dataloader. + NotImplementedError + If the prior type is not supported. + """ if prior == "normal": return StandardPrior() elif prior.startswith("gmm_"): @@ -289,7 +342,7 @@ def _construct_prior(self, prior, prior_init_dataloader=None): n_components = int(prior.split("_")[1]) if prior_init_dataloader is not None: - def preparation_function(prepared_input): + def preparation_function(prepared_input: dict[str, Any]) -> tuple[torch.Tensor, list, dict[str, Any]]: x = prepared_input["encoder_input"] args = [] kwargs = {"cat_full_tensor": prepared_input["cat_full_tensor"]} @@ -310,8 +363,19 @@ def preparation_function(prepared_input): else: raise NotImplementedError() - def _get_inference_input(self, tensors): - """Parse the dictionary to get appropriate args""" + def _get_inference_input(self, tensors: TensorDict) -> dict[str, torch.Tensor | None]: + """Parse the dictionary to get appropriate args. + + Parameters + ---------- + tensors + Dictionary containing tensor data. + + Returns + ------- + dict + Dictionary with parsed input data. + """ x = tensors[REGISTRY_KEYS.X_KEY] cont_covs = tensors.get(REGISTRY_KEYS.CONT_COVS_KEY) @@ -320,7 +384,25 @@ def _get_inference_input(self, tensors): input_dict = {"x": x, "cont_covs": cont_covs, "cat_covs": cat_covs} return input_dict - def _input_pre_processing(self, x, cont_covs=None, cat_covs=None): + def _input_pre_processing( + self, x: torch.Tensor, cont_covs: torch.Tensor | None = None, cat_covs: torch.Tensor | None = None + ) -> dict[str, Any]: + """Pre-process input data for the model. + + Parameters + ---------- + x + Input tensor data. + cont_covs + Continuous covariates. + cat_covs + Categorical covariates. + + Returns + ------- + dict + Dictionary containing pre-processed input data. + """ # log the input to the variational distribution for numerical stability x_, gene_likelihood_additional_info = self.gene_likelihood_module.initial_transformation(x) # Note: this is different from scvi implementation of library size that is log transformed @@ -338,11 +420,26 @@ def _input_pre_processing(self, x, cont_covs=None, cat_covs=None): } @auto_move_data - def inference(self, x, cont_covs=None, cat_covs=None): - """ - High level inference method. + def inference( + self, x: torch.Tensor, cont_covs: torch.Tensor | None = None, cat_covs: torch.Tensor | None = None + ) -> dict[str, Any]: + """High level inference method. Runs the inference (encoder) model. + + Parameters + ---------- + x + Input tensor data. + cont_covs + Continuous covariates. + cat_covs + Categorical covariates. + + Returns + ------- + dict + Dictionary containing inference outputs including latent variables. """ pre_processed_input = self._input_pre_processing(x, cont_covs, cat_covs).copy() x_ = pre_processed_input["encoder_input"] @@ -380,7 +477,21 @@ def inference(self, x, cont_covs=None, cat_covs=None): } return outputs - def _get_generative_input(self, tensors, inference_outputs): + def _get_generative_input(self, tensors: TensorDict, inference_outputs: dict[str, Any]) -> dict[str, Any]: + """Prepare input for the generative model. + + Parameters + ---------- + tensors + Dictionary containing tensor data. + inference_outputs + Outputs from the inference step. + + Returns + ------- + dict + Dictionary containing input for generative model. + """ z = inference_outputs["z"] if self.fully_deterministic: z = inference_outputs["qz_m"] @@ -400,8 +511,34 @@ def _get_generative_input(self, tensors, inference_outputs): return input_dict @auto_move_data - def generative(self, z, library, gene_likelihood_additional_info, cont_covs=None, cat_covs=None): - """Runs the generative model.""" + def generative( + self, + z: torch.Tensor, + library: torch.Tensor, + gene_likelihood_additional_info: Any, + cont_covs: torch.Tensor | None = None, + cat_covs: torch.Tensor | None = None, + ) -> dict[str, Any]: + """Runs the generative model. + + Parameters + ---------- + z + Latent variables. + library + Library size information. + gene_likelihood_additional_info + Additional information for gene likelihood computation. + cont_covs + Continuous covariates. + cat_covs + Categorical covariates. + + Returns + ------- + dict + Dictionary containing generative model outputs. + """ if self.shared_covariate_emb is not None: cat_covs = self.shared_covariate_emb(cat_covs.int()) # form the likelihood @@ -421,12 +558,29 @@ def generative(self, z, library, gene_likelihood_additional_info, cont_covs=None def loss( self, - tensors, - inference_outputs, - generative_outputs, + tensors: TensorDict, + inference_outputs: dict[str, Any], + generative_outputs: dict[str, Any], kl_weight: float = 1.0, - ): - """Loss function.""" + ) -> LossOutput: + """Loss function. + + Parameters + ---------- + tensors + Dictionary containing tensor data. + inference_outputs + Outputs from the inference step. + generative_outputs + Outputs from the generative step. + kl_weight + Weight for KL divergence term. + + Returns + ------- + LossOutput + Loss output object containing various loss components. + """ x = tensors[REGISTRY_KEYS.X_KEY] x_mask = inference_outputs["x_mask"] qz_m = inference_outputs["qz_m"] @@ -462,9 +616,9 @@ def loss( @torch.no_grad() def sample( self, - tensors, - n_samples=1, - library_size=1, + tensors: TensorDict, + n_samples: int = 1, + library_size: int = 1, ) -> torch.Tensor: # Note: Not tested r""" @@ -475,16 +629,16 @@ def sample( Parameters ---------- tensors - Tensors dict + Dictionary containing tensor data. n_samples - Number of required samples for each cell + Number of required samples for each cell. library_size - Library size to scale scamples to + Library size to scale samples to. Returns ------- - x_new - tensor with shape (n_cells, n_genes, n_samples) + torch.Tensor + Tensor with shape (n_cells, n_genes, n_samples). """ inference_kwargs = dict(n_samples=n_samples) # noqa: C408 ( @@ -507,9 +661,21 @@ def sample( @torch.no_grad() @auto_move_data - def marginal_ll(self, tensors: TensorDict, n_mc_samples: int): - # Note: Not tested - """Marginal ll.""" + def marginal_ll(self, tensors: TensorDict, n_mc_samples: int) -> float: + """Compute marginal log-likelihood. + + Parameters + ---------- + tensors + Dictionary containing tensor data. + n_mc_samples + Number of Monte Carlo samples for estimation. + + Returns + ------- + float + Marginal log-likelihood value. + """ sample_batch = tensors[REGISTRY_KEYS.X_KEY] to_sum = torch.zeros(sample_batch.size()[0], n_mc_samples) diff --git a/src/drvi/scvi_tools_based/nn/_base_components.py b/src/drvi/scvi_tools_based/nn/_base_components.py index 4137497..33a8fd0 100644 --- a/src/drvi/scvi_tools_based/nn/_base_components.py +++ b/src/drvi/scvi_tools_based/nn/_base_components.py @@ -1,12 +1,13 @@ import collections import math from collections.abc import Callable, Iterable, Sequence -from typing import Literal +from typing import Any, Literal import torch from scvi.nn._utils import one_hot from torch import nn from torch.distributions import Normal +from torch.nn import functional as F from drvi.nn_modules.embedding import MultiEmbedding from drvi.nn_modules.freezable import FreezableBatchNorm1d, FreezableLayerNorm @@ -15,7 +16,7 @@ from drvi.nn_modules.noise_model import NoiseModel -def _identity(x): +def _identity(x: torch.Tensor) -> torch.Tensor: return x @@ -25,44 +26,46 @@ class FCLayers(nn.Module): Parameters ---------- layers_dim - Number of nodes in layers including input and output dimensions + Number of nodes in layers including input and output dimensions. n_cat_list A list containing, for each category of interest, the number of categories. Each category will be included using a one-hot encoding. dropout_rate - Dropout rate to apply to each of the hidden layers + Dropout rate to apply to each of the hidden layers. split_size - The size of split if input is a 3d tensor otherwise -1 - This parameter is required to handle batch normalization + The size of split if input is a 3d tensor otherwise -1. + This parameter is required to handle batch normalization. reuse_weights - Whether to reuse weights when having multiple splits + Whether to reuse weights when having multiple splits. use_batch_norm - Whether to have `BatchNorm` layers or not + Whether to have `BatchNorm` layers or not. affine_batch_norm - Whether to have affine transformation in `BatchNorm` layers + Whether to have affine transformation in `BatchNorm` layers. use_layer_norm - Whether to have `LayerNorm` layers or not + Whether to have `LayerNorm` layers or not. use_activation - Whether to have layer activation or not + Whether to have layer activation or not. bias - Whether to learn bias in linear layers or not + Whether to learn bias in linear layers or not. inject_covariates - Whether to inject covariates in each layer, or just the first (default). + Whether to inject covariates in each layer, or just the first. activation_fn - Which activation function to use + Which activation function to use. layer_factory - A layer Factory instance to build projection layers based on + A layer Factory instance to build projection layers based on. layers_location An indicator to tell the class where in the architecture these layers reside. covariate_modeling_strategy - The strategy model to consider covariates + The strategy model to consider covariates. + covariate_embs_dim + Dimensions for covariate embeddings when using embedding strategies. """ def __init__( self, layers_dim: Sequence[int], - n_cat_list: Iterable[int] = None, + n_cat_list: Iterable[int] | None = None, dropout_rate: float = 0.1, split_size: int = -1, reuse_weights: bool = True, @@ -72,8 +75,8 @@ def __init__( use_activation: bool = True, bias: bool = True, inject_covariates: bool = True, - activation_fn: nn.Module = nn.ELU, - layer_factory: LayerFactory = None, + activation_fn: type[nn.Module] = nn.ELU, + layer_factory: LayerFactory | None = None, layers_location: Literal["intermediate", "first", "last"] = "intermediate", covariate_modeling_strategy: Literal[ "one_hot", @@ -84,7 +87,7 @@ def __init__( "emb_shared_linear", ] = "one_hot", covariate_embs_dim: Iterable[int] = (), - ): + ) -> None: super().__init__() self.inject_covariates = inject_covariates if covariate_modeling_strategy.endswith("_linear"): @@ -95,16 +98,17 @@ def __init__( self.covariate_vector_modeling = covariate_modeling_strategy layer_factory = layer_factory or FCLayerFactory() - self.n_cat_list = n_cat_list if n_cat_list is not None else [] + self.n_cat_list = list(n_cat_list) if n_cat_list is not None else [] if self.covariate_vector_modeling == "one_hot": covariate_embs_dim = self.n_cat_list else: + covariate_embs_dim = list(covariate_embs_dim) assert len(covariate_embs_dim) == len(self.n_cat_list) self.injectable_layers = [] self.linear_batch_projections = [] - def is_intermediate(i): + def is_intermediate(i: int) -> bool: assert layers_location in ["intermediate", "first", "last"] if layers_location == "first" and i == 0: return False @@ -112,11 +116,11 @@ def is_intermediate(i): return False return True - def inject_into_layer(layer_num): + def inject_into_layer(layer_num: int) -> bool: user_cond = layer_num == 0 or (layer_num > 0 and self.inject_covariates) return user_cond - def get_projection_layer(n_in, n_out, i): + def get_projection_layer(n_in: int, n_out: int, i: int) -> list[nn.Module]: output = [] layer_needs_injection = False if not reuse_weights: @@ -161,7 +165,7 @@ def get_projection_layer(n_in, n_out, i): output.append(layer) return output - def get_normalization_layers(n_out): + def get_normalization_layers(n_out: int) -> list[nn.Module]: output = [] if split_size == -1: if use_batch_norm: @@ -204,13 +208,21 @@ def get_normalization_layers(n_out): ) ) - def set_online_update_hooks(self, previous_n_cats_per_cov, n_cats_per_cov): - """Set online update hooks.""" + def set_online_update_hooks(self, previous_n_cats_per_cov: Sequence[int], n_cats_per_cov: Sequence[int]) -> None: + """Set online update hooks for handling new categories. + + Parameters + ---------- + previous_n_cats_per_cov + Previous number of categories per covariate. + n_cats_per_cov + New number of categories per covariate. + """ if sum(previous_n_cats_per_cov) == sum(n_cats_per_cov): print("Nothing to make hook for!") return - def make_hook_function(weight): + def make_hook_function(weight: torch.Tensor) -> Callable[[torch.Tensor], torch.Tensor]: w_size = weight.size() if weight.dim() == 2: # 2D tensors @@ -249,7 +261,7 @@ def make_hook_function(weight): else: raise NotImplementedError() - def _hook_fn_injectable(grad): + def _hook_fn_injectable(grad: torch.Tensor) -> torch.Tensor: return grad * transfer_mask return _hook_fn_injectable @@ -280,20 +292,20 @@ def _hook_fn_injectable(grad): else: raise NotImplementedError() - def forward(self, x: torch.Tensor, cat_full_tensor: torch.Tensor): + def forward(self, x: torch.Tensor, cat_full_tensor: torch.Tensor | None) -> torch.Tensor: """Forward computation on ``x``. Parameters ---------- x - tensor of values with shape ``(n_in,)`` + Tensor of values with shape ``(batch_size, n_in,)`` or ``(batch_size, n_split, n_in)``. cat_full_tensor - Tensor of category membership(s) for this sample + Tensor of category membership(s) for this sample. Returns ------- - :class:`torch.Tensor` - tensor of shape ``(n_out,)`` + torch.Tensor + Tensor of shape ``(batch_size, n_out,)`` or ``(batch_size, n_split, n_out)``. """ if self.covariate_vector_modeling == "one_hot": concat_list = [] @@ -313,7 +325,7 @@ def forward(self, x: torch.Tensor, cat_full_tensor: torch.Tensor): else: concat_list = [] - def dimension_transformation(t): + def dimension_transformation(t: torch.Tensor) -> torch.Tensor: if x.dim() == t.dim(): return t if x.dim() == 3 and t.dim() == 2: @@ -361,49 +373,47 @@ class Encoder(nn.Module): Parameters ---------- n_input - The dimensionality of the input (data space) + The dimensionality of the input (data space). n_output - The dimensionality of the output (latent space) + The dimensionality of the output (latent space). layers_dim - The number of nodes per hidden layer as a sequence + The number of nodes per hidden layer as a sequence. n_cat_list A list containing the number of categories for each category of interest. Each category will be - included using a one-hot encoding + included using a one-hot encoding. n_continuous_cov - The number of continuous covariates + The number of continuous covariates. inject_covariates - Whether to inject covariates in each layer, or just the first (default). + Whether to inject covariates in each layer, or just the first. use_batch_norm - Whether to use batch norm in layers + Whether to use batch norm in layers. affine_batch_norm - Whether to use affine in batch norms + Whether to use affine in batch norms. use_layer_norm - Whether to use layer norm in layers + Whether to use layer norm in layers. input_dropout_rate - Dropout rate to apply to the input + Dropout rate to apply to the input. dropout_rate - Dropout rate to apply to each of the hidden layers + Dropout rate to apply to each of the hidden layers. distribution - Distribution of z + Distribution of z. var_eps - Minimum value for the variance; - used for numerical stability + Minimum value for the variance; used for numerical stability. var_activation - The activation function to ensure positivity of the variance. Defaults to "exp". + Callable used to ensure positivity of the variance. mean_activation - The activation function at the end of mean encoder. Defaults to "identity". - Possible values are "identity", "relu", "leaky_relu", "leaky_relu_{slope}", "elu", "elu_{min_vaule}". + Callable used to apply activation to the mean. layer_factory - A layer Factory instance for building layers - layers_location - A hint on where the layer resides in the model + A layer Factory instance to build projection layers based on. covariate_modeling_strategy - The strategy model takes to model covariates + The strategy model to consider covariates. + categorical_covariate_dims + Dimensions for covariate embeddings when using embedding strategies. return_dist - Return directly the distribution of z instead of its parameters. + Whether to return the distribution or just the parameters. **kwargs - Keyword args for :class:`~scvi.nn.FCLayers` + Additional keyword arguments. """ def __init__( @@ -411,7 +421,7 @@ def __init__( n_input: int, n_output: int, layers_dim: Sequence[int] = (128,), - n_cat_list: Iterable[int] = None, + n_cat_list: Iterable[int] | None = None, n_continuous_cov: int = 0, inject_covariates: bool = True, use_batch_norm: bool = True, @@ -423,7 +433,7 @@ def __init__( var_eps: float = 1e-4, var_activation: Callable | Literal["exp", "pow2", "2sig"] = "exp", mean_activation: Callable | str = "identity", - layer_factory: LayerFactory = None, + layer_factory: LayerFactory | None = None, covariate_modeling_strategy: Literal[ "one_hot", "emb", @@ -435,9 +445,8 @@ def __init__( categorical_covariate_dims: Sequence[int] = (), return_dist: bool = False, **kwargs, - ): + ) -> None: super().__init__() - self.distribution = distribution self.var_eps = var_eps self.input_dropout = nn.Dropout(p=input_dropout_rate) @@ -528,8 +537,13 @@ def __init__( assert callable(mean_activation) self.mean_activation = mean_activation - def forward(self, x: torch.Tensor, cat_full_tensor: torch.Tensor, cont_full_tensor: torch.Tensor = None): - r"""The forward computation for a single sample. + def forward( + self, + x: torch.Tensor, + cat_full_tensor: torch.Tensor | None, + cont_full_tensor: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[Normal, torch.Tensor]: + r"""Forward computation on ``x``. #. Encodes the data into latent space using the encoder network #. Generates a mean \\( q_m \\) and variance \\( q_v \\) @@ -538,14 +552,17 @@ def forward(self, x: torch.Tensor, cat_full_tensor: torch.Tensor, cont_full_tens Parameters ---------- x - tensor with shape (n_input,) + Tensor with shape (batch_size, n_input). cat_full_tensor - Tensor containing encoding of categorical variables of size n_batch x n_total_cat + Tensor containing encoding of categorical variables of size n_batch x n_total_cat. + cont_full_tensor + Tensor containing continuous covariates. Returns ------- - 3-tuple of :py:class:`torch.Tensor` - tensors of shape ``(n_latent,)`` for mean and var, and sample + tuple + If return_dist=False: (mean, variance, sample) tensors of shape (batch_size, n_latent). + If return_dist=True: (distribution, sample) where distribution is a Normal distribution. """ if cont_full_tensor is not None: @@ -570,42 +587,48 @@ class DecoderDRVI(nn.Module): Parameters ---------- n_input - The dimensionality of the input (latent space) + The dimensionality of the input (latent space). n_output - The dimensionality of the output (data space) - layers_dim - The number of nodes per hidden layer as a sequence + The dimensionality of the output (data space). + gene_likelihood_module + Module defining the noise model for gene expression. n_cat_list A list containing the number of categories for each category of interest. Each category will be - included using a one-hot encoding + included using a one-hot encoding. n_continuous_cov - The number of continuous covariates + The number of continuous covariates. n_split - The number of splits for latent dim + Number of splits in the latent space. split_aggregation - How to aggregate splits in the last layer of the decoder + How to aggregate splits in the last layer of the decoder. split_method - How to make splits. Can be 'split' or 'power', or 'split_map'. + How to make splits: + - "split" : Split the latent space + - "power" : Transform the latent space to n_split vectors of size n_latent + - "split_map" : Split the latent space then map each to latent space using unique transformations reuse_weights - Were to reuse the weights of the decoder layers when using splitting - Possible values are 'everywhere', 'last', 'intermediate', 'nowhere'. + Where to reuse the weights of the decoder layers when using splitting. + layers_dim + The number of nodes per hidden layer as a sequence. dropout_rate - Dropout rate to apply to each of the hidden layers + Dropout rate to apply to each of the hidden layers. inject_covariates - Whether to inject covariates in each layer, or just the first (default). + Whether to inject covariates in each layer, or just the first. use_batch_norm - Whether to use batch norm in layers + Whether to use batch norm in layers. affine_batch_norm - Whether to use affine in batch norms + Whether to use affine in batch norms. use_layer_norm - Whether to use layer norm in layers + Whether to use layer norm in layers. layer_factory - A layer Factory instance for building layers + A layer Factory instance for building layers. covariate_modeling_strategy - The strategy model takes to model covariates + The strategy model takes to model covariates. + categorical_covariate_dims + Dimensions for categorical covariate embeddings. **kwargs - Keyword args for :class:`~scvi.nn.FCLayers`. + Additional keyword arguments. """ def __init__( @@ -613,7 +636,7 @@ def __init__( n_input: int, n_output: int, gene_likelihood_module: NoiseModel, - n_cat_list: Iterable[int] = None, + n_cat_list: Iterable[int] | None = None, n_continuous_cov: int = 0, n_split: int = 1, split_aggregation: Literal["sum", "logsumexp", "max"] = "logsumexp", @@ -625,7 +648,7 @@ def __init__( use_batch_norm: bool = False, affine_batch_norm: bool = True, use_layer_norm: bool = False, - layer_factory: LayerFactory = None, + layer_factory: LayerFactory | None = None, covariate_modeling_strategy: Literal[ "one_hot", "emb", @@ -636,7 +659,7 @@ def __init__( ] = "one_hot", categorical_covariate_dims: Sequence[int] = (), **kwargs, - ): + ) -> None: super().__init__() self.n_output = n_output self.gene_likelihood_module = gene_likelihood_module @@ -724,31 +747,33 @@ def __init__( def forward( self, z: torch.Tensor, - cat_full_tensor: torch.Tensor, - cont_full_tensor: torch.Tensor, + cat_full_tensor: torch.Tensor | None, + cont_full_tensor: torch.Tensor | None, library: torch.Tensor, - gene_likelihood_additional_info: dict, - ): - """The forward computation for a single sample. - - #. Decodes the data from the latent space using the decoder network - #. Returns distribution of expression - #. If ``dispersion != 'gene-cell'`` then value for that param will be ``None`` + gene_likelihood_additional_info: dict[str, Any], + ) -> tuple[Any, dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Forward computation on ``z``. Parameters ---------- - z : - tensor with shape ``(n_input,)`` + z + Tensor of latent values with shape ``(batch_size, n_input)``. cat_full_tensor - Tensor containing encoding of categorical variables of size n_batch x n_total_cat + Tensor containing encoding of categorical variables of size n_batch x n_total_cat. + cont_full_tensor + Tensor of continuous covariate(s) for this sample. library - library size + Library size information. gene_likelihood_additional_info - additional info returned by gene likelihood module + Additional information for gene likelihood computation. Returns ------- - distribution + tuple + (distribution, parameters, original_parameters) where: + - distribution: The gene expression distribution + - parameters: Processed parameters for the distribution + - original_parameters: Raw parameters before processing (for example pooling) """ batch_size = z.shape[0] if self.n_split > 1: @@ -778,6 +803,9 @@ def forward( if self.split_aggregation == "sum": # to get average params[param_name] = param_value.sum(dim=-2) / self.n_split + elif self.split_aggregation == "sum_plus": + # to get average after softplus + params[param_name] = F.softplus(param_value).sum(dim=-2) / self.n_split elif self.split_aggregation == "logsumexp": # to cancel the effect of n_splits params[param_name] = torch.logsumexp(param_value, dim=-2) - math.log(self.n_split) diff --git a/src/drvi/utils/metrics/_aggregation.py b/src/drvi/utils/metrics/_aggregation.py index 50ca850..daf4f7d 100644 --- a/src/drvi/utils/metrics/_aggregation.py +++ b/src/drvi/utils/metrics/_aggregation.py @@ -3,15 +3,163 @@ from scipy.sparse.csgraph import min_weight_full_bipartite_matching -def most_similar_averaging_score(result_matrix): +def most_similar_averaging_score(result_matrix: np.ndarray) -> float: + """Compute the most similar averaging score for disentanglement evaluation. + + This function calculates the average of the maximum scores across all + target variables. It measures how well the each of the targets can be + captured by any of the latent dimensions. + + Parameters + ---------- + result_matrix + Matrix of metric scores with shape (n_dimensions, n_categories). + Each element [i, j] represents the score for dimension i and category j. + + Returns + ------- + float + The average of maximum scores across all categories. + Higher values indicate better disentanglement. + + Notes + ----- + The function computes: + 1. For each target variable, finds the dimension with the highest score + 2. Takes the average of these maximum scores across all target variables + + This approach treats each target variable equally and measures the average + performance of the best-matching dimensions. It's a simple but effective + way to assess overall disentanglement quality. + + **Mathematical formula:** + ``` + score = mean(max(result_matrix[:, j]) for j in range(n_categories)) + ``` + + Examples + -------- + >>> import numpy as np + >>> # Example result matrix (4 dimensions, 3 categories) + >>> result_matrix = np.array( + ... [ + ... [0.9, 0.1, 0.9], + ... [0.1, 0.9, 0.1], + ... [0.1, 0.1, 0.9], + ... [0.1, 0.9, 0.1], + ... ] + ... ) + >>> score = most_similar_averaging_score(result_matrix) + >>> print(f"MSAS score: {score:.3f}") # 0.933 + """ return np.mean(np.max(result_matrix, axis=0)) -def latent_matching_score(result_matrix): +def latent_matching_score(result_matrix: np.ndarray) -> float: + """Compute the latent matching score for disentanglement evaluation. + + This function finds the optimal one-to-one matching between dimensions + and target variables, then computes the average + score of the matched pairs. It ensures each dimension is assigned to + at most one target variable and vice versa. + + Parameters + ---------- + result_matrix + Matrix of metric scores with shape (n_dimensions, n_categories). + Each element [i, j] represents the score for dimension i and category j. + + Returns + ------- + float + The average score of the optimal dimension-category matches. + Higher values indicate better disentanglement. + + Notes + ----- + The function uses the Hungarian algorithm (implemented as minimum weight + bipartite matching) to find the optimal assignment between dimensions + and target variables. + + **Mathematical formula:** + ``` + row_ind, col_ind = max_weight_full_bipartite_matching(result_matrix) + score = mean(result_matrix[row_ind, col_ind]) + ``` + + **Advantages:** + - Ensures each dimension is used only once + - Mathematically rigorous evaluation of disentanglement + + Examples + -------- + >>> import numpy as np + >>> # Example result matrix (4 dimensions, 3 categories) + >>> result_matrix = np.array( + ... [ + ... [0.9, 0.1, 1.0], + ... [0.1, 0.9, 0.1], + ... [0.1, 0.1, 0.9], + ... [0.1, 0.9, 0.1], + ... ] + ... ) + >>> score = latent_matching_score(result_matrix) + >>> print(f"LMS score: {score:.3f}") # 0.900 + """ row_ind, col_ind = min_weight_full_bipartite_matching(csr_matrix(-result_matrix - 1e-10)) return result_matrix[row_ind, col_ind].sum() / result_matrix.shape[1] -def most_similar_gap_score(result_matrix): +def most_similar_gap_score(result_matrix: np.ndarray) -> float: + """Compute the most similar gap score for disentanglement evaluation. + + This function measures the average gap between the best and second-best + scores for each target variable. It quantifies how clearly each target variable is + captured by its best-matching dimension penalized by other dimensions. + + Parameters + ---------- + result_matrix + Matrix of metric scores with shape (n_dimensions, n_categories). + Each element [i, j] represents the score for dimension i and category j. + + Returns + ------- + float + The average gap between best and second-best scores across all target variables. + Higher values indicate clearer dimension-target variable relationships. + + Notes + ----- + The function computes the difference between the highest and second-highest + scores for each target variable, then averages these gaps across all target variables. + + **Algorithm steps:** + 1. Sort scores for each target variable in descending order + 2. Compute gap between first and second highest scores for each target variable + 3. Average the gaps across all target variables + + **Mathematical formula:** + ``` + sorted_scores = sort(result_matrix, axis=0, descending=True) + gaps = sorted_scores[0, :] - sorted_scores[1, :] + score = mean(gaps) + ``` + + Examples + -------- + >>> import numpy as np + >>> # Example result matrix (4 dimensions, 3 categories) + >>> result_matrix = np.array( + ... [ + ... [0.9, 0.1, 1.0], + ... [0.1, 0.9, 0.1], + ... [0.1, 0.1, 0.9], + ... [0.1, 0.9, 0.1], + ... ] + ... ) + >>> score = most_similar_gap_score(result_matrix) + >>> print(f"MSGS score: {score:.3f}") # 0.300 + """ sorted_values = np.sort(result_matrix, axis=0)[::-1, :] return np.mean(sorted_values[0, :] - sorted_values[1, :]) diff --git a/src/drvi/utils/metrics/_benchmark.py b/src/drvi/utils/metrics/_benchmark.py index ea40324..e626d9a 100644 --- a/src/drvi/utils/metrics/_benchmark.py +++ b/src/drvi/utils/metrics/_benchmark.py @@ -30,6 +30,118 @@ class DiscreteDisentanglementBenchmark: + """Benchmark for evaluating discrete disentanglement in latent representations. + + This class provides a comprehensive framework for evaluating how well + latent dimensions capture discrete categorical variables (e.g., cell types, + experimental conditions, biological processes). It supports multiple + evaluation metrics and aggregation methods to provide robust assessment + of disentanglement quality. + + Parameters + ---------- + embed + Latent representations to evaluate. Shape should be (n_samples, n_dimensions). + discrete_target + Discrete categorical target variable. Should contain categorical labels + for each sample. Mutually exclusive with `one_hot_target`. + one_hot_target + One-hot encoded target variable. Shape should be of shape (n_samples, n_categories). + Mutually exclusive with `discrete_target`. + dim_titles + Titles for each latent dimension. If None, will use "dim_0", "dim_1", etc.. + metrics + Metrics to compute for evaluation. Available options: + - "SMI-disc": Discrete mutual information score + - "SPN": Nearest neighbor alignment score + - "ASC": Spearman correlation score + - "SMI-cont": Continuous mutual information score (SMI-cont is not working as expected. More info: https://github.com/scikit-learn/scikit-learn/issues/30772). + aggregation_methods + Methods to aggregate metric scores across dimensions. Available options: + - "LMS": Latent matching score + - "MSAS": Most similar averaging score + - "MSGS": Most similar gap score. + additional_metric_params + Additional parameters to pass to specific metrics. Keys should be metric + names, values should be parameter dictionaries. + + Attributes + ---------- + embed + Copy of the input latent representations. + one_hot_target + One-hot encoded target variable. + dim_titles + Titles for each latent dimension. + metrics + Metrics used for evaluation. + aggregation_methods + Aggregation methods used. + additional_metric_params + Additional parameters for metrics. + results + Raw metric results for each dimension and category. + aggregated_results + Aggregated scores across dimensions. + + Raises + ------ + ValueError + If neither `discrete_target` nor `one_hot_target` is provided. + If both `discrete_target` and `one_hot_target` are provided. + If `discrete_target` is not a pandas Series or numpy array. + If `one_hot_target` is not a pandas DataFrame or numpy array. + + Notes + ----- + The benchmark evaluates disentanglement by measuring how well each latent + dimension captures information about discrete categorical variables. Higher + scores indicate better disentanglement. + + **Available Metrics:** + + - **SMI-disc**: Discrete mutual information between latent dimensions and + categorical targets. Measures how much information each dimension contains + about the categorical variable. + - **SPN**: Nearest neighbor alignment score. Measures how well the nearest + neighbor structure in latent space preserves categorical relationships. + - **ASC**: Spearman correlation score. Measures linear correlation between + latent dimensions and categorical targets (less suitable for discrete targets). + - **SMI-cont**: Continuous mutual information (SMI-cont is not working as expected. + More info: https://github.com/scikit-learn/scikit-learn/issues/30772). + + **Aggregation Methods:** + + - **LMS**: Latent matching score. Finds the optimal matching between latent + dimensions and categories. This aggregation discourages presense of multiple + irrelevant biological processes in a single dimension. + - **MSAS**: Most similar averaging score. Averages scores for the most + similar dimension-category pairs. + - **MSGS**: Most similar gap score. Measures the gap between the best and + second-best matches. This aggregation discourages redundancy. + + Examples + -------- + >>> # Basic usage with discrete targets + >>> import numpy as np + >>> import pandas as pd + >>> # Generate sample data + >>> n_samples, n_dims = 1000, 10 + >>> embed = np.random.randn(n_samples, n_dims) + >>> cell_types = pd.Series(np.random.choice(["A", "B", "C"], n_samples)) + >>> # Create benchmark + >>> benchmark = DiscreteDisentanglementBenchmark( + ... embed, discrete_target=cell_types, metrics=("SMI-disc", "SPN"), aggregation_methods=("LMS", "MSAS") + ... ) + >>> # Run evaluation + >>> benchmark.evaluate() + >>> results = benchmark.get_results() + >>> print(f"LMS-SMI-disc: {results['LMS-SMI-disc']:.3f}") + >>> # With one-hot targets + >>> one_hot = pd.get_dummies(cell_types) + >>> benchmark = DiscreteDisentanglementBenchmark(embed, one_hot_target=one_hot) + """ + version = "v2" def __init__( @@ -82,6 +194,25 @@ def __init__( self.aggregated_results = {} def _compute_metrics(self, embed, one_hot_target, dim_titles=None, metrics=()): + """Compute evaluation metrics for each dimension and category. + + Parameters + ---------- + embed + Latent representations to evaluate. + one_hot_target + One-hot encoded target variable. + dim_titles + Titles for each latent dimension. + metrics + Metrics to compute. + + Returns + ------- + dict + Dictionary with metric names as keys and pandas DataFrames as values. + Each DataFrame has dimensions as rows and categories as columns. + """ if dim_titles is None: dim_titles = [f"dim_{d}" for d in range(embed.shape[1])] @@ -99,6 +230,21 @@ def _compute_metrics(self, embed, one_hot_target, dim_titles=None, metrics=()): @staticmethod def _aggregate_metrics(results, aggregation_methods=()): + """Aggregate metric scores across dimensions. + + Parameters + ---------- + results + Raw metric results from `_compute_metrics`. + aggregation_methods + Aggregation methods to apply. + + Returns + ------- + dict + Dictionary with "{aggregation_method}-{metric_name}" as keys and + aggregated scores as values. + """ aggregated_results = {} for aggregation_method in aggregation_methods: for metric_name in results: @@ -108,6 +254,13 @@ def _aggregate_metrics(results, aggregation_methods=()): return aggregated_results def is_complete(self): + """Check if all metrics and aggregations have been computed. + + Returns + ------- + bool + True if all requested metrics and aggregations are complete. + """ for metric in self.metrics: if metric not in self.results: return False @@ -118,6 +271,30 @@ def is_complete(self): return True def evaluate(self): + """Compute all required metrics and aggregations. + + This method computes any missing metrics and updates the aggregated + results. It's safe to call multiple times - only missing computations + will be performed. + + Notes + ----- + The method performs the following steps: + 1. Identifies any missing metrics that need to be computed + 2. Computes the missing metrics using `_compute_metrics` + 3. Updates the aggregated results using `_aggregate_metrics` + + Aggregation is always performed as it's computationally cheap and + ensures consistency with any new metric results. + + Examples + -------- + >>> # Run evaluation + >>> benchmark.evaluate() + >>> # Check results + >>> results = benchmark.get_results() + >>> print(f"Number of results: {len(results)}") + """ if not self.is_complete(): remaining_metrics = [metric for metric in self.metrics if metric not in self.results] self.results = { @@ -131,6 +308,27 @@ def evaluate(self): } def get_results(self): + """Get aggregated benchmark results. + + Returns + ------- + dict + Dictionary with "{aggregation_method}-{metric_name}" as keys and + aggregated scores as values. + + Notes + ----- + This method returns the final aggregated scores that summarize the + disentanglement performance across all dimensions. Each key follows + the pattern "{aggregation_method}-{metric_name}" (e.g., "LMS-SMI-disc"). + + Examples + -------- + >>> benchmark.evaluate() + >>> results = benchmark.get_results() + >>> for key, value in results.items(): + ... print(f"{key}: {value:.3f}") + """ return { f"{aggregation_method}-{metric}": self.aggregated_results[f"{aggregation_method}-{metric}"] for aggregation_method in self.aggregation_methods @@ -138,9 +336,58 @@ def get_results(self): } def get_results_details(self): + """Get detailed metric results for each dimension and category. + + Returns + ------- + dict + Dictionary with metric names as keys and pandas DataFrames as values. + Each DataFrame shows scores for each dimension (rows) and category (columns). + + Notes + ----- + This method returns the raw metric results before aggregation, allowing + you to examine how each individual dimension performs for each category. + This is useful for identifying which dimensions capture which categories + and for debugging or detailed analysis. + + Examples + -------- + >>> details = benchmark.get_results_details() + >>> smi_scores = details["SMI-disc"] + >>> print(f"SMI-disc shape: {smi_scores.shape}") + >>> print(f"Best dimension for category A: {smi_scores['A'].idxmax()}") + """ return {f"{metric}": self.results[metric] for metric in self.metrics} def save(self, path): + """Save benchmark results to a file. + + Parameters + ---------- + path + File path where to save the benchmark data. + + Notes + ----- + The saved data includes: + - Version information for compatibility + - Raw metric results + - Aggregated results + - Configuration parameters (metrics, aggregation methods, etc.) + + The data is saved using Python's pickle format, which preserves + all object structures and data types. + + Examples + -------- + >>> benchmark.evaluate() + >>> benchmark.save("benchmark_results.pkl") + >>> # Load later + >>> loaded_benchmark = DiscreteDisentanglementBenchmark.load( + ... "benchmark_results.pkl", embed, discrete_target=cell_types + ... ) + """ data = { "version": self.version, "results": self.results, @@ -156,6 +403,49 @@ def save(self, path): @classmethod def load(cls, path, embed, discrete_target=None, one_hot_target=None, metrics=None, aggregation_methods=None): + """Load a saved benchmark instance. + + Parameters + ---------- + path + File path to the saved benchmark data. + embed + Latent representations (must match the original data). + discrete_target + Discrete categorical target variable. + one_hot_target + One-hot encoded target variable. + metrics + Override the metrics from the saved data. + aggregation_methods + Override the aggregation methods from the saved data. + + Returns + ------- + DiscreteDisentanglementBenchmark + Loaded benchmark instance with all results restored. + + Raises + ------ + AssertionError + If the saved version doesn't match the current class version. + FileNotFoundError + If the specified file doesn't exist. + ValueError + If the target data is invalid (same as constructor). + + Notes + ----- + This method creates a new benchmark instance with the same configuration + as the saved one, but allows you to override metrics and aggregation + methods if needed. The embed and target data must be provided to + recreate the instance, but the actual results are loaded from the file. + + Examples + -------- + >>> # Load with same configuration + >>> benchmark = DiscreteDisentanglementBenchmark.load("results.pkl", embed, discrete_target=cell_types) + """ with open(path, "rb") as f: data = pickle.load(f) @@ -179,12 +469,66 @@ def load(cls, path, embed, discrete_target=None, one_hot_target=None, metrics=No @classmethod def load_results(cls, path): + """Load only the aggregated results from a saved benchmark. + + Parameters + ---------- + path + File path to the saved benchmark data. + + Returns + ------- + dict + Dictionary with aggregated results (same format as `get_results()`). + + Notes + ----- + This is a convenience method for quickly accessing just the final + aggregated scores without needing to recreate the full benchmark + instance. Useful when you only need the results for analysis or + comparison. + + Examples + -------- + >>> # Quick access to results + >>> results = DiscreteDisentanglementBenchmark.load_results("results.pkl") + >>> print(f"LMS-SMI-disc score: {results['LMS-SMI-disc']:.3f}") + """ with open(path, "rb") as f: data = pickle.load(f) return data["aggregated_results"] @classmethod def load_results_details(cls, path): + """Load only the detailed results from a saved benchmark. + + Parameters + ---------- + path + File path to the saved benchmark data. + + Returns + ------- + dict + Dictionary with detailed results (same format as `get_results_details()`). + + Notes + ----- + This is a convenience method for quickly accessing the detailed + metric results without needing to recreate the full benchmark + instance. Useful for detailed analysis of dimension-category + relationships. + + Examples + -------- + >>> # Quick access to detailed results + >>> details = DiscreteDisentanglementBenchmark.load_results_details("results.pkl") + >>> smi_scores = details["SMI-disc"] + >>> print(f"Best dimension for each category:") + >>> for col in smi_scores.columns: + ... best_dim = smi_scores[col].idxmax() + ... print(f" {col}: {best_dim}") + """ with open(path, "rb") as f: data = pickle.load(f) return data["results"] diff --git a/src/drvi/utils/metrics/_pairwise.py b/src/drvi/utils/metrics/_pairwise.py index 148d45c..9dc3a6b 100644 --- a/src/drvi/utils/metrics/_pairwise.py +++ b/src/drvi/utils/metrics/_pairwise.py @@ -6,18 +6,104 @@ from sklearn.preprocessing import KBinsDiscretizer -def check_discrete_metric_input(gt_cat_series=None, gt_one_hot=None): +def _check_discrete_metric_input(gt_cat_series=None, gt_one_hot=None): + """Validate input parameters for discrete metric functions. + + Ensures that exactly one of the two input formats is provided for + ground truth categorical data. + + Parameters + ---------- + gt_cat_series + Categorical series with ground truth labels. + gt_one_hot + One-hot encoded ground truth matrix with shape (n_samples, n_categories). + + Raises + ------ + ValueError + If both or neither of gt_cat_series and gt_one_hot are provided. + + Notes + ----- + This function is used internally by all discrete metric functions to + ensure consistent input validation. It prevents ambiguous input scenarios + where both formats might be provided or neither is provided. + """ if gt_cat_series is not None and gt_one_hot is not None: raise ValueError("Only one of gt_cat_series or gt_one_hot should be provided.") if gt_cat_series is None and gt_one_hot is None: raise ValueError("Either gt_cat_series or gt_one_hot must be provided.") -def get_one_hot_encoding(gt_cat_series): +def _get_one_hot_encoding(gt_cat_series: pd.Series) -> np.ndarray: + """Convert categorical series to one-hot encoded matrix. + + Parameters + ---------- + gt_cat_series + Categorical series with ground truth labels. + + Returns + ------- + np.ndarray + One-hot encoded matrix with shape (n_samples, n_categories). + Each row represents a sample, each column represents a category. + + Notes + ----- + This function converts a pandas categorical series to a one-hot encoded + numpy array. The categories are ordered according to the categorical + dtype's categories attribute. + + Examples + -------- + >>> import pandas as pd + >>> gt_series = pd.Series(["A", "B", "A", "C"], dtype="category") + >>> one_hot = _get_one_hot_encoding(gt_series) + >>> print(one_hot) + [[1 0 0] + [0 1 0] + [1 0 0] + [0 0 1]] + """ return np.eye(len(gt_cat_series.cat.categories))[gt_cat_series.cat.codes] -def _nn_alignment_score_per_dim(var_continues, gt_01): +def _nn_alignment_score_per_dim(var_continues: np.ndarray, gt_01: np.ndarray) -> np.ndarray: + """Compute nearest neighbor alignment score for a single continuous variable. + + This function calculates how well the categorical ground truth labels + align with the rightmost nearest neighbor of samples in a continuous variable. + + Parameters + ---------- + var_continues + Continuous variable values with shape (n_samples,). + gt_01 + One-hot encoded ground truth matrix with shape (n_samples, n_categories). + + Returns + ------- + np.ndarray + Alignment scores for each category with shape (n_categories,). + Higher values indicate better alignment between the variable and categories. + + Notes + ----- + The score are adjusted to account for the frequency of the categories. + + **Algorithm:** + 1. Sort samples by the continuous variable + 2. For each category, compute the fraction of adjacent pairs that share + the same category label + 3. Normalize by the expected fraction under random ordering + + **Mathematical formula:** + ``` + alignment = (adjacent_same_category / total_adjacent - category_frequency) / (1 - category_frequency) + ``` + """ order = var_continues.argsort() gt_01 = gt_01[order] alignment = np.clip( @@ -31,9 +117,46 @@ def _nn_alignment_score_per_dim(var_continues, gt_01): return alignment -def nn_alignment_score(all_vars_continues, gt_cat_series=None, gt_one_hot=None): - check_discrete_metric_input(gt_cat_series, gt_one_hot) - gt_01 = get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot +def nn_alignment_score(all_vars_continues: np.ndarray, gt_cat_series=None, gt_one_hot=None) -> np.ndarray: + """Compute nearest neighbor alignment scores for all continuous variables. + + This function calculates how well the categorical ground truth labels + align with the rightmost nearest neighbor of samples in each continuous variable. + + Parameters + ---------- + all_vars_continues + Matrix of continuous variables with shape (n_samples, n_variables). + Each column represents a different continuous variable. + gt_cat_series + Categorical series with ground truth labels. + gt_one_hot + One-hot encoded ground truth matrix with shape (n_samples, n_categories). + + Returns + ------- + np.ndarray + Alignment score matrix with shape (n_variables, n_categories). + Element [i, j] represents the alignment score between variable i + and category j. Higher values indicate better alignment. + + Notes + ----- + The score are adjusted to account for the frequency of the categories. + + Examples + -------- + >>> import numpy as np + >>> import pandas as pd + >>> # Simple example: 3 variables, 2 categories + >>> all_vars = np.array([[1.0, 2.0, 0.5], [2.0, 1.0, 0.8], [3.0, 0.5, 1.2], [0.5, 3.0, 0.9]]) + >>> gt_series = pd.Series(["A", "A", "B", "B"], dtype="category") + >>> scores = nn_alignment_score(all_vars, gt_cat_series=gt_series) + >>> print(scores.shape) # (3, 2) + >>> print(scores) + """ + _check_discrete_metric_input(gt_cat_series, gt_one_hot) + gt_01 = _get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot n_vars = all_vars_continues.shape[1] result = np.zeros([n_vars, gt_01.shape[1]]) @@ -42,16 +165,85 @@ def nn_alignment_score(all_vars_continues, gt_cat_series=None, gt_one_hot=None): return result -def _local_mutual_info_score_per_binary_gt(all_vars_continues, gt_binary): +def _local_mutual_info_score_per_binary_gt(all_vars_continues: np.ndarray, gt_binary: np.ndarray) -> np.ndarray: + """Compute scaled mutual information for a binary ground truth variable. + + Parameters + ---------- + all_vars_continues + Matrix of continuous variables with shape (n_samples, n_variables). + gt_binary + Binary ground truth labels with shape (n_samples,). + + Returns + ------- + np.ndarray + Scaled mutual information scores with shape (n_variables,). + Scores are normalized by the entropy of the binary ground truth. + + Notes + ----- + This function computes mutual information between each continuous variable + and a binary ground truth, then normalizes by the entropy of the ground + truth to obtain scores between 0 and 1. + + This metric is not working as expected. More info: https://github.com/scikit-learn/scikit-learn/issues/30772 + + **Mathematical formula:** + ``` + MI(all_vars_continues, gt_binary) / entropy(gt_binary) + ``` + """ mi_score = mutual_info_classif(all_vars_continues, gt_binary, n_jobs=-1) gt_prob = np.sum(gt_binary == 1) / gt_binary.shape[0] gt_entropy = stats.entropy([gt_prob, 1 - gt_prob]) return mi_score / gt_entropy -def local_mutual_info_score(all_vars_continues, gt_cat_series=None, gt_one_hot=None): - check_discrete_metric_input(gt_cat_series, gt_one_hot) - gt_01 = get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot +def local_mutual_info_score(all_vars_continues: np.ndarray, gt_cat_series=None, gt_one_hot=None) -> np.ndarray: + """Compute local mutual information scores for all variables and categories. + + This function calculates the scaled mutual information between each + continuous variable and each categorical ground truth variable. The scores + are scaled by the entropy of each ground truth category. + + Parameters + ---------- + all_vars_continues + Matrix of continuous variables with shape (n_samples, n_variables). + Each column represents a different continuous variable. + gt_cat_series + Categorical series with ground truth labels. + gt_one_hot + One-hot encoded ground truth matrix with shape (n_samples, n_categories). + + Returns + ------- + np.ndarray + Mutual information score matrix with shape (n_variables, n_categories). + Element [i, j] represents the scaled mutual information between + variable i and category j. Scores range from 0 to 1. + + Notes + ----- + This function calculates the scaled mutual information between each + continuous variable and each categorical ground truth variable. The scores + are scaled by the entropy of each ground truth category. + This metric is not working as expected. More info: https://github.com/scikit-learn/scikit-learn/issues/30772 + + Examples + -------- + >>> import numpy as np + >>> import pandas as pd + >>> # Simple example: 3 variables, 2 categories + >>> all_vars = np.array([[1.0, 2.0, 0.5], [2.0, 1.0, 0.8], [3.0, 0.5, 1.2], [0.5, 3.0, 0.9]]) + >>> gt_series = pd.Series(["A", "A", "B", "B"], dtype="category") + >>> scores = local_mutual_info_score(all_vars, gt_cat_series=gt_series) + >>> print(scores.shape) # (3, 2) + >>> print(scores) + """ + _check_discrete_metric_input(gt_cat_series, gt_one_hot) + gt_01 = _get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot n_vars = all_vars_continues.shape[1] result = np.zeros([n_vars, gt_01.shape[1]]) @@ -60,9 +252,59 @@ def local_mutual_info_score(all_vars_continues, gt_cat_series=None, gt_one_hot=N return result -def discrete_mutual_info_score(all_vars_continues, gt_cat_series=None, gt_one_hot=None, n_bins=10): - check_discrete_metric_input(gt_cat_series, gt_one_hot) - gt_01 = get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot +def discrete_mutual_info_score( + all_vars_continues: np.ndarray, gt_cat_series=None, gt_one_hot=None, n_bins=10 +) -> np.ndarray: + """Compute mutual information scores using discretized continuous variables. + + This function discretizes continuous variables into bins and then computes + mutual information between the discretized variables and categorical ground + truth. This approach can capture non-linear relationships that might be + missed by linear correlation measures. + + Parameters + ---------- + all_vars_continues + Matrix of continuous variables with shape (n_samples, n_variables). + Each column represents a different continuous variable. + gt_cat_series + Categorical series with ground truth labels. + gt_one_hot + One-hot encoded ground truth matrix with shape (n_samples, n_categories). + n_bins + Number of bins to use for discretizing continuous variables. + More bins capture finer details but may be more sensitive to noise. + + Returns + ------- + np.ndarray + Mutual information score matrix with shape (n_variables, n_categories). + Element [i, j] represents the normalized mutual information between + discretized variable i and category j. Scores range from 0 to 1. + + Notes + ----- + This function uses uniform binning to discretize continuous variables, + then computes mutual information between the discretized variables and + categorical ground truth. The scores are normalized by the entropy of + each ground truth category. + + **Advantages over continuous mutual information:** + The continuous mutual information is not working as expected. More info: https://github.com/scikit-learn/scikit-learn/issues/30772 + + Examples + -------- + >>> import numpy as np + >>> import pandas as pd + >>> # Simple example: 3 variables, 2 categories + >>> all_vars = np.array([[1.0, 2.0, 0.5], [2.0, 1.0, 0.8], [3.0, 0.5, 1.2], [0.5, 3.0, 0.9]]) + >>> gt_series = pd.Series(["A", "A", "B", "B"], dtype="category") + >>> scores = discrete_mutual_info_score(all_vars, gt_cat_series=gt_series, n_bins=2) + >>> print(scores.shape) # (3, 2) + >>> print(scores) + """ + _check_discrete_metric_input(gt_cat_series, gt_one_hot) + gt_01 = _get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot discretizer = KBinsDiscretizer(n_bins=n_bins, encode="ordinal", strategy="uniform", random_state=123) all_vars_discrete = discretizer.fit_transform(all_vars_continues) @@ -79,16 +321,101 @@ def discrete_mutual_info_score(all_vars_continues, gt_cat_series=None, gt_one_ho return result -def spearman_correlataion_score(all_vars_continues, gt_cat_series=None, gt_one_hot=None): - check_discrete_metric_input(gt_cat_series, gt_one_hot) - gt_01 = get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot +def spearman_correlataion_score(all_vars_continues: np.ndarray, gt_cat_series=None, gt_one_hot=None) -> np.ndarray: + """Compute Spearman correlation scores between continuous variables and categories. + + This function computes the absolute Spearman correlation coefficients + between each continuous variable and each categorical ground truth + variable (encoded as one-hot). Spearman correlation measures monotonic + relationships and is robust to outliers. + + Parameters + ---------- + all_vars_continues + Matrix of continuous variables with shape (n_samples, n_variables). + Each column represents a different continuous variable. + gt_cat_series + Categorical series with ground truth labels. + gt_one_hot + One-hot encoded ground truth matrix with shape (n_samples, n_categories). + + Returns + ------- + np.ndarray + Absolute Spearman correlation matrix with shape (n_variables, n_categories). + Element [i, j] represents the absolute Spearman correlation between + variable i and category j. Scores range from 0 to 1. + + Notes + ----- + Spearman correlation measures the strength and direction of monotonic + relationships between variables. This function uses absolute values to + focus on the strength of relationships regardless of direction. + + **Limitations:** + Spearman correlaton is not suitable for discrete targets. + More info: https://www.biorxiv.org/content/10.1101/2024.11.06.622266v1.full.pdf lines 985 to 989 + + Examples + -------- + >>> import numpy as np + >>> import pandas as pd + >>> # Simple example: 3 variables, 2 categories + >>> all_vars = np.array([[1.0, 2.0, 0.5], [2.0, 1.0, 0.8], [3.0, 0.5, 1.2], [0.5, 3.0, 0.9]]) + >>> gt_series = pd.Series(["A", "A", "B", "B"], dtype="category") + >>> scores = spearman_correlataion_score(all_vars, gt_cat_series=gt_series) + >>> print(scores.shape) # (3, 2) + >>> print(scores) + """ + _check_discrete_metric_input(gt_cat_series, gt_one_hot) + gt_01 = _get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot n_vars = all_vars_continues.shape[1] result = np.abs(stats.spearmanr(all_vars_continues, gt_01).statistic[:n_vars, n_vars:]) return result -def global_dim_mutual_info_score(all_vars_continues, gt_cat_series): +def global_dim_mutual_info_score(all_vars_continues: np.ndarray, gt_cat_series: pd.Series) -> np.ndarray: + """Compute global mutual information scores for all variables with categorical ground truth. + + This function computes the normalized mutual information between each + continuous variable and the overall categorical ground truth. Unlike + local mutual information, this treats the ground truth as a single + categorical variable rather than separate binary variables. + + Parameters + ---------- + all_vars_continues + Matrix of continuous variables with shape (n_samples, n_variables). + Each column represents a different continuous variable. + gt_cat_series + Categorical series with ground truth labels. + + Returns + ------- + np.ndarray + Global mutual information scores with shape (n_variables,). + Each element represents the normalized mutual information between + a variable and the overall categorical ground truth. + + Notes + ----- + This metric is not used in standard disentanglement analysis but is + provided for completeness. It measures how much information each + variable shares with the overall categorical structure, rather than + with individual categories. + + Examples + -------- + >>> import numpy as np + >>> import pandas as pd + >>> # Simple example: 3 variables, 2 categories + >>> all_vars = np.array([[1.0, 2.0, 0.5], [2.0, 1.0, 0.8], [3.0, 0.5, 1.2], [0.5, 3.0, 0.9]]) + >>> gt_series = pd.Series(["A", "A", "B", "B"], dtype="category") + >>> scores = global_dim_mutual_info_score(all_vars, gt_series) + >>> print(scores.shape) # (3,) + >>> print(scores) + """ # This metric is not used in any analysis, but is provided for completeness. mi_score = mutual_info_classif(all_vars_continues, gt_cat_series) gt_entropy = stats.entropy(pd.Series(gt_cat_series).value_counts(normalize=True, sort=False)) diff --git a/src/drvi/utils/plotting/_interpretability.py b/src/drvi/utils/plotting/_interpretability.py index cf1e349..124730a 100644 --- a/src/drvi/utils/plotting/_interpretability.py +++ b/src/drvi/utils/plotting/_interpretability.py @@ -1,5 +1,6 @@ import itertools from collections.abc import Sequence +from typing import Any import numpy as np import pandas as pd @@ -12,7 +13,38 @@ from drvi.utils.tools.interpretability._latent_traverse import get_dimensions_of_traverse_data -def make_heatmap_groups(ordered_list): +def make_heatmap_groups(ordered_list: list) -> tuple[list[tuple[int, int]], list[Any]]: + """Create group positions and labels for scanpy heatmap visualization of marker genes. + + This helper function processes an ordered list to identify groups of + consecutive identical elements and returns their positions and labels. + It's used to create group annotations for scanpy heatmap plots. + + Parameters + ---------- + ordered_list + List of elements where consecutive identical elements form groups. + + Returns + ------- + tuple[list[tuple[int, int]], list] + A tuple containing: + - List of tuples with (start_index, end_index) for each group + - List of group labels (unique values from ordered_list) + + Notes + ----- + The function uses `itertools.groupby` to identify consecutive groups + of identical elements. Each group is represented by its start and end + indices (inclusive). + + Examples + -------- + >>> # Simple example + >>> groups, labels = make_heatmap_groups(["A", "A", "B", "B", "B", "A"]) + >>> print(f"Groups: {groups}") # [(0, 1), (2, 4), (5, 5)] + >>> print(f"Labels: {labels}") # ['A', 'B', 'A'] + """ n_groups, group_names = zip( *[(len(list(group)), key) for (key, group) in itertools.groupby(ordered_list)], strict=False ) @@ -32,24 +64,85 @@ def differential_vars_heatmap( show: bool = True, **kwargs, ): - """ - Generate a heatmap of differential variables based on traverse data. + """Generate a heatmap of differential variables based on traverse data. + + This function creates a comprehensive heatmap visualization showing how + genes respond to latent dimension traversals. The heatmap displays + stepwise effects across all latent dimensions and genes, with genes + grouped by their maximum effect dimension. Parameters ---------- - - traverse_adata (AnnData): Annotated data object containing traverse data. - - key (str): Key used to access traverse effect data in `traverse_adata.varm`. - - title_col (str): Column name in `traverse_adata.obs` to use as dimension labels. - - score_threshold (float): Threshold value for filtering variables based on the score. - - remove_vanished (bool): Whether to remove variables that have vanished. - - remove_unaffected (bool): Whether to remove variables that have no effect. - - figsize (Optional[Tuple[int, int]]): Size of the figure (width, height). - - show (bool): Whether to show the heatmap. - - **kwargs: Additional keyword arguments to be passed to `sc.pl.heatmap`. + traverse_adata + AnnData object containing traverse data from `traverse_latent` or + `make_traverse_adata`. Must contain differential effect data for the specified key. + key + Key prefix for the differential variables in `traverse_adata.varm`. + Should correspond to a key used in `find_differential_effects` or + `calculate_differential_vars` (e.g., "max_possible", "min_possible", "combined_score"). + title_col + Column name in `traverse_adata.obs` to use as dimension labels. + These titles will be used for axis labels and grouping. + score_threshold + Threshold value for filtering genes based on their maximum effect score. + Only genes with maximum effects above this threshold will be included. + remove_vanished + Whether to remove latent dimensions that have vanished (have no effect). + This helps focus the visualization on meaningful dimensions. + remove_unaffected + Whether to remove genes that have no significant effect (below score_threshold). + When True, only genes with effects above the threshold are shown. + figsize + Size of the figure (width, height) in inches. If None, automatically + calculated based on the number of dimensions. + show + Whether to display the plot. If False, returns the plot object. + **kwargs + Additional keyword arguments passed to `sc.pl.heatmap`. Returns ------- - - None if show is True, otherwise the plot. + matplotlib.axes.Axes or None + The heatmap plot axes if `show=False`, otherwise None. + + Raises + ------ + KeyError + If required data is missing from `traverse_adata`. + ValueError + If the specified key doesn't exist in the AnnData object. + + Notes + ----- + The function performs the following steps: + 1. Calculates maximum effects for each gene in both positive and negative directions + 2. Identifies which dimension has the maximum effect for each gene + 3. Groups genes by their maximum effect dimension + 4. Creates a heatmap showing stepwise effects across all dimensions + 5. Applies filtering based on score threshold and vanished dimensions + + **Visualization Features:** + + - **Color scale**: Red-blue diverging colormap centered at 0 + - **Gene grouping**: Genes are grouped by their maximum effect dimension + - **Dimension ordering**: Dimensions are ordered by their `order` column + - **Gene ordering**: Within each group, genes are ordered by effect magnitude + + **Interpretation:** + + - **Red colors**: Positive effects (increased expression) + - **Blue colors**: Negative effects (decreased expression) + - **Intensity**: Magnitude of the effect + - **Gene groups**: Genes with similar maximum effects are grouped together + + Examples + -------- + >>> # Basic heatmap with combined scores + >>> differential_vars_heatmap(traverse_adata, "combined_score") + >>> # Heatmap with custom parameters + >>> differential_vars_heatmap( + ... traverse_adata, "max_possible", score_threshold=1.0, remove_unaffected=True, figsize=(15, 8) + ... ) """ n_latent, n_steps, n_samples, n_vars = get_dimensions_of_traverse_data(traverse_adata) @@ -166,20 +259,48 @@ def _bar_plot_top_differential_vars( ncols: int = 5, show: bool = True, ): - """ - Plot the top differential variables in a bar plot. + """Plot the top differential variables in a group of bar plots. + + This internal function creates horizontal bar plots showing the top genes + for each latent dimension based on their differential effect scores. Parameters ---------- - plot_info (Sequence[Tuple[str, pd.Series]]): Information about the top differential variables. - dim_subset (Sequence[sre]): List of dimensions to plot in the bar plot. If not specified all dimensions are plotted. - n_top_genes (int, optional): Number of top genes to plot. Defaults to 10. - ncols (int, optional): Number of columns in the plot grid. Defaults to 5. - show (bool, optional): Whether to display the plot. If False, the plot will be returned as a Figure object. Defaults to True. + plot_info + Sequence of tuples containing dimension titles and corresponding gene data. + dim_subset + Subset of dimensions to plot. If None, all dimensions are plotted. + n_top_genes + Number of top genes to show in each plot. + ncols + Number of columns in the subplot grid. + show + Whether to display the plot. If False, returns the figure object. Returns ------- - None if show is True, otherwise the figure. + matplotlib.figure.Figure or None + The figure object if `show=False`, otherwise None. + + Notes + ----- + The function creates a grid of horizontal bar plots, with each subplot + showing the top genes for one latent dimension. Genes are sorted by + their effect scores in descending order. + + **Plot Features:** + + - **Horizontal bars**: Gene names on y-axis, scores on x-axis + - **Color**: Sky blue bars for all genes + - **Grid**: No grid lines for cleaner appearance + - **Layout**: Automatic grid layout based on number of dimensions + + Examples + -------- + >>> # Basic bar plot + >>> _bar_plot_top_differential_vars(plot_info) + >>> # Custom layout + >>> _bar_plot_top_differential_vars(plot_info, n_top_genes=15, ncols=3, show=False) """ if dim_subset is not None: plot_info = dict(plot_info) @@ -224,25 +345,91 @@ def show_top_differential_vars( ncols: int = 5, show: bool = True, ): - """ - Show top differential variables in a bar plot. + """Show top differential variables in a bar plot. + + This function creates a comprehensive visualization of the top differentially + expressed genes for each latent dimension. It generates horizontal bar plots + showing the genes with the highest effect scores for each dimension. Parameters ---------- - traverse_adata (AnnData): Annotated data object containing the variables to be plotted. - key (str): Key to access the traverse effect variables in `traverse_adata.varm`. - title_col (str, optional): Column name in `traverse_adata.obs` that contains the titles for each dimension. Defaults to 'title'. - order_col (str, optional): Column name in `traverse_adata.obs` that specifies the order of dimensions. Defaults to 'order'. Ignored if `dim_subset` is provided. - dim_subset (Sequence[sre]): List of dimensions to plot in the bar plot. If not specified all dimensions are plotted. - gene_symbols (str, optional): Column name in `traverse_adata.var` that contains gene symbols. If provided, gene symbols will be used in the plot instead of gene indices. Defaults to None. - score_threshold (float, optional): Threshold value for gene scores. Only genes with scores above this threshold will be plotted. Defaults to 0. - n_top_genes (int, optional): Number of top genes to plot. Defaults to 10. - ncols (int, optional): Number of columns in the plot grid. Defaults to 5. - show (bool, optional): Whether to display the plot. If False, the plot will be returned as a Figure object. Defaults to True. + traverse_adata + AnnData object containing the differential analysis results from + `calculate_differential_vars`. Must contain differential effect data + for the specified key. + key + Key prefix for the differential variables in `traverse_adata.varm`. + Should correspond to a key used in `find_differential_effects` or + `calculate_differential_vars` (e.g., "max_possible", "min_possible", "combined_score"). + title_col + Column name in `traverse_adata.obs` that contains the titles for each dimension. + These titles will be used as subplot titles. + order_col + Column name in `traverse_adata.obs` that specifies the order of dimensions. + Results will be sorted by this column. Ignored if `dim_subset` is provided. + dim_subset + List of dimensions to plot in the bar plot. If None, all dimensions + with significant effects are plotted. + gene_symbols + Column name in `traverse_adata.var` that contains gene symbols. + If provided, gene symbols will be used in the plot instead of gene indices. + Useful for converting between gene IDs and readable gene names. + score_threshold + Threshold value for gene scores. Only genes with scores above this + threshold will be plotted. + n_top_genes + Number of top genes to plot for each dimension. + ncols + Number of columns in the plot grid. + show + Whether to display the plot. If False, returns the figure object. Returns ------- - None if show is True, otherwise the figure. + matplotlib.figure.Figure or None + The figure object if `show=False`, otherwise None. + + Raises + ------ + KeyError + If required data is missing from `traverse_adata`. + ValueError + If the specified key doesn't exist in the AnnData object. + + Notes + ----- + The function performs the following steps: + 1. Extracts top differential variables using `iterate_on_top_differential_vars` + 2. Filters dimensions based on `dim_subset` if provided + 3. Creates horizontal bar plots for each dimension + 4. Displays top `n_top_genes` genes sorted by their effect scores + + **Visualization Features:** + + - **Gene symbols**: If provided, gene symbols will be used instead of gene indices. + - **Grid layout**: Automatic grid based on number of dimensions and `ncols` + - **Horizontal bars**: Gene names on y-axis, scores on x-axis + - **Color coding**: Sky blue bars for all genes + - **Dimension titles**: Each subplot shows the dimension title + - **Gene ordering**: Genes sorted by effect score (highest first) + + **Interpretation:** + + - **Bar length**: Represents the magnitude of the differential effect + - **Gene position**: Higher bars indicate stronger effects + - **Dimension separation**: Each subplot shows effects for one latent dimension + - **Direction indicators**: Dimension titles include "+" or "-" to indicate effect direction + + Examples + -------- + >>> # Basic visualization with combined scores + >>> show_top_differential_vars(traverse_adata, "combined_score") + >>> # Custom parameters with gene symbols + >>> show_top_differential_vars( + ... traverse_adata, "max_possible", gene_symbols="gene_symbol", score_threshold=1.0, n_top_genes=15, ncols=3 + ... ) + >>> # Subset of dimensions + >>> show_top_differential_vars(traverse_adata, "combined_score", dim_subset=["DR 5+", "DR 12+", "DR 14+"]) """ plot_info = iterate_on_top_differential_vars( traverse_adata, key, title_col, order_col, gene_symbols, score_threshold @@ -265,27 +452,79 @@ def show_differential_vars_scatter_plot( show: bool = True, **kwargs, ): - """ - Show a scatter plot of differential variables conidering multiple criteria. + """Show a scatter plot of differential variables considering multiple criteria. + + This function creates scatter plots comparing different differential effect + (usaully "max_possible" and "min_possible") measures for each latent dimension. + It is color-coded by the combined score. It's useful for understanding how + different analysis methods relate to each other and identifying genes + that show consistent effects across multiple criteria. The top 20 genes + are labeled with their names. Parameters ---------- - - traverse_adata (AnnData): Annotated data object containing the variables to be plotted. - - key_x (str): Key to access the first variable in `traverse_adata.varm`. - - key_y (str): Key to access the second variable in `traverse_adata.varm`. - - key_combined (str): Key to access the combined variable in `traverse_adata.varm`. - - title_col (str, optional): Column name in `traverse_adata.obs` that contains the titles for each dimension. Defaults to 'title'. - - order_col (str, optional): Column name in `traverse_adata.obs` that specifies the order of dimensions. Defaults to 'order'. Ignored if `dim_subset` is provided. - - gene_symbols (str, optional): Column name in `traverse_adata.var` that contains gene symbols. If provided, gene symbols will be used in the plot instead of gene indices. Defaults to None. - - score_threshold (float, optional): Threshold value for gene scores. Only genes with scores above this threshold will be plotted. Defaults to 0. - - dim_subset (Optional[Sequence[str]], optional): Subset of dimensions to plot. If None, all dimensions will be plotted. Defaults to None. - - ncols (int, optional): Number of columns in the plot grid. Defaults to 3. - - show (bool, optional): Whether to display the plot. If False, the plot will be returned as a Figure object. Defaults to True. - - **kwargs: Additional keyword arguments to be passed to the scatter plot. + traverse_adata + AnnData object containing the differential analysis results from + `calculate_differential_vars`. Must contain differential effect data + for all specified keys. + key_x + Key for the x-axis variable in `traverse_adata.varm`. + Typically "max_possible" or "min_possible". + key_y + Key for the y-axis variable in `traverse_adata.varm`. + Typically "min_possible" or "max_possible". + key_combined + Key for the color-coded variable in `traverse_adata.varm`. + Typically "combined_score" for the final combined effect. + title_col + Column name in `traverse_adata.obs` that contains the titles for each dimension. + These titles will be used as subplot titles. + order_col + Column name in `traverse_adata.obs` that specifies the order of dimensions. + Results will be sorted by this column. Ignored if `dim_subset` is provided. + gene_symbols + Column name in `traverse_adata.var` that contains gene symbols. + If provided, gene symbols will be used for point labels instead of gene indices. + score_threshold + Threshold value for gene scores. Only genes with combined scores above + this threshold will be plotted. + dim_subset + Subset of dimensions to plot. If None, all dimensions with significant + effects are plotted. + ncols + Number of columns in the plot grid. + show + Whether to display the plot. If False, returns the figure object. + **kwargs + Additional keyword arguments passed to the scatter plot (e.g., alpha, s for point size). Returns ------- - - None if show is True, otherwise the figure. + matplotlib.figure.Figure or None + The figure object if `show=False`, otherwise None. + + Raises + ------ + KeyError + If required data is missing from `traverse_adata`. + ValueError + If any of the specified keys don't exist in the AnnData object. + + Notes + ----- + The function performs the following steps: + 1. Extracts differential variables for all three keys (x, y, combined) + 2. Creates scatter plots for each dimension comparing the two measures + 3. Color-codes points by the combined score + 4. Labels the top 20 genes by combined score + + **Interpretation:** + + - **X-axis**: Effect measure from `key_x` (e.g., max_possible) + - **Y-axis**: Effect measure from `key_y` (e.g., min_possible) + - **Color**: Combined score from `key_combined` + - **Point position**: Relationship between the two measures + - **Labeled points**: Genes with highest combined scores """ plot_info = {} for key in [key_x, key_y, key_combined]: @@ -353,6 +592,77 @@ def _umap_of_relevant_genes( max_cells_to_plot: int | None = None, **kwargs, ): + """Plot UMAP embeddings for specific latent dimensions together with UMAP embeddings of relevant genes. + + This internal function creates UMAP visualizations showing how genes + associated with specific latent dimensions are expressed across cells. + The latent dimension values are color-coded by the latent dimension values. + The top genes are color-coded by the gene expression. + + Parameters + ---------- + adata + AnnData object containing single-cell data with UMAP coordinates. + Must have UMAP coordinates in `embed.obsm["X_umap"]`. + embed + AnnData object containing latent representations and dimension metadata. + Must have columns in `.var` corresponding to `title_col`. + plot_info + Information about the top differential variables. Each tuple contains + a dimension title and a pandas Series of gene scores. + layer + Layer name in `adata` to use for gene expression visualization. + If None, uses `.X`. + title_col + Column name in `embed.var` that contains dimension titles. + gene_symbols + Column name in `adata.var` that contains gene symbols. + If provided, gene symbols will be used instead of gene indices. + dim_subset + List of dimensions to plot. If None, all dimensions from plot_info are plotted. + n_top_genes + Number of top genes to visualize for each dimension. + max_cells_to_plot + Maximum number of cells to include in the plot. If None, all cells are plotted. + Useful for large datasets to improve performance. + **kwargs + Additional keyword arguments passed to `sc.pl.embedding`. + + Returns + ------- + None + Displays the plots directly. + + Notes + ----- + The function creates two types of visualizations for each dimension: + 1. **Latent dimension values**: Shows how the latent dimension varies across cells + 2. **Top gene expression**: Shows expression patterns of the top genes for that dimension + + **Visualization features:** + + - **UMAP coordinates**: Uses UMAP embedding from the embed object + - **Cell subsetting**: Can limit number of cells for performance + - **Gene labeling**: Shows gene names in plot titles + - **Dimension labeling**: Shows dimension names in plot titles + + Examples + -------- + >>> # Basic UMAP visualization + >>> plot_info = iterate_on_top_differential_vars(traverse_adata, "combined_score") + >>> _umap_of_relevant_genes(adata, embed, plot_info) + >>> # With custom parameters + >>> _umap_of_relevant_genes( + ... adata, + ... embed, + ... plot_info, + ... layer="counts", + ... title_col="title", + ... gene_symbols="gene_symbol", + ... n_top_genes=5, + ... max_cells_to_plot=5000, + ... ) + """ if max_cells_to_plot is not None and adata.n_obs > max_cells_to_plot: adata = sc.pp.subsample(adata, n_obs=max_cells_to_plot, copy=True) @@ -412,11 +722,113 @@ def plot_relevant_genes_on_umap( order_col: str = "order", gene_symbols: str | None = None, score_threshold: float = 0.0, - dim_subset: Sequence[str] = None, + dim_subset: Sequence[str] | None = None, n_top_genes: int = 10, max_cells_to_plot: int | None = None, **kwargs, ): + """Plot relevant genes on UMAP embedding. + + This function creates UMAP visualizations showing how genes associated + with specific latent dimensions are expressed across cells. The latent + dimension values are color-coded by the latent dimension values. The + top genes are color-coded by the gene expression. + + Parameters + ---------- + adata + AnnData object containing single-cell data with gene expression. + This is the original data used for training the model. + embed + AnnData object containing latent representations and dimension metadata. + Must have UMAP coordinates in `embed.obsm["X_umap"]` and dimension + information in `.var` columns. + traverse_adata + AnnData object containing differential analysis results from + `calculate_differential_vars`. Must contain differential effect data + for the specified key. + traverse_adata_key + Key prefix for the differential variables in `traverse_adata.varm`. + Should correspond to a key used in `find_differential_effects` or + `calculate_differential_vars` (e.g., "max_possible", "min_possible", "combined_score"). + layer + Layer name in `adata` to use for gene expression visualization. + If None, uses `.X`. Common options include "counts", "logcounts", etc.. + title_col + Column name in `embed.var` that contains dimension titles. + These titles will be used to match dimensions between objects. + order_col + Column name in `embed.var` that specifies the order of dimensions. + Results will be sorted by this column. Ignored if `dim_subset` is provided. + gene_symbols + Column name in `adata.var` that contains gene symbols. + If provided, gene symbols will be used instead of gene indices. + score_threshold + Threshold value for gene scores. Only genes with scores above this + threshold will be visualized. + dim_subset + List of dimensions to plot. If None, all dimensions with significant + effects are plotted. + n_top_genes + Number of top genes to visualize for each dimension. + max_cells_to_plot + Maximum number of cells to include in the plot. If None, all cells are plotted. + Useful for large datasets to improve performance and reduce memory usage. + **kwargs + Additional keyword arguments passed to `sc.pl.embedding`. + + Returns + ------- + None + Displays the plots directly. + + Raises + ------ + KeyError + If required data is missing from any of the AnnData objects. + ValueError + If the specified key doesn't exist in traverse_adata. + + Notes + ----- + The function performs the following steps: + 1. Extracts top differential variables using `iterate_on_top_differential_vars` + 2. For each dimension, creates two visualizations (I) UMAP of Latent dimension values across cells (II) UMAPs of Expression patterns of top genes for that dimension + + **Interpretation:** + + - **Latent dimension plots**: Show how the dimension varies across cell types + - **Gene expression plots**: Show expression patterns of dimension-specific genes + - **Color intensity**: Indicates magnitude of values/expression + + **Common Use Cases:** + + - **Biological validation**: Verify that latent dimensions capture meaningful biology + - **Gene discovery**: Identify genes associated with specific processes + - **Model interpretation**: Understand what biological processes each dimension represents + - **Quality assessment**: Evaluate the biological relevance of the model + + Examples + -------- + >>> # Basic UMAP visualization with combined scores + >>> plot_relevant_genes_on_umap(adata, embed, traverse_adata, "combined_score") + >>> # With custom parameters + >>> plot_relevant_genes_on_umap( + ... adata, + ... embed, + ... traverse_adata, + ... "max_possible", + ... layer="logcounts", + ... gene_symbols="gene_symbol", + ... score_threshold=1.0, + ... n_top_genes=5, + ... max_cells_to_plot=5000, + ... ) + >>> # Subset of dimensions + >>> plot_relevant_genes_on_umap( + ... adata, embed, traverse_adata, "combined_score", dim_subset=["DR 5+", "DR 12+", "DR 14+"] + ... ) + """ plot_info = iterate_on_top_differential_vars( traverse_adata, traverse_adata_key, title_col, order_col, gene_symbols, score_threshold ) diff --git a/src/drvi/utils/plotting/_latent.py b/src/drvi/utils/plotting/_latent.py index 9b1c7c5..4613b9c 100644 --- a/src/drvi/utils/plotting/_latent.py +++ b/src/drvi/utils/plotting/_latent.py @@ -10,7 +10,32 @@ from drvi.utils.plotting import cmap -def make_balanced_subsample(adata, col, min_count=10): +def make_balanced_subsample(adata: AnnData, col: str, min_count: int = 10) -> AnnData: + """Create a balanced subsample of AnnData based on a categorical column. + + This function creates a balanced subsample by sampling an equal number of cells + from each category in the specified column, ensuring balanced representation. + + Parameters + ---------- + adata + Annotated data object to subsample. + col + Column name in `adata.obs` containing categorical labels for balancing. + min_count + Minimum number of samples per category. If a category has fewer samples + than this, sampling will be done with replacement. + + Returns + ------- + AnnData + Balanced subsample of the input AnnData object. + + Notes + ----- + The function uses a fixed random state (0) for reproducible results. + If a category has fewer samples than `min_count`, sampling is done with replacement. + """ n_sample_per_cond = adata.obs[col].value_counts().min() balanced_sample_index = ( adata.obs.groupby(col) @@ -27,25 +52,63 @@ def plot_latent_dimension_stats( log_scale: bool | Literal["try"] = "try", ncols: int = 5, columns: Sequence[str] = ("reconstruction_effect", "max_value", "mean", "std"), - titles: dict[str, str] = None, + titles: dict[str, str] | None = None, remove_vanished: bool = False, show: bool = True, ): - """ - Plot the statistics of latent dimensions. + """Plot the statistics of latent dimensions. - Args: - - embed (AnnData): The annotated data object containing the latent dimensions. - - figsize (Tuple[int, int], optional): The size of the figure (width, height). Default is (5, 3). - - log_scale (Union[bool, Literal['try']], optional): Whether to use a log scale for the y-axis. If 'try', the log scale is used if the minimum value is greater than 0. Default is 'try'. - - ncols (int, optional): The maximum number of columns in the subplot grid. Default is 5. - - columns (Sequence[str], optional): The columns to plot from the `embed` object. Default is ('reconstruction_effect', 'max_value', 'mean', 'std'). - - titles (Dict[str, str], optional): The titles for each column in the plot. - - show (bool, optional): Whether to display the plot. If False, figure is returned. Default is True. + This function creates line plots showing various statistics of latent dimensions + across their ranking order. It can optionally distinguish between vanished and + non-vanished dimensions. + + Parameters + ---------- + embed + Annotated data object containing the latent dimensions and their statistics + in the `.var` attribute. + figsize + The size of each subplot (width, height) in inches. + log_scale + Whether to use a log scale for the y-axis. If "try", log scale is used + only if the minimum value is greater than 0. + ncols + The maximum number of columns in the subplot grid. + columns + The columns from `embed.var` to plot. These should be numeric columns + containing dimension statistics. + titles + Custom titles for each column in the plot. If None, default titles are used. + remove_vanished + Whether to exclude vanished dimensions from the plot. + show + Whether to display the plot. If False, returns the figure object. Returns ------- - - plt (matplotlib.pyplot module): The matplotlib.pyplot module if show is False. + matplotlib.figure.Figure or None + The matplotlib figure object if `show=False`, otherwise None. + + Notes + ----- + The function expects the following columns in `embed.var`: + - `order`: Ranking of dimensions + - `vanished`: Boolean indicating vanished dimensions + - The columns specified in the `columns` parameter + + If `remove_vanished=False`, a legend is added to distinguish between + vanished (black dots) and non-vanished (blue dots) dimensions. + + Examples + -------- + >>> # Default plot + >>> plot_latent_dimension_stats(embed) + >>> + >>> # Plot basic statistics + >>> plot_latent_dimension_stats(embed, columns=["reconstruction_effect", "max_value"]) + >>> # Plot with custom titles and log scale + >>> titles = {"reconstruction_effect": "Reconstruction Impact", "max_value": "Max Activation"} + >>> plot_latent_dimension_stats(embed, titles=titles, log_scale=True) """ if titles is None: titles = { @@ -130,11 +193,11 @@ def plot_latent_dimension_stats( def plot_latent_dims_in_umap( - embed, + embed: AnnData, title_col: str = "title", - additional_columns=(), + additional_columns: Sequence[str] = (), max_cells_to_plot: int | None = None, - order_col="order", + order_col: str = "order", dim_subset: Sequence[str] | None = None, directional: bool = False, remove_vanished: bool = True, @@ -143,28 +206,73 @@ def plot_latent_dims_in_umap( show: bool = True, **kwargs, ): - """ - Plot the latent dimensions of a UMAP embedding. - - Args: - embed (AnnData): Annotated data object containing the UMAP embedding. - title_col (str, optional): Name of the column in `embed.var` to use as titles for each dimension. - If None, default titles will be used. - additional_columns (tuple, optional): Additional columns to plot alongside the latent dimensions. - max_cells_to_plot (int, optional): Maximum number of cells to plot. If the number of cells in `embed` - is greater than `max_cells_to_plot`, a subsample will be taken. - order_col (str, optional): The column in the `embed.var` DataFrame to use for ordering the dimensions. Ignored if `dim_subset` is provided. - dim_subset (Sequence[str], optional): The subset of dimensions to plot. Defaults to None. - directional (bool, optional): Consider + and - directions as two separate dimensions. Default is False. - remove_vanished (bool, optional): Whether to remove the vanished dimensions from the plot. Default is True. - rearrange_titles (bool, optional): Whether to rearrange the titles to bottom right of the plot. Default is True. - color_bar_rescale_ratio (float, optional): The ratio to rescale the height of colorbars. Default is 1.0. - show (bool, optional): Whether to display the plot. If False, the plot is returned. Default is True. - **kwargs: Additional keyword arguments to be passed to `sc.pl.umap`. + """Plot the latent dimensions on a UMAP embedding. + + This function creates UMAP plots for each latent dimension, showing how cells + are distributed in the UMAP space based on their values for each dimension. + It can optionally handle directional dimensions and subsample cells for performance. + + Parameters + ---------- + embed + Annotated data object containing the UMAP embedding in `.obsm['X_umap']` + and latent dimensions in `.X`. + title_col + Name of the column in `embed.var` to use as titles for each dimension. + If None, default titles will be used. + additional_columns + Additional columns from `embed.obs` to plot alongside the latent dimensions. + max_cells_to_plot + Maximum number of cells to plot. If the number of cells in `embed` + is greater than this value, a subsample will be taken. + order_col + The column in `embed.var` to use for ordering the dimensions. + Ignored if `dim_subset` is provided. + dim_subset + The subset of dimensions to plot. If provided, overrides `order_col`. + directional + Whether to consider positive and negative directions as separate dimensions. + If True, creates separate plots for + and - directions. + remove_vanished + Whether to remove vanished dimensions from the plot. + rearrange_titles + Whether to rearrange titles to the bottom right of each plot. + color_bar_rescale_ratio + Ratio to rescale the height of colorbars. + show + Whether to display the plot. If False, returns the figure object. + **kwargs + Additional keyword arguments passed to `sc.pl.umap`. Returns ------- - matplotlib.figure.Figure: The UMAP plot figure if show is False. + matplotlib.figure.Figure or None + The UMAP plot figure if `show=False`, otherwise None. + + Raises + ------ + ValueError + If required columns (`order_col` or "vanished") are not found in `embed.var`. + + Notes + ----- + The function expects the following columns in `embed.var`: + - `order_col`: For ordering dimensions + - `title_col`: For dimension titles + - `vanished`: Boolean indicating vanished dimensions (if `remove_vanished=True`) + - `min`, `max`: For setting color scale limits + + When `directional=True`, the function creates separate plots for positive + and negative directions, effectively doubling the number of plots. + + Examples + -------- + >>> # Basic UMAP plot of latent dimensions + >>> plot_latent_dims_in_umap(embed) + >>> # Plot with directional dimensions and custom subset + >>> plot_latent_dims_in_umap(embed, directional=True, dim_subset=["DR 1", "DR 2"]) + >>> # Plot with additional metadata columns + >>> plot_latent_dims_in_umap(embed, additional_columns=["cell_type", "batch"]) """ if order_col not in embed.var: raise ValueError( @@ -258,25 +366,72 @@ def plot_latent_dims_in_heatmap( show: bool = True, **kwargs, ): - """ - Plot the latent dimensions in a heatmap. + """Plot the latent dimensions in a heatmap. + + This function creates a heatmap showing the values of latent dimensions + across different categories. It can optionally create balanced subsamples + and sort dimensions based on categorical differences. Parameters ---------- - embed (AnnData): The annotated data object containing the latent dimensions. - categorical_column (str): The column in the `embed.obs` DataFrame that represents the categorical variable. - title_col (str, optional): The column in the `embed.var` DataFrame to use as the title for each dimension. Defaults to None. - sort_by_categorical (bool, optional): Whether to sort the dimensions based on the categorical variable. Defaults to True. - make_balanced (bool, optional): Whether to make a balanced subsample of the data based on the categorical variable. Defaults to True. - order_col (str, optional): The column in the `embed.var` DataFrame to use for ordering the dimensions. Discarded if sort_by_categorical is set. Defaults to 'order'. - remove_vanished (bool, optional): Whether to remove the vanished dimensions from the plot. Default is True. - figsize (Tuple[int, int], optional): The size of the figure. Defaults to None. - show (bool, optional): Whether to show the plot. Defaults to True. - **kwargs: Additional keyword arguments to be passed to the `sc.pl.heatmap` function. + embed + Annotated data object containing the latent dimensions in `.X` + and categorical metadata in `.obs`. + categorical_column + The column in `embed.obs` that represents the categorical variable + for grouping cells. + title_col + The column in `embed.var` to use as titles for each dimension. + If None, uses the dimension indices. + sort_by_categorical + Whether to sort dimensions based on their maximum absolute values + within each category. If True, `order_col` is ignored. + make_balanced + Whether to create a balanced subsample of the data based on the + categorical variable using `make_balanced_subsample`. + order_col + The column in `embed.var` to use for ordering the dimensions. + Ignored if `sort_by_categorical=True`. + remove_vanished + Whether to remove vanished dimensions from the plot. + figsize + The size of the figure (width, height) in inches. + If None, automatically calculated based on number of categories. + show + Whether to display the plot. If False, returns the plot object. + **kwargs + Additional keyword arguments passed to `sc.pl.heatmap`. Returns ------- - plot if Show is not False. + matplotlib.axes.Axes or None + The heatmap axes if `show=False`, otherwise None. + + Raises + ------ + ValueError + If required columns (`order_col` or "vanished") are not found in `embed.var`. + + Notes + ----- + The function expects the following columns in `embed.var`: + - `order_col`: For ordering dimensions (if `sort_by_categorical=False`) + - `title_col`: For dimension titles + - `vanished`: Boolean indicating vanished dimensions (if `remove_vanished=True`) + + If `figsize=None`, the figure height is automatically calculated as + `len(unique_categories) / 6` to accommodate all categories. + + The heatmap uses a red-blue color map centered at 0, with no dendrogram. + + Examples + -------- + >>> # Basic heatmap of latent dimensions by cell type + >>> plot_latent_dims_in_heatmap(embed, categorical_column="cell_type") + >>> # Heatmap with balanced sampling and custom sorting + >>> plot_latent_dims_in_heatmap(embed, categorical_column="condition", sort_by_categorical=True, make_balanced=True) + >>> # Heatmap with custom figure size + >>> plot_latent_dims_in_heatmap(embed, categorical_column="batch", figsize=(12, 8)) """ if order_col is not None and order_col not in embed.var: raise ValueError( diff --git a/src/drvi/utils/tools/_latent.py b/src/drvi/utils/tools/_latent.py index 5697be4..391f514 100644 --- a/src/drvi/utils/tools/_latent.py +++ b/src/drvi/utils/tools/_latent.py @@ -8,19 +8,55 @@ def set_latent_dimension_stats( model: DRVI, embed: AnnData, inplace: bool = True, - vanished_threshold=0.1, + vanished_threshold: float = 0.1, ): - """ - Set the latent dimension statistics of a DRVI model into var of an embedding anndata. + """Set the latent dimension statistics of a DRVI embedding into var of an AnnData. + + This function computes and stores various statistics for each latent dimension + in the embedding AnnData object. It calculates reconstruction effects, ordering, + and basic statistical measures (mean, std, min, max) for each dimension. Parameters ---------- model - DRVI model object. + DRVI model object that has been trained and can compute reconstruction effects. embed - latent representation of the model. + AnnData object containing the latent representation (embedding) of the model. + The latent dimensions should be in the `.X` attribute. inplace - Whether to modify the input AnnData object or return a new one. + Whether to modify the input AnnData object in-place or return a new copy. + vanished_threshold + Threshold for determining if a latent dimension has "vanished" (become inactive). + Dimensions with max absolute values below this threshold are marked as vanished. + + Returns + ------- + AnnData or None + If `inplace=True` (default), modifies the input AnnData object and returns None. + If `inplace=False`, returns a new AnnData object with the statistics added. + + Notes + ----- + The function adds the following columns to `embed.var`: + + - `original_dim_id`: Original dimension indices + - `reconstruction_effect`: Reconstruction effect scores from the DRVI model + - `order`: Ranking of dimensions by reconstruction effect (descending) + - `max_value`: Maximum absolute value across all cells for each dimension + - `mean`: Mean value across all cells for each dimension + - `min`: Minimum value across all cells for each dimension + - `max`: Maximum value across all cells for each dimension + - `std`: Standard deviation of absolute values across all cells for each dimension + - `title`: Dimension titles in format "DR {order+1}" + - `vanished`: Boolean indicating if dimension is considered "vanished" (max_value < threshold) + + Examples + -------- + >>> # Assuming you have a trained DRVI model and latent representation + >>> latent_adata = model.get_latent_representation(adata, return_anndata=True) + >>> set_latent_dimension_stats(model, latent_adata, inplace=True) + >>> # Now latent_adata.var contains all the statistics + >>> print(latent_adata.var[["order", "reconstruction_effect", "vanished"]].head()) """ if not inplace: embed = embed.copy() diff --git a/src/drvi/utils/tools/interpretability/_differential_vars.py b/src/drvi/utils/tools/interpretability/_differential_vars.py index 94bbaf1..a8d492a 100644 --- a/src/drvi/utils/tools/interpretability/_differential_vars.py +++ b/src/drvi/utils/tools/interpretability/_differential_vars.py @@ -15,7 +15,102 @@ def find_differential_effects( key_added: str = "effect", add_to_counts: float = 0.1, relax_max_by: float = 0.0, -): +) -> None: + """Find differential effects in latent space traversal data. + + This function analyzes the differential effects between control and effect + conditions in traversal data to identify genes that respond to latent + dimension changes. It supports two methods for calculating differential effects. + + The "max_possible" method is the simple log-fold-change effect between effect and control. + The "min_possible" method is more conservative and considers the maximum possible effect of + other dimensions to normalize the log-fold-change effect of the current dimension. + + Parameters + ---------- + traverse_adata + AnnData object created by `traverse_latent` or `make_traverse_adata`. + Must contain `.layers['control']` and `.layers['effect']`. + method + Method for calculating differential effects: + - "max_possible": Simple log-fold-change between effect and control conditions. + This is the direct difference in log-space and represents the maximum + possible effect a latent dimension can have on gene expression. + - "min_possible": Conservative estimate that normalizes the effect by considering + the maximum possible effects from other dimensions. Although the effect in count space + is deterministic, the effect in log-space is not. This accounts for that, so changes in + one dimension may constrain the possible change in other dimensions. + key_added + Prefix for the keys added to `traverse_adata.uns` and `traverse_adata.varm`. + Results will be stored with keys like `{key_added}_traverse_effect_stepwise`. + add_to_counts + Small value added to counts to avoid log(0) issues in log-space calculations. + This pseudo-count ensures numerical stability when computing log-fold-changes. + relax_max_by + Relaxation factor for the maximum possible effect calculation + (only used with "min_possible" method). + + Returns + ------- + None + Results are stored in `traverse_adata`: + - `.uns[f"{key_added}_traverse_effect_stepwise"]`: Stepwise effects for each + latent dimension, step, and gene (shape: n_latent × n_steps × n_vars) + - `.varm[f"{key_added}_traverse_effect_pos"]`: Maximum positive effects per + dimension and gene (DataFrame with genes as rows, dimensions as columns) + - `.varm[f"{key_added}_traverse_effect_neg"]`: Maximum negative effects per + dimension and gene (DataFrame with genes as rows, dimensions as columns) + - `.uns[f"{key_added}_traverse_effect_pos_dim_ids"]`: Array of dimension IDs + for positive effects + - `.uns[f"{key_added}_traverse_effect_neg_dim_ids"]`: Array of dimension IDs + for negative effects + + Raises + ------ + AssertionError + If `method` is not one of the allowed values. + ValueError + If required data is missing from `traverse_adata` (e.g., missing layers). + KeyError + If required columns are missing from `traverse_adata.obs`. + + Notes + ----- + The function performs the following steps: + 1. Calculates differential effects using the specified method: + - "max_possible": Direct difference between effect and control conditions + - "min_possible": Normalized difference that considers the maximum possible + effect from other dimensions + 2. Identifies maximum effects in positive and negative directions + 3. Stores results in the AnnData object for further analysis + + **Method Details:** + + - **max_possible**: Computes the direct log-fold-change between effect and control + conditions. This represents the maximum possible effect a latent dimension can + have on gene expression, assuming no constraints from other dimensions. + + - **min_possible**: More conservative approach that normalizes the effect by + considering the maximum possible effects from other dimensions: + ``` + normalized_effect = log(exp(effect) + pseudo_count + baseline) - baseline + ``` + where baseline is the maximum possible effect from other dimensions. + + Examples + -------- + >>> # Using max_possible method (default) + >>> find_differential_effects(traverse_adata, method="max_possible") + >>> + >>> # Using min_possible method with custom parameters + >>> find_differential_effects( + ... traverse_adata, method="min_possible", key_added="conservative", add_to_counts=0.05, relax_max_by=0.1 + ... ) + >>> # Access results + >>> stepwise_effects = traverse_adata.uns["effect_traverse_effect_stepwise"] + >>> positive_effects = traverse_adata.varm["effect_traverse_effect_pos"] + >>> negative_effects = traverse_adata.varm["effect_traverse_effect_neg"] + """ assert method in ["max_possible", "min_possible"] # Reorder the traverse_adata to original order @@ -103,7 +198,79 @@ def combine_differential_effects( keys: list[str], key_added: str, combine_function: callable, -): +) -> None: + """Combine differential effects from multiple analyses. + + This function combines differential effects calculated using different methods + or parameters to create a unified score. It applies a custom combination + function to merge the stepwise effects and recalculates the positive/negative + direction effects. + + Parameters + ---------- + traverse_adata + AnnData object containing differential effects from `find_differential_effects`. + Must have `.uns[f"{key}_traverse_effect_stepwise"]` for each key in `keys`. + keys + List of keys corresponding to existing differential effect analyses. + Each key should have been used in a previous call to `find_differential_effects`. + key_added + Prefix for the keys added to store the combined results. + Results will be stored with keys like `{key_added}_traverse_effect_stepwise`. + combine_function + Function to combine the stepwise effects. Should take multiple arrays + (one for each key) as positional arguments and return a single combined + array with the same shape as the input arrays. + + Returns + ------- + None + Combined results are stored in `traverse_adata`: + - `.uns[f"{key_added}_traverse_effect_stepwise"]`: Combined stepwise effects + (shape: n_latent × n_steps × n_vars) + - `.varm[f"{key_added}_traverse_effect_pos"]`: Combined positive direction effects + (DataFrame with genes as rows, dimensions as columns) + - `.varm[f"{key_added}_traverse_effect_neg"]`: Combined negative direction effects + (DataFrame with genes as rows, dimensions as columns) + - `.uns[f"{key_added}_traverse_effect_pos_dim_ids"]`: Array of dimension IDs + for positive effects + - `.uns[f"{key_added}_traverse_effect_neg_dim_ids"]`: Array of dimension IDs + for negative effects + + Raises + ------ + KeyError + If any key in `keys` is missing from `traverse_adata.uns`. + ValueError + If the combine_function returns an array with unexpected shape. + + Notes + ----- + The function performs the following steps: + 1. Reorders the data to original order for consistent processing + 2. Extracts stepwise effects for each key in `keys` + 3. Applies the `combine_function` to merge the effects + 4. Recalculates positive and negative direction effects from the combined data + 5. Stores the results with the new `key_added` prefix + + The `combine_function` should be designed to work with the specific types of + differential effects being combined. Common approaches include: + - Taking the maximum/minimum of multiple analyses + - Computing weighted averages + - Applying logical operations (AND/OR) for binary effects + + Examples + -------- + >>> # Combine max_possible and min_possible effects using maximum + >>> def combine_max_min(max_effects, min_effects): + ... return np.maximum(max_effects, min_effects) + >>> combine_differential_effects( + ... traverse_adata, + ... keys=["max_possible", "min_possible"], + ... key_added="combined", + ... combine_function=combine_max_min, + ... ) + """ # Reorder the traverse_adata to original order original_traverse_adata = traverse_adata traverse_adata = traverse_adata[traverse_adata.obs.sort_values(["original_order"]).index].copy() @@ -143,7 +310,73 @@ def combine_differential_effects( ].columns = original_traverse_adata.varm[f"{key_added}_traverse_effect_{effect_sign}"].columns.astype(str) -def calculate_differential_vars(traverse_adata: AnnData, **kwargs): +def calculate_differential_vars(traverse_adata: AnnData, **kwargs) -> None: + """Calculate differential variables based on a combination of max_possible and min_possible effects. + + This function performs a comprehensive differential variable analysis by + calculating effects using both "max_possible" and "min_possible" methods, + then combining them into a unified score that considers both approaches. + + Parameters + ---------- + traverse_adata + AnnData object created by `traverse_latent` or `make_traverse_adata`. + Must contain `.layers['control']` and `.layers['effect']`. + **kwargs + Additional keyword arguments passed to `find_differential_effects`. + + Returns + ------- + None + Results are stored in `traverse_adata`: + - `.uns["max_possible_traverse_effect_stepwise"]`: Max possible stepwise effects + - `.uns["min_possible_traverse_effect_stepwise"]`: Min possible stepwise effects + - `.uns["combined_score_traverse_effect_stepwise"]`: Combined stepwise effects + - `.varm["max_possible_traverse_effect_pos/neg"]`: Max possible direction effects + - `.varm["min_possible_traverse_effect_pos/neg"]`: Min possible direction effects + - `.varm["combined_score_traverse_effect_pos/neg"]`: Combined direction effects + + Notes + ----- + The function performs the following steps: + 1. Calculates differential effects using "max_possible" method + 2. Calculates differential effects using "min_possible" method + 3. Combines the results using a custom scoring function + + **Scoring Logic:** + + The combined score is designed to identify genes that show consistent + differential effects across both calculation methods. The filtering criteria are: + + 1. **Base threshold**: max_possible >= 1.0 (ensures substantial effect) + 2. **Relative thresholds**: Either: + - max_possible > 50% of its maximum across all dimensions and steps, OR + - min_possible > 10% of its maximum across all dimensions and steps + 3. **Final score**: For genes passing the filters, score = max_possible × min_possible + + This approach provides a robust way to identify genes that show consistent + differential effects across both calculation methods, reducing false positives + while maintaining sensitivity to real biological effects. + + **Biological Interpretation:** + + - **High combined scores**: Genes with strong specific effects + - **High max_possible, low min_possible**: These genes are usually shared with another program, + with the other program having a larger effect. + - **Low max_possible, high min_possible**: Not possible. + - **Low scores in both**: Genes with weak or inconsistent effects + + Examples + -------- + >>> # Basic differential variable calculation + >>> calculate_differential_vars(traverse_adata) + >>> # With custom parameters + >>> calculate_differential_vars(traverse_adata, add_to_counts=0.05, relax_max_by=0.1) + >>> # Access results + >>> max_effects = traverse_adata.varm["max_possible_traverse_effect_pos"] + >>> min_effects = traverse_adata.varm["min_possible_traverse_effect_pos"] + >>> combined_effects = traverse_adata.varm["combined_score_traverse_effect_pos"] + """ print("Finding differential variables per latent dimension ...") find_differential_effects(traverse_adata, method="max_possible", key_added="max_possible", **kwargs) find_differential_effects(traverse_adata, method="min_possible", key_added="min_possible", **kwargs) @@ -169,28 +402,97 @@ def combine_function(min_possible, max_possible): def get_split_effects( model: DRVI, embed: AnnData, - n_steps=10 * 2, - n_samples=100, - traverse_kwargs=None, - de_kwargs=None, -): - """ - Get the split effect of a latent dimension. + n_steps: int = 20, + n_samples: int = 100, + traverse_kwargs: dict | None = None, + de_kwargs: dict | None = None, +) -> AnnData: + """Get split effects by performing latent traversal and differential analysis. + + This is a high-level function that combines latent space traversal with + differential variable analysis. It performs the complete pipeline from + generating traversal data to calculating differential effects. Parameters ---------- model - DRVI model object. + Trained DRVI model for decoding latent representations. embed - latent representation of the model. + AnnData object containing latent dimension statistics in `.var`. + Must have columns: `original_dim_id`, `min`, `max`, `std`, `title`, `vanished`, `order`. n_steps - Number of steps. + Number of steps in the traversal. Must be even (half negative, half positive). n_samples - Number of samples. + Number of samples to generate for each step. traverse_kwargs - Additional arguments passed to `traverse_the_latent`. - kwargs - Additional arguments passed to `find_differential_vars`. + Additional arguments passed to `traverse_latent`. Common options include: + - `copy_adata_var_info`: Whether to copy variable information + - `noise_formula`: Custom noise generation function + - `max_noise_std`: Maximum noise standard deviation. + de_kwargs + Additional arguments passed to `calculate_differential_vars`. Common options include: + - `add_to_counts`: Pseudo-count for log calculations + - `relax_max_by`: Relaxation factor for min_possible method. + + Returns + ------- + AnnData + AnnData object containing both traversal data and differential analysis results. + Includes all outputs from `traverse_latent` and `calculate_differential_vars`: + + **Traversal Data:** + - `.X`: Difference between effect and control conditions + - `.layers['control']`: Control condition gene expression + - `.layers['effect']`: Effect condition gene expression + - `.obs`: Metadata including dim_id, sample_id, step_id, span_value, title, vanished, order + + **Differential Analysis Results:** + - `.uns["max_possible_traverse_effect_stepwise"]`: Max possible stepwise effects + - `.uns["min_possible_traverse_effect_stepwise"]`: Min possible stepwise effects + - `.uns["combined_score_traverse_effect_stepwise"]`: Combined stepwise effects + - `.varm["*_traverse_effect_pos/neg"]`: Direction-specific effects for all methods + + Raises + ------ + ValueError + If required columns are missing from `embed.var`. + KeyError + If required data is missing from the model or embed objects. + + Notes + ----- + This function is a convenience wrapper that performs the complete analysis + pipeline in one call. It's equivalent to: + + ```python + traverse_adata = traverse_latent(model, embed, n_steps, n_samples, **traverse_kwargs) + calculate_differential_vars(traverse_adata, **de_kwargs) + return traverse_adata + ``` + + **Functionality:** + + 1. **Traversal Generation**: Creates systematic traversals through latent space + 2. **Decoding**: Converts latent traversals to gene expression predictions + 3. **Differential Analysis**: Calculates effects using max_possible and min_possible methods + 4. **Combination**: Creates unified scores for robust gene identification + + **Use Cases:** + + - **Gene Discovery**: Identify genes associated with specific latent dimensions + - **Pathway Analysis**: Understand biological processes captured by latent factors + - **Model Validation**: Verify that latent dimensions have interpretable biological meaning + - **Comparative Analysis**: Compare effects across different models or conditions + + Examples + -------- + >>> # Basic split effects analysis + >>> split_effects = get_split_effects(model, embed) + >>> # With custom parameters + >>> split_effects = get_split_effects(model, embed, n_steps=30, n_samples=50) + >>> # Access results + >>> combined_effects = split_effects.varm["combined_score_traverse_effect_pos"] + >>> stepwise_effects = split_effects.uns["combined_score_traverse_effect_stepwise"] """ if traverse_kwargs is None: traverse_kwargs = {} @@ -210,24 +512,83 @@ def iterate_on_top_differential_vars( order_col: str = "order", gene_symbols: str | None = None, score_threshold: float = 0.0, -): - """ - Make an iterator of the top differential variables per latent dimension. +) -> list[tuple[str, pd.Series]]: + """Create an iterator of top differential variables per latent dimension. + + This function processes differential analysis results to create an organized + list of top differentially expressed genes for each latent dimension, + sorted by their effect scores and organized by dimension. Parameters ---------- traverse_adata - Anndata object with the split effects. + AnnData object with differential analysis results from `calculate_differential_vars`. + Must contain differential effect data for the specified `key`. key - Key to the differential variables in traverse_adata. + Key prefix for the differential variables in `traverse_adata`. + Should correspond to a key used in `find_differential_effects` or `calculate_differential_vars`. + Common value: "combined_score". title_col - Title columns defining title of the dimensions. + Column name in `traverse_adata.obs` containing dimension titles. + These titles will be used in the output dimension names. order_col - Order column that defines latent dimension orders. + Column name in `traverse_adata.obs` containing dimension ordering. gene_symbols - Column with the gene symbols in traverse_adata.var. + Column name in `traverse_adata.var` containing gene symbols. + If None, uses the index of `traverse_adata.var` (usually gene IDs). + Useful for converting between gene IDs and readable gene names. score_threshold - A threshold to filter the differential variables. + Minimum score threshold to include genes in the results. + Only genes with scores above this threshold will be included. + + Returns + ------- + list[tuple[str, pd.Series]] + List of tuples, where each tuple contains: + - str: Dimension title with direction indicator (e.g., "Cell Cycle+", "Cell Cycle-") + - pd.Series: Series of gene scores for that dimension/direction, sorted descending + + The list is sorted by dimension order, with each dimension appearing at most twice + (once for positive effects, once for negative effects). + + Raises + ------ + KeyError + If required columns or differential effect data are missing. + ValueError + If the specified key doesn't exist in the AnnData object. + + Notes + ----- + The function performs the following steps: + 1. Extracts positive and negative differential effects for the specified key + 2. Maps gene names to symbols if `gene_symbols` is provided + 3. Filters genes by score threshold + 4. Organizes results by dimension and direction (positive/negative) + 5. Returns a list sorted by dimension order + + **Output Structure:** + + Each dimension appears twice in the results - once for positive effects + and once for negative effects. The direction is indicated by "+" or "-" + appended to the dimension title. + + Only dimensions with at least one gene above the threshold are included. + + Examples + -------- + >>> # Basic iteration over top differential variables + >>> top_vars = iterate_on_top_differential_vars(traverse_adata, "combined_score") + >>> for dim_title, gene_scores in top_vars: + ... print(f"{dim_title}: {len(gene_scores)} genes") + ... print(f"Top genes: {gene_scores.head().index.tolist()}") + >>> # With custom parameters and gene symbols + >>> top_vars = iterate_on_top_differential_vars( + ... traverse_adata, "max_possible", gene_symbols="gene_symbol", score_threshold=1.0 + ... ) + >>> # Create a summary of results + >>> for dim_title, gene_scores in top_vars: + ... print(f"{dim_title}: {gene_scores.head().index.tolist()}") """ df_pos = traverse_adata.varm[f"{key}_traverse_effect_pos"].copy() df_neg = traverse_adata.varm[f"{key}_traverse_effect_neg"].copy() diff --git a/src/drvi/utils/tools/interpretability/_latent_traverse.py b/src/drvi/utils/tools/interpretability/_latent_traverse.py index 6312ea1..8e288bb 100644 --- a/src/drvi/utils/tools/interpretability/_latent_traverse.py +++ b/src/drvi/utils/tools/interpretability/_latent_traverse.py @@ -8,14 +8,72 @@ def iterate_dimensions( - # n_latent: int, latent_dims: np.ndarray, latent_min: np.ndarray, latent_max: np.ndarray, n_steps: int = 10 * 2, n_samples: int = 100, -): - assert n_steps % 2 == 0 +) -> AnnData: + """Generate systematic traversal data for latent dimensions. + + This function creates a systematic grid of latent space traversals by + generating combinations of dimension IDs, step IDs, and sample IDs. + It creates a sparse matrix representation of the traversal vectors. + + Parameters + ---------- + latent_dims + Array of latent dimension indices to traverse. + latent_min + Minimum values for each latent dimension (typically negative). + Must have same length as `latent_dims`. + latent_max + Maximum values for each latent dimension (typically positive). + Must have same length as `latent_dims`. + n_steps + Number of steps in the traversal. Must be even (half negative, half positive). + n_samples + Number of samples to generate for each step. + + Returns + ------- + AnnData + AnnData object containing the traversal vectors in `.X` and metadata in `.obs`. + The `.obs` contains: + - `original_order`: Original index of each row + - `dim_id`: Which latent dimension this row corresponds to + - `sample_id`: Sample identifier (0 to n_samples-1) + - `step_id`: Step identifier (0 to n_steps-1) + - `span_value`: The actual latent value for this step + + Raises + ------ + AssertionError + If `n_steps` is not even. + + Notes + ----- + The function creates a systematic traversal where: + - For each latent dimension, it generates n_steps steps + - Each step has n_samples samples + - The first half of steps go from latent_min to 0 + - The second half of steps go from 0 to latent_max + - The result is a sparse matrix of shape (n_latent * n_steps * n_samples, n_latent) + + The traversal values are linearly interpolated between the min/max bounds, + ensuring smooth coverage of the latent space for each dimension. + + Examples + -------- + >>> # Basic traversal + >>> latent_dims = np.array([0, 1, 2]) + >>> latent_min = np.array([-2, -1, -3]) + >>> latent_max = np.array([2, 1, 3]) + >>> traverse_data = iterate_dimensions(latent_dims, latent_min, latent_max) + >>> print(f"Shape: {traverse_data.X.shape}") + >>> print(f"Unique dimensions: {traverse_data.obs['dim_id'].unique()}") + """ + assert n_steps % 2 == 0, "n_steps must be even" # Sometimes it is negligibly not like below # assert np.all(latent_min <= 0) & np.all(latent_max >= 0) @@ -73,7 +131,74 @@ def make_traverse_adata( max_noise_std: float = 0.2, copy_adata_var_info: bool = True, **kwargs, -): +) -> AnnData: + """Create traversal AnnData by decoding latent space traversals. + + This function generates systematic traversals through the latent space + and decodes them to observe the effects on gene expression. It creates + both control (baseline) and effect (traversal) conditions. + + Parameters + ---------- + model + Trained DRVI model for decoding latent representations. + embed + AnnData object containing latent dimension statistics in `.var`. + Must have columns: `original_dim_id`, `min`, `max`, `std`. + n_steps + Number of steps in the traversal. Must be even (half negative, half positive). + n_samples + Number of samples to generate for each step. + noise_formula + Function to compute noise standard deviation from dimension std values. + Should take a numpy array and return a numpy array. + max_noise_std + Maximum allowed noise standard deviation. + copy_adata_var_info + Whether to copy variable information from the original model's AnnData. + **kwargs + Additional keyword arguments passed to `model.decode_latent_samples`. + + Returns + ------- + AnnData + AnnData object containing the traversal results with the following structure: + - `.X`: Difference between effect and control conditions (main result) + - `.layers['control']`: Control condition gene expression (noise only) + - `.layers['effect']`: Effect condition gene expression (noise + traversal) + - `.obsm['control_latent']`: Control condition latent values (noise only) + - `.obsm['effect_latent']`: Effect condition latent values (noise + traversal) + - `.obsm['cat_covs']`: Categorical covariates (if model uses them) + - `.obs['lib_size']`: Library size for each sample (set to 10,000) + - `.obs`: Metadata including dim_id, sample_id, step_id, span_value + + Raises + ------ + NotImplementedError + If the model has continuous covariates (not yet supported). + ValueError + If required columns are missing from `embed.var`. + + Notes + ----- + The function performs the following steps: + 1. Generates systematic traversal vectors using `iterate_dimensions` + 2. Creates baseline noise for each sample based on dimension std values + 3. Generates random categorical covariates for individual samples if the model uses them + 4. Sets library size to 10,000 for all samples (typical for single-cell data) + 5. Decodes both control (noise only) and effect (noise + traversal) conditions + 6. Returns the difference between effect and control as the main data + + The traversal goes from negative to positive values for each dimension, + allowing observation of how each dimension affects gene expression. + + The control condition contains only noise, while the effect condition + contains noise plus the systematic traversal. The difference shows + the pure effect of the latent dimension traversal on gene expression. + + The noise is generated based on the standard deviation of each latent dimension, + scaled by the `noise_formula` function and clipped to `max_noise_std`. + """ # Generate random delta vectors for each dimension span_adata = iterate_dimensions( latent_dims=embed.var["original_dim_id"].values, @@ -141,7 +266,45 @@ def make_traverse_adata( return traverse_adata -def get_dimensions_of_traverse_data(traverse_adata: AnnData): +def get_dimensions_of_traverse_data(traverse_adata: AnnData) -> tuple[int, int, int, int]: + """Get the dimensions of traversal data. + + This function extracts the key dimensions from a traversal AnnData object + created by `make_traverse_adata` or `traverse_latent`. + + Parameters + ---------- + traverse_adata + AnnData object created by traversal functions. + + Returns + ------- + tuple[int, int, int, int] + A tuple containing: + - n_latent: Number of latent dimensions traversed + - n_steps: Number of steps in each traversal + - n_samples: Number of samples per step + - n_vars: Number of variables (genes) in the data + + Raises + ------ + KeyError + If required columns are missing from `traverse_adata.obs`. + + Notes + ----- + The function expects the following columns in `traverse_adata.obs`: + - `dim_id`: Latent dimension identifier + - `step_id`: Step identifier within each dimension + - `sample_id`: Sample identifier within each step + + Examples + -------- + >>> # Get dimensions of traversal data + >>> n_latent, n_steps, n_samples, n_vars = get_dimensions_of_traverse_data(traverse_adata) + >>> print(f"Traversed {n_latent} dimensions with {n_steps} steps and {n_samples} samples") + >>> print(f"Data has {n_vars} variables") + """ # Get the number of latent dimensions, steps, samples, and vars n_latent = traverse_adata.obs["dim_id"].nunique() n_steps = traverse_adata.obs["step_id"].nunique() @@ -153,11 +316,74 @@ def get_dimensions_of_traverse_data(traverse_adata: AnnData): def traverse_latent( model: DRVI, embed: AnnData, - n_steps=10 * 2, - n_samples=100, - copy_adata_var_info=True, + n_steps: int = 10 * 2, + n_samples: int = 100, + copy_adata_var_info: bool = True, **kwargs, -): +) -> AnnData: + """Perform latent space traversal and enrich with metadata. + + This function generates systematic traversals through the latent space + and decodes them to observe the effects on gene expression. It creates + both control (baseline) and effect (traversal) conditions. + Additionally, it enriches the data with dimension-specific metadata + like titles, vanished status, and ordering. + + Parameters + ---------- + model + Trained DRVI model for decoding latent representations. + embed + AnnData object containing latent dimension statistics in `.var`. + Must have columns: `original_dim_id`, `min`, `max`, `std`, `title`, `vanished`, `order`. + n_steps + Number of steps in the traversal. Must be even. + n_samples + Number of samples to generate for each step. + copy_adata_var_info + Whether to copy variable information from the original model's AnnData. + **kwargs + Additional keyword arguments passed to `make_traverse_adata`. + + Returns + ------- + AnnData + AnnData object containing the traversal results with enriched metadata. + In addition to the structure returned by `make_traverse_adata`, the `.obs` + also contains: + - `title`: Dimension titles from `embed.var['title']` + - `vanished`: Vanished status from `embed.var['vanished']` + - `order`: Dimension ordering from `embed.var['order']` + + Raises + ------ + ValueError + If required columns are missing from `embed.var`. + + Notes + ----- + The function performs the following steps: + 1. Generates systematic traversal vectors using `iterate_dimensions` + 2. Creates baseline noise for each sample based on dimension std values + 3. Generates random categorical covariates for individual samples if the model uses them + 4. Decodes both control (noise only) and effect (noise + traversal) conditions + 5. Returns the difference between effect and control as the main data + 6. Enriches the traversal data with dimension-specific information like titles, vanished status, and ordering. + + The function expects the following columns in `embed.var`: + - `original_dim_id`: Original dimension indices + - `min`, `max`, `std`: Dimension statistics + - `title`: Human-readable dimension titles + - `vanished`: Boolean indicating vanished dimensions + - `order`: Dimension ordering + + Examples + -------- + >>> # Basic traversal + >>> traverse_data = traverse_latent(model, embed) + >>> # Traversal with custom parameters + >>> traverse_data = traverse_latent(model, embed, n_steps=30, n_samples=50) + """ if "original_dim_id" not in embed.var: raise ValueError( 'Column "original_dim_id" not found in `embed.var`. Please run `set_latent_dimension_stats` to set vanished status.'