diffbio 0.1.0__py3-none-any.whl

diffbio/operators/_masked_gene_transformer.py

@@ -0,0 +1,118 @@
+"""Shared scaffolding for masked-gene transformer operators.
+
+These helpers centralize the common single-cell transformer setup used by the
+foundation-model and imputation operators so both paths share one encoder
+construction and one mask/input preparation flow.
+"""
+
+from __future__ import annotations
+
+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, PyTree
+
+from diffbio.operators.foundation_models.transformer_encoder import (
+    TransformerSequenceEncoder,
+    TransformerSequenceEncoderConfig,
+)
+
+
+@dataclass(frozen=True)
+class MaskedGeneTransformerConfigBase(OperatorConfig):
+    """Shared config fields for masked-gene transformer operators."""
+
+    n_genes: int = 2000
+    hidden_dim: int = 128
+    num_layers: int = 2
+    num_heads: int = 4
+    mask_ratio: float = 0.15
+    dropout_rate: float = 0.1
+
+    def __post_init__(self) -> None:
+        """Default masked-gene operators to sampled stochastic execution."""
+        object.__setattr__(self, "stochastic", True)
+        if self.stream_name is None:
+            object.__setattr__(self, "stream_name", "sample")
+        super().__post_init__()
+
+
+def build_masked_gene_transformer_encoder(
+    config: MaskedGeneTransformerConfigBase,
+    *,
+    rngs: nnx.Rngs,
+) -> TransformerSequenceEncoder:
+    """Build the shared token-embedding transformer encoder contract."""
+    encoder_config = TransformerSequenceEncoderConfig(
+        hidden_dim=config.hidden_dim,
+        num_layers=config.num_layers,
+        num_heads=config.num_heads,
+        intermediate_dim=4 * config.hidden_dim,
+        max_length=config.n_genes,
+        input_embedding_type="token_embedding",
+        vocab_size=config.n_genes,
+        dropout_rate=config.dropout_rate,
+        pooling="mean",
+    )
+    return TransformerSequenceEncoder(encoder_config, rngs=rngs)
+
+
+def build_masked_gene_mask(
+    *,
+    random_params: Any,
+    mask_ratio: float,
+    n_genes: int,
+) -> Array:
+    """Build a per-gene binary mask for masked-gene transformer operators."""
+    if random_params is not None and mask_ratio > 0:
+        noise = jax.random.uniform(random_params, (n_genes,))
+        return (noise < mask_ratio).astype(jnp.float32)
+    return jnp.zeros(n_genes, dtype=jnp.float32)
+
+
+def prepare_masked_gene_batch(
+    data: PyTree,
+    *,
+    random_params: Any,
+    mask_ratio: float,
+) -> tuple[Array, Array, Array]:
+    """Extract counts, int32 gene IDs, and the shared masking vector."""
+    counts = data["counts"]
+    gene_ids = jnp.asarray(data["gene_ids"], dtype=jnp.int32)
+    mask = build_masked_gene_mask(
+        random_params=random_params,
+        mask_ratio=mask_ratio,
+        n_genes=int(counts.shape[1]),
+    )
+    return counts, gene_ids, mask
+
+
+class MaskedGeneTransformerOperatorMixin:
+    """Mixin for operators built on the shared masked-gene transformer flow."""
+
+    config: MaskedGeneTransformerConfigBase
+
+    def generate_random_params(
+        self,
+        rng: jax.Array,
+        data_shapes: PyTree,
+    ) -> jax.Array:
+        """Return the RNG key used for reproducible masking inside apply."""
+        del data_shapes
+        return rng
+
+    def prepare_masked_gene_batch(
+        self,
+        data: PyTree,
+        random_params: Any,
+    ) -> tuple[Array, Array, Array]:
+        """Prepare shared masked-gene inputs for per-cell `vmap` execution."""
+        return prepare_masked_gene_batch(
+            data,
+            random_params=random_params,
+            mask_ratio=self.config.mask_ratio,
+        )

diffbio/operators/_transformer_validation.py

