Add Shapley attribution and redundancy/nonlinear validation cases to the failure harness

scripts/run_failure_analysis.py

@@ -129,10 +129,15 @@ def _run_predictions(predictions_path: str, out: str) -> None:
 
     def predict(present):
         if present not in preds:
-            raise SystemExit(f"no scored condition covers the modality subset {sorted(present)}")
+            # KeyError, not SystemExit: the harness degrades gracefully when an optional
+            # subset (the all-masked one Shapley probes) has no scored condition behind it.
+            raise KeyError(f"no scored condition covers the modality subset {sorted(present)}")
         return preds[present]
 
-    report = analyze_modality_failure(["echo", "ecg"], lvef, ef, predict)
+    try:
+        report = analyze_modality_failure(["echo", "ecg"], lvef, ef, predict)
+    except KeyError as e:
+        raise SystemExit(e.args[0]) from None
     _emit(
         report,
         out,

scripts/validate_harness.py

@@ -5,6 +5,9 @@
 (1) Multi-seed recovery (echo/ECG): does the harness reliably recover planted echo-dominance and
     the modality-specific silent-failure asymmetry across seeds?
 (2) N-modality generality (echo/ECG/labs): does it produce a genuine 3x3 complementarity matrix?
+(3) Redundancy stress case: two equally strong modalities plus one weak one -- leave-one-out
+    reads the redundant pair as worthless, Shapley splits the credit between them.
+(4) Nonlinear probe: the two-modality recovery repeated with a small MLP instead of Ridge.
 """
 
 from __future__ import annotations
@@ -18,16 +21,17 @@
     analyze_modality_failure,
     make_synthetic_modalities,
     make_synthetic_multimodal,
+    masked_mlp_predict_fn,
     masked_ridge_predict_fn,
     stratified_train_mask,
 )
 from primed_ai.failure.plots import plot_complementarity_matrix
 
 
-def _analyze(emb, lvef, ef, seed):
+def _analyze(emb, lvef, ef, seed, probe=masked_ridge_predict_fn):
     train = stratified_train_mask(ef, 0.7, seed=seed + 1000)
     test = ~train
-    predict_full = masked_ridge_predict_fn(emb, lvef, train)
+    predict_full = probe(emb, lvef, train)
     return analyze_modality_failure(
         {m: emb[m][test] for m in emb},
         lvef[test],
@@ -92,6 +96,32 @@ def threemodal(out: str, figdir: str) -> dict:
     return d
 
 
+def redundancy(out: str) -> dict:
+    emb, lvef, ef = make_synthetic_modalities({"m1": 2.0, "m2": 2.0, "weak": 0.5}, n=750, seed=0)
+    d = _analyze(emb, lvef, ef, 0).to_dict()
+    c = d["complementarity"]
+    Path(out).mkdir(parents=True, exist_ok=True)
+    (Path(out) / "validation_redundancy.json").write_text(json.dumps(d, indent=2))
+    print("=== redundancy stress case (m1 == m2 strong, weak) ===")
+    print("leave-one-out :", c["marginal_value"])
+    print("shapley       :", c["shapley_value"])
+    return d
+
+
+def nonlinear(out: str) -> dict:
+    emb, lvef, ef, _ = make_synthetic_multimodal(seed=0)
+    d = _analyze(emb, lvef, ef, 0, probe=masked_mlp_predict_fn).to_dict()
+    c = d["complementarity"]
+    Path(out).mkdir(parents=True, exist_ok=True)
+    (Path(out) / "validation_nonlinear.json").write_text(json.dumps(d, indent=2))
+    print("=== nonlinear probe (MLP) on the planted two-modality data ===")
+    print("leave-one-out :", c["marginal_value"])
+    print("shapley       :", c["shapley_value"])
+    print("winners       :", c["per_example_winners"])
+    print("silent rates  :", {k: v["silent_rate"] for k, v in d["dropout"].items()})
+    return d
+
+
 def main() -> None:
     ap = argparse.ArgumentParser(description=__doc__)
     ap.add_argument("--seeds", type=int, default=12)
@@ -100,6 +130,8 @@ def main() -> None:
     a = ap.parse_args()
     multiseed(a.seeds, a.out)
     threemodal(a.out, a.figdir)
+    redundancy(a.out)
+    nonlinear(a.out)
 
 
 if __name__ == "__main__":

src/primed_ai/failure/__init__.py

@@ -14,6 +14,7 @@
 from .demo import (
     make_synthetic_modalities,
     make_synthetic_multimodal,
+    masked_mlp_predict_fn,
     masked_ridge_predict_fn,
     stratified_train_mask,
 )
@@ -25,6 +26,7 @@
     "CATEGORIES",
     "make_synthetic_multimodal",
     "make_synthetic_modalities",
+    "masked_mlp_predict_fn",
     "masked_ridge_predict_fn",
     "stratified_train_mask",
 ]

src/primed_ai/failure/core.py

@@ -6,8 +6,9 @@
 
 1. **Failure taxonomy** -- per-example category (``correct`` / ``imprecise`` / ``critical``)
    under full-modality inference and under each missing-modality condition.
-2. **Complementarity** -- each modality's leave-one-out marginal value, a complementarity
-   matrix over modality subsets, and per-example attribution ("which modality wins").
+2. **Complementarity** -- each modality's leave-one-out marginal value, its exact Shapley
+   value over modality coalitions, a complementarity matrix over modality subsets, and
+   per-example attribution ("which modality wins").
 3. **Dropout profile (loud vs. silent)** -- when a modality is missing at inference, does the
    model fail *loudly* (output shifts / sits near the decision boundary -> monitorable) or
    *silently* (confident, stable-looking output that is now clinically wrong)?
@@ -21,6 +22,7 @@
 
 from dataclasses import dataclass, field
 from itertools import combinations
+from math import factorial
 from typing import Callable
 
 import numpy as np
@@ -124,7 +126,23 @@ def analyze_modality_failure(
         needed.add(allm - {m})
     for a, b in combinations(modalities, 2):
         needed.add(frozenset({a, b}))
-    preds = {s: np.asarray(predict_fn(s), dtype=float) for s in needed}
+    # Exact Shapley attribution needs every coalition, the empty set (all branches
+    # masked) included. 2^N calls are cheap at the N this harness sees; past the cap
+    # we keep leave-one-out attribution only rather than blow up predict_fn calls.
+    shapley_ok = len(modalities) <= 8
+    if shapley_ok:
+        for k in range(len(modalities) + 1):
+            for combo in combinations(modalities, k):
+                needed.add(frozenset(combo))
+    preds = {}
+    for s in needed:
+        try:
+            preds[s] = np.asarray(predict_fn(s), dtype=float)
+        except Exception:
+            if s:
+                raise
+            # predict_fn that cannot mask everything: drop Shapley, keep the rest
+            shapley_ok = False
     full = preds[allm]
     err_full = _abs_err(full, y)
 
@@ -157,11 +175,31 @@ def _condition(pred) -> dict:
             row.append(round(float(_abs_err(preds[subset], y).mean()), 4))
         matrix.append(row)
 
+    # Exact Shapley value of each modality over the coalition game v(S) = -MAE(f(S)):
+    # the weighted mean MAE drop from adding m across all coalitions, so positive phi
+    # is error the modality removes. Unlike leave-one-out, two redundant modalities
+    # split the credit instead of both reading as worthless.
+    shapley = None
+    if shapley_ok:
+        mae_of = {s: float(_abs_err(p, y).mean()) for s, p in preds.items()}
+        nm = len(modalities)
+        shapley = {}
+        for m in modalities:
+            rest = [x for x in modalities if x != m]
+            phi = 0.0
+            for k in range(nm):
+                w = factorial(k) * factorial(nm - k - 1) / factorial(nm)
+                for combo in combinations(rest, k):
+                    s = frozenset(combo)
+                    phi += w * (mae_of[s] - mae_of[s | {m}])
+            shapley[m] = round(phi, 4)
+
     best_solo = min(solo_mae.values())
     complementarity = {
         "full_mae": round(float(err_full.mean()), 4),
         "solo_mae": solo_mae,
         "marginal_value": marginal_value,  # how much each modality is worth (LOO)
+        "shapley_value": shapley,  # coalition-fair attribution; None past the subset cap
         "fusion_gain_vs_best_solo": round(best_solo - float(err_full.mean()), 4),
         "per_example_winners": winners,  # which modality "wins" overall
         "matrix": {"modalities": modalities, "values": matrix},

src/primed_ai/failure/demo.py

@@ -50,14 +50,12 @@ def make_synthetic_multimodal(
     return embeddings, lvef, ef_le_40, planted
 
 
-def masked_ridge_predict_fn(embeddings: dict, lvef, train_mask, alpha: float = 1.0):
-    """Train one Ridge model on the full concatenated embeddings; mask absent modalities at
+def _masked_predict_fn(model, embeddings: dict, lvef, train_mask):
+    """Fit ``model`` on the full concatenated embeddings; mask absent modalities at
     inference by zeroing their columns. This mimics a *single deployed multimodal model*
     evaluated under missing inputs (no per-condition retraining), which is exactly what the
     dropout analysis assumes.
     """
-    from sklearn.linear_model import Ridge
-
     mods = list(embeddings)
     cols, off = {}, 0
     for m in mods:
@@ -67,7 +65,7 @@ def masked_ridge_predict_fn(embeddings: dict, lvef, train_mask, alpha: float = 1
     X = np.concatenate([np.asarray(embeddings[m], dtype=float) for m in mods], axis=1)
     y = np.asarray(lvef, dtype=float)
     train_mask = np.asarray(train_mask, dtype=bool)
-    model = Ridge(alpha=alpha).fit(X[train_mask], y[train_mask])
+    model.fit(X[train_mask], y[train_mask])
 
     def predict_fn(present: frozenset) -> np.ndarray:
         Xm = X.copy()
@@ -79,6 +77,26 @@ def predict_fn(present: frozenset) -> np.ndarray:
     return predict_fn
 
 
+def masked_ridge_predict_fn(embeddings: dict, lvef, train_mask, alpha: float = 1.0):
+    """One Ridge model on the full concatenation, absent modalities zeroed at inference."""
+    from sklearn.linear_model import Ridge
+
+    return _masked_predict_fn(Ridge(alpha=alpha), embeddings, lvef, train_mask)
+
+
+def masked_mlp_predict_fn(embeddings: dict, lvef, train_mask, seed: int = 0):
+    """Nonlinear counterpart of :func:`masked_ridge_predict_fn`: one small MLP trained on the
+    full concatenation. Exists so the harness's recovery checks are not conditional on a
+    linear probe.
+    """
+    from sklearn.neural_network import MLPRegressor
+
+    model = MLPRegressor(
+        hidden_layer_sizes=(64,), max_iter=2000, random_state=seed, early_stopping=True
+    )
+    return _masked_predict_fn(model, embeddings, lvef, train_mask)
+
+
 def make_synthetic_modalities(
     strengths: dict, n: int = 600, dim: int = 16, seed: int = 0, exclusive_frac: float = 0.12
 ):

tests/test_failure.py

@@ -67,6 +67,61 @@ def test_report_is_json_serializable():
     assert d["complementarity"]["matrix"]["modalities"] == ["echo", "ecg"]
 
 
+def test_shapley_efficiency_and_closed_form():
+    # exact Shapley over v(S) = -MAE(f(S)): the values must sum to the full-vs-empty MAE
+    # gap, and at N = 2 each phi has a closed form over the four subset MAEs
+    emb, lvef, ef, _ = make_synthetic_multimodal(n=600, seed=0)
+    train = stratified_train_mask(ef, 0.7, seed=12345)
+    test = ~train
+    pf = masked_ridge_predict_fn(emb, lvef, train)
+    report = analyze_modality_failure(
+        {m: emb[m][test] for m in emb}, lvef[test], ef[test], lambda present: pf(present)[test]
+    )
+    shap = report.complementarity["shapley_value"]
+    y = report.predictions["__labels__"]
+
+    def mae(key):
+        return float(np.abs(report.predictions[key] - y).mean())
+
+    assert sum(shap.values()) == pytest.approx(mae("") - mae("ecg,echo"), abs=1e-3)
+    closed_echo = 0.5 * ((mae("") - mae("echo")) + (mae("ecg") - mae("ecg,echo")))
+    assert shap["echo"] == pytest.approx(closed_echo, abs=1e-3)
+    assert shap["echo"] > shap["ecg"]  # planted dominance survives the attribution change
+
+
+def test_shapley_credits_redundant_modalities_beyond_loo():
+    # two equally strong modalities: leave-one-out under-credits each (the twin covers),
+    # Shapley counts the coalitions where the twin is absent and pays both more
+    from primed_ai.failure import make_synthetic_modalities
+
+    emb, lvef, ef = make_synthetic_modalities({"m1": 2.0, "m2": 2.0, "weak": 0.5}, n=750, seed=0)
+    train = stratified_train_mask(ef, 0.7, seed=1000)
+    test = ~train
+    pf = masked_ridge_predict_fn(emb, lvef, train)
+    c = analyze_modality_failure(
+        {m: emb[m][test] for m in emb}, lvef[test], ef[test], lambda present: pf(present)[test]
+    ).complementarity
+    loo, shap = c["marginal_value"], c["shapley_value"]
+    assert shap["m1"] > loo["m1"] and shap["m2"] > loo["m2"]
+    assert abs(shap["m1"] - shap["m2"]) < 1.0  # symmetric plants get symmetric credit
+    assert shap["weak"] < min(shap["m1"], shap["m2"])
+
+
+def test_shapley_none_when_empty_set_unsupported():
+    rng = np.random.default_rng(0)
+    y = rng.uniform(15, 75, 40)
+    gate = y <= 40
+
+    def predict_fn(present: frozenset) -> np.ndarray:
+        if not present:
+            raise ValueError("cannot mask everything")
+        return y + rng.normal(0, len(present), 40)
+
+    report = analyze_modality_failure(["echo", "ecg"], y, gate, predict_fn)
+    assert report.complementarity["shapley_value"] is None
+    assert report.complementarity["marginal_value"]  # the rest of the report still works
+
+
 def test_classify_taxonomy_categories():
     y = np.array([35.0, 55.0, 50.0])
     pred = np.array(