diffbio 0.1.0__py3-none-any.whl

diffbio/losses/statistical_losses.py

@@ -0,0 +1,345 @@
+"""Statistical loss functions for differentiable bioinformatics.
+
+This module provides differentiable implementations of statistical loss
+functions commonly used in bioinformatics applications.
+
+Includes:
+- zinb_negative_log_likelihood: Zero-Inflated Negative Binomial NLL
+- NegativeBinomialLoss: For count data modeling (scRNA-seq, RNA-seq)
+- VAELoss: ELBO loss for variational autoencoders
+- HMMLikelihoodLoss: Negative log-likelihood for HMM sequence models
+"""
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jaxtyping import Array, Float, Int
+
+from diffbio.constants import EPSILON
+from diffbio.core import soft_ops
+
+
+def zinb_negative_log_likelihood(
+    counts: Float[Array, "... n_genes"],
+    log_rate: Float[Array, "... n_genes"],
+    log_theta: Float[Array, "... n_genes"],
+    pi_logit: Float[Array, "... n_genes"],
+) -> Float[Array, ""]:
+    """Zero-Inflated Negative Binomial negative log-likelihood.
+
+    Uses the scVI-style logit-space formulation for numerical stability.
+    All log-sigmoid terms are computed via softplus, avoiding explicit
+    materialisation of ``sigmoid(pi_logit)`` which is unstable when
+    ``pi_logit`` is large positive (pi near 1, so 1-pi near 0).
+
+    Key identities::
+
+        log(sigmoid(pi))   = -softplus(-pi)
+        log(1-sigmoid(pi)) = -softplus(pi)
+
+    ZINB: ``P(x) = sigmoid(pi) * delta_0(x) + (1 - sigmoid(pi)) * NB(x; mu, theta)``
+    where ``mu = exp(log_rate)``, ``theta = exp(log_theta)``.
+
+    Args:
+        counts: Observed counts.
+        log_rate: Log mean parameter from decoder.
+        log_theta: Log dispersion parameter.
+        pi_logit: Logit of zero-inflation probability.
+
+    Returns:
+        Negative log-likelihood (scalar, summed over all elements).
+    """
+    mu = jnp.exp(log_rate)
+    theta = jnp.exp(jnp.clip(log_theta, -10.0, 10.0))
+    eps = EPSILON
+
+    # Log-space sigmoid computations (numerically stable)
+    softplus_pi = jax.nn.softplus(-pi_logit)  # = -log(sigmoid(pi_logit))
+    log_theta_mu = jnp.log(theta + mu + eps)
+
+    # NB(0) in log-space combined with dropout logit
+    pi_theta_log = -pi_logit + theta * (jnp.log(theta + eps) - log_theta_mu)
+
+    # Case x == 0: log[sigmoid(pi) + (1 - sigmoid(pi)) * NB(0)]
+    case_zero = jax.nn.softplus(pi_theta_log) - softplus_pi
+
+    # Case x > 0: log[(1 - sigmoid(pi)) * NB(x)]
+    case_nonzero = (
+        -softplus_pi
+        + pi_theta_log
+        + counts * (jnp.log(mu + eps) - log_theta_mu)
+        + jax.scipy.special.gammaln(counts + theta)
+        - jax.scipy.special.gammaln(theta)
+        - jax.scipy.special.gammaln(counts + 1.0)
+    )
+
+    is_zero = soft_ops.less(counts, jnp.array(eps), softness=0.01)
+    log_prob = is_zero * case_zero + (1.0 - is_zero) * case_nonzero
+
+    return -jnp.sum(log_prob)
+
+
+class NegativeBinomialLoss(nnx.Module):
+    """Negative binomial log-likelihood loss for count data.
+
+    The negative binomial distribution is parameterized by mean (mu) and
+    dispersion (theta), suitable for overdispersed count data like RNA-seq.
+
+    NB(x | mu, theta) = Gamma(x + theta) / (Gamma(theta) * Gamma(x + 1))
+                        * (theta / (theta + mu))^theta
+                        * (mu / (theta + mu))^x
+
+    Args:
+        eps: Small constant for numerical stability.
+        rngs: Flax NNX random number generators.
+
+    Example:
+        ```python
+        loss_fn = NegativeBinomialLoss(rngs=nnx.Rngs(42))
+        loss = loss_fn(counts, mu, theta)
+        ```
+    """
+
+    def __init__(
+        self,
+        eps: float = 1e-8,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize the NB loss.
+
+        Args:
+            eps: Numerical stability constant.
+            rngs: Random number generators (not used, for API consistency).
+        """
+        super().__init__()
+        self.eps = eps
+
+    def __call__(
+        self,
+        counts: Float[Array, "batch genes"],
+        mu: Float[Array, "batch genes"],
+        theta: Float[Array, "genes"],
+    ) -> Float[Array, ""]:
+        """Compute negative binomial negative log-likelihood.
+
+        Args:
+            counts: Observed counts.
+            mu: Predicted mean.
+            theta: Dispersion parameter (per gene).
+
+        Returns:
+            Mean negative log-likelihood (scalar).
+        """
+        # Ensure positive values
+        mu = jnp.maximum(mu, self.eps)
+        theta = jnp.maximum(theta, self.eps)
+
+        # Log-likelihood of NB distribution
+        # log P(x | mu, theta) = log Gamma(x + theta) - log Gamma(theta) - log Gamma(x + 1)
+        #                       + theta * log(theta / (theta + mu))
+        #                       + x * log(mu / (theta + mu))
+
+        log_theta_mu = jnp.log(theta[None, :] + mu + self.eps)
+
+        ll = (
+            jax.scipy.special.gammaln(counts + theta[None, :])
+            - jax.scipy.special.gammaln(theta[None, :])
+            - jax.scipy.special.gammaln(counts + 1)
+            + theta[None, :] * (jnp.log(theta[None, :] + self.eps) - log_theta_mu)
+            + counts * (jnp.log(mu + self.eps) - log_theta_mu)
+        )
+
+        # Return mean negative log-likelihood
+        return -jnp.mean(ll)
+
+
+class VAELoss(nnx.Module):
+    """Variational autoencoder ELBO loss.
+
+    Combines reconstruction loss with KL divergence regularization:
+    ELBO = E[log p(x|z)] - KL(q(z|x) || p(z))
+
+    For Gaussian encoder and prior:
+    KL = -0.5 * sum(1 + log(var) - mean^2 - var)
+
+    Args:
+        kl_weight: Weight for KL divergence term (beta-VAE).
+        reconstruction_type: Type of reconstruction loss ("mse" or "bce").
+        rngs: Flax NNX random number generators.
+
+    Example:
+        ```python
+        loss_fn = VAELoss(kl_weight=1.0, rngs=nnx.Rngs(42))
+        loss = loss_fn(x, x_recon, mean, logvar)
+        ```
+    """
+
+    def __init__(
+        self,
+        kl_weight: float = 1.0,
+        reconstruction_type: str = "mse",
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize the VAE loss.
+
+        Args:
+            kl_weight: Weight for KL term.
+            reconstruction_type: "mse" or "bce".
+            rngs: Random number generators (not used, for API consistency).
+        """
+        super().__init__()
+        self.kl_weight = kl_weight
+        self.reconstruction_type = reconstruction_type
+
+    def __call__(
+        self,
+        x: Float[Array, "batch features"],
+        x_recon: Float[Array, "batch features"],
+        mean: Float[Array, "batch latent"],
+        logvar: Float[Array, "batch latent"],
+    ) -> Float[Array, ""]:
+        """Compute VAE ELBO loss.
+
+        Args:
+            x: Original input.
+            x_recon: Reconstructed input.
+            mean: Encoder mean.
+            logvar: Encoder log-variance.
+
+        Returns:
+            Negative ELBO (scalar).
+        """
+        # Reconstruction loss
+        if self.reconstruction_type == "mse":
+            recon_loss = jnp.mean((x - x_recon) ** 2)
+        else:  # bce
+            x_recon = jax.nn.sigmoid(x_recon)
+            recon_loss = -jnp.mean(
+                x * jnp.log(x_recon + 1e-8) + (1 - x) * jnp.log(1 - x_recon + 1e-8)
+            )
+
+        # KL divergence: -0.5 * sum(1 + log(var) - mean^2 - var)
+        kl_loss = -0.5 * jnp.mean(1 + logvar - mean**2 - jnp.exp(logvar))
+
+        # Total loss
+        return recon_loss + self.kl_weight * kl_loss
+
+
+class HMMLikelihoodLoss(nnx.Module):
+    """HMM negative log-likelihood loss.
+
+    Computes the negative log-likelihood of sequences under a Hidden
+    Markov Model using the forward algorithm with logsumexp for stability.
+
+    Args:
+        n_states: Number of hidden states.
+        n_emissions: Number of emission symbols.
+        rngs: Flax NNX random number generators.
+
+    Example:
+        ```python
+        loss_fn = HMMLikelihoodLoss(n_states=3, n_emissions=4, rngs=rnx.Rngs(42))
+        nll = loss_fn(observations)
+        ```
+    """
+
+    def __init__(
+        self,
+        n_states: int,
+        n_emissions: int,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize the HMM loss.
+
+        Args:
+            n_states: Number of hidden states.
+            n_emissions: Number of emission symbols.
+            rngs: Random number generators.
+        """
+        super().__init__()
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        self.n_states = n_states
+        self.n_emissions = n_emissions
+
+        key = rngs.params()
+        k1, k2, k3 = jax.random.split(key, 3)
+
+        # Learnable parameters (log-space for stability)
+        # Initial state distribution
+        self.log_initial = nnx.Param(jax.random.normal(k1, (n_states,)) * 0.1)
+
+        # Transition matrix (log probabilities)
+        self.log_transitions = nnx.Param(jax.random.normal(k2, (n_states, n_states)) * 0.1)
+
+        # Emission matrix (log probabilities)
+        self.log_emissions = nnx.Param(jax.random.normal(k3, (n_states, n_emissions)) * 0.1)
+
+    def _get_log_initial(self) -> Float[Array, "n_states"]:
+        """Get normalized log initial distribution."""
+        return jax.nn.log_softmax(self.log_initial[...])
+
+    def _get_log_transitions(self) -> Float[Array, "n_states n_states"]:
+        """Get normalized log transition matrix."""
+        return jax.nn.log_softmax(self.log_transitions[...], axis=-1)
+
+    def _get_log_emissions(self) -> Float[Array, "n_states n_emissions"]:
+        """Get normalized log emission matrix."""
+        return jax.nn.log_softmax(self.log_emissions[...], axis=-1)
+
+    def _forward_single(
+        self,
+        observations: Int[Array, "seq_len"],
+    ) -> Float[Array, ""]:
+        """Forward algorithm for a single sequence.
+
+        Args:
+            observations: Integer-encoded observations.
+
+        Returns:
+            Log-likelihood of the sequence.
+        """
+        log_init = self._get_log_initial()
+        log_trans = self._get_log_transitions()
+        log_emit = self._get_log_emissions()
+
+        # Initialize with first observation
+        log_alpha = log_init + log_emit[:, observations[0]]
+
+        # Forward pass
+        def step(log_alpha, obs):
+            # log_alpha: (n_states,)
+            # Transition: log_alpha[i] + log_trans[i, j] for all j
+            log_alpha_expanded = log_alpha[:, None] + log_trans  # (n_states, n_states)
+            log_alpha_new = jax.scipy.special.logsumexp(log_alpha_expanded, axis=0)
+            # Add emission
+            log_alpha_new = log_alpha_new + log_emit[:, obs]
+            return log_alpha_new, None
+
+        log_alpha, _ = jax.lax.scan(step, log_alpha, observations[1:])
+
+        # Final log-likelihood
+        return jax.scipy.special.logsumexp(log_alpha)
+
+    def __call__(
+        self,
+        observations: Int[Array, "batch seq_len"],
+    ) -> Float[Array, ""]:
+        """Compute mean negative log-likelihood over batch.
+
+        Args:
+            observations: Batch of integer-encoded sequences.
+
+        Returns:
+            Mean negative log-likelihood (scalar).
+        """
+        # Compute log-likelihood for each sequence
+        log_probs = jax.vmap(self._forward_single)(observations)
+
+        # Return mean negative log-likelihood
+        return -jnp.mean(log_probs)

diffbio/operators/__init__.py

@@ -0,0 +1,60 @@
+"""Differentiable bioinformatics operators for DiffBio.
+
+This module provides differentiable operators for common bioinformatics
+operations such as quality filtering, sequence alignment, variant calling,
+epigenomics analysis, and RNA-seq processing.
+
+All operators extend Datarax's OperatorModule for seamless integration.
+"""
+
+from diffbio.operators.quality_filter import (
+    DifferentiableQualityFilter,
+    QualityFilterConfig,
+)
+
+# Import submodules for convenient access
+from diffbio.operators import alignment
+from diffbio.operators import assembly
+from diffbio.operators import crispr
+from diffbio.operators import drug_discovery
+from diffbio.operators import epigenomics
+from diffbio.operators import foundation_models
+from diffbio.operators import mapping
+from diffbio.operators import metabolomics
+from diffbio.operators import molecular_dynamics
+from diffbio.operators import multiomics
+from diffbio.operators import normalization
+from diffbio.operators import population
+from diffbio.operators import preprocessing
+from diffbio.operators import protein
+from diffbio.operators import rna_structure
+from diffbio.operators import rnaseq
+from diffbio.operators import singlecell
+from diffbio.operators import statistical
+from diffbio.operators import variant
+
+__all__ = [
+    # Core quality filter
+    "DifferentiableQualityFilter",
+    "QualityFilterConfig",
+    # Submodules
+    "alignment",
+    "assembly",
+    "crispr",
+    "drug_discovery",
+    "epigenomics",
+    "foundation_models",
+    "mapping",
+    "metabolomics",
+    "molecular_dynamics",
+    "multiomics",
+    "normalization",
+    "population",
+    "preprocessing",
+    "protein",
+    "rna_structure",
+    "rnaseq",
+    "singlecell",
+    "statistical",
+    "variant",
+]

diffbio/operators/_count_vae.py

@@ -0,0 +1,197 @@
+"""Shared count-VAE building blocks for DiffBio operators."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from artifex.generative_models.core.base import MLP
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, Float
+
+from diffbio.losses.statistical_losses import zinb_negative_log_likelihood
+from diffbio.utils.nn_utils import ensure_rngs
+
+
+class CountVAEBackboneMixin:
+    """Shared encoder/decoder backbone for count-based VAEs."""
+
+    encoder_backbone: MLP | None
+    fc_mean: nnx.Linear
+    fc_logvar: nnx.Linear
+    decoder_backbone: MLP | None
+    fc_output: nnx.Linear
+    n_genes: int
+    stream_name: Any
+
+    def _init_count_vae_backbone(
+        self,
+        *,
+        n_inputs: int,
+        latent_dim: int,
+        hidden_dims: list[int],
+        n_outputs: int,
+        rngs: nnx.Rngs | None,
+    ) -> None:
+        """Initialise the shared count-VAE encoder and decoder layers."""
+        safe_rngs = ensure_rngs(rngs)
+
+        encoder_hidden_dims = list(hidden_dims)
+        decoder_hidden_dims = list(reversed(hidden_dims))
+
+        if encoder_hidden_dims:
+            self.encoder_backbone = MLP(
+                hidden_dims=encoder_hidden_dims,
+                in_features=n_inputs,
+                activation="relu",
+                output_activation="relu",
+                use_batch_norm=False,
+                rngs=safe_rngs,
+            )
+            encoder_out_dim = encoder_hidden_dims[-1]
+        else:
+            self.encoder_backbone = None
+            encoder_out_dim = n_inputs
+
+        self.fc_mean = nnx.Linear(
+            in_features=encoder_out_dim,
+            out_features=latent_dim,
+            rngs=safe_rngs,
+        )
+        self.fc_logvar = nnx.Linear(
+            in_features=encoder_out_dim,
+            out_features=latent_dim,
+            rngs=safe_rngs,
+        )
+
+        if decoder_hidden_dims:
+            self.decoder_backbone = MLP(
+                hidden_dims=decoder_hidden_dims,
+                in_features=latent_dim,
+                activation="relu",
+                output_activation="relu",
+                use_batch_norm=False,
+                rngs=safe_rngs,
+            )
+            decoder_out_dim = decoder_hidden_dims[-1]
+        else:
+            self.decoder_backbone = None
+            decoder_out_dim = latent_dim
+
+        self.fc_output = nnx.Linear(
+            in_features=decoder_out_dim,
+            out_features=n_outputs,
+            rngs=safe_rngs,
+        )
+
+    def _init_count_vae_operator(
+        self,
+        *,
+        config: Any,
+        rngs: nnx.Rngs | None,
+    ) -> nnx.Rngs:
+        """Initialise shared count-VAE operator state and return safe RNGs."""
+        safe_rngs = ensure_rngs(rngs)
+        self.n_genes = config.n_genes
+        self.stream_name = nnx.static(config.stream_name)
+        self._init_count_vae_backbone(
+            n_inputs=config.n_genes,
+            latent_dim=config.latent_dim,
+            hidden_dims=config.hidden_dims,
+            n_outputs=config.n_genes,
+            rngs=safe_rngs,
+        )
+        return safe_rngs
+
+    def encode(
+        self,
+        counts: Float[Array, "... n_genes"],
+    ) -> tuple[Float[Array, "... latent_dim"], Float[Array, "... latent_dim"]]:
+        """Encode count vectors to latent Gaussian parameters."""
+        x = jnp.log1p(counts)
+        if self.encoder_backbone is not None:
+            encoded: jax.Array = self.encoder_backbone(x)
+            x = encoded
+        mean = self.fc_mean(x)
+        logvar = jnp.clip(self.fc_logvar(x), -10.0, 10.0)
+        return mean, logvar
+
+    def decode_hidden(
+        self,
+        z: Float[Array, "... latent_dim"],
+    ) -> Float[Array, "... hidden_dim"]:
+        """Decode latent vectors to the shared decoder hidden representation."""
+        if self.decoder_backbone is None:
+            return z
+        decoded: jax.Array = self.decoder_backbone(z)
+        return decoded
+
+    def decode_rates(
+        self,
+        z: Float[Array, "... latent_dim"],
+    ) -> Float[Array, "... n_outputs"]:
+        """Decode latent vectors directly to output log-rates."""
+        return self.fc_output(self.decode_hidden(z))
+
+
+class CountVAEBackbone(CountVAEBackboneMixin, nnx.Module):
+    """Concrete test harness for the shared count-VAE backbone."""
+
+    def __init__(
+        self,
+        *,
+        n_inputs: int,
+        latent_dim: int,
+        hidden_dims: list[int],
+        n_outputs: int,
+        rngs: nnx.Rngs | None = None,
+    ) -> None:
+        """Initialise a standalone shared count-VAE backbone."""
+        super().__init__()
+        self._init_count_vae_backbone(
+            n_inputs=n_inputs,
+            latent_dim=latent_dim,
+            hidden_dims=hidden_dims,
+            n_outputs=n_outputs,
+            rngs=rngs,
+        )
+
+
+class CountReconstructionMixin:
+    """Shared reconstruction losses for count-based VAE operators."""
+
+    @staticmethod
+    def _poisson_nll(
+        counts: Float[Array, "... n_genes"],
+        log_rate: Float[Array, "... n_genes"],
+    ) -> Float[Array, ""]:
+        """Compute Poisson negative log-likelihood."""
+        rate = jnp.exp(log_rate)
+        return jnp.sum(rate - counts * log_rate)
+
+    @staticmethod
+    def _zinb_nll(
+        counts: Float[Array, "... n_genes"],
+        log_rate: Float[Array, "... n_genes"],
+        log_theta: Float[Array, "... n_genes"],
+        pi_logit: Float[Array, "... n_genes"],
+    ) -> Float[Array, ""]:
+        """Compute Zero-Inflated Negative Binomial negative log-likelihood."""
+        return zinb_negative_log_likelihood(counts, log_rate, log_theta, pi_logit)
+
+    def reconstruction_loss(
+        self,
+        counts: Float[Array, "... n_genes"],
+        decode_output: dict[str, jax.Array],
+    ) -> Float[Array, ""]:
+        """Compute reconstruction loss for Poisson or ZINB decoder outputs."""
+        log_rate = decode_output["log_rate"]
+        if "log_theta" in decode_output:
+            return self._zinb_nll(
+                counts,
+                log_rate,
+                decode_output["log_theta"],
+                decode_output["pi_logit"],
+            )
+        return self._poisson_nll(counts, log_rate)

diffbio/operators/_loss_balancing.py

@@ -0,0 +1,65 @@
+"""Shared scalar-loss balancing helpers for DiffBio operators."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from flax import nnx
+from jaxtyping import Array, Float
+from opifex.core.physics.gradnorm import GradNormBalancer
+
+from diffbio.utils.nn_utils import ensure_rngs
+
+
+def combine_scalar_losses(
+    losses: Mapping[str, Float[Array, ""]],
+    *,
+    use_gradnorm: bool,
+    rngs: nnx.Rngs | None = None,
+) -> Float[Array, ""]:
+    """Combine scalar losses with optional GradNorm-based balancing.
+
+    Args:
+        losses: Named scalar losses to combine.
+        use_gradnorm: Whether to balance losses with ``GradNormBalancer``.
+        rngs: Optional random generators used when constructing GradNorm.
+
+    Returns:
+        Combined scalar loss.
+
+    Raises:
+        ValueError: If *losses* is empty.
+    """
+    if not losses:
+        msg = "losses must contain at least one scalar loss"
+        raise ValueError(msg)
+
+    loss_values = list(losses.values())
+    if use_gradnorm:
+        balancer = GradNormBalancer(
+            num_losses=len(loss_values),
+            rngs=ensure_rngs(rngs),
+        )
+        return balancer(loss_values)
+
+    total_loss = loss_values[0]
+    for loss_value in loss_values[1:]:
+        total_loss = total_loss + loss_value
+    return total_loss
+
+
+class LossBalancingMixin:
+    """Reusable operator mixin exposing ``compute_balanced_loss``."""
+
+    config: Any
+
+    def compute_balanced_loss(
+        self,
+        losses: Mapping[str, Float[Array, ""]],
+    ) -> Float[Array, ""]:
+        """Combine operator loss terms using the config's GradNorm flag."""
+        return combine_scalar_losses(
+            losses,
+            use_gradnorm=bool(getattr(self.config, "use_gradnorm", False)),
+        )