diffbio 0.1.0__py3-none-any.whl

diffbio/operators/alignment/smith_waterman.py

@@ -0,0 +1,261 @@
+"""Smooth Smith-Waterman alignment operator.
+
+This module provides a differentiable implementation of the Smith-Waterman
+local alignment algorithm using the logsumexp relaxation (SMURF-style).
+
+Key technique: Replace max with logsumexp and argmax with softmax
+to enable gradient flow through the alignment computation.
+
+Reference:
+    Petti et al. "End-to-end learning of multiple sequence alignments with
+    differentiable Smith-Waterman." Bioinformatics 39(1):btac724, 2023.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any, NamedTuple
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jaxtyping import Array, Float, PyTree
+
+from diffbio.configs import TemperatureConfig
+from diffbio.constants import DEFAULT_GAP_EXTEND, DEFAULT_GAP_OPEN
+from diffbio.core.base_operators import TemperatureOperator
+from diffbio.utils.nn_utils import init_learnable_param
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class SmithWatermanConfig(TemperatureConfig):
+    """Configuration for SmoothSmithWaterman.
+
+    Attributes:
+        temperature: Temperature for logsumexp smoothing.
+            Lower = sharper (closer to hard max), Higher = smoother.
+        gap_open: Penalty for opening a gap.
+        gap_extend: Penalty for extending a gap.
+    """
+
+    cacheable: bool = True
+    gap_open: float = DEFAULT_GAP_OPEN
+    gap_extend: float = DEFAULT_GAP_EXTEND
+
+
+class AlignmentResult(NamedTuple):
+    """Result of a smooth alignment.
+
+    Attributes:
+        score: The soft alignment score.
+        alignment_matrix: The DP matrix H[i,j] of shape (len1+1, len2+1).
+        soft_alignment: Soft alignment matrix showing position correspondences.
+    """
+
+    score: Float[Array, ""]
+    alignment_matrix: Float[Array, "len1_plus1 len2_plus1"]
+    soft_alignment: Float[Array, "len1 len2"]
+
+
+class SmoothSmithWaterman(TemperatureOperator):
+    """Differentiable Smith-Waterman local alignment.
+
+    This operator implements a smooth version of the Smith-Waterman algorithm
+    where max operations are replaced with logsumexp, enabling gradient flow
+    through the alignment computation.
+
+    The smoothness is controlled by the temperature parameter:
+    - temperature -> 0: Approaches hard max (standard Smith-Waterman)
+    - temperature -> inf: Uniform averaging
+
+    Inherits from TemperatureOperator to get:
+
+    - Learnable temperature parameter management
+    - soft_max() method using logsumexp relaxation
+
+    Args:
+        config: SmithWatermanConfig with alignment parameters.
+        scoring_matrix: Scoring matrix for matches/mismatches.
+        rngs: Flax NNX random number generators (optional).
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = SmithWatermanConfig(temperature=1.0)
+        scoring = create_dna_scoring_matrix(match=2.0, mismatch=-1.0)
+        aligner = SmoothSmithWaterman(config, scoring_matrix=scoring)
+        result = aligner.align(seq1, seq2)
+        print(result.score)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: SmithWatermanConfig,
+        scoring_matrix: Array,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the smooth Smith-Waterman aligner.
+
+        Args:
+            config: Alignment configuration.
+            scoring_matrix: Scoring matrix (alphabet_size, alphabet_size).
+            rngs: Random number generators (optional).
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        # Domain-specific learnable parameters (temperature managed by base class)
+        self.scoring_matrix = nnx.Param(scoring_matrix)
+        self.gap_open = init_learnable_param(config.gap_open)
+        self.gap_extend = init_learnable_param(config.gap_extend)
+
+    def _compute_score_matrix(
+        self,
+        seq1: Float[Array, "len1 alphabet"],
+        seq2: Float[Array, "len2 alphabet"],
+    ) -> Float[Array, "len1 len2"]:
+        """Compute pairwise scoring matrix between sequences.
+
+        Args:
+            seq1: First sequence, one-hot encoded (len1, alphabet_size).
+            seq2: Second sequence, one-hot encoded (len2, alphabet_size).
+
+        Returns:
+            Score matrix of shape (len1, len2).
+        """
+        # S[i,j] = seq1[i] @ scoring_matrix @ seq2[j].T
+        # Using einsum for clarity
+        scoring = self.scoring_matrix[...]
+        return jnp.einsum("ia,ab,jb->ij", seq1, scoring, seq2)
+
+    def align(
+        self,
+        seq1: Float[Array, "len1 alphabet"],
+        seq2: Float[Array, "len2 alphabet"],
+    ) -> AlignmentResult:
+        """Perform smooth Smith-Waterman local alignment.
+
+        Args:
+            seq1: First sequence, one-hot encoded (len1, alphabet_size).
+            seq2: Second sequence, one-hot encoded (len2, alphabet_size).
+
+        Returns:
+            AlignmentResult with score, alignment matrix, and soft alignment.
+        """
+        len1, len2 = seq1.shape[0], seq2.shape[0]
+
+        # Compute pairwise scores
+        score_matrix = self._compute_score_matrix(seq1, seq2)
+
+        # Gap penalties
+        gap_open = self.gap_open[...]
+        gap_extend = self.gap_extend[...]
+        # Simplified: use linear gap penalty (gap_open + gap_extend per position)
+        gap_penalty = gap_open + gap_extend
+
+        # Initialize DP matrices
+        # H[i,j] = alignment score ending at seq1[i-1], seq2[j-1]
+        # Using scan for efficient JAX computation
+
+        # Initialize H matrix with zeros
+        H = jnp.zeros((len1 + 1, len2 + 1))
+
+        # Fill the DP matrix row by row
+        # Note: fori_loop body signature is (i, carry) -> carry
+        def fill_row(i, H):
+            """Fill row i+1 of the DP matrix."""
+
+            def cell_update(H_prev_col, j_idx):
+                """Update single cell H[i+1, j+1]."""
+                j = j_idx.astype(jnp.int32)
+                s = score_matrix[i, j]
+                diag = H[i, j] + s
+                up = H[i, j + 1] + gap_penalty
+                left = H_prev_col + gap_penalty
+
+                # Use inherited soft_max from TemperatureOperator
+                candidates = jnp.stack([jnp.array(0.0), diag, up, left], axis=-1)
+                h_new = self.soft_max(candidates, axis=-1)
+                return h_new, h_new
+
+            # Scan across columns
+            _, new_row = jax.lax.scan(
+                cell_update, jnp.array(0.0), jnp.arange(len2, dtype=jnp.int32)
+            )
+
+            # Update H matrix
+            H = H.at[i + 1, 1:].set(new_row)
+            return H
+
+        # Fill all rows using fori_loop for efficiency
+        H = jax.lax.fori_loop(0, len1, fill_row, H)
+
+        # Compute final score as smooth max over all positions
+        # (local alignment can end anywhere)
+        final_candidates = jnp.stack([jnp.array(0.0), jnp.max(H)], axis=-1)
+        final_score = self.soft_max(final_candidates, axis=-1)
+
+        # Compute soft alignment (position correspondence probabilities)
+        # This is the softmax of the DP matrix (excluding borders)
+        H_inner = H[1:, 1:]  # (len1, len2)
+        temp = self._temperature  # Use property from TemperatureOperator
+        soft_alignment = jax.nn.softmax(H_inner.flatten() / temp).reshape(len1, len2)
+
+        return AlignmentResult(
+            score=final_score,
+            alignment_matrix=H,
+            soft_alignment=soft_alignment,
+        )
+
+    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 alignment to sequence pair data.
+
+        This method implements the OperatorModule interface for batch processing.
+        It expects data containing two sequences and returns alignment results.
+
+        Note: Output preserves input keys for Datarax vmap compatibility,
+        while adding alignment result keys.
+
+        Args:
+            data: Dictionary containing:
+                - "seq1": First sequence, one-hot encoded (len1, alphabet_size)
+                - "seq2": Second sequence, one-hot encoded (len2, alphabet_size)
+            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 input sequences plus alignment results
+                  (score, alignment_matrix, soft_alignment)
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        seq1 = data["seq1"]
+        seq2 = data["seq2"]
+
+        # Perform alignment
+        result = self.align(seq1, seq2)
+
+        # Build output data - preserve input keys for Datarax vmap compatibility
+        transformed_data = {
+            "seq1": seq1,
+            "seq2": seq2,
+            "score": result.score,
+            "alignment_matrix": result.alignment_matrix,
+            "soft_alignment": result.soft_alignment,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/alignment/soft_msa.py

@@ -0,0 +1,419 @@
+"""Soft progressive multiple sequence alignment operator.
+
+This module provides differentiable multiple sequence alignment using
+soft operations throughout the alignment process.
+
+Key technique: Uses neural network sequence encoders to compute pairwise
+similarities, builds a soft guide tree, and performs progressive profile
+alignment with soft gap handling.
+
+Applications: Multiple sequence alignment for homology detection, phylogenetic
+analysis, and protein family characterization.
+
+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 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 (
+    ARTIFEX_GELU_MLP_KWARGS,
+    ARTIFEX_GELU_NO_OUTPUT_MLP_KWARGS,
+    ensure_rngs,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class SoftProgressiveMSAConfig(OperatorConfig):
+    """Configuration for SoftProgressiveMSA.
+
+    Attributes:
+        max_seq_length: Maximum sequence length.
+        hidden_dim: Hidden dimension for neural networks.
+        num_layers: Number of encoder layers.
+        alphabet_size: Size of sequence alphabet (4 for DNA, 20 for protein).
+        temperature: Temperature for softmax operations.
+        gap_open_penalty: Gap opening penalty.
+        gap_extend_penalty: Gap extension penalty.
+    """
+
+    max_seq_length: int = 100
+    hidden_dim: int = 64
+    num_layers: int = 2
+    alphabet_size: int = 4
+    temperature: float = 1.0
+    gap_open_penalty: float = -10.0
+    gap_extend_penalty: float = -1.0
+
+    def __post_init__(self) -> None:
+        """Validate soft MSA configuration."""
+        super().__post_init__()
+        if self.num_layers < 1:
+            raise ValueError("SoftProgressiveMSAConfig.num_layers must be at least 1.")
+
+
+class SequenceEncoder(nnx.Module):
+    """Encoder for biological sequences."""
+
+    def __init__(
+        self,
+        alphabet_size: int,
+        hidden_dim: int,
+        num_layers: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize the sequence encoder.
+
+        Args:
+            alphabet_size: Size of input alphabet.
+            hidden_dim: Hidden dimension.
+            num_layers: Number of layers.
+            rngs: Random number generators.
+        """
+        super().__init__()
+        self.backbone = MLP(
+            hidden_dims=[hidden_dim] * num_layers,
+            in_features=alphabet_size,
+            rngs=rngs,
+            **ARTIFEX_GELU_MLP_KWARGS,
+        )
+
+        # Output projection for sequence embedding
+        self.output_proj = nnx.Linear(
+            in_features=hidden_dim,
+            out_features=hidden_dim,
+            rngs=rngs,
+        )
+
+    def __call__(
+        self,
+        sequence: Float[Array, "seq_len alphabet_size"],
+    ) -> Float[Array, "hidden_dim"]:
+        """Encode a sequence to a fixed-size embedding.
+
+        Args:
+            sequence: One-hot encoded sequence.
+
+        Returns:
+            Sequence embedding vector.
+        """
+        backbone_output = self.backbone(sequence)
+        if isinstance(backbone_output, tuple):
+            raise TypeError("Soft MSA sequence backbone must return a single tensor.")
+
+        # Global average pooling to get fixed-size embedding
+        embedding = jnp.mean(backbone_output, axis=0)  # (hidden_dim,)
+        embedding = self.output_proj(embedding)
+
+        return embedding
+
+
+class ProfileBuilder(nnx.Module):
+    """Builds alignment profiles from sequences."""
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        alphabet_size: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize the profile builder.
+
+        Args:
+            hidden_dim: Hidden dimension.
+            alphabet_size: Alphabet size.
+            rngs: Random number generators.
+        """
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+        self.alphabet_size = alphabet_size
+
+        self.backbone = MLP(
+            hidden_dims=[hidden_dim, alphabet_size],
+            in_features=alphabet_size,
+            rngs=rngs,
+            **ARTIFEX_GELU_NO_OUTPUT_MLP_KWARGS,
+        )
+
+    def __call__(
+        self,
+        sequences: Float[Array, "n_seqs seq_len alphabet_size"],
+        weights: Float[Array, "n_seqs"],
+    ) -> Float[Array, "seq_len alphabet_size"]:
+        """Build a profile from weighted sequences.
+
+        Args:
+            sequences: Stack of aligned sequences.
+            weights: Weights for each sequence.
+
+        Returns:
+            Profile (position-specific scoring matrix).
+        """
+        # Weighted average of sequences
+        weights = weights / (jnp.sum(weights) + 1e-8)
+        profile = jnp.einsum("n,nla->la", weights, sequences)
+
+        # Refine profile
+        refinement = self.backbone(profile)
+        if isinstance(refinement, tuple):
+            raise TypeError("Soft MSA profile backbone must return a single tensor.")
+        profile = profile + 0.1 * refinement  # Small residual update
+
+        # Normalize to valid probability distribution
+        profile = jax.nn.softmax(profile, axis=-1)
+
+        return profile
+
+
+class SoftProgressiveMSA(TemperatureOperator):
+    """Differentiable progressive multiple sequence alignment.
+
+    This operator performs multiple sequence alignment using soft
+    operations that maintain gradient flow throughout the process.
+
+    Algorithm:
+    1. Encode each sequence to get embeddings
+    2. Compute pairwise distances from embeddings
+    3. Build soft guide tree from distances
+    4. Progressive alignment following guide tree order
+    5. Build consensus profile
+
+    Inherits from TemperatureOperator to get:
+
+    - _temperature property for temperature-controlled smoothing
+    - soft_max() for logsumexp-based smooth maximum
+    - soft_argmax() for soft position selection
+
+    Args:
+        config: SoftProgressiveMSAConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = SoftProgressiveMSAConfig(max_seq_length=100)
+        msa = SoftProgressiveMSA(config, rngs=nnx.Rngs(42))
+        data = {"sequences": seqs}  # (n_seqs, seq_len, alphabet_size)
+        result, state, meta = msa.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: SoftProgressiveMSAConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the soft progressive MSA operator.
+
+        Args:
+            config: MSA configuration.
+            rngs: Random number generators for initialization.
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        rngs = ensure_rngs(rngs)
+
+        self.hidden_dim = config.hidden_dim
+        # Temperature is now managed by TemperatureOperator via self._temperature
+        self.alphabet_size = config.alphabet_size
+
+        # Sequence encoder for computing pairwise similarities
+        self.seq_encoder = SequenceEncoder(
+            alphabet_size=config.alphabet_size,
+            hidden_dim=config.hidden_dim,
+            num_layers=config.num_layers,
+            rngs=rngs,
+        )
+
+        # Profile builder for progressive alignment
+        self.profile_builder = ProfileBuilder(
+            hidden_dim=config.hidden_dim,
+            alphabet_size=config.alphabet_size,
+            rngs=rngs,
+        )
+
+        # Alignment scoring
+        self.align_score = nnx.Linear(
+            in_features=config.alphabet_size * 2,
+            out_features=1,
+            rngs=rngs,
+        )
+
+    def _compute_pairwise_distances(
+        self,
+        sequences: Float[Array, "n_seqs seq_len alphabet_size"],
+    ) -> Float[Array, "n_seqs n_seqs"]:
+        """Compute pairwise distances between sequences.
+
+        Args:
+            sequences: Input sequences.
+
+        Returns:
+            Pairwise distance matrix.
+        """
+        n_seqs = sequences.shape[0]
+
+        # Encode all sequences
+        embeddings = jax.vmap(self.seq_encoder)(sequences)  # (n_seqs, hidden_dim)
+
+        # Compute pairwise distances (negative cosine similarity)
+        norms = jnp.linalg.norm(embeddings, axis=-1, keepdims=True) + 1e-8
+        normalized = embeddings / norms
+
+        # Cosine similarity -> distance
+        similarities = jnp.einsum("ih,jh->ij", normalized, normalized)
+        distances = 1.0 - similarities
+
+        # Zero diagonal
+        distances = distances * (1.0 - jnp.eye(n_seqs))
+
+        return distances
+
+    def _soft_align_pair(
+        self,
+        seq1: Float[Array, "len1 alphabet_size"],
+        seq2: Float[Array, "len2 alphabet_size"],
+    ) -> tuple[
+        Float[Array, "max_len alphabet_size"],
+        Float[Array, "max_len alphabet_size"],
+        Float[Array, ""],
+    ]:
+        """Perform soft pairwise alignment.
+
+        Args:
+            seq1: First sequence.
+            seq2: Second sequence.
+
+        Returns:
+            Tuple of (aligned_seq1, aligned_seq2, alignment_score).
+        """
+        len1, len2 = seq1.shape[0], seq2.shape[0]
+        max_len = max(len1, len2)
+
+        # Compute position-wise similarity scores
+        # (len1, alphabet) x (len2, alphabet) -> (len1, len2)
+        match_scores = jnp.einsum("ia,ja->ij", seq1, seq2)
+
+        # Soft alignment via attention
+        # Each position in seq1 attends to positions in seq2
+        # Use inherited _temperature property from TemperatureOperator
+        attn_weights = jax.nn.softmax(match_scores / self._temperature, axis=-1)
+
+        # Soft-aligned seq2 based on seq1 positions
+        aligned_seq2_to_seq1 = jnp.einsum("ij,ja->ia", attn_weights, seq2)
+
+        # Pad to max_len
+        pad1 = max_len - len1
+
+        aligned1 = jnp.pad(seq1, ((0, pad1), (0, 0)))
+        aligned2 = jnp.pad(aligned_seq2_to_seq1, ((0, pad1), (0, 0)))
+
+        # Compute alignment score
+        alignment_score = jnp.mean(match_scores * attn_weights)
+
+        return aligned1, aligned2, alignment_score
+
+    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 progressive MSA.
+
+        Args:
+            data: Dictionary containing:
+                - "sequences": Input sequences (n_seqs, seq_len, alphabet_size)
+            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:
+
+                    - "sequences": Original sequences
+                    - "aligned_sequences": Soft-aligned sequences
+                    - "pairwise_distances": Guide tree distances
+                    - "alignment_scores": Pairwise alignment scores
+                    - "consensus_profile": Consensus profile
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        sequences = data["sequences"]
+        n_seqs = sequences.shape[0]
+
+        # Step 1: Compute pairwise distances for guide tree
+        pairwise_distances = self._compute_pairwise_distances(sequences)
+
+        # Step 2: Progressive alignment
+        # For simplicity, align all sequences to the first one
+        # (full progressive alignment would follow guide tree)
+        aligned_sequences = []
+        alignment_scores = []
+
+        # Use first sequence as anchor
+        anchor = sequences[0]
+        aligned_sequences.append(anchor)
+
+        for i in range(1, n_seqs):
+            aligned1, aligned2, score = self._soft_align_pair(anchor, sequences[i])
+            # Keep the aligned version of sequences[i]
+            aligned_sequences.append(aligned2[: sequences[i].shape[0]])
+            alignment_scores.append(score)
+
+        # Stack aligned sequences (pad to same length)
+        max_len = max(s.shape[0] for s in aligned_sequences)
+        padded_aligned = []
+        for seq in aligned_sequences:
+            pad_len = max_len - seq.shape[0]
+            padded = jnp.pad(seq, ((0, pad_len), (0, 0)))
+            padded_aligned.append(padded)
+
+        aligned_stack = jnp.stack(padded_aligned, axis=0)
+
+        # Step 3: Build consensus profile
+        uniform_weights = jnp.ones(n_seqs) / n_seqs
+        consensus = self.profile_builder(aligned_stack, uniform_weights)
+
+        # Alignment scores matrix
+        scores_matrix = jnp.zeros((n_seqs, n_seqs))
+        if alignment_scores:
+            scores_array = jnp.array(alignment_scores)
+            # Fill first row/column with computed scores
+            scores_matrix = scores_matrix.at[0, 1:].set(scores_array)
+            scores_matrix = scores_matrix.at[1:, 0].set(scores_array)
+
+        transformed_data = {
+            "sequences": sequences,
+            "aligned_sequences": aligned_stack,
+            "pairwise_distances": pairwise_distances,
+            "alignment_scores": scores_matrix,
+            "consensus_profile": consensus,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/assembly/__init__.py

@@ -0,0 +1,27 @@
+"""Assembly operators for differentiable genome assembly.
+
+This module provides graph neural network-based approaches to
+assembly graph traversal that enable gradient flow through the
+assembly process.
+
+- GNNAssemblyNavigator: Message passing GNN for soft edge selection
+- DifferentiableMetagenomicBinner: VAMB-style VAE for metagenomic binning
+"""
+
+from diffbio.operators.assembly.gnn_assembly import (
+    GNNAssemblyNavigator,
+    GNNAssemblyNavigatorConfig,
+)
+from diffbio.operators.assembly.metagenomic_binning import (
+    DifferentiableMetagenomicBinner,
+    MetagenomicBinnerConfig,
+    create_metagenomic_binner,
+)
+
+__all__ = [
+    "GNNAssemblyNavigator",
+    "GNNAssemblyNavigatorConfig",
+    "DifferentiableMetagenomicBinner",
+    "MetagenomicBinnerConfig",
+    "create_metagenomic_binner",
+]