@@ -0,0 +1,50 @@
+"""Shared validation for transformer-style operator configs."""
+
+
+def validate_transformer_encoder_shape(
+    *,
+    hidden_dim: int,
+    num_layers: int,
+    num_heads: int,
+    intermediate_dim: int,
+    max_length: int,
+    dropout_rate: float,
+) -> None:
+    """Validate common transformer encoder hyperparameters."""
+    if hidden_dim <= 0:
+        raise ValueError("hidden_dim must be positive.")
+    if num_layers <= 0:
+        raise ValueError("num_layers must be positive.")
+    if num_heads <= 0:
+        raise ValueError("num_heads must be positive.")
+    if hidden_dim % num_heads != 0:
+        raise ValueError("hidden_dim must be divisible by num_heads.")
+    if intermediate_dim <= 0:
+        raise ValueError("intermediate_dim must be positive.")
+    if max_length <= 0:
+        raise ValueError("max_length must be positive.")
+    if not 0.0 <= dropout_rate < 1.0:
+        raise ValueError("dropout_rate must be in [0.0, 1.0).")
+
+
+class TransformerEncoderShapeValidationMixin:
+    """Mixin that validates shared transformer encoder dimensions."""
+
+    hidden_dim: int
+    num_layers: int
+    num_heads: int
+    intermediate_dim: int
+    max_length: int
+    dropout_rate: float
+
+    def __post_init__(self) -> None:
+        """Run base config validation and shared transformer checks."""
+        super().__post_init__()
+        validate_transformer_encoder_shape(
+            hidden_dim=self.hidden_dim,
+            num_layers=self.num_layers,
+            num_heads=self.num_heads,
+            intermediate_dim=self.intermediate_dim,
+            max_length=self.max_length,
+            dropout_rate=self.dropout_rate,
+        )

diffbio/operators/alignment/__init__.py

@@ -0,0 +1,51 @@
+"""Differentiable sequence alignment operators.
+
+This module provides differentiable implementations of sequence alignment
+algorithms including smooth Smith-Waterman for local alignment,
+profile HMM search for domain detection, and soft progressive MSA.
+"""
+
+from diffbio.operators.alignment.profile_hmm import (
+    ProfileHMMConfig,
+    ProfileHMMSearch,
+)
+from diffbio.operators.alignment.scoring import (
+    PROTEIN_ALPHABET,
+    ScoringMatrix,
+    create_dna_scoring_matrix,
+    create_rna_scoring_matrix,
+    get_blosum62,
+    get_dna_simple,
+    get_rna_simple,
+)
+from diffbio.operators.alignment.smith_waterman import (
+    AlignmentResult,
+    SmithWatermanConfig,
+    SmoothSmithWaterman,
+)
+from diffbio.operators.alignment.soft_msa import (
+    SoftProgressiveMSA,
+    SoftProgressiveMSAConfig,
+)
+
+
+__all__ = [
+    # Profile HMM
+    "ProfileHMMConfig",
+    "ProfileHMMSearch",
+    # Scoring
+    "PROTEIN_ALPHABET",
+    "ScoringMatrix",
+    "create_dna_scoring_matrix",
+    "create_rna_scoring_matrix",
+    "get_blosum62",
+    "get_dna_simple",
+    "get_rna_simple",
+    # Smith-Waterman
+    "AlignmentResult",
+    "SmithWatermanConfig",
+    "SmoothSmithWaterman",
+    # Soft MSA
+    "SoftProgressiveMSA",
+    "SoftProgressiveMSAConfig",
+]

diffbio/operators/alignment/profile_hmm.py

