diffbio 0.1.0__py3-none-any.whl

diffbio/operators/population/ancestry_estimation.py

@@ -0,0 +1,274 @@
+"""Differentiable Ancestry Estimation Operator.
+
+This module implements a Neural ADMIXTURE-style differentiable ancestry estimator
+using an autoencoder architecture. The model learns to decompose individual
+genotypes into ancestry proportions from K ancestral populations.
+
+Reference:
+    Dias et al. (2022). "Neural ADMIXTURE: A Neural Network Approach for
+    Fast and Accurate Estimation of Population Structure."
+    https://github.com/AI-sandbox/neural-admixture
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from flax import nnx
+
+from diffbio.configs import TemperatureConfig
+from diffbio.core.base_operators import TemperatureOperator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class AncestryEstimatorConfig(TemperatureConfig):
+    """Configuration for DifferentiableAncestryEstimator.
+
+    Attributes:
+        n_snps: Number of SNP markers in genotype input.
+        n_populations: Number of ancestral populations (K).
+        hidden_dims: Hidden layer dimensions for encoder.
+        dropout_rate: Dropout rate for regularization.
+    """
+
+    n_snps: int = 10000
+    n_populations: int = 5
+    hidden_dims: tuple[int, ...] = (128, 64)
+    dropout_rate: float = 0.1
+
+
+class DifferentiableAncestryEstimator(TemperatureOperator):
+    """Neural ADMIXTURE-style differentiable ancestry estimator.
+
+    This operator uses an autoencoder architecture to estimate ancestry
+    proportions from genotype data. The encoder maps genotypes to a latent
+    representation, which is then transformed to ancestry proportions via
+    temperature-controlled softmax. The decoder reconstructs genotypes
+    from ancestry proportions, enabling unsupervised learning.
+
+    The model follows the ADMIXTURE generative model:
+        G_ij = sum_k Q_ik * P_kj
+    Where:
+        - G is the genotype matrix (individuals x SNPs)
+        - Q is the ancestry proportion matrix (individuals x K populations)
+        - P is the population allele frequency matrix (K x SNPs)
+
+    Attributes:
+        config: Operator configuration.
+        backbone: Shared encoder MLP, or ``None`` when ``hidden_dims`` is empty.
+        ancestry_head: Linear layer for ancestry proportions.
+        population_frequencies: Learnable population allele frequencies (P matrix).
+
+    Example:
+        ```python
+        from diffbio.operators.population import (
+            DifferentiableAncestryEstimator,
+            AncestryEstimatorConfig,
+        )
+        config = AncestryEstimatorConfig(n_snps=1000, n_populations=5)
+        estimator = DifferentiableAncestryEstimator(config, rngs=nnx.Rngs(42))
+        data = {"genotypes": genotype_matrix}  # (n_samples, n_snps)
+        result, _, _ = estimator.apply(data, {}, None)
+        ancestry = result["ancestry_proportions"]  # (n_samples, K)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: AncestryEstimatorConfig,
+        *,
+        rngs: nnx.Rngs,
+    ) -> None:
+        """Initialize the ancestry estimator.
+
+        Args:
+            config: Operator configuration.
+            rngs: Flax NNX random number generators.
+        """
+        super().__init__(config, rngs=rngs)
+
+        self.config = config
+
+        if config.hidden_dims:
+            self.backbone = MLP(
+                hidden_dims=list(config.hidden_dims),
+                in_features=config.n_snps,
+                activation="relu",
+                dropout_rate=config.dropout_rate,
+                output_activation="relu",
+                use_batch_norm=False,
+                rngs=rngs,
+            )
+            latent_dim = config.hidden_dims[-1]
+        else:
+            self.backbone = None
+            latent_dim = config.n_snps
+
+        # Ancestry proportion head
+        self.ancestry_head = nnx.Linear(latent_dim, config.n_populations, rngs=rngs)
+
+        # Learnable population allele frequencies (P matrix)
+        # Shape: (n_populations, n_snps)
+        # Initialized with small random values
+        self.population_frequencies = nnx.Param(
+            jnp.abs(
+                jnp.array(
+                    nnx.initializers.normal(0.1)(
+                        rngs.params(),
+                        (config.n_populations, config.n_snps),
+                    )
+                )
+            )
+            + 0.01
+        )
+
+    def encode(self, genotypes: jnp.ndarray) -> jnp.ndarray:
+        """Encode genotypes to latent representation.
+
+        Args:
+            genotypes: Genotype matrix of shape (n_samples, n_snps).
+                Values should be 0, 1, or 2 representing allele counts.
+
+        Returns:
+            Latent representation of shape ``(n_samples, hidden_dims[-1])`` when
+            hidden layers are configured, otherwise the original genotype matrix.
+        """
+        if self.backbone is None:
+            return genotypes
+        latent: jnp.ndarray = self.backbone(genotypes)
+        return latent
+
+    def compute_ancestry(self, latent: jnp.ndarray) -> jnp.ndarray:
+        """Compute ancestry proportions from latent representation.
+
+        Args:
+            latent: Latent representation of shape (n_samples, hidden_dims[-1]).
+
+        Returns:
+            Ancestry proportions of shape (n_samples, n_populations).
+            Each row sums to 1 and all values are non-negative.
+        """
+        # Get raw ancestry logits
+        logits = self.ancestry_head(latent)
+
+        # Apply temperature-controlled softmax
+        temperature = jnp.maximum(self._temperature, 1e-6)
+        proportions = nnx.softmax(logits / temperature, axis=-1)
+
+        return proportions
+
+    def decode(self, ancestry: jnp.ndarray) -> jnp.ndarray:
+        """Decode ancestry proportions to reconstructed genotypes.
+
+        Following the ADMIXTURE model: G = Q @ P
+        Where Q is ancestry proportions and P is population frequencies.
+
+        Args:
+            ancestry: Ancestry proportions of shape (n_samples, n_populations).
+
+        Returns:
+            Reconstructed genotypes of shape (n_samples, n_snps).
+            Values represent expected allele counts (continuous 0-2).
+        """
+        # Get population frequencies normalized to [0, 1]
+        pop_freqs = nnx.sigmoid(self.population_frequencies[...])
+
+        # Reconstruct: G = Q @ P, scaled to 0-2 range
+        # ancestry: (n_samples, n_populations)
+        # pop_freqs: (n_populations, n_snps)
+        reconstructed = 2.0 * jnp.matmul(ancestry, pop_freqs)
+
+        return reconstructed
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
+        """Apply ancestry estimation to genotype data.
+
+        Args:
+            data: Dictionary containing:
+                - "genotypes": Genotype matrix (n_samples, n_snps) with values 0/1/2.
+            state: Per-element state (passed through).
+            metadata: Optional metadata (passed through).
+            random_params: Random parameters for stochastic operations.
+            stats: Optional statistics dictionary.
+
+        Returns:
+            Tuple of (transformed_data, state, metadata) where transformed_data
+            contains:
+
+                - "genotypes": Original genotype matrix.
+                - "ancestry_proportions": Estimated ancestry (n_samples, K).
+                - "reconstructed": Reconstructed genotypes (n_samples, n_snps).
+                - "latent": Latent representation.
+        """
+        genotypes = data["genotypes"]
+
+        # Encode genotypes to latent space
+        latent = self.encode(genotypes)
+
+        # Compute ancestry proportions
+        ancestry = self.compute_ancestry(latent)
+
+        # Decode to reconstructed genotypes
+        reconstructed = self.decode(ancestry)
+
+        # Build output
+        output = {
+            **data,
+            "ancestry_proportions": ancestry,
+            "reconstructed": reconstructed,
+            "latent": latent,
+        }
+
+        return output, state, metadata
+
+
+def create_ancestry_estimator(
+    n_snps: int,
+    n_populations: int,
+    hidden_dims: tuple[int, ...] = (128, 64),
+    temperature: float = 1.0,
+    dropout_rate: float = 0.1,
+    seed: int = 42,
+) -> DifferentiableAncestryEstimator:
+    """Factory function to create an ancestry estimator.
+
+    Args:
+        n_snps: Number of SNP markers.
+        n_populations: Number of ancestral populations (K).
+        hidden_dims: Hidden layer dimensions for encoder.
+        temperature: Softmax temperature for ancestry proportions.
+        dropout_rate: Dropout rate for regularization.
+        seed: Random seed for initialization.
+
+    Returns:
+        Configured DifferentiableAncestryEstimator instance.
+
+    Example:
+        ```python
+        estimator = create_ancestry_estimator(
+            n_snps=10000,
+            n_populations=5,
+        )
+        result, _, _ = estimator.apply({"genotypes": data}, {}, None)
+        ```
+    """
+    config = AncestryEstimatorConfig(
+        n_snps=n_snps,
+        n_populations=n_populations,
+        hidden_dims=hidden_dims,
+        temperature=temperature,
+        dropout_rate=dropout_rate,
+    )
+
+    return DifferentiableAncestryEstimator(config, rngs=nnx.Rngs(seed))

diffbio/operators/preprocessing/__init__.py

@@ -0,0 +1,76 @@
+"""Differentiable preprocessing operators for bioinformatics sequences.
+
+This module provides preprocessing operators that maintain gradient flow
+for end-to-end trainable pipelines.
+
+Operators:
+    SoftAdapterRemoval: Differentiable adapter trimming using soft alignment
+    DifferentiableDuplicateWeighting: Probabilistic duplicate weighting
+    SoftErrorCorrection: Neural network-based error correction
+
+Factories:
+    wrap_probabilistic: Wrap any preprocessing operator in a datarax
+        ProbabilisticOperator for random augmentation during training.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from diffbio.operators.preprocessing.adapter_removal import (
+    AdapterRemovalConfig,
+    SoftAdapterRemoval,
+)
+from diffbio.operators.preprocessing.duplicate_filter import (
+    DifferentiableDuplicateWeighting,
+    DuplicateWeightingConfig,
+)
+from diffbio.operators.preprocessing.error_correction import (
+    ErrorCorrectionConfig,
+    SoftErrorCorrection,
+)
+
+if TYPE_CHECKING:
+    from datarax.core.operator import OperatorModule
+    from datarax.operators import ProbabilisticOperator
+
+
+def wrap_probabilistic(
+    operator: OperatorModule,
+    probability: float = 0.5,
+) -> ProbabilisticOperator:
+    """Wrap a preprocessing operator in a ProbabilisticOperator.
+
+    The wrapped operator is applied with the given probability during
+    each forward pass, enabling random augmentation during training.
+    When not applied, the input data passes through unchanged.
+
+    Args:
+        operator: A preprocessing operator to wrap.
+        probability: Probability of applying the operator (0.0 to 1.0).
+
+    Returns:
+        A ProbabilisticOperator wrapping the given operator.
+
+    Example:
+        >>> adapter_remover = SoftAdapterRemoval(AdapterRemovalConfig(), rngs=rngs)
+        >>> prob_remover = wrap_probabilistic(adapter_remover, probability=0.8)
+    """
+    from datarax.operators import (  # noqa: PLC0415
+        ProbabilisticOperator,
+        ProbabilisticOperatorConfig,
+    )
+
+    config = ProbabilisticOperatorConfig(operator=operator, probability=probability)
+    return ProbabilisticOperator(config)
+
+
+__all__ = [
+    "AdapterRemovalConfig",
+    "DifferentiableDuplicateWeighting",
+    "DuplicateWeightingConfig",
+    "ErrorCorrectionConfig",
+    "SoftAdapterRemoval",
+    "SoftErrorCorrection",
+    "wrap_probabilistic",
+]

diffbio/operators/preprocessing/adapter_removal.py

@@ -0,0 +1,311 @@
+"""Differentiable adapter removal operator.
+
+This module provides a soft adapter removal operator that uses differentiable
+alignment to find adapter sequences and applies soft trimming to maintain
+gradient flow.
+
+Key technique: Use SmoothSmithWaterman for adapter matching, then apply
+sigmoid-weighted soft trimming based on the match position.
+
+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
+from diffbio.sequences.dna import encode_dna_string
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class AdapterRemovalConfig(OperatorConfig):
+    """Configuration for SoftAdapterRemoval.
+
+    Attributes:
+        adapter_sequence: Adapter sequence to remove (default: Illumina universal).
+        temperature: Temperature for soft matching and trimming.
+            Lower = sharper trimming, Higher = smoother.
+        match_threshold: Minimum alignment score ratio to consider a match.
+        min_overlap: Minimum overlap length to consider adapter presence.
+    """
+
+    adapter_sequence: str = "AGATCGGAAGAG"  # Illumina universal adapter
+    temperature: float = 1.0
+    learnable_temperature: bool = True
+    match_threshold: float = 0.5
+    min_overlap: int = 6
+
+
+class SoftAdapterRemoval(TemperatureOperator):
+    """Differentiable adapter removal for sequencing reads.
+
+    This operator performs soft adapter trimming using a differentiable
+    approach. It finds potential adapter matches at the 3' end of reads
+    and applies sigmoid-weighted trimming that maintains gradient flow.
+
+    The algorithm:
+    1. Compute soft alignment scores between sequence suffix and adapter
+    2. Find the soft trim position using weighted position averaging
+    3. Apply sigmoid-weighted retention to each position
+
+    Args:
+        config: AdapterRemovalConfig with adapter parameters.
+        rngs: Flax NNX random number generators (optional).
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = AdapterRemovalConfig(adapter_sequence="AGATCGGAAGAG")
+        remover = SoftAdapterRemoval(config)
+        data = {"sequence": encoded_seq, "quality_scores": quality}
+        result, state, meta = remover.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: AdapterRemovalConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the soft adapter removal operator.
+
+        Args:
+            config: Adapter removal configuration.
+            rngs: Random number generators (optional).
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        # Encode adapter sequence as one-hot
+        adapter_onehot = encode_dna_string(config.adapter_sequence)
+        self.adapter = nnx.Param(adapter_onehot)
+
+        # Learnable parameters
+        # Temperature is managed by TemperatureOperator via self._temperature
+        self.match_threshold = nnx.Param(jnp.array(config.match_threshold))
+
+        # Scoring matrix for DNA alignment (match=2, mismatch=-1)
+        scoring = jnp.array(
+            [
+                [2.0, -1.0, -1.0, -1.0],  # A
+                [-1.0, 2.0, -1.0, -1.0],  # C
+                [-1.0, -1.0, 2.0, -1.0],  # G
+                [-1.0, -1.0, -1.0, 2.0],  # T
+            ]
+        )
+        self.scoring_matrix = nnx.Param(scoring)
+
+        # Store config values
+        self.min_overlap = config.min_overlap
+
+    def _compute_suffix_adapter_scores(
+        self,
+        sequence: Float[Array, "length alphabet"],
+    ) -> Float[Array, "length"]:
+        """Compute adapter match scores for each suffix position.
+
+        For each position i, compute how well the suffix starting at i
+        matches the adapter prefix. This finds where the adapter might start.
+
+        Args:
+            sequence: One-hot encoded sequence (length, 4).
+
+        Returns:
+            Scores for adapter match starting at each position (length,).
+        """
+        seq_len = sequence.shape[0]
+        adapter = self.adapter[...]
+        adapter_len = adapter.shape[0]
+        scoring = self.scoring_matrix[...]
+
+        # For each starting position, compute match score with adapter
+        def score_at_position(start_pos: Array | int) -> Float[Array, ""]:
+            """Compute alignment score for suffix starting at start_pos."""
+            # Length of overlap
+            overlap_len = jnp.minimum(seq_len - start_pos, adapter_len)
+
+            # Extract overlapping regions
+            # Use dynamic slicing with masking for differentiability
+            positions = jnp.arange(adapter_len)
+            mask = positions < overlap_len
+
+            # Compute scores for each position in the overlap
+            def score_position(pos: Array | int) -> Float[Array, ""]:
+                seq_pos = start_pos + pos
+                # Handle out-of-bounds with zeros
+                valid = (seq_pos < seq_len) & (pos < adapter_len)
+                seq_base = jnp.where(
+                    valid,
+                    jax.lax.dynamic_slice(sequence, (seq_pos, 0), (1, 4)).squeeze(0),
+                    jnp.zeros(4),
+                )
+                adapter_base = jnp.where(
+                    valid,
+                    jax.lax.dynamic_slice(adapter, (pos, 0), (1, 4)).squeeze(0),
+                    jnp.zeros(4),
+                )
+                # Score = seq @ scoring @ adapter.T
+                score = jnp.einsum("a,ab,b->", seq_base, scoring, adapter_base)
+                return jnp.where(valid, score, 0.0)
+
+            # Sum scores across overlap positions
+            scores = jax.vmap(score_position)(jnp.arange(adapter_len))
+            total_score = jnp.sum(scores * mask)
+
+            # Normalize by maximum possible score for this overlap length
+            max_score = 2.0 * overlap_len  # Perfect match score
+            normalized = jnp.where(
+                overlap_len >= self.min_overlap,
+                total_score / jnp.maximum(max_score, 1.0),
+                0.0,
+            )
+            return normalized
+
+        # Compute scores for all suffix positions
+        scores = jax.vmap(score_at_position)(jnp.arange(seq_len))
+        return scores
+
+    def _compute_soft_trim_position(
+        self,
+        adapter_scores: Float[Array, "length"],
+    ) -> Float[Array, ""]:
+        """Compute soft trim position using weighted average.
+
+        Uses softmax over adapter scores to compute a soft trim position.
+        High adapter match scores at position i suggest trimming should
+        start there.
+
+        Args:
+            adapter_scores: Adapter match scores at each position.
+
+        Returns:
+            Soft trim position (continuous value).
+        """
+        seq_len = adapter_scores.shape[0]
+        temp = self._temperature
+        threshold = self.match_threshold[...]
+
+        # Apply threshold - only consider positions with good matches
+        thresholded_scores = nnx.relu(adapter_scores - threshold)
+
+        # Soft position selection using softmax
+        # Add small epsilon for numerical stability
+        weights = jax.nn.softmax(thresholded_scores / temp + 1e-10)
+
+        # Weighted average of positions
+        positions = jnp.arange(seq_len, dtype=jnp.float32)
+        soft_position = jnp.sum(weights * positions)
+
+        # If no adapter found (all scores below threshold), return seq_len
+        has_adapter = jnp.any(adapter_scores > threshold)
+        soft_position = jnp.where(has_adapter, soft_position, float(seq_len))
+
+        return soft_position
+
+    def _apply_soft_trimming(
+        self,
+        sequence: Float[Array, "length alphabet"],
+        quality_scores: Float[Array, "length"],
+        soft_trim_pos: Float[Array, ""],
+    ) -> tuple[Float[Array, "length alphabet"], Float[Array, "length"]]:
+        """Apply soft trimming to sequence and quality scores.
+
+        Uses sigmoid to create smooth retention weights based on position
+        relative to the trim point. Positions before trim_pos have high
+        retention, positions after have low retention.
+
+        Args:
+            sequence: One-hot encoded sequence.
+            quality_scores: Quality scores for each position.
+            soft_trim_pos: Soft trim position.
+
+        Returns:
+            Tuple of (weighted_sequence, weighted_quality).
+        """
+        seq_len = sequence.shape[0]
+        temp = self._temperature
+
+        # Create position-based retention weights
+        # retention = sigmoid((trim_pos - position) / temperature)
+        positions = jnp.arange(seq_len, dtype=jnp.float32)
+        retention_weights = soft_ops.greater(soft_trim_pos, positions, softness=temp)
+
+        # Apply weights
+        weighted_sequence = sequence * retention_weights[:, None]
+        weighted_quality = quality_scores * retention_weights
+
+        return weighted_sequence, weighted_quality
+
+    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 soft adapter removal to sequence data.
+
+        This method finds potential adapter sequences and applies
+        differentiable soft trimming to remove them while maintaining
+        gradient flow.
+
+        Args:
+            data: Dictionary containing:
+                - "sequence": One-hot encoded sequence (length, alphabet_size)
+                - "quality_scores": Phred quality scores (length,)
+            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:
+
+                    - "sequence": Soft-trimmed sequence
+                    - "quality_scores": Soft-trimmed quality scores
+                    - "adapter_score": Maximum adapter match score
+                    - "trim_position": Soft trim position
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        sequence = data["sequence"]
+        quality_scores = data["quality_scores"]
+
+        # Compute adapter match scores for each suffix position
+        adapter_scores = self._compute_suffix_adapter_scores(sequence)
+
+        # Find soft trim position
+        soft_trim_pos = self._compute_soft_trim_position(adapter_scores)
+
+        # Apply soft trimming
+        trimmed_sequence, trimmed_quality = self._apply_soft_trimming(
+            sequence, quality_scores, soft_trim_pos
+        )
+
+        # Build output data
+        transformed_data = {
+            "sequence": trimmed_sequence,
+            "quality_scores": trimmed_quality,
+            "adapter_score": jnp.max(adapter_scores),
+            "trim_position": soft_trim_pos,
+        }
+
+        return transformed_data, state, metadata