splicecraft_splice.py

"""splicecraft_splice — plant splice-site scoring (layer L0).

Answers one question: **would a plant's spliceosome plausibly use this GT as a
donor, or this AG as an acceptor?** That matters when designing an intronless
transgene, because a codon-optimized CDS can accidentally spell a strong
splice site, and the plant will then splice a chunk out of the message. The
protein disappears and the DNA looks perfect — one of the nastier silent
failure modes in plant expression.

**Why not a regex.** `GT` and `AG` each occur roughly every 16 bp by chance, so
a pattern match flags essentially the entire CDS and forces meaningless edits.
The dinucleotide is necessary but nowhere near sufficient: real sites carry
information across a ~9 nt (donor) or ~23 nt (acceptor) window. Scoring that
window is what separates "a GT" from "a usable donor".

**The model.** Position weight matrices, scored as a log-odds ratio of two
EMPIRICAL models:

    score = sum_i log2( P_real(base_i at position i) / P_decoy(base_i at i) )

The null model is not uniform base composition. It is real GT/AG dinucleotides
that sit *inside* verified introns and are therefore demonstrably NOT used as
boundaries — exactly the population the scrubber is trying to discriminate
against. A composition-based background would make every GT look informative
and inflate every score.

Positive: a score above 0 means the window looks more like a real site than
like an unused one.

**Clade-matched training matters.** The shipped matrices are trained on
verified introns from a reference plant genome in the same clade as the
intended host. Intron recognition is NOT uniform across plants — some clades
lean hard on intron UA-richness, others tolerate GC-richer introns — so a
matrix trained on a mismatched clade carries the wrong bias into the design
and will mis-rank sites. Retrain with `_splice_train` if a better-matched set
arrives.

**Thresholds are sensitivity-anchored,** not arbitrary. The default is the 5th
percentile of REAL site scores: a candidate scoring above it is as strong as
95% of genuine splice sites. That is an interpretable claim ("this is as good a
donor as a real one") rather than an arbitrary cutoff, and the decoy
false-positive rate at that threshold is recorded alongside it so the flagging
burden is known up front.

Layer L0: imports only logging. Pure functions over strings.
"""
from __future__ import annotations

from math import log2

from splicecraft_logging import _log

# Window geometry, matching the MaxEntScan convention so scores are comparable
# to published work.
#   donor    9 nt = exon[-3:] + intron[:6],  GT at window offsets 3..4
#   acceptor 23 nt = intron[-20:] + exon[:3], AG at window offsets 18..19
_SPLICE_DONOR_LEN = 9
_SPLICE_DONOR_GT = 3
_SPLICE_ACCEPTOR_LEN = 23
_SPLICE_ACCEPTOR_AG = 18

_SPLICE_BASES = "ACGT"

# Pseudocount as a FRACTION of the sample, not a flat count. This matters more
# than it looks. The real and decoy sets differ in size by ~18x, so a flat
# Laplace +1 smooths them by different amounts: at an invariant position (the
# GT / AG anchor, where every sequence in BOTH classes carries the same base) a
# base that never occurs would score
#     log2( (1/N_real) / (1/N_decoy) ) = log2(N_decoy/N_real) ~ +4.2 bits
# — i.e. an IMPOSSIBLE base would outscore the correct one. Smoothing each
# class by the same proportion makes the log-odds at an invariant position
# exactly 0, which is the truthful answer: a position that is constant in both
# populations cannot discriminate between them.
_SPLICE_PSEUDO_FRAC = 0.001

# ── Trained matrices ───────────────────────────────────────────────────────
# Generated by `_splice_train` (see outputs/tools/train_splice_model.py in the
# thaumatin project for the reproducible pipeline). Each is a list of
# per-position {base: log2-odds} dicts. Provenance + calibration live in
# _SPLICE_MODEL_META so a stale matrix is visible rather than silent.
_SPLICE_DONOR_PWM: "list[dict[str, float]]" = []
_SPLICE_ACCEPTOR_PWM: "list[dict[str, float]]" = []
_SPLICE_MODEL_META: dict = {}


