diffbio 0.1.0__py3-none-any.whl

diffbio/operators/singlecell/ambient_removal.py

@@ -0,0 +1,333 @@
+"""Differentiable ambient RNA removal operator.
+
+This module provides CellBender-style ambient RNA removal using a VAE
+architecture that learns to separate cell-intrinsic from ambient signal.
+
+Key technique: Uses variational autoencoder to model cell-specific expression
+and ambient contamination fraction, enabling decontamination with uncertainty
+quantification.
+
+Applications: Removing ambient RNA contamination from single-cell RNA-seq data,
+improving cell type identification and differential expression analysis.
+
+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
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from datarax.core.config import OperatorConfig
+from flax import nnx
+from jaxtyping import Array, Float, PyTree
+
+from diffbio.configs import apply_stochastic_sampling_defaults
+from diffbio.core.base_operators import EncoderDecoderOperator
+from diffbio.utils.nn_utils import ARTIFEX_GELU_MLP_KWARGS
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class AmbientRemovalConfig(OperatorConfig):
+    """Configuration for DifferentiableAmbientRemoval.
+
+    Attributes:
+        n_genes: Number of genes in expression profiles.
+        latent_dim: Dimension of latent space.
+        hidden_dims: Hidden layer dimensions for encoder/decoder.
+        ambient_prior: Prior probability of ambient contamination.
+        temperature: Temperature for softmax operations.
+    """
+
+    n_genes: int = 2000
+    latent_dim: int = 64
+    hidden_dims: list[int] = field(default_factory=lambda: [256, 128])
+    ambient_prior: float = 0.01
+    temperature: float = 1.0
+
+    def __post_init__(self) -> None:
+        """Set stochastic defaults and validate."""
+        apply_stochastic_sampling_defaults(self)
+        super().__post_init__()
+        if not self.hidden_dims:
+            raise ValueError(
+                "AmbientRemovalConfig.hidden_dims must contain at least one hidden dimension."
+            )
+
+
+class AmbientEncoder(nnx.Module):
+    """Encoder network for ambient removal VAE."""
+
+    def __init__(
+        self,
+        n_genes: int,
+        hidden_dims: list[int],
+        latent_dim: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize the encoder.
+
+        Args:
+            n_genes: Number of input genes.
+            hidden_dims: Hidden layer dimensions.
+            latent_dim: Latent space dimension.
+            rngs: Random number generators.
+        """
+        super().__init__()
+        self.backbone = MLP(
+            hidden_dims=hidden_dims,
+            in_features=n_genes,
+            rngs=rngs,
+            **ARTIFEX_GELU_MLP_KWARGS,
+        )
+
+        # Latent projections
+        self.mean_proj = nnx.Linear(in_features=hidden_dims[-1], out_features=latent_dim, rngs=rngs)
+        self.logvar_proj = nnx.Linear(
+            in_features=hidden_dims[-1], out_features=latent_dim, rngs=rngs
+        )
+
+        # Contamination fraction projection
+        self.contamination_proj = nnx.Linear(in_features=hidden_dims[-1], out_features=1, rngs=rngs)
+
+    def __call__(
+        self,
+        counts: Float[Array, "n_cells n_genes"],
+    ) -> tuple[
+        Float[Array, "n_cells latent_dim"],
+        Float[Array, "n_cells latent_dim"],
+        Float[Array, "n_cells"],
+    ]:
+        """Encode counts to latent space.
+
+        Args:
+            counts: Input count matrix.
+
+        Returns:
+            Tuple of (mean, logvar, contamination_fraction).
+        """
+        # Log-normalize for encoder input
+        backbone_output = self.backbone(jnp.log1p(counts))
+        if isinstance(backbone_output, tuple):
+            raise TypeError("AmbientEncoder backbone must return a single tensor output.")
+
+        # Latent parameters
+        mean = self.mean_proj(backbone_output)
+        logvar = self.logvar_proj(backbone_output)
+
+        # Contamination fraction (bounded 0-1)
+        contamination = jax.nn.sigmoid(self.contamination_proj(backbone_output)).squeeze(-1)
+
+        return mean, logvar, contamination
+
+
+class AmbientDecoder(nnx.Module):
+    """Decoder network for ambient removal VAE."""
+
+    def __init__(
+        self,
+        latent_dim: int,
+        hidden_dims: list[int],
+        n_genes: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize the decoder.
+
+        Args:
+            latent_dim: Latent space dimension.
+            hidden_dims: Hidden layer dimensions (reversed from encoder).
+            n_genes: Number of output genes.
+            rngs: Random number generators.
+        """
+        super().__init__()
+        decoder_hidden_dims = list(reversed(hidden_dims))
+        self.backbone = MLP(
+            hidden_dims=decoder_hidden_dims,
+            in_features=latent_dim,
+            rngs=rngs,
+            **ARTIFEX_GELU_MLP_KWARGS,
+        )
+
+        # Output projection (log-rate for Poisson/NB)
+        self.output_proj = nnx.Linear(
+            in_features=decoder_hidden_dims[-1], out_features=n_genes, rngs=rngs
+        )
+
+    def __call__(
+        self,
+        z: Float[Array, "n_cells latent_dim"],
+    ) -> Float[Array, "n_cells n_genes"]:
+        """Decode latent to gene expression rates.
+
+        Args:
+            z: Latent representation.
+
+        Returns:
+            Log-rate parameters for gene expression.
+        """
+        backbone_output = self.backbone(z)
+        if isinstance(backbone_output, tuple):
+            raise TypeError("AmbientDecoder backbone must return a single tensor output.")
+
+        # Output log-rates
+        log_rate = self.output_proj(backbone_output)
+
+        return log_rate
+
+
+class DifferentiableAmbientRemoval(EncoderDecoderOperator):
+    """Differentiable ambient RNA removal using VAE.
+
+    This operator removes ambient RNA contamination from single-cell
+    count data using a variational autoencoder that models both
+    cell-intrinsic expression and ambient contamination.
+
+    Algorithm:
+    1. Encode counts to latent space + contamination fraction
+    2. Sample latent (reparameterization trick)
+    3. Decode to cell-intrinsic expression rate
+    4. Compute decontaminated counts by subtracting ambient contribution
+
+    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: AmbientRemovalConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = AmbientRemovalConfig(n_genes=2000)
+        remover = DifferentiableAmbientRemoval(config, rngs=nnx.Rngs(42))
+        data = {"counts": counts, "ambient_profile": ambient}
+        result, state, meta = remover.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: AmbientRemovalConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the ambient removal operator.
+
+        Args:
+            config: Ambient removal configuration.
+            rngs: Random number generators for initialization.
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        self.ambient_prior = config.ambient_prior
+        self.stochastic = nnx.static(config.stochastic)
+
+        # Encoder
+        self.encoder = AmbientEncoder(
+            n_genes=config.n_genes,
+            hidden_dims=config.hidden_dims,
+            latent_dim=config.latent_dim,
+            rngs=rngs,
+        )
+
+        # Decoder
+        self.decoder = AmbientDecoder(
+            latent_dim=config.latent_dim,
+            hidden_dims=config.hidden_dims,
+            n_genes=config.n_genes,
+            rngs=rngs,
+        )
+
+    # reparameterize() is inherited from EncoderDecoderOperator
+    # kl_divergence() is inherited from EncoderDecoderOperator
+    # elbo_loss() is inherited from EncoderDecoderOperator
+
+    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 ambient RNA removal.
+
+        Args:
+            data: Dictionary containing:
+                - "counts": Raw count matrix (n_cells, n_genes)
+                - "ambient_profile": Ambient expression profile (n_genes,)
+            state: Element state (passed through unchanged)
+            metadata: Element metadata (passed through unchanged)
+            random_params: Random key for stochastic sampling
+            stats: Not used
+
+        Returns:
+            Tuple of (transformed_data, state, metadata):
+                - transformed_data contains:
+
+                    - "counts": Original counts
+                    - "ambient_profile": Original ambient profile
+                    - "decontaminated_counts": Decontaminated counts
+                    - "contamination_fraction": Estimated contamination per cell
+                    - "latent": Latent representation
+                    - "latent_mean": Mean of latent distribution
+                    - "latent_logvar": Log variance of latent distribution
+                    - "reconstructed": Reconstructed expression
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        counts = data["counts"]
+        ambient_profile = data["ambient_profile"]
+
+        # Encode
+        mean, logvar, contamination = self.encoder(counts)
+
+        # Sample latent using inherited reparameterize (uses self.rngs)
+        z = self.reparameterize(mean, logvar)
+
+        # Decode to cell-intrinsic expression rate
+        log_rate = self.decoder(z)
+        cell_rate = jax.nn.softplus(log_rate)  # Non-negative
+
+        # Compute total counts per cell for scaling
+        total_counts = jnp.sum(counts, axis=-1, keepdims=True)
+
+        # Model: observed = (1 - contamination) * cell + contamination * ambient * total
+        # Decontaminated = observed - contamination * ambient * total
+        ambient_contribution = contamination[:, None] * ambient_profile[None, :] * total_counts
+        decontaminated = jnp.maximum(counts - ambient_contribution, 0.0)
+
+        # Reconstructed expression (for loss computation)
+        cell_contribution = (1.0 - contamination[:, None]) * cell_rate
+        reconstructed = cell_contribution + ambient_contribution
+
+        transformed_data = {
+            "counts": counts,
+            "ambient_profile": ambient_profile,
+            "decontaminated_counts": decontaminated,
+            "contamination_fraction": contamination,
+            "latent": z,
+            "latent_mean": mean,
+            "latent_logvar": logvar,
+            "reconstructed": reconstructed,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/singlecell/archetypes.py

@@ -0,0 +1,191 @@
+"""Differentiable archetypal analysis for single-cell data.
+
+Implements PCHA (Principal Convex Hull Analysis, Morup & Hansen 2012) as a
+differentiable autoencoder with softmax bottleneck.  Each cell is represented
+as a temperature-controlled convex combination of learnable archetype
+prototypes.
+
+Algorithm:
+    1. Encode cells to archetype weight space via MLP.
+    2. Apply temperature-scaled softmax to enforce simplex constraints.
+    3. Reconstruct cells as the convex combination ``weights @ archetypes``.
+
+Inherits from ``TemperatureOperator`` to get temperature-controlled smoothing.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+from artifex.generative_models.core.base import MLP
+from datarax.core.config import OperatorConfig
+from flax import nnx
+from jaxtyping import Array, Float, PyTree
+
+from diffbio.core.base_operators import TemperatureOperator
+from diffbio.utils.nn_utils import ensure_rngs, get_rng_key
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "ArchetypalAnalysisConfig",
+    "DifferentiableArchetypalAnalysis",
+]
+
+
+@dataclass(frozen=True)
+class ArchetypalAnalysisConfig(OperatorConfig):
+    """Configuration for DifferentiableArchetypalAnalysis.
+
+    Attributes:
+        n_genes: Number of input genes (features per cell).
+        n_archetypes: Number of archetype prototypes to learn.
+        hidden_dim: Hidden dimension for the encoder MLP.
+        temperature: Softmax temperature (lower = sharper assignments).
+        learnable_temperature: Whether temperature is a learnable parameter.
+    """
+
+    n_genes: int = 2000
+    n_archetypes: int = 5
+    hidden_dim: int = 64
+    temperature: float = 1.0
+    learnable_temperature: bool = False
+
+
+class DifferentiableArchetypalAnalysis(TemperatureOperator):
+    """Differentiable archetypal analysis with softmax simplex constraints.
+
+    Each cell is encoded into archetype weight space via an MLP, then
+    temperature-controlled softmax produces simplex weights.  The
+    reconstruction is the convex combination of learnable archetype
+    prototypes, enabling end-to-end gradient-based optimisation.
+
+    Inherits from ``TemperatureOperator`` to get:
+
+    - ``_temperature`` property for temperature-controlled smoothing
+    - ``soft_max()`` for logsumexp-based smooth maximum
+
+    Args:
+        config: ArchetypalAnalysisConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        import jax.numpy as jnp
+        config = ArchetypalAnalysisConfig(n_genes=2000, n_archetypes=5)
+        op = DifferentiableArchetypalAnalysis(config, rngs=nnx.Rngs(0))
+        data = {"counts": jnp.ones((100, 2000))}
+        result, state, meta = op.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: ArchetypalAnalysisConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Initialize the archetypal analysis operator.
+
+        Args:
+            config: Archetypal analysis configuration.
+            rngs: Random number generators for weight initialisation.
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        rngs = ensure_rngs(rngs)
+
+        self.encoder_layers = MLP(
+            hidden_dims=[config.hidden_dim],
+            in_features=config.n_genes,
+            activation="relu",
+            output_activation="relu",
+            use_batch_norm=False,
+            rngs=rngs,
+        )
+        self.projection = nnx.Linear(
+            in_features=config.hidden_dim,
+            out_features=config.n_archetypes,
+            rngs=rngs,
+        )
+
+        # Learnable archetype prototypes (n_archetypes, n_genes)
+        key = get_rng_key(rngs, "params", fallback_seed=1)
+        init_archetypes = jax.random.normal(key, (config.n_archetypes, config.n_genes)) * 0.1
+        self.archetypes = nnx.Param(init_archetypes)
+
+    def encode(
+        self,
+        counts: Float[Array, "n_cells n_genes"],
+    ) -> Float[Array, "n_cells n_archetypes"]:
+        """Encode cells to simplex weights over archetypes.
+
+        Args:
+            counts: Cell-by-gene count matrix.
+
+        Returns:
+            Simplex weights of shape ``(n_cells, n_archetypes)``.
+        """
+        hidden: jax.Array = self.encoder_layers(counts)
+        logits = self.projection(hidden)
+        weights = jax.nn.softmax(logits / self._temperature, axis=-1)
+        return weights
+
+    def reconstruct(
+        self,
+        weights: Float[Array, "n_cells n_archetypes"],
+    ) -> Float[Array, "n_cells n_genes"]:
+        """Reconstruct cells as convex combinations of archetypes.
+
+        Args:
+            weights: Simplex weights per cell.
+
+        Returns:
+            Reconstructed cell-by-gene matrix.
+        """
+        return weights @ self.archetypes[...]
+
+    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 archetypal analysis to a cell-by-gene count matrix.
+
+        Args:
+            data: Dictionary containing:
+                - ``"counts"``: Cell-by-gene matrix ``(n_cells, n_genes)``
+            state: Element state (passed through unchanged).
+            metadata: Element metadata (passed through unchanged).
+            random_params: Not used.
+            stats: Not used.
+
+        Returns:
+            Tuple of ``(transformed_data, state, metadata)`` where
+            ``transformed_data`` contains:
+
+            - ``"counts"``: Original count matrix
+            - ``"archetype_weights"``: Simplex weights ``(n_cells, n_archetypes)``
+            - ``"archetypes"``: Archetype prototypes ``(n_archetypes, n_genes)``
+            - ``"reconstructed"``: Reconstructed counts ``(n_cells, n_genes)``
+        """
+        counts = data["counts"]
+
+        weights = self.encode(counts)
+        reconstructed = self.reconstruct(weights)
+
+        transformed_data = {
+            **data,
+            "archetype_weights": weights,
+            "archetypes": self.archetypes[...],
+            "reconstructed": reconstructed,
+        }
+
+        return transformed_data, state, metadata