diffbio 0.1.0__py3-none-any.whl

diffbio/operators/statistical/hmm.py

@@ -0,0 +1,234 @@
+"""Differentiable Hidden Markov Model operator.
+
+This module provides a differentiable implementation of the HMM
+forward algorithm using logsumexp for numerical stability.
+
+Key technique: Use logsumexp instead of direct probability multiplication
+to maintain numerical stability and enable gradient flow.
+
+Applications: Gene finding, chromatin state annotation, profile search.
+
+Inherits from HMMOperator to get:
+
+- forward_pass() for likelihood computation
+- forward_backward_posteriors() for posterior computation
+- get_log_transition_matrix(), get_log_emission_matrix(),
+  get_log_initial_distribution() for parameter access
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from flax import nnx
+from jaxtyping import Array, Float, Int, PyTree
+
+from diffbio.core.base_operators import HMMOperator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class HMMConfig(OperatorConfig):
+    """Configuration for DifferentiableHMM.
+
+    Attributes:
+        num_states: Number of hidden states.
+        num_emissions: Number of possible emissions (e.g., 4 for DNA).
+        temperature: Temperature for softmax operations.
+        learnable_transitions: Whether transition probabilities are learnable.
+        learnable_emissions: Whether emission probabilities are learnable.
+    """
+
+    num_states: int = 3
+    num_emissions: int = 4
+    temperature: float = 1.0
+    learnable_transitions: bool = True
+    learnable_emissions: bool = True
+
+
+class DifferentiableHMM(HMMOperator):
+    """Differentiable Hidden Markov Model.
+
+    This operator implements the HMM forward algorithm with differentiable
+    operations, enabling gradient-based learning of transition and emission
+    parameters.
+
+    The forward algorithm computes P(observations | model) using dynamic
+    programming with logsumexp for numerical stability:
+
+    alpha[t, j] = sum_i(alpha[t-1, i] * A[i,j]) * B[j, o_t]
+
+    In log space:
+    log_alpha[t, j] = logsumexp_i(log_alpha[t-1, i] + log_A[i,j]) + log_B[j, o_t]
+
+    Inherits from HMMOperator to get:
+
+    - forward_pass() for likelihood computation
+    - forward_backward_posteriors() for posterior computation
+    - get_log_transition_matrix(), get_log_emission_matrix(),
+      get_log_initial_distribution() for parameter access
+
+    Args:
+        config: HMMConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = HMMConfig(num_states=3, num_emissions=4)
+        hmm = DifferentiableHMM(config, rngs=nnx.Rngs(42))
+        data = {"observations": jnp.array([0, 1, 2, 3])}
+        result, state, meta = hmm.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: HMMConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the HMM operator.
+
+        Args:
+            config: HMM configuration.
+            rngs: Random number generators for initialization.
+            name: Optional operator name.
+        """
+        # HMMOperator handles parameter initialization
+        super().__init__(config, rngs=rngs, name=name)
+
+    # get_log_transition_matrix() is inherited from HMMOperator
+    # get_log_emission_matrix() is inherited from HMMOperator
+    # get_log_initial_distribution() is inherited from HMMOperator
+
+    def forward(
+        self,
+        observations: Int[Array, "seq_len"],
+    ) -> Float[Array, ""]:
+        """Compute log probability of observations using forward algorithm.
+
+        Delegates to inherited forward_pass() from HMMOperator.
+
+        Args:
+            observations: Integer-encoded observations (seq_len,).
+
+        Returns:
+            Log probability of the observation sequence.
+        """
+        return self.forward_pass(observations)
+
+    def forward_soft(
+        self,
+        observations: Float[Array, "seq_len n_emissions"],
+    ) -> Float[Array, ""]:
+        """Compute log probability with soft (probabilistic) observations.
+
+        This variant accepts soft observations (probability distributions
+        over emissions) for fully differentiable operation.
+
+        Args:
+            observations: Soft observations (seq_len, n_emissions).
+
+        Returns:
+            Log probability of the observation sequence.
+        """
+        log_trans = self.get_log_transition_matrix()
+        log_emit = self.get_log_emission_matrix()
+        log_init = self.get_log_initial_distribution()
+
+        # Soft emission: sum over emissions weighted by observation probs
+        # log P(o_t | state) = logsumexp(log_emit + log(o_t))
+        def soft_emission(obs: Float[Array, "n_emissions"]) -> Float[Array, "n_states"]:
+            """Compute log emission probability for soft observation vector."""
+            # obs is (n_emissions,), log_emit is (n_states, n_emissions)
+            log_obs = jnp.log(obs + 1e-10)
+            return jax.scipy.special.logsumexp(log_emit + log_obs, axis=1)
+
+        # Initialize
+        log_alpha = log_init + soft_emission(observations[0])
+
+        # Forward pass
+        def forward_step(
+            log_alpha: Float[Array, "n_states"], obs: Float[Array, "n_emissions"]
+        ) -> tuple[Float[Array, "n_states"], None]:
+            """Advance the forward variable by one time step."""
+            log_alpha_expanded = log_alpha[:, None]
+            log_alpha_new = jax.scipy.special.logsumexp(log_alpha_expanded + log_trans, axis=0)
+            log_alpha_new = log_alpha_new + soft_emission(obs)
+            return log_alpha_new, None
+
+        log_alpha, _ = jax.lax.scan(forward_step, log_alpha, observations[1:])
+
+        log_prob = jax.scipy.special.logsumexp(log_alpha)
+
+        return log_prob
+
+    def forward_backward(
+        self,
+        observations: Int[Array, "seq_len"],
+    ) -> Float[Array, "seq_len num_states"]:
+        """Compute state posteriors using forward-backward algorithm.
+
+        Delegates to inherited forward_backward_posteriors() from HMMOperator.
+
+        Args:
+            observations: Integer-encoded observations.
+
+        Returns:
+            State posteriors P(state_t | observations) for each position.
+        """
+        return self.forward_backward_posteriors(observations)
+
+    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 HMM to observation sequence.
+
+        This method computes the log-likelihood and state posteriors
+        for a given observation sequence.
+
+        Args:
+            data: Dictionary containing:
+                - "observations": Integer-encoded observations (seq_len,)
+            state: Element state (passed through unchanged)
+            metadata: Element metadata (passed through unchanged)
+            random_params: Not used (deterministic operator)
+            stats: Not used
+
+        Returns:
+            Tuple of (transformed_data, state, metadata):
+                - transformed_data contains:
+
+                    - "observations": Original observations
+                    - "log_likelihood": Log probability of sequence
+                    - "state_posteriors": P(state | observations) at each position
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        observations = data["observations"]
+
+        # Compute log likelihood
+        log_likelihood = self.forward(observations)
+
+        # Compute state posteriors
+        state_posteriors = self.forward_backward(observations)
+
+        # Build output data
+        transformed_data = {
+            "observations": observations,
+            "log_likelihood": log_likelihood,
+            "state_posteriors": state_posteriors,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/statistical/nb_glm.py

@@ -0,0 +1,272 @@
+"""Differentiable Negative Binomial GLM operator.
+
+This module provides a differentiable implementation of the negative
+binomial generalized linear model for differential expression analysis,
+inspired by DESeq2.
+
+Key technique: Parameterize the NB mean through a log-linear model
+and estimate dispersion parameters per gene.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+from jaxtyping import Array, Float, PyTree
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class NBGLMConfig(OperatorConfig):
+    """Configuration for DifferentiableNBGLM.
+
+    Attributes:
+        n_features: Number of features (genes).
+        n_covariates: Number of covariates in design matrix.
+        estimate_dispersion: Whether to estimate dispersion parameters.
+    """
+
+    n_features: int = 2000
+    n_covariates: int = 2
+    estimate_dispersion: bool = True
+
+
+class DifferentiableNBGLM(OperatorModule):
+    """Differentiable Negative Binomial GLM for differential expression.
+
+    This operator implements a negative binomial GLM where:
+    - log(mu) = X @ beta (design matrix @ coefficients)
+    - P(count | mu, dispersion) = NB(count; mu, dispersion)
+
+    Gradients flow through both the coefficient (beta) and dispersion
+    parameters, enabling end-to-end learning.
+
+    The negative binomial distribution is parameterized as:
+    - mean = mu
+    - variance = mu + mu^2 / dispersion
+
+    Args:
+        config: NBGLMConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = NBGLMConfig(n_features=2000, n_covariates=2)
+        glm = DifferentiableNBGLM(config, rngs=nnx.Rngs(42))
+        data = {"counts": counts, "design": design_row, "size_factor": sf}
+        result, state, meta = glm.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: NBGLMConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the NB GLM operator.
+
+        Args:
+            config: NB GLM 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.n_features = config.n_features
+        self.n_covariates = config.n_covariates
+
+        # Initialize coefficients (beta)
+        # Shape: (n_covariates, n_features)
+        key = rngs.params()
+        init_beta = jax.random.normal(key, (config.n_covariates, config.n_features)) * 0.1
+        self.beta = nnx.Param(init_beta)
+
+        # Initialize log dispersion parameters
+        # Shape: (n_features,)
+        # Start with dispersion = 1.0 (log_dispersion = 0)
+        self.log_dispersion = nnx.Param(jnp.zeros(config.n_features))
+
+    def get_coefficients(self) -> Float[Array, "n_covariates n_features"]:
+        """Get coefficient matrix.
+
+        Returns:
+            Coefficient matrix beta (n_covariates, n_features).
+        """
+        return self.beta[...]
+
+    def get_dispersion(self) -> Float[Array, "n_features"]:
+        """Get dispersion parameters.
+
+        Returns:
+            Dispersion parameters (n_features,), always positive.
+        """
+        # Use softplus to ensure positivity
+        return jax.nn.softplus(self.log_dispersion[...]) + 1e-4
+
+    def predict_mean(
+        self,
+        design: Float[Array, "n_covariates"],
+        size_factor: Float[Array, ""],
+    ) -> Float[Array, "n_features"]:
+        """Predict mean expression for a sample.
+
+        Args:
+            design: Design matrix row for this sample.
+            size_factor: Library size normalization factor.
+
+        Returns:
+            Predicted mean expression (n_features,).
+        """
+        beta = self.get_coefficients()
+
+        # log(mu) = design @ beta
+        # design is (n_covariates,), beta is (n_covariates, n_features)
+        log_mu = jnp.dot(design, beta)
+
+        # Add size factor effect
+        log_mu = log_mu + jnp.log(size_factor + 1e-8)
+
+        # Exponentiate to get mean
+        mu = jnp.exp(log_mu)
+
+        return mu
+
+    def negative_binomial_log_prob(
+        self,
+        counts: Float[Array, "n_features"],
+        design: Float[Array, "n_covariates"],
+        size_factor: Float[Array, ""],
+    ) -> Float[Array, ""]:
+        """Compute negative binomial log probability.
+
+        Uses the parameterization where:
+        - variance = mu + mu^2 / dispersion
+
+        Args:
+            counts: Observed counts.
+            design: Design matrix row.
+            size_factor: Size factor for normalization.
+
+        Returns:
+            Total log probability (scalar).
+        """
+        mu = self.predict_mean(design, size_factor)
+        dispersion = self.get_dispersion()
+
+        # Negative binomial log probability
+        # NB(k; mu, r) where r = dispersion
+        # log P(k) = log(Gamma(k + r)) - log(Gamma(k + 1)) - log(Gamma(r))
+        #          + r * log(r / (r + mu)) + k * log(mu / (r + mu))
+
+        r = dispersion
+        k = counts
+
+        # Compute log probability using the NB PMF
+        # Use jax.scipy.special functions for numerical stability
+        log_prob = (
+            jax.scipy.special.gammaln(k + r)
+            - jax.scipy.special.gammaln(k + 1)
+            - jax.scipy.special.gammaln(r)
+            + r * jnp.log(r / (r + mu + 1e-8))
+            + k * jnp.log(mu / (r + mu + 1e-8) + 1e-8)
+        )
+
+        # Sum over features
+        total_log_prob = jnp.sum(log_prob)
+
+        return total_log_prob
+
+    def batch_log_likelihood(
+        self,
+        counts: Float[Array, "n_samples n_features"],
+        design: Float[Array, "n_samples n_covariates"],
+        size_factors: Float[Array, "n_samples"],
+    ) -> Float[Array, ""]:
+        """Compute log likelihood for a batch of samples.
+
+        Args:
+            counts: Count matrix (n_samples, n_features).
+            design: Design matrix (n_samples, n_covariates).
+            size_factors: Size factors (n_samples,).
+
+        Returns:
+            Total log likelihood.
+        """
+
+        def sample_log_prob(args):
+            c, d, s = args
+            return self.negative_binomial_log_prob(c, d, s)
+
+        log_probs = jax.vmap(sample_log_prob)((counts, design, size_factors))
+        return jnp.sum(log_probs)
+
+    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 NB GLM to count data.
+
+        This method computes the log likelihood and predicted mean
+        for a given sample.
+
+        Args:
+            data: Dictionary containing:
+                - "counts": Gene counts (n_features,)
+                - "design": Design matrix row (n_covariates,)
+                - "size_factor": Library size factor (scalar)
+            state: Element state (passed through unchanged)
+            metadata: Element metadata (passed through unchanged)
+            random_params: Not used (deterministic operator)
+            stats: Not used
+
+        Returns:
+            Tuple of (transformed_data, state, metadata):
+                - transformed_data contains:
+
+                    - "counts": Original counts
+                    - "log_likelihood": Log probability of counts
+                    - "predicted_mean": Predicted expression
+                    - "dispersion": Dispersion parameters
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        counts = data["counts"]
+        design = data["design"]
+        size_factor = data["size_factor"]
+
+        # Compute log likelihood
+        log_likelihood = self.negative_binomial_log_prob(counts, design, size_factor)
+
+        # Compute predicted mean
+        predicted_mean = self.predict_mean(design, size_factor)
+
+        # Get dispersion
+        dispersion = self.get_dispersion()
+
+        # Build output data
+        transformed_data = {
+            "counts": counts,
+            "design": design,
+            "size_factor": size_factor,
+            "log_likelihood": log_likelihood,
+            "predicted_mean": predicted_mean,
+            "dispersion": dispersion,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/variant/__init__.py

@@ -0,0 +1,64 @@
+"""Variant calling operators for differentiable variant detection.
+
+This module provides differentiable components for variant calling:
+- DifferentiablePileup: Generates pileup from aligned reads
+- DeepVariantStylePileup: DeepVariant-style multi-channel pileup images
+- VariantClassifier: MLP-based variant classifier
+- CellTypeAwareVariantClassifier: Cell-type-weighted variant classifier
+- CNNVariantClassifier: CNN-based pileup image classifier (DeepVariant-style)
+- SoftVariantQualityFilter: GMM-based quality filtering (VQSR-style)
+- DifferentiableCNVSegmentation: Attention-based CNV detection
+"""
+
+from diffbio.operators.variant.classifier import (
+    CellTypeAwareVariantClassifier,
+    CellTypeAwareVariantClassifierConfig,
+    VariantClassifier,
+    VariantClassifierConfig,
+)
+from diffbio.operators.variant.cnn_classifier import (
+    CNNVariantClassifier,
+    CNNVariantClassifierConfig,
+)
+from diffbio.operators.variant.cnv_segmentation import (
+    CNVSegmentationConfig,
+    DifferentiableCNVSegmentation,
+    EnhancedCNVSegmentation,
+    EnhancedCNVSegmentationConfig,
+)
+from diffbio.operators.variant.deepvariant_pileup import (
+    DeepVariantPileupConfig,
+    DeepVariantStylePileup,
+)
+from diffbio.operators.variant.pileup import DifferentiablePileup, PileupConfig
+from diffbio.operators.variant.quality_recalibration import (
+    SoftVariantQualityFilter,
+    VariantQualityFilterConfig,
+)
+
+__all__ = [
+    # Pileup
+    "DifferentiablePileup",
+    "PileupConfig",
+    # DeepVariant-style Pileup
+    "DeepVariantStylePileup",
+    "DeepVariantPileupConfig",
+    # MLP Classifier
+    "VariantClassifier",
+    "VariantClassifierConfig",
+    # Cell-Type-Aware Classifier
+    "CellTypeAwareVariantClassifier",
+    "CellTypeAwareVariantClassifierConfig",
+    # CNN Classifier
+    "CNNVariantClassifier",
+    "CNNVariantClassifierConfig",
+    # Quality Filter
+    "SoftVariantQualityFilter",
+    "VariantQualityFilterConfig",
+    # CNV Segmentation
+    "CNVSegmentationConfig",
+    "DifferentiableCNVSegmentation",
+    # Enhanced CNV Segmentation
+    "EnhancedCNVSegmentation",
+    "EnhancedCNVSegmentationConfig",
+]