def _splice_count_matrix(windows, width: int) -> "list[dict[str, float]]":
    """Pseudocounted base frequencies per position."""
    cols = [{b: 0.0 for b in _SPLICE_BASES} for _ in range(width)]
    used = 0
    for w in windows:
        if len(w) != width:
            continue
        w = w.upper()
        if set(w) - set(_SPLICE_BASES):
            continue
        used += 1
        for i, base in enumerate(w):
            cols[i][base] += 1.0
    if not used:
        raise ValueError(f"no usable training windows of width {width}")
    # Proportional smoothing — see _SPLICE_PSEUDO_FRAC.
    pseudo = max(_SPLICE_PSEUDO_FRAC * used, 1e-9) / len(_SPLICE_BASES)
    for col in cols:
        for base in col:
            col[base] += pseudo
    for col in cols:
        total = sum(col.values())
        for base in col:
            col[base] /= total
    return cols


def _splice_train(real, decoy, width: int) -> "list[dict[str, float]]":
    """Build a per-position log2-odds matrix from real vs decoy windows.

    Both models are estimated from data; the decoy set defines the null. See
    the module docstring for why that beats a composition background."""
    p_real = _splice_count_matrix(real, width)
    p_decoy = _splice_count_matrix(decoy, width)
    return [{b: log2(p_real[i][b] / p_decoy[i][b]) for b in _SPLICE_BASES}
            for i in range(width)]


def _splice_score_window(window: str, pwm: "list[dict[str, float]]") -> "float | None":
    """Log2-odds score for one window. None when it cannot be scored (wrong
    length, or an ambiguity code) — never a silent 0.0, which would read as
    'perfectly neutral site' rather than 'no answer'."""
    if not pwm or len(window) != len(pwm):
        return None
    window = window.upper()
    total = 0.0
    for i, base in enumerate(window):
        col = pwm[i]
        if base not in col:
            return None
        total += col[base]
    return total


def _splice_score_donor(seq: str, gt_pos: int) -> "float | None":
    """Score the donor whose GT begins at `gt_pos` (0-based, forward strand).

    Returns None when there is no `GT` there. The anchor is deliberately NOT
    part of the score: the decoy population is GT-anchored too, so those two
    positions carry zero discriminative information and contribute 0 bits. The
    model's job is to rank GTs against each other, not to decide whether a
    dinucleotide is a GT — so that check lives here, as a precondition. Without
    it a `CAG|ATAAGT` window would score identically to `CAG|GTAAGT` and a
    caller scoring arbitrary offsets would get a confident number for a site
    that cannot splice."""
    if seq[gt_pos:gt_pos + 2].upper() != "GT":
        return None
    start = gt_pos - _SPLICE_DONOR_GT
    if start < 0 or start + _SPLICE_DONOR_LEN > len(seq):
        return None
    return _splice_score_window(seq[start:start + _SPLICE_DONOR_LEN],
                                _SPLICE_DONOR_PWM)


def _splice_score_acceptor(seq: str, ag_pos: int) -> "float | None":
    """Score the acceptor whose AG begins at `ag_pos` (0-based, forward).

    Returns None when there is no `AG` there — same precondition rationale as
    `_splice_score_donor`."""
    if seq[ag_pos:ag_pos + 2].upper() != "AG":
        return None
    start = ag_pos - _SPLICE_ACCEPTOR_AG
    if start < 0 or start + _SPLICE_ACCEPTOR_LEN > len(seq):
        return None
    return _splice_score_window(seq[start:start + _SPLICE_ACCEPTOR_LEN],
                                _SPLICE_ACCEPTOR_PWM)


def _splice_threshold(kind: str) -> "float | None":
    meta = _SPLICE_MODEL_META.get(kind) or {}
    return meta.get("threshold")


