analyses/deep/15_calibration_fix.py

#!/usr/bin/env python
"""Fix uncertainty calibration via post-hoc temperature scaling.

The raw posterior has 27% coverage for 95% CI. This script:
1. Learns a temperature T on synthetic data
2. Scales posterior variance by T^2
3. Shows improved calibration on held-out synthetic + real data
"""
from _common import *
from scptr.deep.synthetic import generate_kinetic_data
from scipy.optimize import minimize_scalar

OUT = output_dir("15_calibration_fix")


def compute_coverage(gamma_true, gamma_mean, gamma_var, level=0.95, temperature=1.0):
    """Compute CI coverage at given temperature."""
    z = stats.norm.ppf(0.5 + level / 2)
    std = np.sqrt(np.clip(gamma_var * temperature**2, 1e-10, None))
    inside = (gamma_true >= gamma_mean - z * std) & (gamma_true <= gamma_mean + z * std)
    return float(inside.mean())


def find_temperature(gamma_true, gamma_mean, gamma_var, target=0.95):
    """Find temperature that gives target coverage."""
    def loss(log_t):
        t = np.exp(log_t)
        cov = compute_coverage(gamma_true, gamma_mean, gamma_var, target, t)
        return (cov - target) ** 2

    result = minimize_scalar(loss, bounds=(-2, 5), method="bounded")
    return np.exp(result.x)


