Upload explainability_engine (1).py

explainability_engine (1).py

@@ -0,0 +1,556 @@
+"""
+explainability_engine.py
+========================
+Extract ALL internal explainability signals from each of the three models.
+No signal is simplified or omitted.
+
+Splice model signals:
+  - probability
+  - conv3 activation norm vector (99,)
+  - mutation-centered activation peak
+  - splice aura distance (donor / acceptor)
+  - counterfactual delta (all alternative bases)
+  - feature ablation response (splice / region / mutation groups)
+  - risk tier classification
+
+V4 model signals:
+  - probability
+  - importance head vector (via conv3 hook — identical architecture)
+  - mutation-centered importance density
+
+Classic model signals:
+  - probability
+  - importance head output (scalar)
+  - region importance (exon / intron)
+  - conv3 activation norm vector (99,)
+"""
+
+from __future__ import annotations
+import logging
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+import torch
+
+from model_loader import (
+    MutationPredictorCNN_v2,
+    MutationPredictorCNN_v4,
+    MutationPredictorClassic,
+    ModelRegistry,
+    encode_for_v2,
+    encode_for_v4,
+    find_mutation_pos,
+    ALL_BASES,
+    MUT_TYPES,
+)
+
+logger = logging.getLogger("mutation_xai.xai")
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Shared helpers
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def _conv3_activation_norm(model: torch.nn.Module, x: torch.Tensor,
+                            forward_fn) -> np.ndarray:
+    """
+    Register a forward hook on model.conv3, run forward_fn(x), return
+    L2-normalised per-position activation norm vector of shape (99,).
+    """
+    activations: dict = {}
+
+    def _hook(module, inp, out):
+        activations["conv3"] = out.detach()
+
+    hook = model.conv3.register_forward_hook(_hook)
+    try:
+        with torch.no_grad():
+            forward_fn(x)
+    finally:
+        hook.remove()
+
+    act = activations.get("conv3")
+    if act is None:
+        return np.zeros(99)
+
+    # act shape: (1, 256, 99)
+    norm = act.squeeze(0).norm(dim=0).numpy()   # (99,)
+    if norm.max() > 0:
+        norm = norm / norm.max()
+    return norm
+
+
+def _gradient_attribution(model: torch.nn.Module, enc: torch.Tensor,
+                            forward_fn_grad) -> np.ndarray:
+    """
+    Compute input-gradient attribution for the sequence portion.
+    Returns normalised per-position attribution of shape (99,).
+    """
+    model.eval()
+    enc_leaf = enc.clone().detach().requires_grad_(True)
+    logit = forward_fn_grad(enc_leaf)
+    model.zero_grad()
+    logit.backward()
+    grad = enc_leaf.grad
+    if grad is None:
+        return np.zeros(99)
+    seq_grad = grad[:1089].view(99, 11)
+    attr = seq_grad.abs().norm(dim=1).detach().numpy()
+    if attr.max() > 0:
+        attr = attr / attr.max()
+    return attr
+
+
+def _mutation_peak_ratio(profile: np.ndarray, mutation_pos: int) -> float:
+    """
+    peak_signal / mean_signal, where peak_signal is the profile value at
+    mutation_pos.  Returns 0.0 if mutation_pos < 0 or mean == 0.
+    """
+    if mutation_pos < 0 or mutation_pos >= len(profile):
+        return 0.0
+    mean_sig = float(profile.mean())
+    if mean_sig == 0:
+        return 0.0
+    return float(profile[mutation_pos]) / mean_sig
+
+
+def _signal_concentration_index(profile: np.ndarray, mutation_pos: int,
+                                  window: int = 10) -> float:
+    """
+    Fraction of total activation energy within ±window of mutation_pos.
+    Ranges 0–1; 1.0 = perfectly concentrated.
+    """
+    if mutation_pos < 0:
+        return 0.0
+    total = float(profile.sum())
+    if total == 0:
+        return 0.0
+    lo = max(0, mutation_pos - window)
+    hi = min(len(profile), mutation_pos + window + 1)
+    local = float(profile[lo:hi].sum())
+    return local / total
+
+
+def _splice_distances(ref_seq: str, mutation_pos: int):
+    """
+    Scan ref_seq for GT (donor) and AG (acceptor) dinucleotides.
+    Returns (dist_donor, dist_acceptor, nearest_donor_pos, nearest_acceptor_pos).
+    Any value may be None if no site found.
+    """
+    seq = (ref_seq.upper() + "N" * 99)[:99]
+    donors, acceptors = [], []
+    for i in range(len(seq) - 1):
+        if seq[i:i+2] == "GT": donors.append(i)
+        if seq[i:i+2] == "AG": acceptors.append(i)
+
+    if mutation_pos < 0:
+        return None, None, None, None
+
+    dist_d = nearest_d = None
+    dist_a = nearest_a = None
+
+    if donors:
+        pairs = sorted([(abs(mutation_pos - p), p) for p in donors])
+        dist_d, nearest_d = pairs[0]
+    if acceptors:
+        pairs = sorted([(abs(mutation_pos - p), p) for p in acceptors])
+        dist_a, nearest_a = pairs[0]
+
+    return dist_d, dist_a, nearest_d, nearest_a
+
+
+def _classify_splice_risk(distance: Optional[int]) -> str:
+    if distance is None:   return "UNKNOWN"
+    if distance <= 2:      return "CRITICAL SPLICE SITE"
+    if distance <= 8:      return "SPLICE REGION"
+    return "NON-SPLICE"
+
+
+def _classify_risk_tier(prob: float) -> tuple[str, str]:
+    if prob >= 0.90: return "PATHOGENIC",          "Very high confidence"
+    if prob >= 0.70: return "LIKELY PATHOGENIC",   "High confidence"
+    if prob >= 0.50: return "POSSIBLY PATHOGENIC", "Moderate confidence"
+    if prob >= 0.20: return "LIKELY BENIGN",        "Low pathogenic signal"
+    return           "BENIGN",                     "Very low pathogenic signal"
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Signal dataclasses
+# ═══════════════════════════════════════════════════════════════════════════════
+
+@dataclass
+class SpliceSignals:
+    probability:              float
+    risk_tier:                str
+    tier_desc:                str
+    conv3_norm:               np.ndarray          # (99,)
+    gradient_attribution:     np.ndarray          # (99,)
+    mutation_pos:             int
+    mutation_peak_ratio:      float
+    signal_concentration:     float
+    imp_score:                float               # importance_head output
+    region_imp:               np.ndarray          # (2,) [exon, intron]
+    splice_imp:               np.ndarray          # (3,) [donor, acc, region]
+    dist_donor:               Optional[int]
+    dist_acceptor:            Optional[int]
+    nearest_donor:            Optional[int]
+    nearest_acceptor:         Optional[int]
+    splice_risk_donor:        str
+    splice_risk_acceptor:     str
+    counterfactual:           dict                # all-base CF results
+    ablation:                 dict                # feature ablation deltas
+    splice_aura_score:        float               # proximity-weighted splice signal
+
+
+@dataclass
+class V4Signals:
+    probability:          float
+    conv3_norm:           np.ndarray    # (99,)
+    gradient_attribution: np.ndarray   # (99,)
+    mutation_pos:         int
+    mutation_peak_ratio:  float
+    signal_concentration: float
+
+
+@dataclass
+class ClassicSignals:
+    probability:      float
+    conv3_norm:       np.ndarray   # (99,)
+    importance_head:  float        # scalar importance_head output
+    region_imp:       np.ndarray   # (2,) [exon, intron]
+    mutation_pos:     int
+    mutation_peak_ratio: float
+    signal_concentration: float
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# ① Extract Splice Signals
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def extract_splice_signals(model: MutationPredictorCNN_v2,
+                            ref_seq: str, mut_seq: str,
+                            exon_flag: int, intron_flag: int) -> SpliceSignals:
+    enc = encode_for_v2(ref_seq, mut_seq, exon_flag, intron_flag)
+
+    # ── base forward pass ────────────────────────────────────────────────────
+    with torch.no_grad():
+        x = enc.unsqueeze(0)
+        logit, imp_t, r_imp_t, s_imp_t = model(x)
+        prob      = float(torch.sigmoid(logit).item())
+        imp_score = float(imp_t.item())
+        region_imp= r_imp_t[0].numpy()
+        splice_imp= s_imp_t[0].numpy()
+
+    tier, tier_desc = _classify_risk_tier(prob)
+    mutation_pos    = find_mutation_pos(ref_seq, mut_seq)
+
+    # ── conv3 activation norm ────────────────────────────────────────────────
+    def _fwd(x_in):
+        return model(x_in.unsqueeze(0))
+
+    conv3_norm = _conv3_activation_norm(
+        model, enc,
+        lambda x: model(x.unsqueeze(0))
+    )
+
+    # ── gradient attribution ─────────────────────────────────────────────────
+    def _fwd_grad(leaf: torch.Tensor):
+        logit_g, _, _, _ = model(leaf.unsqueeze(0))
+        return logit_g
+
+    grad_attr = _gradient_attribution(model, enc, _fwd_grad)
+
+    # ── mutation-peak derived metrics ────────��────────────────────────────────
+    mpr = _mutation_peak_ratio(conv3_norm, mutation_pos)
+    sci = _signal_concentration_index(conv3_norm, mutation_pos)
+
+    # ── splice distances ─────────────────────────────────────────────────────
+    dist_d, dist_a, nearest_d, nearest_a = _splice_distances(ref_seq, mutation_pos)
+    risk_d = _classify_splice_risk(dist_d)
+    risk_a = _classify_splice_risk(dist_a)
+
+    # ── splice aura score — proximity-weighted composite ────────────────────
+    def _proximity_weight(dist):
+        if dist is None: return 0.0
+        if dist <= 2:    return 1.0
+        if dist <= 8:    return 0.5
+        return 0.1
+
+    aura = (
+        _proximity_weight(dist_d) * float(splice_imp[0]) +
+        _proximity_weight(dist_a) * float(splice_imp[1]) +
+        float(splice_imp[2]) * 0.3
+    ) / 1.6   # normalise to ~[0,1]
+    aura = float(np.clip(aura, 0.0, 1.0))
+
+    # ── counterfactual analysis ───────────────────────────────────────────────
+    cf = _counterfactual_splice(model, ref_seq, mut_seq, mutation_pos,
+                                 exon_flag, intron_flag, prob)
+
+    # ── feature ablation ─────────────────────────────────────────────────────
+    abl = _ablation_splice(model, enc, prob)
+
+    return SpliceSignals(
+        probability=prob, risk_tier=tier, tier_desc=tier_desc,
+        conv3_norm=conv3_norm, gradient_attribution=grad_attr,
+        mutation_pos=mutation_pos,
+        mutation_peak_ratio=mpr, signal_concentration=sci,
+        imp_score=imp_score, region_imp=region_imp, splice_imp=splice_imp,
+        dist_donor=dist_d, dist_acceptor=dist_a,
+        nearest_donor=nearest_d, nearest_acceptor=nearest_a,
+        splice_risk_donor=risk_d, splice_risk_acceptor=risk_a,
+        counterfactual=cf, ablation=abl,
+        splice_aura_score=aura,
+    )
+
+
+def _counterfactual_splice(model: MutationPredictorCNN_v2,
+                             ref_seq: str, mut_seq: str,
+                             mutation_pos: int, exon_flag: int,
+                             intron_flag: int, orig_prob: float) -> dict:
+    if mutation_pos < 0 or mutation_pos >= len(ref_seq):
+        return {"error": "mutation position not detected",
+                "original_probability": orig_prob}
+
+    ref_base = ref_seq[mutation_pos].upper()
+    results  = []
+
+    for alt in ALL_BASES:
+        if alt == ref_base:
+            continue
+        alt_mut = ref_seq[:mutation_pos] + alt + ref_seq[mutation_pos+1:]
+        enc_cf  = encode_for_v2(ref_seq, alt_mut, exon_flag, intron_flag)
+        with torch.no_grad():
+            logit_cf, _, _, _ = model(enc_cf.unsqueeze(0))
+            p = float(torch.sigmoid(logit_cf).item())
+        results.append({"mutation": f"{ref_base}>{alt}", "alt_base": alt,
+                         "probability": round(p, 4)})
+
+    all_probs = [r["probability"] for r in results] + [orig_prob]
+    return {
+        "original_probability":   round(orig_prob, 4),
+        "ref_base":               ref_base,
+        "table":                  sorted(results, key=lambda x: x["probability"], reverse=True),
+        "max_probability":        round(max(all_probs), 4),
+        "min_probability":        round(min(all_probs), 4),
+        "probability_range":      round(max(all_probs) - min(all_probs), 4),
+        "counterfactual_delta":   round(abs(max(all_probs) - min(all_probs)), 4),
+    }
+
+
+def _ablation_splice(model: MutationPredictorCNN_v2,
+                      enc: torch.Tensor, prob_base: float) -> dict:
+    def _prob(e):
+        with torch.no_grad():
+            logit, _, _, _ = model(e.unsqueeze(0))
+            return float(torch.sigmoid(logit).item())
+
+    enc_no_splice = enc.clone(); enc_no_splice[1103:1106] = 0.0
+    enc_no_region = enc.clone(); enc_no_region[1101:1103] = 0.0
+    enc_no_mut    = enc.clone(); enc_no_mut[1089:1101]    = 0.0
+    enc_no_seq    = enc.clone(); enc_no_seq[:1089]        = 0.0
+
+    d_splice = round(abs(prob_base - _prob(enc_no_splice)), 4)
+    d_region = round(abs(prob_base - _prob(enc_no_region)), 4)
+    d_mut    = round(abs(prob_base - _prob(enc_no_mut)),    4)
+    d_seq    = round(abs(prob_base - _prob(enc_no_seq)),    4)
+
+    total = d_splice + d_region + d_mut + d_seq
+    def _pct(v): return round(v / total * 100, 1) if total > 0 else 0.0
+
+    return {
+        "baseline_probability":  round(prob_base, 4),
+        "splice_delta":          d_splice, "splice_pct": _pct(d_splice),
+        "region_delta":          d_region, "region_pct": _pct(d_region),
+        "mutation_delta":        d_mut,    "mutation_pct": _pct(d_mut),
+        "sequence_delta":        d_seq,    "sequence_pct": _pct(d_seq),
+        "dominant_feature": max(
+            [("Splice features", d_splice), ("Region flags", d_region),
+             ("Mutation type",   d_mut),    ("Sequence context", d_seq)],
+            key=lambda x: x[1]
+        )[0],
+    }
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# ② Extract V4 Signals
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def extract_v4_signals(model: MutationPredictorCNN_v4,
+                        ref_seq: str, mut_seq: str,
+                        exon_flag: int, intron_flag: int) -> V4Signals:
+    seq_t, mut_oh, region_t, splice_t = encode_for_v4(ref_seq, mut_seq,
+                                                       exon_flag, intron_flag)
+    # ── base forward ─────────────────────────────────────────────────────────
+    with torch.no_grad():
+        logit = model(seq_t, mut_oh, region_t, splice_t)
+        prob  = float(torch.sigmoid(logit).item())
+
+    mutation_pos = find_mutation_pos(ref_seq, mut_seq)
+
+    # ── conv3 activation norm ────────────────────────────────────────────────
+    def _fwd_v4(seq_in):
+        return model(seq_in, mut_oh, region_t, splice_t)
+
+    conv3_norm = _conv3_activation_norm(
+        model, seq_t.squeeze(0),
+        lambda x: model(x.unsqueeze(0), mut_oh, region_t, splice_t)
+    )
+
+    # ── gradient attribution — through sequence tensor only ──────────────────
+    model.eval()
+    seq_leaf = seq_t.clone().detach().requires_grad_(True)
+    logit_g  = model(seq_leaf, mut_oh, region_t, splice_t)
+    model.zero_grad()
+    logit_g.backward()
+    grad = seq_leaf.grad  # (1, 11, 99)
+    if grad is not None:
+        # L2 norm per position across 11 channels
+        grad_attr = grad.squeeze(0).abs().norm(dim=0).numpy()   # (99,)
+        if grad_attr.max() > 0:
+            grad_attr = grad_attr / grad_attr.max()
+    else:
+        grad_attr = np.zeros(99)
+
+    mpr = _mutation_peak_ratio(conv3_norm, mutation_pos)
+    sci = _signal_concentration_index(conv3_norm, mutation_pos)
+
+    return V4Signals(
+        probability=prob,
+        conv3_norm=conv3_norm, gradient_attribution=grad_attr,
+        mutation_pos=mutation_pos,
+        mutation_peak_ratio=mpr, signal_concentration=sci,
+    )
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# ③ Extract Classic Signals
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def extract_classic_signals(model: MutationPredictorClassic,
+                              ref_seq: str, mut_seq: str,
+                              exon_flag: int, intron_flag: int) -> ClassicSignals:
+    enc = encode_for_v2(ref_seq, mut_seq, exon_flag, intron_flag)
+
+    # ── base forward ─────────────────────────────────────────────────────────
+    with torch.no_grad():
+        x = enc.unsqueeze(0)
+        logit, imp_t, r_imp_t = model(x)
+        prob      = float(torch.sigmoid(logit).item())
+        imp_score = float(imp_t.item())
+        region_imp= r_imp_t[0].numpy()
+
+    mutation_pos = find_mutation_pos(ref_seq, mut_seq)
+
+    # ── conv3 activation norm ────────────────────────────────────────────────
+    conv3_norm = _conv3_activation_norm(
+        model, enc,
+        lambda x: model(x.unsqueeze(0))
+    )
+
+    mpr = _mutation_peak_ratio(conv3_norm, mutation_pos)
+    sci = _signal_concentration_index(conv3_norm, mutation_pos)
+
+    return ClassicSignals(
+        probability=prob,
+        conv3_norm=conv3_norm,
+        importance_head=imp_score,
+        region_imp=region_imp,
+        mutation_pos=mutation_pos,
+        mutation_peak_ratio=mpr,
+        signal_concentration=sci,
+    )
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Cross-model analysis
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def compute_cross_model_analysis(splice: SpliceSignals,
+                                  v4: V4Signals,
+                                  classic: ClassicSignals) -> dict:
+    """
+    Compute all five XAI Engine metrics and cross-model locality score.
+    """
+
+    # 1. Mutation Peak Ratio — average across models
+    mpr_avg = float(np.mean([
+        splice.mutation_peak_ratio,
+        v4.mutation_peak_ratio,
+        classic.mutation_peak_ratio,
+    ]))
+
+    # 2. Counterfactual magnitude — from splice model (has full CF data)
+    cf_mag = float(splice.counterfactual.get("counterfactual_delta", 0.0))
+
+    # 3. Cross-model locality score
+    #    Are activation peaks aligned across models?
+    #    Compute correlation of all three conv3_norm profiles.
+    profiles = [splice.conv3_norm, v4.conv3_norm, classic.conv3_norm]
+    cors = []
+    for i in range(len(profiles)):
+        for j in range(i+1, len(profiles)):
+            a, b = profiles[i], profiles[j]
+            if a.std() > 0 and b.std() > 0:
+                cors.append(float(np.corrcoef(a, b)[0, 1]))
+            else:
+                cors.append(0.0)
+    cross_locality = float(np.clip(np.mean(cors), -1.0, 1.0))
+
+    # 4. Signal concentration index — average across models
+    sci_avg = float(np.mean([
+        splice.signal_concentration,
+        v4.signal_concentration,
+        classic.signal_concentration,
+    ]))
+
+    # 5. Explainability Strength Score (0–1)
+    mpr_norm = float(np.clip(mpr_avg / 3.0, 0.0, 1.0))   # >3× peak = full score
+    cf_norm  = float(np.clip(cf_mag,  0.0, 1.0))
+    loc_norm = float(np.clip((cross_locality + 1.0) / 2.0, 0.0, 1.0))
+
+    ess = (0.35 * mpr_norm + 0.35 * cf_norm + 0.30 * loc_norm)
+    ess = float(np.clip(ess, 0.0, 1.0))
+
+    # Activation pattern type
+    peak = float(np.max(splice.conv3_norm))
+    if peak > 0:
+        above_half  = int(np.sum(splice.conv3_norm > 0.5 * peak))
+        above_tenth = int(np.sum(splice.conv3_norm > 0.1 * peak))
+    else:
+        above_half = above_tenth = 0
+
+    if above_half <= 5:
+        pattern = "Sharp"
+    elif above_half <= 25:
+        pattern = "Broad"
+    else:
+        pattern = "Flat"
+
+    # Per-model probability agreement
+    probs = [splice.probability, v4.probability, classic.probability]
+    prob_std = float(np.std(probs))
+
+    return {
+        "mutation_peak_ratio":          round(mpr_avg, 4),
+        "counterfactual_magnitude":     round(cf_mag, 4),
+        "cross_model_locality_score":   round(cross_locality, 4),
+        "signal_concentration_index":   round(sci_avg, 4),
+        "explainability_strength_score": round(ess, 4),
+        "activation_pattern_type":      pattern,
+        "prob_std":                     round(prob_std, 4),
+        "model_agreement":              _agreement_level(prob_std),
+        # raw profiles for plotting
+        "_splice_norm":   splice.conv3_norm,
+        "_v4_norm":       v4.conv3_norm,
+        "_classic_norm":  classic.conv3_norm,
+        "_splice_grad":   splice.gradient_attribution,
+        "_v4_grad":       v4.gradient_attribution,
+    }
+
+
+def _agreement_level(std: float) -> str:
+    if std < 0.05:  return "Strong"
+    if std < 0.12:  return "Moderate"
+    return "Weak"