def _splice_scan(seq: str, *, kind: str = "both",
                 donor_threshold: "float | None" = None,
                 acceptor_threshold: "float | None" = None,
                 both_strands: bool = False) -> "list[dict]":
    """Every cryptic splice site in `seq` scoring at or above threshold.

    Returns ``[{kind, position, score, threshold, window, strand}]`` sorted by
    descending score. `position` is the 0-based offset of the GT (donor) or AG
    (acceptor) dinucleotide on the FORWARD strand.

    `both_strands` is off by default and should usually stay off: an mRNA is
    single-stranded and only the sense strand is spliced, so a reverse-strand
    "site" in a CDS is not a splicing hazard. It is available for scanning
    regions whose transcribed orientation is unknown or bidirectional.

    Raises RuntimeError when the model has not been trained — a scan that
    silently returned [] would read as 'clean' when it means 'not checked'."""
    if not _SPLICE_DONOR_PWM or not _SPLICE_ACCEPTOR_PWM:
        raise RuntimeError(
            "splice model is not trained — _SPLICE_DONOR_PWM/_SPLICE_ACCEPTOR_PWM "
            "are empty. An untrained scan cannot report 'no sites found'.")
    seq = (seq or "").upper()
    if kind not in ("both", "donor", "acceptor"):
        raise ValueError("kind must be 'both', 'donor' or 'acceptor'")
    if donor_threshold is None:
        donor_threshold = _splice_threshold("donor")
    if acceptor_threshold is None:
        acceptor_threshold = _splice_threshold("acceptor")
    # `_splice_threshold` returns None when the metadata carries no calibrated
    # cutoff. That is a MODEL fault, not a scan input fault, and it would
    # otherwise reach the `score >= threshold` comparisons below and die on
    # `float >= None` — an opaque TypeError from deep inside the loop. Same
    # class as the untrained-PWM guard above: refuse loudly and say why.
    if donor_threshold is None or acceptor_threshold is None:
        raise RuntimeError(
            "splice model has no calibrated threshold — the model metadata is "
            "missing 'threshold' for donor and/or acceptor. Retrain the model, "
            "or pass explicit donor_threshold / acceptor_threshold.")

    strands = [("+", seq)]
    if both_strands:
        comp = str.maketrans("ACGT", "TGCA")
        strands.append(("-", seq.translate(comp)[::-1]))

    hits: "list[dict]" = []
    for strand, s in strands:
        n = len(s)
        for i in range(n - 1):
            pair = s[i:i + 2]
            if kind in ("both", "donor") and pair == "GT":
                score = _splice_score_donor(s, i)
                if score is not None and score >= donor_threshold:
                    pos = i if strand == "+" else n - i - 2
                    hits.append({"kind": "donor", "position": pos,
                                 "score": round(score, 3),
                                 "threshold": donor_threshold,
                                 "window": s[i - _SPLICE_DONOR_GT:
                                             i - _SPLICE_DONOR_GT + _SPLICE_DONOR_LEN],
                                 "strand": strand})
            if kind in ("both", "acceptor") and pair == "AG":
                score = _splice_score_acceptor(s, i)
                if score is not None and score >= acceptor_threshold:
                    pos = i if strand == "+" else n - i - 2
                    hits.append({"kind": "acceptor", "position": pos,
                                 "score": round(score, 3),
                                 "threshold": acceptor_threshold,
                                 "window": s[i - _SPLICE_ACCEPTOR_AG:
                                             i - _SPLICE_ACCEPTOR_AG + _SPLICE_ACCEPTOR_LEN],
                                 "strand": strand})
    hits.sort(key=lambda h: -h["score"])
    return hits


def _splice_pair_risk(seq: str, **kw) -> "list[dict]":
    """Donor/acceptor PAIRS that could excise a real cryptic intron.

    A lone strong donor is much less dangerous than a donor followed, at a
    plausible intron distance, by a strong acceptor — that pair is what
    actually removes sequence from the message. Plant introns are short
    (training-set median ~150 nt), so the window checked is 60–3000 nt.

    Returns ``[{donor, acceptor, intron_len, combined_score, in_frame_loss}]``
    sorted by combined score. `in_frame_loss` flags an excision that is a
    multiple of 3: those are the truly insidious ones, since they delete
    residues without frameshifting and the product may still be detected by an
    antibody while being biologically wrong."""
    hits = _splice_scan(seq, **kw)
    donors = sorted([h for h in hits if h["kind"] == "donor"],
                    key=lambda h: h["position"])
    acceptors = sorted([h for h in hits if h["kind"] == "acceptor"],
                       key=lambda h: h["position"])
    pairs = []
    for d in donors:
        for a in acceptors:
            span = a["position"] + 2 - d["position"]
            if 60 <= span <= 3000:
                pairs.append({
                    "donor": d, "acceptor": a, "intron_len": span,
                    "combined_score": round(d["score"] + a["score"], 3),
                    "in_frame_loss": span % 3 == 0,
                })
    pairs.sort(key=lambda p: -p["combined_score"])
    return pairs


