diffbio 0.1.0__py3-none-any.whl

diffbio/operators/variant/quality_recalibration.py

@@ -0,0 +1,274 @@
+"""Soft Variant Quality Filter for VQSR-style recalibration.
+
+This module provides a differentiable variant quality filter using
+a Gaussian Mixture Model for scoring and sigmoid-based soft filtering.
+
+Key technique: Differentiable GMM enables end-to-end learning of
+quality distributions, with sigmoid thresholds maintaining gradients.
+
+Applications: VQSR-style variant filtering, quality score recalibration.
+
+Inherits from TemperatureOperator to get:
+
+- _temperature property for temperature-controlled smoothing
+- soft_max() for logsumexp-based smooth maximum
+- soft_argmax() for soft position selection
+"""
+
+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, PyTree
+
+from diffbio.core import soft_ops
+from diffbio.core.base_operators import TemperatureOperator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class VariantQualityFilterConfig(OperatorConfig):
+    """Configuration for SoftVariantQualityFilter.
+
+    Attributes:
+        n_components: Number of GMM components.
+        n_features: Number of variant features.
+        threshold: Quality score threshold for filtering.
+        temperature: Temperature for softmax/sigmoid operations.
+    """
+
+    n_components: int = 3
+    n_features: int = 4  # depth, qual, strand_bias, mapq
+    threshold: float = 0.5
+    temperature: float = 1.0
+
+
+class SoftVariantQualityFilter(TemperatureOperator):
+    """Differentiable variant quality filter using GMM.
+
+    This operator implements VQSR-style variant quality recalibration
+    using a learnable Gaussian Mixture Model. Variants are scored
+    by their likelihood under the GMM, and soft filtering is applied
+    via sigmoid thresholds.
+
+    Algorithm:
+    1. Compute GMM component responsibilities (E-step style)
+    2. Score variants by weighted log-likelihood
+    3. Apply sigmoid threshold for soft filtering
+
+    Args:
+        config: VariantQualityFilterConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = VariantQualityFilterConfig(n_components=3)
+        filter_op = SoftVariantQualityFilter(config, rngs=nnx.Rngs(42))
+        data = {"variant_features": features}  # (n_variants, n_features)
+        result, state, meta = filter_op.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: VariantQualityFilterConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the quality filter.
+
+        Args:
+            config: Filter 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_components = config.n_components
+        self.n_features = config.n_features
+        self.threshold = config.threshold
+        # Temperature is managed by TemperatureOperator via self._temperature
+
+        # Initialize GMM parameters
+        # Component means: (n_components, n_features)
+        key = rngs.params()
+        init_means = jax.random.normal(key, (config.n_components, config.n_features)) * 0.5
+        self.means = nnx.Param(init_means)
+
+        # Component log variances (diagonal covariance): (n_components, n_features)
+        key = rngs.params()
+        init_log_var = jax.random.normal(key, (config.n_components, config.n_features)) * 0.1
+        self.log_variances = nnx.Param(init_log_var)
+
+        # Component mixing weights (unnormalized): (n_components,)
+        key = rngs.params()
+        noise = jax.random.normal(key, (config.n_components,)) * 0.1
+        init_weights = jnp.ones(config.n_components) + noise
+        self.log_mixing_weights = nnx.Param(init_weights)
+
+        # Learned quality score projection
+        self.quality_proj = nnx.Linear(config.n_features, 1, rngs=rngs)
+
+    def get_mixing_weights(self) -> Float[Array, "n_components"]:
+        """Get normalized mixing weights.
+
+        Returns:
+            Mixing weights summing to 1.
+        """
+        return jax.nn.softmax(self.log_mixing_weights[...] / self._temperature)
+
+    def get_variances(self) -> Float[Array, "n_components n_features"]:
+        """Get positive variances from log parameters.
+
+        Returns:
+            Variance values.
+        """
+        return jnp.exp(self.log_variances[...])
+
+    def compute_component_log_probs(
+        self,
+        features: Float[Array, "n_variants n_features"],
+    ) -> Float[Array, "n_variants n_components"]:
+        """Compute log probability under each GMM component.
+
+        Args:
+            features: Variant feature vectors.
+
+        Returns:
+            Log probability for each variant under each component.
+        """
+        means = self.means[...]  # (n_components, n_features)
+        variances = self.get_variances()  # (n_components, n_features)
+
+        # Expand for broadcasting
+        # features: (n_variants, 1, n_features)
+        # means: (1, n_components, n_features)
+        features_exp = features[:, None, :]
+        means_exp = means[None, :, :]
+        variances_exp = variances[None, :, :]
+
+        # Gaussian log probability (diagonal covariance)
+        # log p(x | mu, sigma^2) = -0.5 * sum((x - mu)^2 / sigma^2 + log(2*pi*sigma^2))
+        diff_sq = (features_exp - means_exp) ** 2
+        log_prob = -0.5 * jnp.sum(
+            diff_sq / variances_exp + jnp.log(2 * jnp.pi * variances_exp), axis=-1
+        )  # (n_variants, n_components)
+
+        return log_prob
+
+    def compute_responsibilities(
+        self,
+        features: Float[Array, "n_variants n_features"],
+    ) -> Float[Array, "n_variants n_components"]:
+        """Compute component responsibilities (soft assignments).
+
+        Args:
+            features: Variant feature vectors.
+
+        Returns:
+            Responsibility matrix (probability each variant belongs to each component).
+        """
+        log_probs = self.compute_component_log_probs(features)  # (n_variants, n_components)
+        mixing_weights = self.get_mixing_weights()  # (n_components,)
+
+        # log P(x, z) = log P(x | z) + log P(z)
+        log_joint = log_probs + jnp.log(mixing_weights + 1e-10)
+
+        # Responsibilities via softmax (with temperature)
+        responsibilities = jax.nn.softmax(log_joint / self._temperature, axis=-1)
+
+        return responsibilities
+
+    def compute_quality_scores(
+        self,
+        features: Float[Array, "n_variants n_features"],
+    ) -> Float[Array, "n_variants"]:
+        """Compute quality scores for variants.
+
+        Combines GMM likelihood with learned projection.
+
+        Args:
+            features: Variant feature vectors.
+
+        Returns:
+            Quality scores in [0, 1] range.
+        """
+        # GMM-based score
+        log_probs = self.compute_component_log_probs(features)  # (n_variants, n_components)
+        mixing_weights = self.get_mixing_weights()  # (n_components,)
+
+        # Total log likelihood under mixture
+        log_likelihood = jax.scipy.special.logsumexp(
+            log_probs + jnp.log(mixing_weights + 1e-10), axis=-1
+        )  # (n_variants,)
+
+        # Learned quality projection
+        learned_quality = self.quality_proj(features).squeeze(-1)  # (n_variants,)
+
+        # Combine and normalize to [0, 1]
+        combined_score = log_likelihood + learned_quality
+        quality_scores = jax.nn.sigmoid(combined_score)
+
+        return quality_scores
+
+    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 quality filtering to variants.
+
+        Args:
+            data: Dictionary containing:
+                - "variant_features": Feature vectors (n_variants, n_features)
+            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):
+                - transformed_data contains:
+
+                    - "variant_features": Original features
+                    - "quality_scores": Computed quality scores [0, 1]
+                    - "filter_weights": Soft filter weights [0, 1]
+                    - "component_probs": GMM component responsibilities
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        features = data["variant_features"]
+
+        # Compute quality scores
+        quality_scores = self.compute_quality_scores(features)
+
+        # Soft filter weights using sigmoid threshold
+        filter_weights = soft_ops.greater(
+            quality_scores, self.threshold, softness=self._temperature
+        )
+
+        # Component responsibilities
+        component_probs = self.compute_responsibilities(features)
+
+        # Build output data
+        transformed_data = {
+            "variant_features": features,
+            "quality_scores": quality_scores,
+            "filter_weights": filter_weights,
+            "component_probs": component_probs,
+        }
+
+        return transformed_data, state, metadata

diffbio/pipelines/__init__.py

@@ -0,0 +1,65 @@
+"""Pre-built differentiable pipeline templates.
+
+This module provides ready-to-use pipeline templates for common bioinformatics
+workflows such as variant calling, preprocessing, differential expression,
+and single-cell analysis.
+"""
+
+from diffbio.pipelines.differential_expression import (
+    DEPipelineConfig,
+    DifferentialExpressionPipeline,
+)
+from diffbio.pipelines.enhanced_variant_calling import (
+    EnhancedVariantCallingPipeline,
+    EnhancedVariantCallingPipelineConfig,
+    create_enhanced_variant_calling_pipeline,
+)
+from diffbio.pipelines.preprocessing import (
+    PreprocessingPipeline,
+    PreprocessingPipelineConfig,
+    create_preprocessing_pipeline,
+)
+from diffbio.pipelines.single_cell import (
+    SingleCellPipeline,
+    SingleCellPipelineConfig,
+    create_single_cell_pipeline,
+)
+from diffbio.pipelines.perturbation import (
+    PerturbationPipeline,
+    PerturbationPipelineConfig,
+    PerturbationPipelineResult,
+)
+from diffbio.pipelines.variant_calling import (
+    VariantCallingPipeline,
+    VariantCallingPipelineConfig,
+    create_cnn_variant_pipeline,
+    create_variant_calling_pipeline,
+)
+
+
+__all__ = [
+    # Differential Expression
+    "DEPipelineConfig",
+    "DifferentialExpressionPipeline",
+    # Enhanced Variant Calling
+    "EnhancedVariantCallingPipeline",
+    "EnhancedVariantCallingPipelineConfig",
+    "create_enhanced_variant_calling_pipeline",
+    # Preprocessing
+    "PreprocessingPipeline",
+    "PreprocessingPipelineConfig",
+    "create_preprocessing_pipeline",
+    # Single-Cell Analysis
+    "SingleCellPipeline",
+    "SingleCellPipelineConfig",
+    "create_single_cell_pipeline",
+    # Perturbation
+    "PerturbationPipeline",
+    "PerturbationPipelineConfig",
+    "PerturbationPipelineResult",
+    # Variant calling
+    "VariantCallingPipeline",
+    "VariantCallingPipelineConfig",
+    "create_variant_calling_pipeline",
+    "create_cnn_variant_pipeline",
+]

diffbio/pipelines/differential_expression.py

@@ -0,0 +1,279 @@
+"""Differentiable differential expression analysis pipeline.
+
+This module implements an end-to-end differentiable differential expression
+pipeline inspired by DESeq2, with negative binomial modeling and size factor
+normalization.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+
+from diffbio.core import soft_ops
+from diffbio.operators.statistical.nb_glm import DifferentiableNBGLM, NBGLMConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class DEPipelineConfig(OperatorConfig):
+    """Configuration for differential expression pipeline.
+
+    Attributes:
+        n_genes: Number of genes to analyze.
+        n_conditions: Number of conditions (covariates) in design matrix.
+        alpha: Significance threshold for differential expression.
+        use_size_factors: Whether to compute and use size factors.
+    """
+
+    n_genes: int = 1000
+    n_conditions: int = 2
+    alpha: float = 0.05
+    use_size_factors: bool = True
+
+
+class DifferentialExpressionPipeline(OperatorModule):
+    """End-to-end differentiable differential expression analysis.
+
+    This pipeline implements a DESeq2-style analysis with:
+    1. Size factor normalization (median-of-ratios)
+    2. Negative binomial GLM fitting
+    3. Wald test for significance
+    4. Multiple testing correction (soft approximation)
+
+    All steps maintain gradient flow for end-to-end learning.
+
+    Example:
+        ```python
+        config = DEPipelineConfig(
+            n_genes=5000,
+            n_conditions=2,
+        )
+        pipeline = DifferentialExpressionPipeline(config, rngs=rngs)
+
+        data = {
+            "counts": count_matrix,  # (n_samples, n_genes)
+            "design": design_matrix,  # (n_samples, n_conditions)
+        }
+        result, state, metadata = pipeline.apply(data, {}, None)
+        lfc = result["log_fold_change"]
+        pvals = result["p_values"]
+        significant = result["significant"]
+        ```
+    """
+
+    def __init__(self, config: DEPipelineConfig, *, rngs: nnx.Rngs | None = None):
+        """Initialize the differential expression pipeline.
+
+        Args:
+            config: Configuration for the pipeline.
+            rngs: Random number generators for initialization.
+        """
+        super().__init__(config, rngs=rngs)
+        self.config = config
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        # Initialize the NB GLM
+        nb_config = NBGLMConfig(
+            n_features=config.n_genes,
+            n_covariates=config.n_conditions,
+            estimate_dispersion=True,
+            stream_name=config.stream_name,
+        )
+        self.nb_glm = DifferentiableNBGLM(nb_config, rngs=rngs)
+
+    def _compute_size_factors(self, counts: jax.Array) -> jax.Array:
+        """Compute size factors using median-of-ratios method (DESeq2 style).
+
+        Args:
+            counts: Count matrix of shape (n_samples, n_genes).
+
+        Returns:
+            Size factors of shape (n_samples,).
+        """
+        # Compute geometric mean per gene (reference sample)
+        # Add pseudocount for numerical stability
+        log_counts = jnp.log(counts + 1)
+        geo_mean_log = jnp.mean(log_counts, axis=0)
+
+        # Compute ratios to geometric mean
+        log_ratios = log_counts - geo_mean_log[None, :]
+
+        # Size factor = median of ratios for each sample
+        size_factors = jnp.exp(soft_ops.median(log_ratios, axis=1, softness=0.1))
+
+        # Normalize to have geometric mean of 1
+        size_factors = size_factors / jnp.exp(jnp.mean(jnp.log(size_factors + 1e-8)))
+
+        return size_factors
+
+    def _compute_wald_statistic(
+        self,
+        beta: jax.Array,
+        dispersion: jax.Array,
+        design: jax.Array,
+        size_factors: jax.Array,
+    ) -> tuple[jax.Array, jax.Array]:
+        """Compute Wald test statistics for differential expression.
+
+        The Wald statistic tests H0: beta[coef_idx] = 0.
+
+        Args:
+            beta: Coefficient matrix of shape (n_conditions, n_genes).
+            dispersion: Dispersion parameters of shape (n_genes,).
+            design: Design matrix of shape (n_samples, n_conditions).
+            size_factors: Size factors of shape (n_samples,).
+
+        Returns:
+            Tuple of (wald_statistic, standard_error).
+        """
+        # Compute predicted means
+        log_mu = jnp.dot(design, beta)  # (n_samples, n_genes)
+        mu = jnp.exp(log_mu) * size_factors[:, None]
+
+        # Compute variance of NB distribution
+        # Var = mu + mu^2 / dispersion
+        variance = mu + jnp.square(mu) / (dispersion[None, :] + 1e-8)
+
+        # Fisher information for beta (approximate)
+        # I = X^T W X where W = diag(mu^2 / variance)
+        weights = jnp.square(mu) / (variance + 1e-8)
+
+        # Standard error of beta[1] (treatment effect)
+        # SE = sqrt(diag((X^T W X)^{-1}))
+        # Simplified: use diagonal approximation
+
+        # For the treatment coefficient (index 1), approximate SE
+        design_sq = jnp.square(design[:, 1:2])  # Treatment column
+        weighted_design = jnp.sum(weights * design_sq, axis=0)
+        se = 1.0 / jnp.sqrt(weighted_design + 1e-8)
+
+        # Wald statistic for treatment coefficient
+        wald_stat = beta[1, :] / (se + 1e-8)
+
+        return wald_stat, se
+
+    def _wald_to_pvalue(self, wald_stat: jax.Array) -> jax.Array:
+        """Convert Wald statistic to p-value using soft normal CDF.
+
+        Args:
+            wald_stat: Wald statistics of shape (n_genes,).
+
+        Returns:
+            Two-sided p-values of shape (n_genes,).
+        """
+        # Two-sided p-value using standard normal
+        # p = 2 * (1 - Phi(|z|))
+        # Use jax.scipy.stats.norm.sf for survival function
+
+        # Soft approximation using sigmoid for differentiability
+        # Approximate normal CDF: Phi(x) ≈ sigmoid(1.7 * x)
+        abs_z = jnp.abs(wald_stat)
+        p_one_sided = 1.0 - jax.nn.sigmoid(1.7 * abs_z)
+        p_values = 2.0 * p_one_sided
+
+        # Clamp to [0, 1]
+        p_values = soft_ops.clip(p_values, 0.0, 1.0, softness=0.1)
+
+        return p_values
+
+    def _soft_significance(
+        self, p_values: jax.Array, alpha: float, temperature: float = 0.1
+    ) -> jax.Array:
+        """Compute soft significance indicator.
+
+        Args:
+            p_values: P-values of shape (n_genes,).
+            alpha: Significance threshold.
+            temperature: Temperature for sigmoid smoothing.
+
+        Returns:
+            Soft significance indicators of shape (n_genes,).
+        """
+        # Soft thresholding: significant if p < alpha
+        return soft_ops.less(p_values, alpha, softness=temperature)
+
+    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 differential expression analysis.
+
+        Args:
+            data: Dictionary containing:
+                - 'counts': Count matrix of shape (n_samples, n_genes)
+                - 'design': Design matrix of shape (n_samples, n_conditions)
+            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:
+
+                - 'counts': Original count matrix
+                - 'design': Original design matrix
+                - 'size_factors': Computed size factors
+                - 'predicted_mean': Predicted mean expression
+                - 'log_fold_change': Log2 fold change estimates
+                - 'wald_statistic': Wald test statistics
+                - 'standard_error': Standard errors
+                - 'p_values': P-values for differential expression
+                - 'significant': Soft significance indicators
+        """
+        del random_params, stats  # Unused
+
+        counts = data["counts"]
+        design = data["design"]
+
+        # Compute size factors
+        if self.config.use_size_factors:
+            size_factors = self._compute_size_factors(counts)
+        else:
+            size_factors = jnp.ones(counts.shape[0])
+
+        # Fit NB GLM for each sample
+        # The NB GLM expects single samples, so we process in batch
+        beta = self.nb_glm.beta[...]
+        dispersion = jnp.exp(self.nb_glm.log_dispersion[...])
+
+        # Compute predicted means
+        log_mu = jnp.dot(design, beta)
+        predicted_mean = jnp.exp(log_mu) * size_factors[:, None]
+
+        # Compute Wald statistics
+        wald_stat, se = self._compute_wald_statistic(beta, dispersion, design, size_factors)
+
+        # Convert to p-values
+        p_values = self._wald_to_pvalue(wald_stat)
+
+        # Compute log fold change (treatment coefficient in log2 scale)
+        log_fold_change = beta[1, :] / jnp.log(2)
+
+        # Soft significance
+        significant = self._soft_significance(p_values, self.config.alpha)
+
+        output_data = {
+            **data,
+            "size_factors": size_factors,
+            "predicted_mean": predicted_mean,
+            "log_fold_change": log_fold_change,
+            "wald_statistic": wald_stat,
+            "standard_error": se,
+            "p_values": p_values,
+            "significant": significant,
+        }
+
+        return output_data, state, metadata