dee/core/align.py

"""Pairwise sequence alignment — Needleman–Wunsch (global) and Smith–Waterman
(local). numpy + stdlib only; no Biopython dependency so it's trivially testable.

Identity-based scoring (match / mismatch + linear gap) — enough for the
"compare two sequences" use case: visualise percent identity, mismatches, and
gaps for DNA or protein. O(n·m) DP, so we cap |a|·|b| to keep it snappy.
"""
from __future__ import annotations

from typing import Any, Dict

import numpy as np

# |a| * |b| ceiling (~1225 × 1225). Above this we ask the user to align a
# shorter region — the pure-Python traceback + DP fill stays sub-second below it.
MAX_CELLS = 1_500_000


def _clean(s: str) -> str:
    return "".join(c for c in (s or "").upper() if not c.isspace())


def align(a: str, b: str, *, mode: str = "global",
          match: int = 2, mismatch: int = -1, gap: int = -2) -> Dict[str, Any]:
    """Align sequences `a` and `b`.

    mode: 'global' (Needleman–Wunsch) or 'local' (Smith–Waterman).
    Returns aligned strings, a match midline, percent identity, score, gaps.
    """
    a = _clean(a)
    b = _clean(b)
    if not a or not b:
        raise ValueError("Both sequences must be non-empty.")
    if len(a) * len(b) > MAX_CELLS:
        raise ValueError(
            f"Sequences too large to align in-browser (|a|×|b| > {MAX_CELLS:,}). "
            "Align a shorter region."
        )

    local = (mode == "local")
    n, m = len(a), len(b)
    H = np.zeros((n + 1, m + 1), dtype=np.int32)
    # traceback codes: 0 diag, 1 up (gap in b), 2 left (gap in a), 3 stop
    T = np.zeros((n + 1, m + 1), dtype=np.int8)

    if not local:
        H[:, 0] = np.arange(n + 1) * gap
        H[0, :] = np.arange(m + 1) * gap
        T[1:, 0] = 1
        T[0, 1:] = 2

    best_score, best_i, best_j = 0, 0, 0
    for i in range(1, n + 1):
        ai = a[i - 1]
        Hi, Hprev, Ti = H[i], H[i - 1], T[i]
        for j in range(1, m + 1):
            sc = match if ai == b[j - 1] else mismatch
            cell = Hprev[j - 1] + sc
            t = 0
            up = Hprev[j] + gap
            if up > cell:
                cell, t = up, 1
            left = Hi[j - 1] + gap
            if left > cell:
                cell, t = left, 2
            if local and cell < 0:
                cell, t = 0, 3
            Hi[j] = cell
            Ti[j] = t
            if local and cell > best_score:
                best_score, best_i, best_j = cell, i, j

    if local:
        score, i, j = best_score, best_i, best_j
    else:
        score, i, j = int(H[n, m]), n, m

    out_a, out_b = [], []
    while i > 0 or j > 0:
        t = int(T[i, j])
        if local and (H[i, j] == 0 or t == 3):
            break
        if t == 0:
            out_a.append(a[i - 1]); out_b.append(b[j - 1]); i -= 1; j -= 1
        elif t == 1:
            out_a.append(a[i - 1]); out_b.append("-"); i -= 1
        elif t == 2:
            out_a.append("-"); out_b.append(b[j - 1]); j -= 1
        else:
            break
    out_a.reverse(); out_b.reverse()
    aa, bb = "".join(out_a), "".join(out_b)

    cols = len(aa)
    matches = sum(1 for x, y in zip(aa, bb) if x == y and x != "-")
    identity = round(100.0 * matches / cols, 1) if cols else 0.0
    gaps = aa.count("-") + bb.count("-")
    midline = "".join(
        "|" if (x == y and x != "-") else (" " if (x == "-" or y == "-") else ".")
        for x, y in zip(aa, bb)
    )
    return {
        "mode": "local" if local else "global",
        "score": int(score),
        "identity": identity,
        "length": cols,
        "matches": matches,
        "gaps": gaps,
        "aligned_a": aa,
        "aligned_b": bb,
        "midline": midline,
    }