diffbio 0.1.0__py3-none-any.whl

diffbio/operators/normalization/umap.py

@@ -0,0 +1,261 @@
+"""Differentiable UMAP dimensionality reduction.
+
+This module implements a differentiable version of UMAP (Uniform Manifold
+Approximation and Projection) for dimensionality reduction with end-to-end
+gradient flow.
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+
+from diffbio.constants import DISTANCE_MASK_SENTINEL
+from diffbio.core.graph_utils import (
+    compute_fuzzy_membership,
+    compute_pairwise_distances,
+    symmetrize_graph,
+)
+
+
+@dataclass(frozen=True)
+class UMAPConfig(OperatorConfig):
+    """Configuration for differentiable UMAP.
+
+    Attributes:
+        n_components: Number of dimensions in the embedding.
+        n_neighbors: Number of neighbors for local structure preservation.
+        metric: Distance metric ('euclidean' or 'cosine').
+        input_features: Number of input features (required for initialization).
+        hidden_dim: Hidden dimension for projection network.
+        stream_name: Name of the data stream to process.
+    """
+
+    n_components: int = 2
+    n_neighbors: int = 15
+    metric: str = "euclidean"
+    input_features: int = 64
+    hidden_dim: int = 32
+
+    def __post_init__(self) -> None:
+        """Validate the supported UMAP configuration surface."""
+        super().__post_init__()
+
+        if self.n_components <= 0:
+            raise ValueError(f"n_components must be positive, got {self.n_components}")
+        if self.n_neighbors <= 0:
+            raise ValueError(f"n_neighbors must be positive, got {self.n_neighbors}")
+        if self.input_features <= 0:
+            raise ValueError(f"input_features must be positive, got {self.input_features}")
+        if self.hidden_dim <= 0:
+            raise ValueError(f"hidden_dim must be positive, got {self.hidden_dim}")
+        if self.metric not in {"euclidean", "cosine"}:
+            raise ValueError(f"metric must be 'euclidean' or 'cosine', got '{self.metric}'")
+
+
+class ParametricUMAPHead(nnx.Module):
+    """Parametric UMAP embedding head with learnable similarity curve."""
+
+    def __init__(self, config: UMAPConfig, *, rngs: nnx.Rngs):
+        """Initialize the parametric embedding head."""
+        super().__init__()
+        self.curve_params = nnx.Param(jnp.array([1.929, 0.7915], dtype=jnp.float32))
+        self.projection_backbone = MLP(
+            hidden_dims=[config.hidden_dim, config.n_components],
+            in_features=config.input_features,
+            activation="relu",
+            output_activation=None,
+            use_batch_norm=False,
+            rngs=rngs,
+        )
+
+    def project(self, features: jax.Array) -> jax.Array:
+        """Project high-dimensional features to low-dimensional embedding."""
+        embedding = self.projection_backbone(features)
+        if isinstance(embedding, tuple):
+            raise TypeError("ParametricUMAPHead projection backbone must return a single tensor.")
+        return embedding
+
+    def curve_coefficients(self) -> tuple[jax.Array, jax.Array]:
+        """Return positive low-dimensional similarity curve coefficients."""
+        a = jnp.abs(self.curve_params[0]) + 1e-6
+        b = jnp.abs(self.curve_params[1]) + 1e-6
+        return a, b
+
+
+class DifferentiableUMAP(OperatorModule):
+    """Differentiable UMAP for dimensionality reduction.
+
+    This operator implements a simplified differentiable version of UMAP that
+    learns a low-dimensional embedding while preserving local structure.
+
+    The UMAP loss function is:
+        L = sum_edges [p_ij * log(q_ij) + (1 - p_ij) * log(1 - q_ij)]
+
+    where:
+        - p_ij is the high-dimensional similarity (fuzzy set membership)
+        - q_ij is the low-dimensional similarity
+
+    This implementation uses a parametric approach with learnable curve
+    parameters (a, b) for the low-dimensional similarity function.
+
+    Example:
+        ```python
+        config = UMAPConfig(
+            n_components=2,
+            n_neighbors=15,
+        )
+        umap = DifferentiableUMAP(config, rngs=rngs)
+
+        data = {"features": high_dim_data}  # (n_samples, n_features)
+        result, state, metadata = umap.apply(data, {}, None)
+        embedding = result["embedding"]  # (n_samples, n_components)
+        ```
+    """
+
+    def __init__(self, config: UMAPConfig, *, rngs: nnx.Rngs | None = None):
+        """Initialize the differentiable UMAP.
+
+        Args:
+            config: Configuration for UMAP.
+            rngs: Random number generators for initialization.
+        """
+        super().__init__(config, rngs=rngs)
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        self.embedding_head = ParametricUMAPHead(config, rngs=rngs)
+
+    def _project(self, features: jax.Array) -> jax.Array:
+        """Project high-dimensional features to low-dimensional embedding.
+
+        Args:
+            features: Input features of shape (n_samples, n_features).
+
+        Returns:
+            Embedding of shape (n_samples, n_components).
+        """
+        return self.embedding_head.project(features)
+
+    def _compute_high_dim_similarities(self, features: jax.Array) -> jax.Array:
+        """Compute high-dimensional fuzzy set membership (p_ij).
+
+        Uses k-nearest neighbors and Gaussian kernel with local bandwidth.
+        Delegates to reusable graph utilities in ``diffbio.core.graph_utils``.
+
+        Args:
+            features: Input features of shape (n_samples, n_features).
+
+        Returns:
+            Similarity matrix of shape (n_samples, n_samples).
+        """
+        n_samples = features.shape[0]
+        n_neighbors = min(self.config.n_neighbors, n_samples - 1)
+
+        distances = compute_pairwise_distances(features, metric=self.config.metric)
+        distances = distances + jnp.eye(n_samples) * DISTANCE_MASK_SENTINEL
+
+        p_ij = compute_fuzzy_membership(distances, k=n_neighbors)
+        return symmetrize_graph(p_ij)
+
+    def _compute_low_dim_similarities(self, embedding: jax.Array) -> jax.Array:
+        """Compute low-dimensional similarities (q_ij).
+
+        Uses the UMAP student-t like kernel:
+            q(d) = 1 / (1 + a * d^(2b))
+
+        Args:
+            embedding: Low-dimensional embedding of shape (n_samples, n_components).
+
+        Returns:
+            Similarity matrix of shape (n_samples, n_samples).
+        """
+        n_samples = embedding.shape[0]
+
+        # Compute pairwise distances in embedding space
+        diff = embedding[:, None, :] - embedding[None, :, :]
+        distances_sq = jnp.sum(diff**2, axis=-1)
+
+        # UMAP similarity kernel
+        a, b = self.embedding_head.curve_coefficients()
+
+        q_ij = 1.0 / (1.0 + a * jnp.power(distances_sq + 1e-8, b))
+
+        # Set diagonal to 0
+        q_ij = q_ij * (1 - jnp.eye(n_samples))
+
+        return q_ij
+
+    def _compute_umap_loss(self, p_ij: jax.Array, q_ij: jax.Array) -> jax.Array:
+        """Compute UMAP cross-entropy loss.
+
+        Args:
+            p_ij: High-dimensional similarities.
+            q_ij: Low-dimensional similarities.
+
+        Returns:
+            Scalar loss value.
+        """
+        # Cross-entropy: -sum(p * log(q) + (1-p) * log(1-q))
+        eps = 1e-8
+        q_ij = jnp.clip(q_ij, eps, 1 - eps)
+
+        # Attractive term
+        attractive = -jnp.sum(p_ij * jnp.log(q_ij))
+
+        # Repulsive term
+        repulsive = -jnp.sum((1 - p_ij) * jnp.log(1 - q_ij))
+
+        return attractive + repulsive
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict | None,
+        random_params: dict | None = None,
+        stats: dict | None = None,
+    ) -> tuple[dict, dict, dict | None]:
+        """Apply UMAP dimensionality reduction.
+
+        Args:
+            data: Dictionary containing:
+                - 'features': High-dimensional features of shape (n_samples, n_features)
+            state: Operator state dictionary.
+            metadata: Optional metadata dictionary.
+            random_params: Optional random parameters (unused).
+            stats: Optional statistics dictionary (unused).
+
+        Returns:
+            Tuple of (output_data, state, metadata) where output_data contains:
+
+                - 'features': Original high-dimensional features
+                - 'embedding': Low-dimensional embedding
+                - 'high_dim_similarities': Fuzzy set memberships (p_ij)
+                - 'low_dim_similarities': Embedding similarities (q_ij)
+        """
+        del random_params, stats  # Unused
+
+        features = data["features"]
+
+        # Project to low-dimensional space
+        embedding = self._project(features)
+
+        # Compute similarities
+        p_ij = self._compute_high_dim_similarities(features)
+        q_ij = self._compute_low_dim_similarities(embedding)
+
+        output_data = {
+            **data,
+            "embedding": embedding,
+            "high_dim_similarities": p_ij,
+            "low_dim_similarities": q_ij,
+        }
+
+        return output_data, state, metadata

diffbio/operators/normalization/vae_normalizer.py

@@ -0,0 +1,258 @@
+"""VAE-based count normalization operator.
+
+This module provides a variational autoencoder for normalizing gene
+expression count data, inspired by scVI (Lopez et al., 2018).
+
+Key technique: Learn a latent representation of cell state while
+modeling count data with a configurable likelihood (Poisson or ZINB).
+
+The ZINB (Zero-Inflated Negative Binomial) likelihood is particularly
+suited for single-cell RNA-seq data which exhibits both overdispersion
+and excess zeros (dropout events).
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+import jax
+import jax.numpy as jnp
+from artifex.generative_models.core.losses.divergence import gaussian_kl_divergence
+from datarax.core.config import OperatorConfig
+from flax import nnx
+from jaxtyping import Array, Float, PyTree
+
+from diffbio.core.base_operators import EncoderDecoderOperator
+from diffbio.operators._count_vae import CountReconstructionMixin, CountVAEBackboneMixin
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class VAENormalizerConfig(OperatorConfig):
+    """Configuration for VAENormalizer.
+
+    Attributes:
+        latent_dim: Dimension of latent space.
+        hidden_dims: Hidden layer dimensions for encoder/decoder.
+        n_genes: Number of genes (input/output dimension).
+        use_batch_correction: Whether to include batch effects.
+        likelihood: Likelihood model for reconstruction loss.
+            'poisson' for standard Poisson NLL, 'zinb' for
+            Zero-Inflated Negative Binomial.
+    """
+
+    latent_dim: int = 10
+    hidden_dims: list[int] = field(default_factory=lambda: [128, 64])
+    n_genes: int = 2000
+    use_batch_correction: bool = False
+    likelihood: Literal["poisson", "zinb"] = "poisson"
+
+    def __post_init__(self) -> None:
+        """Set stochastic defaults for VAE sampling and validate."""
+        object.__setattr__(self, "stochastic", True)
+        if self.stream_name is None:
+            object.__setattr__(self, "stream_name", "sample")
+        super().__post_init__()
+
+
+class VAENormalizer(CountReconstructionMixin, CountVAEBackboneMixin, EncoderDecoderOperator):
+    """Variational autoencoder for count normalization.
+
+    This operator learns a low-dimensional latent representation of
+    single-cell gene expression data while accounting for technical
+    factors like library size.
+
+    The model:
+    - Encoder: counts -> latent (mean, logvar)
+    - Reparameterization: z = mean + exp(0.5 * logvar) * epsilon
+    - Decoder: z -> gene expression rates (and optionally dispersion/dropout)
+
+    Supports two likelihood models:
+    - Poisson: Simple count model (default)
+    - ZINB: Zero-Inflated Negative Binomial for overdispersed data
+      with excess zeros, as used in scVI
+
+    Inherits from EncoderDecoderOperator to get:
+
+    - reparameterize() for sampling with reparameterization trick
+    - kl_divergence() for KL from standard normal
+    - elbo_loss() for combining reconstruction and KL losses
+
+    Args:
+        config: VAENormalizerConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = VAENormalizerConfig(n_genes=2000, latent_dim=10)
+        normalizer = VAENormalizer(config, rngs=nnx.Rngs(42))
+        data = {"counts": counts, "library_size": lib_size}
+        result, state, meta = normalizer.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: VAENormalizerConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Initialize the VAE normalizer.
+
+        Args:
+            config: VAE configuration.
+            rngs: Random number generators for initialization and sampling.
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+        safe_rngs = self._init_count_vae_operator(config=config, rngs=rngs)
+
+        decoder_out_dim = config.hidden_dims[0] if config.hidden_dims else config.latent_dim
+
+        # ZINB-specific decoder heads
+        if config.likelihood == "zinb":
+            self.fc_log_theta = nnx.Linear(
+                in_features=decoder_out_dim, out_features=config.n_genes, rngs=safe_rngs
+            )
+            self.fc_pi_logit = nnx.Linear(
+                in_features=decoder_out_dim, out_features=config.n_genes, rngs=safe_rngs
+            )
+
+    def decode(
+        self,
+        z: Float[Array, "latent_dim"],
+        library_size: Float[Array, ""],
+    ) -> dict[str, jax.Array]:
+        """Decode latent representation to gene expression parameters.
+
+        Args:
+            z: Latent representation.
+            library_size: Total counts (library size) for normalization.
+
+        Returns:
+            Dictionary with keys:
+                - 'log_rate': Log rates for each gene (always present).
+                - 'log_theta': Log dispersion parameter (ZINB only).
+                - 'pi_logit': Dropout logit for zero inflation (ZINB only).
+        """
+        x = self.decode_hidden(z)
+
+        # Output layer (log rates, normalized by library size)
+        log_rate = self.fc_output(x)
+
+        result: dict[str, jax.Array] = {}
+
+        # ZINB heads are computed before adding library size
+        if self.config.likelihood == "zinb":
+            result["log_theta"] = self.fc_log_theta(x)
+            result["pi_logit"] = self.fc_pi_logit(x)
+
+        # Add library size effect (log scale)
+        log_rate = log_rate + jnp.log(library_size + 1e-8)
+        result["log_rate"] = log_rate
+
+        return result
+
+    # kl_divergence() is inherited from EncoderDecoderOperator
+
+    def compute_elbo_loss(
+        self,
+        counts: Float[Array, "n_genes"],
+        library_size: Float[Array, ""],
+    ) -> Float[Array, ""]:
+        """Compute negative ELBO loss.
+
+        Uses gaussian_kl_divergence from artifex for the KL term.
+
+        Args:
+            counts: Gene expression counts.
+            library_size: Total counts.
+
+        Returns:
+            Negative ELBO (reconstruction loss + KL divergence).
+        """
+        # Encode
+        mean, logvar = self.encode(counts)
+
+        # Sample latent using inherited reparameterize (uses self.rngs)
+        z = self.reparameterize(mean, logvar)
+
+        # Decode (returns dict)
+        decode_output = self.decode(z, library_size)
+
+        # Reconstruction loss
+        recon_loss = self.reconstruction_loss(counts, decode_output)
+
+        # KL divergence via artifex (handles unbatched 1-D inputs with sum)
+        kl = gaussian_kl_divergence(mean, logvar, reduction="sum")
+
+        return recon_loss + kl
+
+    def apply(
+        self,
+        data: PyTree,
+        state: PyTree,
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
+        """Apply VAE normalization to count data.
+
+        This method encodes the counts to latent space, samples a
+        latent representation, and decodes to normalized expression.
+
+        Args:
+            data: Dictionary containing:
+                - "counts": Gene expression counts (n_genes,)
+                - "library_size": Total counts for the cell
+            state: Element state (passed through unchanged)
+            metadata: Element metadata (passed through unchanged)
+            random_params: Optional random parameters (not used)
+            stats: Not used
+
+        Returns:
+            Tuple of (transformed_data, state, metadata):
+                - transformed_data contains:
+
+                    - "counts": Original counts
+                    - "normalized": Normalized expression
+                    - "latent_z": Sampled latent representation
+                    - "latent_mean": Mean of latent distribution
+                    - "latent_logvar": Log variance of latent distribution
+                    - "log_rate": Decoded log rates
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        counts = data["counts"]
+        library_size = data["library_size"]
+
+        # Encode to latent distribution
+        mean, logvar = self.encode(counts)
+
+        # Sample from latent distribution using inherited reparameterize
+        # (uses self.rngs from EncoderDecoderOperator)
+        z = self.reparameterize(mean, logvar)
+
+        # Decode to gene expression rates (returns dict)
+        decode_output = self.decode(z, library_size)
+        log_rate = decode_output["log_rate"]
+
+        # Compute normalized expression (rate normalized by library size)
+        # This is the "denoised" expression
+        normalized = jnp.exp(log_rate - jnp.log(library_size + 1e-8))
+
+        # Build output data
+        transformed_data = {
+            "counts": counts,
+            "library_size": library_size,
+            "normalized": normalized,
+            "latent_z": z,
+            "latent_mean": mean,
+            "latent_logvar": logvar,
+            "log_rate": log_rate,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/population/__init__.py

@@ -0,0 +1,17 @@
+"""Population genetics operators.
+
+This module provides differentiable operators for population genetics analysis
+including ancestry estimation, phasing, and imputation.
+"""
+
+from diffbio.operators.population.ancestry_estimation import (
+    AncestryEstimatorConfig,
+    DifferentiableAncestryEstimator,
+    create_ancestry_estimator,
+)
+
+__all__ = [
+    "AncestryEstimatorConfig",
+    "DifferentiableAncestryEstimator",
+    "create_ancestry_estimator",
+]