def main():
    set_figure_style()

    # ── Learn temperature on synthetic data ───────────────────────────
    print("=" * 60)
    print("Learning calibration temperature on synthetic data")
    print("=" * 60)

    adata_train, truth_train = generate_kinetic_data(n_cells=1500, n_genes=100, seed=0)
    adata_test, truth_test = generate_kinetic_data(n_cells=1500, n_genes=100, seed=42)

    torch.set_num_threads(4)

    # Fit on training synthetic
    scptr.deep.fit_deepptr(adata_train, verbose=False, **DEEP_HP)
    # Fit on test synthetic (separate model)
    scptr.deep.fit_deepptr(adata_test, verbose=False, **{**DEEP_HP, "seed": 42})

    # Learn T on training set
    T = find_temperature(truth_train["gamma"], adata_train.layers["gamma"],
                          adata_train.layers["gamma_var"])
    print(f"  Learned temperature: T = {T:.4f}")

    # Evaluate on test set (held-out)
    cov_raw = compute_coverage(truth_test["gamma"], adata_test.layers["gamma"],
                                adata_test.layers["gamma_var"])
    cov_cal = compute_coverage(truth_test["gamma"], adata_test.layers["gamma"],
                                adata_test.layers["gamma_var"], temperature=T)
    print(f"  Test coverage (raw):        {cov_raw:.4f}")
    print(f"  Test coverage (calibrated): {cov_cal:.4f}")

    # Full calibration curve
    levels = [0.50, 0.60, 0.70, 0.80, 0.90, 0.95, 0.99]
    raw_covs = [compute_coverage(truth_test["gamma"], adata_test.layers["gamma"],
                                  adata_test.layers["gamma_var"], l) for l in levels]
    cal_covs = [compute_coverage(truth_test["gamma"], adata_test.layers["gamma"],
                                  adata_test.layers["gamma_var"], l, T) for l in levels]

    print(f"\n  {'Level':>8} {'Raw':>8} {'Calibrated':>12}")
    for l, r, c in zip(levels, raw_covs, cal_covs):
        print(f"  {l:>8.2f} {r:>8.4f} {c:>12.4f}")

    # ── Apply to real data ────────────────────────────────────────────
    print(f"\n{'=' * 60}")
    print("Applying calibration to pancreas")
    print("=" * 60)

    adata_real = scptr.datasets.pancreas()
    scptr.pp.filter_genes(adata_real)
    scptr.pp.normalize_layers(adata_real)
    scptr.pp.neighbors(adata_real, n_neighbors=30)
    scptr.pp.smooth_layers(adata_real)
    scptr.tl.estimate_beta(adata_real)
    adata_real = select_top_genes(adata_real, n_top=300)
    from scipy.sparse import issparse
    for key in ("spliced", "unspliced"):
        if key in adata_real.layers and issparse(adata_real.layers[key]):
            adata_real.layers[key] = np.asarray(adata_real.layers[key].todense())

    torch.set_num_threads(4)
    scptr.deep.fit_deepptr(adata_real, verbose=False, **DEEP_HP)

    # Calibrated variance
    gamma_var_cal = adata_real.layers["gamma_var"] * T**2
    adata_real.layers["gamma_var_calibrated"] = gamma_var_cal

    # Show calibrated uncertainty is more useful for gene filtering
    _, hl_human = load_halflife_refs()
    gamma_med = np.median(adata_real.layers["gamma"], axis=0)
    gamma_cv_raw = np.sqrt(np.median(adata_real.layers["gamma_var"], axis=0)) / (gamma_med + 1e-8)
    gamma_cv_cal = np.sqrt(np.median(gamma_var_cal, axis=0)) / (gamma_med + 1e-8)

    g, h, names = match_halflife(adata_real, hl_human)
    name_to_idx = {n: i for i, n in enumerate(adata_real.var_names)}
    cv_matched = np.array([gamma_cv_cal[name_to_idx[n]] for n in names])

    r_all, _ = stats.spearmanr(g, h)
    # Filter by calibrated CV
    mask_50 = cv_matched <= np.percentile(cv_matched, 50)
    mask_25 = cv_matched <= np.percentile(cv_matched, 25)
    r_50, _ = stats.spearmanr(g[mask_50], h[mask_50]) if mask_50.sum() > 10 else (np.nan, None)
    r_25, _ = stats.spearmanr(g[mask_25], h[mask_25]) if mask_25.sum() > 10 else (np.nan, None)

    print(f"  Calibrated uncertainty filtering:")
    print(f"    All genes:      r={r_all:.4f} (n={len(g)})")
    print(f"    Bottom 50% CV:  r={r_50:.4f} (n={mask_50.sum()})")
    print(f"    Bottom 25% CV:  r={r_25:.4f} (n={mask_25.sum()})")

    results = {
        "temperature": float(T),
        "test_coverage_raw": cov_raw,
        "test_coverage_calibrated": cov_cal,
        "calibration_curve": [{"level": l, "raw": r, "calibrated": c}
                               for l, r, c in zip(levels, raw_covs, cal_covs)],
        "real_filtering": {
            "all": {"r": float(r_all), "n": len(g)},
            "bottom_50pct": {"r": float(r_50), "n": int(mask_50.sum())},
            "bottom_25pct": {"r": float(r_25), "n": int(mask_25.sum())},
        },
    }
    save_json(results, "calibration_fix", OUT)

    # Figure
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    # Calibration curve
    axes[0].plot(levels, raw_covs, "o-", label="Raw", color="gray")
    axes[0].plot(levels, cal_covs, "o-", label=f"Calibrated (T={T:.2f})", color="darkorange")
    axes[0].plot([0, 1], [0, 1], "k--", alpha=0.3, label="Perfect")
    axes[0].set_xlabel("Nominal coverage")
    axes[0].set_ylabel("Actual coverage")
    axes[0].set_title("CI Calibration (held-out synthetic)")
    axes[0].legend()

    # Uncertainty filtering improvement
    bars = axes[1].bar(["All", "Bottom\n50% CV", "Bottom\n25% CV"],
                       [abs(r_all), abs(r_50), abs(r_25)],
                       color=["gray", "steelblue", "darkorange"], alpha=0.7)
    axes[1].set_ylabel("|r| with half-life")
    axes[1].set_title("Calibrated uncertainty filtering (pancreas)")

    fig.suptitle(f"Post-hoc calibration (T={T:.2f})", y=1.02)
    fig.tight_layout()
    save_fig(fig, "calibration_fix", OUT)


if __name__ == "__main__":
    main()