@@ -0,0 +1,350 @@
+"""Profile HMM search operator for HMMER-style sequence alignment.
+
+This module provides a differentiable implementation of profile HMM
+search, enabling gradient-based learning of profile parameters.
+
+Key technique: Use the forward algorithm with logsumexp for
+differentiable profile-sequence alignment scoring.
+
+Applications: Protein domain detection, remote homology search.
+
+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.base_operators import TemperatureOperator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ProfileHMMConfig(OperatorConfig):
+    """Configuration for ProfileHMMSearch.
+
+    Attributes:
+        profile_length: Length of the profile (number of match states).
+        alphabet_size: Size of sequence alphabet (20 for protein, 4 for DNA).
+        temperature: Temperature for softmax operations.
+        learnable_profile: Whether profile parameters are learnable.
+    """
+
+    cacheable: bool = True
+    profile_length: int = 100
+    alphabet_size: int = 20  # Amino acids by default
+    temperature: float = 1.0
+    learnable_profile: bool = True
+
+
+class ProfileHMMSearch(TemperatureOperator):
+    """Profile HMM search with differentiable scoring.
+
+    This operator implements a simplified profile HMM with match, insert,
+    and delete states. The forward algorithm computes the alignment score
+    differentiably using logsumexp.
+
+    Profile HMM structure (per position):
+    - Match state: emits according to position-specific distribution
+    - Insert state: emits according to background distribution
+    - Delete state: silent (no emission)
+
+    Transitions:
+    - M->M, M->I, M->D (from match)
+    - I->M, I->I (from insert)
+    - D->M, D->D (from delete)
+
+    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: ProfileHMMConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = ProfileHMMConfig(profile_length=100, alphabet_size=20)
+        profiler = ProfileHMMSearch(config, rngs=nnx.Rngs(42))
+        data = {"sequence": one_hot_sequence}
+        result, state, meta = profiler.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: ProfileHMMConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the profile HMM operator.
+
+        Args:
+            config: Profile HMM 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.profile_length = config.profile_length
+        self.alphabet_size = config.alphabet_size
+        # Temperature is now managed by TemperatureOperator via self._temperature
+
+        # Initialize match emissions (profile_length, alphabet_size)
+        key = rngs.params()
+        init_match = jax.random.normal(key, (config.profile_length, config.alphabet_size)) * 0.1
+        self.log_match_emissions = nnx.Param(init_match)
+
+        # Initialize insert emissions (profile_length, alphabet_size)
+        # Insert states use near-uniform distribution
+        key = rngs.params()
+        init_insert = jax.random.normal(key, (config.profile_length, config.alphabet_size)) * 0.01
+        self.log_insert_emissions = nnx.Param(init_insert)
+
+        # Initialize transition parameters
+        # For each position: [M->M, M->I, M->D, I->M, I->I, D->M, D->D]
+        key = rngs.params()
+        # Default: prefer M->M transitions
+        init_trans = jnp.zeros((config.profile_length, 7))
+        init_trans = init_trans.at[:, 0].set(2.0)  # M->M bias
+        init_trans = init_trans.at[:, 3].set(1.0)  # I->M bias
+        init_trans = init_trans.at[:, 5].set(1.0)  # D->M bias
+        init_trans = init_trans + jax.random.normal(key, init_trans.shape) * 0.1
+        self.log_transitions = nnx.Param(init_trans)
+
+    def get_match_emissions(self) -> Float[Array, "profile_length alphabet_size"]:
+        """Get normalized match emission probabilities.
+
+        Returns:
+            Log match emission probabilities.
+        """
+        return jax.nn.log_softmax(self.log_match_emissions[...] / self._temperature, axis=1)
+
+    def get_insert_emissions(self) -> Float[Array, "profile_length alphabet_size"]:
+        """Get normalized insert emission probabilities.
+
+        Returns:
+            Log insert emission probabilities.
+        """
+        return jax.nn.log_softmax(self.log_insert_emissions[...] / self._temperature, axis=1)
+
+    def get_transitions(self) -> dict[str, Float[Array, "profile_length"]]:
+        """Get normalized transition probabilities.
+
+        Returns:
+            Dictionary with log transition probabilities for each type.
+        """
+        trans = self.log_transitions[...]
+
+        # Normalize M->* transitions (positions 0, 1, 2)
+        log_m_trans = jax.nn.log_softmax(trans[:, :3] / self._temperature, axis=1)
+
+        # Normalize I->* transitions (positions 3, 4)
+        log_i_trans = jax.nn.log_softmax(trans[:, 3:5] / self._temperature, axis=1)
+
+        # Normalize D->* transitions (positions 5, 6)
+        log_d_trans = jax.nn.log_softmax(trans[:, 5:7] / self._temperature, axis=1)
+
+        return {
+            "m_to_m": log_m_trans[:, 0],
+            "m_to_i": log_m_trans[:, 1],
+            "m_to_d": log_m_trans[:, 2],
+            "i_to_m": log_i_trans[:, 0],
+            "i_to_i": log_i_trans[:, 1],
+            "d_to_m": log_d_trans[:, 0],
+            "d_to_d": log_d_trans[:, 1],
+        }
+
+    def score_sequence(
+        self,
+        sequence: Float[Array, "seq_len alphabet_size"],
+    ) -> Float[Array, ""]:
+        """Score a sequence against the profile using forward algorithm.
+
+        Computes log P(sequence | profile) using dynamic programming.
+
+        Args:
+            sequence: One-hot encoded sequence.
+
+        Returns:
+            Log probability (alignment score).
+        """
+        seq_len = sequence.shape[0]
+        log_match = self.get_match_emissions()
+        log_insert = self.get_insert_emissions()
+        trans = self.get_transitions()
+
+        # DP matrices: (seq_pos, profile_pos)
+        # We use a simplified forward algorithm
+
+        # Emission scores for each sequence position at each profile position
+        # match_scores[i, j] = log P(seq[i] | Match[j])
+        match_scores = jnp.einsum("sa,pa->sp", sequence, jnp.exp(log_match))
+        match_scores = jnp.log(match_scores + 1e-10)
+
+        # Insert scores per position (use position-specific insert emissions)
+        insert_scores = jnp.einsum("sa,pa->sp", sequence, jnp.exp(log_insert))
+        insert_scores = jnp.log(insert_scores + 1e-10)
+
+        # Initialize DP
+        # log_M[j] = log probability of being in Match state j
+        # log_I[j] = log probability of being in Insert state j
+        # log_D[j] = log probability of being in Delete state j
+
+        neg_inf = -1e10
+
+        # Initial state: can start at any match state with decreasing probability
+        # or go through leading deletes
+        log_M = jnp.full(self.profile_length, neg_inf)
+        log_I = jnp.full(self.profile_length, neg_inf)
+        log_D = jnp.full(self.profile_length, neg_inf)
+
+        # Start: emit first sequence position at first match state
+        log_M = log_M.at[0].set(match_scores[0, 0])
+        log_I = log_I.at[0].set(insert_scores[0, 0])
+
+        # Forward pass
+        def forward_step(carry, seq_idx):
+            log_M, log_I, log_D = carry
+            obs_match = match_scores[seq_idx]
+            obs_insert = insert_scores[seq_idx]
+
+            # New M states: can come from M, I, or D at previous profile position
+            # M[j] <- M[j-1] + M->M + emit[j]
+            #      <- I[j-1] + I->M + emit[j]
+            #      <- D[j-1] + D->M + emit[j]
+            new_log_M = jnp.full(self.profile_length, neg_inf)
+
+            # From previous M (M[j-1] -> M[j] for j=1..L-1)
+            from_M = log_M[:-1] + trans["m_to_m"][:-1]
+            # From previous I (I[j-1] -> M[j] for j=1..L-1)
+            from_I = log_I[:-1] + trans["i_to_m"][:-1]
+            # From previous D (D[j-1] -> M[j] for j=1..L-1)
+            from_D = log_D[:-1] + trans["d_to_m"][:-1]
+
+            # Combine and add emission
+            combined = jax.scipy.special.logsumexp(jnp.stack([from_M, from_I, from_D]), axis=0)
+            new_log_M = new_log_M.at[1:].set(combined + obs_match[1:])
+
+            # Can also start fresh at position 0
+            new_log_M = new_log_M.at[0].set(
+                jax.scipy.special.logsumexp(jnp.array([new_log_M[0], obs_match[0]]))
+            )
+
+            # New I states: can come from M or I at same profile position
+            from_M_to_I = log_M + trans["m_to_i"]
+            from_I_to_I = log_I + trans["i_to_i"]
+            new_log_I = (
+                jax.scipy.special.logsumexp(jnp.stack([from_M_to_I, from_I_to_I]), axis=0)
+                + obs_insert
+            )
+
+            # D states don't emit, handled separately
+            # For simplicity, we skip explicit D state tracking in emissions
+            # and just allow gaps through the transition structure
+            new_log_D = log_D  # Simplified: D states updated via M transitions
+
+            return (new_log_M, new_log_I, new_log_D), None
+
+        # Scan over sequence positions (skip first, handled in init)
+        (final_M, final_I, final_D), _ = jax.lax.scan(
+            forward_step, (log_M, log_I, log_D), jnp.arange(1, seq_len)
+        )
+
+        # Final score: sum over all final states
+        score = jax.scipy.special.logsumexp(jnp.concatenate([final_M, final_I]))
+
+        return score
+
+    def compute_posteriors(
+        self,
+        sequence: Float[Array, "seq_len alphabet_size"],
+    ) -> Float[Array, "seq_len profile_length 3"]:
+        """Compute state posteriors (simplified).
+
+        Returns soft alignment between sequence and profile positions.
+
+        Args:
+            sequence: One-hot encoded sequence.
+
+        Returns:
+            Posterior probabilities for M/I/D states at each position.
+        """
+        seq_len = sequence.shape[0]
+        log_match = self.get_match_emissions()
+
+        # Simplified: just compute match scores as posteriors
+        match_scores = jnp.einsum("sa,pa->sp", sequence, jnp.exp(log_match))
+
+        # Normalize across profile positions for each sequence position
+        posteriors = jax.nn.softmax(match_scores / self._temperature, axis=1)
+
+        # Expand to include I/D placeholder dimensions
+        # Shape: (seq_len, profile_length, 3) where dim 2 is [M, I, D]
+        full_posteriors = jnp.zeros((seq_len, self.profile_length, 3))
+        full_posteriors = full_posteriors.at[:, :, 0].set(posteriors)
+
+        return full_posteriors
+
+    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 profile HMM search to sequence.
+
+        Args:
+            data: Dictionary containing:
+                - "sequence": One-hot encoded sequence (seq_len, 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:
+
+                    - "sequence": Original sequence
+                    - "score": Profile alignment score
+                    - "state_posteriors": Soft state assignments
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        sequence = data["sequence"]
+
+        # Compute alignment score
+        score = self.score_sequence(sequence)
+
+        # Compute state posteriors
+        state_posteriors = self.compute_posteriors(sequence)
+
+        # Build output data
+        transformed_data = {
+            "sequence": sequence,
+            "score": score,
+            "state_posteriors": state_posteriors,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/alignment/scoring.py

@@ -0,0 +1,127 @@
+"""Scoring matrices for sequence alignment.
+
+This module provides pre-defined scoring matrices and utilities for
+creating custom scoring matrices for DNA, RNA, and protein alignment.
+"""
+
+import functools
+import logging
+from typing import NamedTuple
+
+import jax.numpy as jnp
+from jaxtyping import Array, Float
+
+logger = logging.getLogger(__name__)
+
+
+class ScoringMatrix(NamedTuple):
+    """Scoring matrix with metadata.
+
+    Attributes:
+        matrix: The scoring matrix array.
+        alphabet: The alphabet string (e.g., "ACGT" for DNA).
+        name: Optional name for the matrix.
+    """
+
+    matrix: Float[Array, "alphabet alphabet"]
+    alphabet: str
+    name: str = ""
+
+
+def create_dna_scoring_matrix(
+    match: float = 2.0,
+    mismatch: float = -1.0,
+) -> Float[Array, "4 4"]:
+    """Create a simple DNA scoring matrix.
+
+    Args:
+        match: Score for matching nucleotides (diagonal).
+        mismatch: Score for mismatching nucleotides (off-diagonal).
+
+    Returns:
+        4x4 scoring matrix for DNA (A, C, G, T order).
+    """
+    # Create identity matrix scaled by match score
+    identity = jnp.eye(4) * match
+    # Create off-diagonal with mismatch score
+    off_diag = (jnp.ones((4, 4)) - jnp.eye(4)) * mismatch
+    return identity + off_diag
+
+
+def create_rna_scoring_matrix(
+    match: float = 2.0,
+    mismatch: float = -1.0,
+) -> Float[Array, "4 4"]:
+    """Create a simple RNA scoring matrix.
+
+    Args:
+        match: Score for matching nucleotides (diagonal).
+        mismatch: Score for mismatching nucleotides (off-diagonal).
+
+    Returns:
+        4x4 scoring matrix for RNA (A, C, G, U order).
+    """
+    # RNA uses same simple scoring as DNA
+    return create_dna_scoring_matrix(match, mismatch)
+
+
+@functools.cache
+def get_dna_simple() -> Float[Array, "4 4"]:
+    """Get pre-defined DNA scoring matrix (simple match/mismatch).
+
+    Returns:
+        4x4 scoring matrix with match=2.0, mismatch=-1.0.
+    """
+    return create_dna_scoring_matrix(match=2.0, mismatch=-1.0)
+
+
+@functools.cache
+def get_rna_simple() -> Float[Array, "4 4"]:
+    """Get pre-defined RNA scoring matrix (simple match/mismatch).
+
+    Returns:
+        4x4 scoring matrix with match=2.0, mismatch=-1.0.
+    """
+    return create_rna_scoring_matrix(match=2.0, mismatch=-1.0)
+
+
+# BLOSUM62 scoring matrix for proteins (20 amino acids)
+# Standard BLOSUM62 matrix values
+# Amino acid order: A, R, N, D, C, Q, E, G, H, I, L, K, M, F, P, S, T, W, Y, V
+_BLOSUM62_VALUES = [
+    # A   R   N   D   C   Q   E   G   H   I   L   K   M   F   P   S   T   W   Y   V
+    [4, -1, -2, -2, 0, -1, -1, 0, -2, -1, -1, -1, -1, -2, -1, 1, 0, -3, -2, 0],  # A
+    [-1, 5, 0, -2, -3, 1, 0, -2, 0, -3, -2, 2, -1, -3, -2, -1, -1, -3, -2, -3],  # R
+    [-2, 0, 6, 1, -3, 0, 0, 0, 1, -3, -3, 0, -2, -3, -2, 1, 0, -4, -2, -3],  # N
+    [-2, -2, 1, 6, -3, 0, 2, -1, -1, -3, -4, -1, -3, -3, -1, 0, -1, -4, -3, -3],  # D
+    [0, -3, -3, -3, 9, -3, -4, -3, -3, -1, -1, -3, -1, -2, -3, -1, -1, -2, -2, -1],  # C
+    [-1, 1, 0, 0, -3, 5, 2, -2, 0, -3, -2, 1, 0, -3, -1, 0, -1, -2, -1, -2],  # Q
+    [-1, 0, 0, 2, -4, 2, 5, -2, 0, -3, -3, 1, -2, -3, -1, 0, -1, -3, -2, -2],  # E
+    [0, -2, 0, -1, -3, -2, -2, 6, -2, -4, -4, -2, -3, -3, -2, 0, -2, -2, -3, -3],  # G
+    [-2, 0, 1, -1, -3, 0, 0, -2, 8, -3, -3, -1, -2, -1, -2, -1, -2, -2, 2, -3],  # H
+    [-1, -3, -3, -3, -1, -3, -3, -4, -3, 4, 2, -3, 1, 0, -3, -2, -1, -3, -1, 3],  # I
+    [-1, -2, -3, -4, -1, -2, -3, -4, -3, 2, 4, -2, 2, 0, -3, -2, -1, -2, -1, 1],  # L
+    [-1, 2, 0, -1, -3, 1, 1, -2, -1, -3, -2, 5, -1, -3, -1, 0, -1, -3, -2, -2],  # K
+    [-1, -1, -2, -3, -1, 0, -2, -3, -2, 1, 2, -1, 5, 0, -2, -1, -1, -1, -1, 1],  # M
+    [-2, -3, -3, -3, -2, -3, -3, -3, -1, 0, 0, -3, 0, 6, -4, -2, -2, 1, 3, -1],  # F
+    [-1, -2, -2, -1, -3, -1, -1, -2, -2, -3, -3, -1, -2, -4, 7, -1, -1, -4, -3, -2],  # P
+    [1, -1, 1, 0, -1, 0, 0, 0, -1, -2, -2, 0, -1, -2, -1, 4, 1, -3, -2, -2],  # S
+    [0, -1, 0, -1, -1, -1, -1, -2, -2, -1, -1, -1, -1, -2, -1, 1, 5, -2, -2, 0],  # T
+    [-3, -3, -4, -4, -2, -2, -3, -2, -2, -3, -2, -3, -1, 1, -4, -3, -2, 11, 2, -3],  # W
+    [-2, -2, -2, -3, -2, -1, -2, -3, 2, -1, -1, -2, -1, 3, -3, -2, -2, 2, 7, -1],  # Y
+    [0, -3, -3, -3, -1, -2, -2, -3, -3, 3, 1, -2, 1, -1, -2, -2, 0, -3, -1, 4],  # V
+]
+
+
+@functools.cache
+def get_blosum62() -> Float[Array, "20 20"]:
+    """Get BLOSUM62 scoring matrix for protein alignment.
+
+    Returns:
+        20x20 scoring matrix for 20 standard amino acids.
+    """
+    return jnp.array(_BLOSUM62_VALUES, dtype=jnp.float32)
+
+
+# Amino acid alphabet for BLOSUM62
+PROTEIN_ALPHABET = "ARNDCQEGHILKMFPSTWYV"