def _splice_model_summary() -> dict:
    """What model is loaded, and how it was calibrated."""
    return {
        "trained": bool(_SPLICE_DONOR_PWM and _SPLICE_ACCEPTOR_PWM),
        "donor_width": len(_SPLICE_DONOR_PWM),
        "acceptor_width": len(_SPLICE_ACCEPTOR_PWM),
        **_SPLICE_MODEL_META,
    }


def _splice_load_model(donor_pwm, acceptor_pwm, meta) -> None:
    """Install a trained model (used by the vendored block below and by tests
    that want a deliberately tiny model)."""
    global _SPLICE_DONOR_PWM, _SPLICE_ACCEPTOR_PWM, _SPLICE_MODEL_META
    _SPLICE_DONOR_PWM = donor_pwm
    _SPLICE_ACCEPTOR_PWM = acceptor_pwm
    _SPLICE_MODEL_META = meta
    _log.info("splice model loaded: donor %d nt, acceptor %d nt, source %s",
              len(donor_pwm), len(acceptor_pwm), meta.get("source", "?"))


def _splice_consensus(pwm: "list[dict[str, float]]", *,
                      min_bits: float = 0.15) -> str:
    """Consensus encoded by the matrix — uppercase where the position carries
    real discriminative signal, `.` where it does not.

    **The GT / AG anchor positions read as `.`, and that is correct.** The
    decoy set is itself GT- and AG-anchored (unused GT/AG dinucleotides inside
    real introns), so at those positions both classes are constant, the
    log-odds is ~0 for every base, and `max()` would return an arbitrary
    letter. Printing that letter as if it were a consensus is actively
    misleading — it looks like the model failed to learn `GT` when in fact
    `GT` carries no information *relative to the population being
    discriminated against*.

    What the model must get right is the INFORMATIVE positions. For a correctly
    trained donor that is `CAG` on the exon side and `AAGT` at intron
    +3..+6 (the `CAG|GTAAGT` consensus); for the acceptor, a long pyrimidine
    tract and `G` at exon +1. Those are what `test_splice.py` asserts."""
    out = []
    for col in pwm:
        best = max(col, key=lambda b: col[b])
        spread = col[best] - min(col.values())
        out.append(best if spread >= min_bits else ".")
    return "".join(out)


def _splice_position_bits(pwm: "list[dict[str, float]]") -> "list[float]":
    """Per-position discriminative range (max - min log-odds), in bits.

    Near-zero identifies the anchor positions; the profile as a whole is what
    to inspect when a retrained model behaves oddly."""
    return [round(max(c.values()) - min(c.values()), 3) for c in pwm]


# ── Load the vendored model at import. Optional by design: the scoring code is
# useful with a caller-supplied model (tests, a retrained species-specific
# matrix), and a missing vendored block must not make the module unimportable.
# `_splice_scan` raises rather than returning [] when nothing is loaded, so an
# absent model can never be mistaken for a clean result.
try:                                    # pragma: no cover - import-time wiring
    from splicecraft_splice_model import (
        ACCEPTOR_PWM as _VENDORED_ACCEPTOR,
        DONOR_PWM as _VENDORED_DONOR,
        SPLICE_MODEL_META as _VENDORED_META,
    )
except ImportError:                     # pragma: no cover
    _log.warning("no vendored splice model found — _splice_scan will raise "
                 "until a model is loaded via _splice_load_model")
else:                                   # pragma: no cover
    _splice_load_model(_VENDORED_DONOR, _VENDORED_ACCEPTOR, _VENDORED_META)