diffbio 0.1.0__py3-none-any.whl

diffbio/losses/alignment_losses.py

@@ -0,0 +1,222 @@
+"""Alignment loss functions for differentiable sequence alignment.
+
+This module provides loss functions for training differentiable alignment
+models, including alignment score losses, soft edit distance, and
+alignment consistency losses for multi-sequence alignment.
+"""
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jaxtyping import Array, Float
+
+
+class AlignmentScoreLoss(nnx.Module):
+    """Loss function based on alignment quality score.
+
+    Computes a loss that measures how well the alignment captures
+    the similarity between two sequences. Lower loss indicates
+    better alignment of similar positions.
+
+    The loss computes the weighted sum of position-wise mismatches,
+    where weights come from the soft alignment matrix.
+
+    Args:
+        rngs: Flax NNX random number generators.
+    """
+
+    def __init__(self, *, rngs: nnx.Rngs | None = None):
+        """Initialize alignment score loss.
+
+        Args:
+            rngs: Random number generators (optional).
+        """
+        super().__init__()
+
+    def __call__(
+        self,
+        seq1: Float[Array, "len1 alphabet"],
+        seq2: Float[Array, "len2 alphabet"],
+        alignment: Float[Array, "len1 len2"],
+    ) -> Float[Array, ""]:
+        """Compute alignment score loss.
+
+        Args:
+            seq1: First sequence, soft one-hot encoded (len1, alphabet).
+            seq2: Second sequence, soft one-hot encoded (len2, alphabet).
+            alignment: Soft alignment matrix where alignment[i,j] indicates
+                      probability of aligning position i to position j.
+
+        Returns:
+            Scalar loss value. Lower is better alignment.
+        """
+        # Compute position-wise similarity: seq1[i] dot seq2[j]
+        # Higher similarity when same nucleotide
+        similarity = jnp.einsum("ia,ja->ij", seq1, seq2)
+
+        # Weight by alignment probabilities
+        # High alignment probability * high similarity = good
+        weighted_similarity = jnp.sum(alignment * similarity)
+
+        # Convert to loss (negate similarity, normalize)
+        max_possible = jnp.sum(alignment)  # If all positions perfectly matched
+        loss = 1.0 - (weighted_similarity / jnp.maximum(max_possible, 1e-8))
+
+        return loss
+
+
+class SoftEditDistanceLoss(nnx.Module):
+    """Differentiable approximation of edit distance.
+
+    Computes a soft version of edit distance between two sequences
+    that allows gradient flow. Uses the relationship between
+    edit distance and alignment scores.
+
+    The edit distance is approximated as the complement of the
+    optimal alignment score, scaled appropriately.
+
+    Args:
+        normalize: Whether to normalize by sequence length.
+        temperature: Temperature for soft minimum operations.
+        rngs: Flax NNX random number generators.
+    """
+
+    def __init__(
+        self,
+        normalize: bool = False,
+        temperature: float = 0.1,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize soft edit distance loss.
+
+        Args:
+            normalize: Whether to normalize by sequence length.
+            temperature: Temperature for softmax operations. Lower values
+                give sharper approximation of true edit distance.
+                Default 0.1 works well for one-hot encoded sequences.
+            rngs: Random number generators (optional).
+        """
+        super().__init__()
+        self.normalize = normalize
+        self.temperature = nnx.Param(jnp.array(temperature))
+
+    def __call__(
+        self,
+        seq1: Float[Array, "len1 alphabet"],
+        seq2: Float[Array, "len2 alphabet"],
+    ) -> Float[Array, ""]:
+        """Compute soft edit distance between sequences.
+
+        Args:
+            seq1: First sequence, soft one-hot encoded (len1, alphabet).
+            seq2: Second sequence, soft one-hot encoded (len2, alphabet).
+
+        Returns:
+            Scalar soft edit distance. 0 for identical sequences.
+        """
+        len1, len2 = seq1.shape[0], seq2.shape[0]
+
+        # Compute position-wise similarity matrix
+        # similarity[i,j] = probability that seq1[i] matches seq2[j]
+        similarity = jnp.einsum("ia,ja->ij", seq1, seq2)
+
+        temp = self.temperature[...]
+
+        # Use soft assignment to find best match per position
+        # Softmax over similarities gives assignment weights
+        soft_assign_row = jax.nn.softmax(similarity / temp, axis=1)
+        soft_assign_col = jax.nn.softmax(similarity / temp, axis=0)
+
+        # Compute expected similarity under soft assignment
+        # For identical sequences: assignment concentrates on diagonal (sim=1)
+        # For different sequences: assignment spread out (sim=0 everywhere)
+        expected_sim_row = jnp.sum(soft_assign_row * similarity, axis=1)
+        expected_sim_col = jnp.sum(soft_assign_col * similarity, axis=0)
+
+        # Total match score = sum of best similarities per position
+        match_score_row = jnp.sum(expected_sim_row)
+        match_score_col = jnp.sum(expected_sim_col)
+
+        # Distance = unmatched positions
+        # Average of (len - match_score) from both perspectives
+        dist_row = len1 - match_score_row
+        dist_col = len2 - match_score_col
+        total_distance = (dist_row + dist_col) / 2.0
+
+        # Ensure non-negative (numerical precision)
+        total_distance = jnp.maximum(total_distance, 0.0)
+
+        # Add length difference penalty
+        length_penalty = jnp.abs(len1 - len2).astype(jnp.float32)
+        total_distance = total_distance + length_penalty
+
+        if self.normalize:
+            # Normalize by total length
+            total_distance = total_distance / (len1 + len2)
+
+        return total_distance
+
+
+class AlignmentConsistencyLoss(nnx.Module):
+    """Loss for enforcing transitivity in multi-sequence alignments.
+
+    For three sequences A, B, C with pairwise alignments:
+    - A->B (align_ab)
+    - B->C (align_bc)
+    - A->C (align_ac)
+
+    The alignments are consistent if: align_ac ≈ align_ab @ align_bc
+
+    This loss penalizes violations of this transitivity property,
+    which is important for producing coherent multiple sequence alignments.
+
+    Args:
+        rngs: Flax NNX random number generators.
+    """
+
+    def __init__(self, *, rngs: nnx.Rngs | None = None):
+        """Initialize alignment consistency loss.
+
+        Args:
+            rngs: Random number generators (optional).
+        """
+        super().__init__()
+
+    def __call__(
+        self,
+        align_ab: Float[Array, "len_a len_b"],
+        align_bc: Float[Array, "len_b len_c"],
+        align_ac: Float[Array, "len_a len_c"],
+    ) -> Float[Array, ""]:
+        """Compute alignment consistency loss.
+
+        Args:
+            align_ab: Soft alignment from sequence A to B.
+            align_bc: Soft alignment from sequence B to C.
+            align_ac: Soft alignment from sequence A to C.
+
+        Returns:
+            Scalar loss measuring transitivity violation.
+        """
+        # Compute expected A->C alignment through B
+        # align_ac_expected[i,k] = sum_j align_ab[i,j] * align_bc[j,k]
+        align_ac_expected = jnp.matmul(align_ab, align_bc)
+
+        # Normalize to make it a proper probability distribution
+        align_ac_expected = align_ac_expected / jnp.maximum(
+            jnp.sum(align_ac_expected, axis=1, keepdims=True), 1e-8
+        )
+
+        # Compute KL divergence between expected and actual A->C alignment
+        # KL(actual || expected) = sum(actual * log(actual / expected))
+        eps = 1e-8
+        kl_div = jnp.sum(align_ac * jnp.log((align_ac + eps) / (align_ac_expected + eps)))
+
+        # Also compute reverse KL for symmetry
+        kl_div_reverse = jnp.sum(
+            align_ac_expected * jnp.log((align_ac_expected + eps) / (align_ac + eps))
+        )
+
+        # Return symmetric KL (Jensen-Shannon style)
+        return (kl_div + kl_div_reverse) / 2.0

diffbio/losses/biological_regularization.py

@@ -0,0 +1,288 @@
+"""Biological regularization losses for differentiable bioinformatics.
+
+This module provides regularization losses that help prevent adversarial
+optimization of differentiable bioinformatics components. These losses
+encourage biologically plausible sequences and alignments.
+
+Reference:
+    Petti et al. (2023) observed that purely differentiable alignment can
+    produce biologically implausible solutions without proper regularization.
+"""
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+from flax import nnx
+from jaxtyping import Array, Float
+
+from diffbio.core import soft_ops
+
+
+@dataclass(frozen=True)
+class BiologicalRegularizationConfig:
+    """Configuration for biological regularization losses.
+
+    Attributes:
+        gc_content_weight: Weight for GC content regularization.
+        gap_pattern_weight: Weight for gap pattern regularization.
+        complexity_weight: Weight for sequence complexity loss.
+        target_gc_content: Target GC content (typically 0.4-0.6).
+        target_gc_tolerance: Tolerance around target GC content.
+    """
+
+    gc_content_weight: float = 1.0
+    gap_pattern_weight: float = 1.0
+    complexity_weight: float = 1.0
+    target_gc_content: float = 0.5
+    target_gc_tolerance: float = 0.2
+
+
+class GCContentRegularization(nnx.Module):
+    """Regularization loss for GC content.
+
+    Penalizes sequences with GC content far from biological norms.
+    For most organisms, GC content ranges from 25% to 75%.
+
+    Args:
+        target_gc: Target GC content (default 0.5 for balanced).
+        tolerance: Tolerance around target before penalizing.
+        rngs: Flax NNX random number generators.
+    """
+
+    def __init__(
+        self,
+        target_gc: float = 0.5,
+        tolerance: float = 0.2,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize GC content regularization.
+
+        Args:
+            target_gc: Target GC content.
+            tolerance: Tolerance around target.
+            rngs: Random number generators (optional).
+        """
+        super().__init__()
+        self.target_gc = nnx.Param(jnp.array(target_gc))
+        self.tolerance = nnx.Param(jnp.array(tolerance))
+
+    def __call__(
+        self,
+        sequence: Float[Array, "length alphabet"],
+    ) -> Float[Array, ""]:
+        """Compute GC content regularization loss.
+
+        Args:
+            sequence: Soft one-hot encoded sequence (length, alphabet_size).
+                     Assumes alphabet order: A, C, G, T (indices 0, 1, 2, 3).
+
+        Returns:
+            Scalar loss penalizing deviation from target GC content.
+        """
+        # GC content = sum of C and G probabilities
+        # C is index 1, G is index 2
+        gc_content = jnp.mean(sequence[:, 1] + sequence[:, 2])
+
+        # Compute deviation from target
+        target = self.target_gc[...]
+        tolerance = self.tolerance[...]
+
+        # Soft penalty: quadratic beyond tolerance
+        deviation = jnp.abs(gc_content - target)
+        excess = soft_ops.relu(deviation - tolerance, softness=0.1)
+
+        return excess**2
+
+
+class GapPatternRegularization(nnx.Module):
+    """Regularization loss for gap patterns in alignments.
+
+    Penalizes unrealistic gap patterns such as:
+    - Very long consecutive gaps
+    - Many scattered small gaps
+
+    Args:
+        max_gap_length: Maximum expected gap length before penalizing.
+        rngs: Flax NNX random number generators.
+    """
+
+    def __init__(
+        self,
+        max_gap_length: int = 10,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize gap pattern regularization.
+
+        Args:
+            max_gap_length: Maximum expected gap length.
+            rngs: Random number generators (optional).
+        """
+        super().__init__()
+        self.max_gap_length = max_gap_length
+
+    def __call__(
+        self,
+        alignment_weights: Float[Array, "len1 len2"],
+    ) -> Float[Array, ""]:
+        """Compute gap pattern regularization loss.
+
+        Args:
+            alignment_weights: Soft alignment matrix where entry (i,j)
+                             indicates probability of aligning position i to j.
+
+        Returns:
+            Scalar loss penalizing unrealistic gap patterns.
+        """
+        # Compute row-wise and column-wise alignment strengths
+        row_aligned = soft_ops.max(alignment_weights, axis=1, softness=0.1)
+        col_aligned = soft_ops.max(alignment_weights, axis=0, softness=0.1)
+
+        # Penalize positions with very low alignment probability (gaps)
+        # Using smooth measure of "gappiness"
+        row_gap_penalty = jnp.mean(1.0 - row_aligned)
+        col_gap_penalty = jnp.mean(1.0 - col_aligned)
+
+        # Also penalize non-monotonic alignments (jumps)
+        # A good alignment should roughly follow the diagonal
+        len1, len2 = alignment_weights.shape
+        expected_diag = jnp.linspace(0, len2 - 1, len1)
+
+        # Compute weighted average position for each row
+        positions = jnp.arange(len2)
+        weighted_pos = jnp.sum(alignment_weights * positions[None, :], axis=1) / jnp.maximum(
+            jnp.sum(alignment_weights, axis=1), 1e-8
+        )
+
+        # Penalize deviation from expected diagonal progression
+        diag_penalty = jnp.mean((weighted_pos - expected_diag) ** 2) / (len2**2)
+
+        return row_gap_penalty + col_gap_penalty + diag_penalty
+
+
+class SequenceComplexityLoss(nnx.Module):
+    """Regularization loss for sequence complexity.
+
+    Penalizes low-complexity sequences that might arise from adversarial
+    optimization (e.g., all-A sequences, repetitive patterns).
+
+    Uses entropy as a measure of complexity.
+
+    Args:
+        min_entropy: Minimum expected entropy per position.
+        rngs: Flax NNX random number generators.
+    """
+
+    def __init__(
+        self,
+        min_entropy: float = 1.0,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize sequence complexity loss.
+
+        Args:
+            min_entropy: Minimum expected entropy.
+            rngs: Random number generators (optional).
+        """
+        super().__init__()
+        self.min_entropy = nnx.Param(jnp.array(min_entropy))
+
+    def __call__(
+        self,
+        sequence: Float[Array, "length alphabet"],
+    ) -> Float[Array, ""]:
+        """Compute sequence complexity loss.
+
+        Args:
+            sequence: Soft one-hot encoded sequence (length, alphabet_size).
+
+        Returns:
+            Scalar loss penalizing low-complexity sequences.
+        """
+        # Compute per-position entropy
+        # Add small epsilon for numerical stability
+        eps = 1e-8
+        entropy = -jnp.sum(sequence * jnp.log(sequence + eps), axis=-1)
+
+        # Average entropy across positions
+        avg_entropy = jnp.mean(entropy)
+
+        # Penalize if entropy is below minimum
+        min_ent = self.min_entropy[...]
+        deficit = jnp.maximum(min_ent - avg_entropy, 0.0)
+
+        return deficit**2
+
+
+class BiologicalPlausibilityLoss(nnx.Module):
+    """Combined biological plausibility regularization.
+
+    Combines multiple regularization terms to encourage biologically
+    plausible sequences and alignments during differentiable optimization.
+
+    Args:
+        config: BiologicalRegularizationConfig with weights and targets.
+        rngs: Flax NNX random number generators.
+    """
+
+    def __init__(
+        self,
+        config: BiologicalRegularizationConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize combined biological plausibility loss.
+
+        Args:
+            config: Configuration with weights and targets.
+            rngs: Random number generators (optional).
+        """
+        super().__init__()
+        self.config = config
+
+        # Initialize component losses
+        self.gc_loss = GCContentRegularization(
+            target_gc=config.target_gc_content,
+            tolerance=config.target_gc_tolerance,
+            rngs=rngs,
+        )
+        self.complexity_loss = SequenceComplexityLoss(
+            min_entropy=1.0,
+            rngs=rngs,
+        )
+
+    def __call__(
+        self,
+        sequence: Float[Array, "length alphabet"],
+        alignment_weights: Float[Array, "len1 len2"] | None = None,
+    ) -> Float[Array, ""]:
+        """Compute combined biological plausibility loss.
+
+        Args:
+            sequence: Soft one-hot encoded sequence.
+            alignment_weights: Optional soft alignment matrix.
+
+        Returns:
+            Scalar combined regularization loss.
+        """
+        total_loss = jnp.array(0.0)
+
+        # GC content regularization
+        if self.config.gc_content_weight > 0:
+            gc_loss = self.gc_loss(sequence)
+            total_loss = total_loss + self.config.gc_content_weight * gc_loss
+
+        # Sequence complexity regularization
+        if self.config.complexity_weight > 0:
+            complexity_loss = self.complexity_loss(sequence)
+            total_loss = total_loss + self.config.complexity_weight * complexity_loss
+
+        # Gap pattern regularization (if alignment provided)
+        if alignment_weights is not None and self.config.gap_pattern_weight > 0:
+            gap_loss_fn = GapPatternRegularization(rngs=None)
+            gap_loss = gap_loss_fn(alignment_weights)
+            total_loss = total_loss + self.config.gap_pattern_weight * gap_loss
+
+        return total_loss

diffbio/losses/metric_losses.py

@@ -0,0 +1,139 @@
+"""Differentiable metric-based loss functions.
+
+This module provides differentiable approximations of evaluation metrics
+so they can be used as training objectives with gradient-based optimizers,
+plus exact evaluation metrics backed by calibrax.
+
+Includes:
+- DifferentiableAUROC: Sigmoid-approximated Area Under the ROC Curve (training)
+- ExactAUROC: Trapezoidal-rule AUROC via calibrax (evaluation)
+"""
+
+import jax
+import jax.numpy as jnp
+from calibrax.metrics.functional.classification import roc_auc
+from flax import nnx
+from jaxtyping import Array, Float
+
+
+class DifferentiableAUROC(nnx.Module):
+    """Differentiable approximation of the Area Under the ROC Curve.
+
+    This is a smooth training surrogate.  For exact AUROC evaluation use
+    :class:`ExactAUROC`, which delegates to calibrax's trapezoidal-rule
+    implementation.
+
+    Approximates AUROC by replacing the hard indicator in the Wilcoxon-Mann-Whitney
+    statistic with a sigmoid function, making it fully differentiable and
+    JIT-compatible.
+
+    For every (positive, negative) pair the hard AUROC checks whether the
+    positive score exceeds the negative score.  This module replaces that
+    indicator with ``sigmoid((pos - neg) / temperature)``, yielding a smooth
+    surrogate whose gradient can drive optimisation.
+
+    Args:
+        temperature: Controls sharpness of the sigmoid approximation.
+            Lower values approach the hard indicator; higher values give
+            smoother gradients.  Default 1.0.
+
+    Example:
+        ```python
+        auroc_loss = DifferentiableAUROC(temperature=1.0)
+        predictions = jnp.array([0.9, 0.8, 0.1, 0.2])
+        labels = jnp.array([1.0, 1.0, 0.0, 0.0])
+        value = auroc_loss(predictions, labels)
+        ```
+    """
+
+    def __init__(self, temperature: float = 1.0) -> None:
+        """Initialise the differentiable AUROC loss.
+
+        Args:
+            temperature: Sigmoid temperature.  Lower values produce a
+                sharper (closer to hard) approximation.
+        """
+        super().__init__()
+        self.temperature = nnx.Param(jnp.array(temperature))
+
+    def __call__(
+        self,
+        predictions: Float[Array, " n"],
+        labels: Float[Array, " n"],
+    ) -> Float[Array, ""]:
+        """Compute the differentiable AUROC approximation.
+
+        Args:
+            predictions: Model output scores, shape ``(n,)``.
+            labels: Binary ground-truth labels (0 or 1), shape ``(n,)``.
+
+        Returns:
+            Scalar AUROC approximation in ``[0, 1]``.
+        """
+        temp = self.temperature[...]
+
+        pos_mask = labels == 1  # (n,)
+        neg_mask = labels == 0  # (n,)
+
+        n_pos = jnp.sum(pos_mask)
+        n_neg = jnp.sum(neg_mask)
+
+        # Pairwise differences: pos_scores[:, None] - neg_scores[None, :]
+        # Built via broadcasting with masks to stay JIT-compatible.
+        # pos_vals[i] = predictions[i] where label==1, else 0
+        pos_vals = jnp.where(pos_mask, predictions, 0.0)  # (n,)
+        neg_vals = jnp.where(neg_mask, predictions, 0.0)  # (n,)
+
+        # Outer difference over all (i, j) pairs
+        diffs = pos_vals[:, None] - neg_vals[None, :]  # (n, n)
+
+        # Mask to select only valid (positive_i, negative_j) pairs
+        pair_mask = pos_mask[:, None] & neg_mask[None, :]  # (n, n)
+
+        sigmoid_diffs = jax.nn.sigmoid(diffs / temp)  # (n, n)
+
+        # Mean over valid pairs
+        auroc = jnp.sum(sigmoid_diffs * pair_mask) / jnp.maximum(n_pos * n_neg, 1.0)
+
+        return auroc
+
+
+class ExactAUROC(nnx.Module):
+    """Exact AUROC metric using calibrax's trapezoidal-rule implementation.
+
+    Delegates to :func:`calibrax.metrics.functional.classification.roc_auc`
+    to compute the exact Area Under the ROC Curve via threshold-sweep and
+    the trapezoidal rule.
+
+    Use this for evaluation; use :class:`DifferentiableAUROC` for training
+    (the sorting-based trapezoidal rule has zero gradients w.r.t. predictions
+    because ``argsort`` is not differentiable).
+
+    Example:
+        ```python
+        exact = ExactAUROC()
+        predictions = jnp.array([0.9, 0.8, 0.1, 0.2])
+        labels = jnp.array([1.0, 1.0, 0.0, 0.0])
+        value = exact(predictions, labels)  # 1.0
+        ```
+    """
+
+    def __init__(self) -> None:
+        """Initialise the exact AUROC metric (no learnable parameters)."""
+        super().__init__()
+
+    def __call__(
+        self,
+        predictions: Float[Array, " n"],
+        labels: Float[Array, " n"],
+    ) -> Float[Array, ""]:
+        """Compute the exact AUROC via calibrax.
+
+        Args:
+            predictions: Model output scores, shape ``(n,)``.
+            labels: Binary ground-truth labels (0 or 1), shape ``(n,)``.
+
+        Returns:
+            Scalar AUROC in ``[0, 1]``.
+        """
+        return roc_auc(predictions, labels)