Blame for run_trace_benchmark.py

#!/usr/bin/env python3
"""TRACE: regulatory-guided cubic-Hermite temporal CFM benchmark.

Input NPZ must contain:
  X         cells x genes expression matrix
  t         biological time for each cell
  var_names gene names aligned to X columns
  used_scale (optional bool; if true, retained-only gene scaling is applied before PCA)

For each interior holdout h:
  - exclude h from PCA/scalers/activity map/OT/path/flow training;
  - source is the previous retained time;
  - compare linear CFM, PCHIP cubic CFM, TRACE, and matched permutations;
  - save held-out population metrics.

This script is dataset-generic: it does not hard-code Gupta days.
"""
from __future__ import annotations

import argparse
import csv
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import ot
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split

import benchmark_table as BT
from biochime_regulon_core import (
    RegulonScorer,
    fit_regulon_bio_flow,
    fit_ridge_map,
    permute_tf_channels,
    permute_within_day,
)


def floats(s):
    if s is None or not str(s).strip():
        return []
    return [float(x) for x in str(s).split(",") if x.strip()]


def sub(A, B, mx, seed):
    rng = np.random.default_rng(seed)
    if len(A) > mx:
        A = A[rng.choice(len(A), mx, replace=False)]
    if len(B) > mx:
        B = B[rng.choice(len(B), mx, replace=False)]
    return A, B


def w1(A, B, mx=1500, seed=0):
    A, B = sub(A, B, mx, seed)
    C = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(-1))
    return float(max(ot.emd2(ot.unif(len(A)), ot.unif(len(B)), C, numItermax=2_000_000), 0.0))


def w2(A, B, mx=1500, seed=0):
    A, B = sub(A, B, mx, seed)
    C = ((A[:, None, :] - B[None, :, :]) ** 2).sum(-1)
    return float(np.sqrt(max(ot.emd2(ot.unif(len(A)), ot.unif(len(B)), C, numItermax=2_000_000), 0.0)))


def mmd(A, B, mx=1000, seed=0):
    A, B = sub(A, B, mx, seed)
    Z = np.vstack([A, B])
    D2 = ((Z[:, None, :] - Z[None, :, :]) ** 2).sum(-1)
    med = np.median(D2[D2 > 0]) if np.any(D2 > 0) else 1.0
    gamma = 1.0 / (med + 1e-9)
    K = np.exp(-gamma * D2)
    n = len(A)
    return float(K[:n, :n].mean() + K[n:, n:].mean() - 2 * K[:n, n:].mean())


def mean_rmse(A, B):
    return float(np.sqrt(np.mean((A.mean(0) - B.mean(0)) ** 2)))


def write_rows(path, rows):
    if not rows:
        return
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]))
        w.writeheader(); w.writerows(rows)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mat", required=True)
    ap.add_argument("--prior", required=True)
    ap.add_argument("--holdouts", default="", help="comma list; default=all interior timepoints")
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--steps", type=int, default=4000)
    ap.add_argument("--ncomp", type=int, default=30)
    ap.add_argument("--batch", type=int, default=128)
    ap.add_argument("--ridge", type=float, default=1.0, help="ridge for TF activity -> latent state map")
    ap.add_argument("--min_targets", type=int, default=5)
    ap.add_argument("--min_coverage", type=float, default=0.15)
    ap.add_argument("--max_tfs", type=int, default=0)
    ap.add_argument("--candidate_source", choices=["taiji", "all"], default="all")
    ap.add_argument("--edge_weight_mode", choices=["sign", "sign_resources", "weight"], default="sign",
                    help="Paper default: sign-only TF-target edges")
    ap.add_argument("--rho", type=float, default=1.0)
    ap.add_argument("--tag", default="trace")
    args = ap.parse_args()
    if not 0.0 <= args.rho <= 1.0:
        ap.error("--rho must be between 0 and 1")
    Path(args.tag).parent.mkdir(parents=True, exist_ok=True)

    dev = "cuda" if torch.cuda.is_available() else "cpu"
    M = np.load(args.mat, allow_pickle=True)
    for k in ["X", "t", "var_names"]:
        if k not in M:
            raise ValueError(f"{args.mat} lacks required key '{k}'")

    X = M["X"].astype(np.float64)
    t = M["t"].astype(float)
    var_names = np.asarray([str(x) for x in M["var_names"]])
    use_scale = bool(M["used_scale"]) if "used_scale" in M else False
    days = np.sort(np.unique(t))
    if len(days) < 3:
        raise ValueError(f"Need >=3 temporal marginals, found {list(days)}")

    requested = floats(args.holdouts)
    holdouts = requested if requested else [float(x) for x in days[1:-1]]
    for h in holdouts:
        if not np.any(np.isclose(days, h)) or np.isclose(h, days.min()) or np.isclose(h, days.max()):
            raise ValueError(f"holdout {h} must be an interior observed time; days={list(days)}")

    tn = {float(d): float((d - days.min()) / (days.max() - days.min())) for d in days}
    print(
        f"device={dev} X={X.shape} days={list(days)} holdouts={holdouts} "
        f"folds={args.folds} steps={args.steps}", flush=True
    )

    records = []
    coverage_written = False

    for h in holdouts:
        h = float(h)
        kept = [float(d) for d in days if not np.isclose(d, h)]
        lo = max(d for d in kept if d < h)
        kept_idx = np.where(np.isin(np.round(t, 8), np.round(kept, 8)))[0]

        # Retained-only expression scaling/PCA.
        Xs = X.copy()
        gm = np.zeros(X.shape[1], dtype=np.float64)
        gs = np.ones(X.shape[1], dtype=np.float64)
        if use_scale:
            gm = X[kept_idx].mean(0)
            gs = X[kept_idx].std(0)
            gs[gs < 1e-8] = 1.0
            Xs = np.clip((X - gm) / gs, -10.0, 10.0)

        nc = min(args.ncomp, Xs.shape[1], len(kept_idx) - 1)
        pca = PCA(n_components=nc, svd_solver="full").fit(Xs[kept_idx])
        Zall = pca.transform(Xs).astype(np.float32)
        byday_z = {float(d): Zall[np.isclose(t, d)] for d in days}
        byday_x = {float(d): X[np.isclose(t, d)] for d in days}
        target_z = byday_z[h]
        source_z = byday_z[lo]

        print(f"\n[holdout {h:g}] retained={kept} source={lo:g} target_n={len(target_z)}", flush=True)

        for fold in range(args.folds):
            tr_idx = {}
            for d in kept:
                idx = np.where(np.isclose(t, d))[0]
                aa, _ = train_test_split(idx, test_size=0.2, random_state=fold, shuffle=True)
                tr_idx[d] = aa

            # Latent scaling on retained training cells only.
            zpool = np.vstack([Zall[tr_idx[d]] for d in kept])
            zmu, zsd = zpool.mean(0), zpool.std(0) + 1e-8
            zscale = lambda U: ((U - zmu) / zsd).astype(np.float32)
            zunscale = lambda U: U * zsd + zmu
            cl_z = {tn[d]: zscale(Zall[tr_idx[d]]) for d in kept}

            # TF activities; only the scaler is fitted from retained training cells.
            scorer = RegulonScorer(
                args.prior, var_names,
                min_targets=args.min_targets,
                min_coverage=args.min_coverage,
                candidate_source=args.candidate_source,
                max_tfs=args.max_tfs,
                edge_weight_mode=args.edge_weight_mode,
            )
            xpool = np.vstack([X[tr_idx[d]] for d in kept])
            scorer.fit_gene_scaler(xpool)
            raw_a = {d: scorer.score(X[tr_idx[d]]) for d in kept}
            apool = np.vstack([raw_a[d] for d in kept])
            amu, asd = apool.mean(0), apool.std(0)
            asd[asd < 1e-8] = 1.0
            cl_a = {tn[d]: ((raw_a[d] - amu) / asd).astype(np.float32) for d in kept}

            if not coverage_written:
                covpath = f"{args.tag}_tf_coverage.csv"
                write_rows(covpath, scorer.coverage_rows())
                coverage_written = True
                print(f"[TF prior] {len(scorer.tf_names)} regulons -> {covpath}")
                print("[TFs] " + ", ".join(scorer.tf_names))

            # Heldout TF activities are evaluation-only.
            real_a = ((scorer.score(byday_x[h]) - amu) / asd).astype(np.float32)

            map_times = sorted(cl_z)
            Zmap = np.vstack([cl_z[q] for q in map_times])
            Amap = np.vstack([cl_a[q] for q in map_times])
            # TRACE maps activity derivatives to PCA space: Z ~= A@W+b.
            # Freeze it on the TRUE activities so cell-/TF-perm controls change only the regulatory
            # information, never a re-fitted map (a re-fit would undo a global TF permutation exactly).
            W_bio_true, b_bio_true = fit_ridge_map(Amap, Zmap, lam=args.ridge)

            # Same frozen maps for true and controls: only regulatory information is changed.
            a_cellperm = permute_within_day(cl_a, 10000 + fold)
            a_tfperm = permute_tf_channels(cl_a, 20000 + fold)

            methods = [
                ("Linear_CFM", "linear", None),
                ("Cubic_CFM", "pchip", None),
                ("TRACE", "trace", (cl_a, None)),
                ("TRACE_cellperm", "trace", (a_cellperm, (W_bio_true, b_bio_true))),
                ("TRACE_tfperm", "trace", (a_tfperm, (W_bio_true, b_bio_true))),
            ]

            for name, kind, payload in methods:
                if kind in {"linear", "pchip"}:
                    field = BT.fit_flow(cl_z, sorted(cl_z), kind, dev, args.steps, fold, batch=args.batch)
                else:
                    bio_clouds, bio_map = payload
                    field, _ = fit_regulon_bio_flow(
                        cl_z, bio_clouds, sorted(cl_z), dev, args.steps, fold,
                        rho=args.rho, batch=args.batch, ridge=args.ridge,
                        ridge_map=bio_map, shape_project=True,
                    )

                pred_scaled = BT.ode(
                    field,
                    torch.tensor(zscale(source_z), dtype=torch.float32, device=dev),
                    tn[lo], tn[h], dev, N=50,
                )
                pred_z = zunscale(pred_scaled)

                pred_xs = pca.inverse_transform(pred_z)
                pred_x = pred_xs * gs + gm if use_scale else pred_xs
                pred_a = ((scorer.score(pred_x) - amu) / asd).astype(np.float32)

                rec = {
                    "holdout": h,
                    "fold": fold,
                    "method": name,
                    "W1": w1(pred_z, target_z, seed=fold),
                    "W2": w2(pred_z, target_z, seed=fold),
                    "MMD": mmd(pred_z, target_z, seed=fold),
                    "TF_W2": w2(pred_a, real_a, seed=fold),
                    "TF_mean_RMSE": mean_rmse(pred_a, real_a),
                    "n_tfs": len(scorer.tf_names),
                }
                records.append(rec)

                print(
                    f"  fold {fold} {name:24s} W2={rec['W2']:.3f} MMD={rec['MMD']:.4f}",
                    flush=True,
                )

    results_path = f"{args.tag}_results.csv"
    write_rows(results_path, records)

    groups = defaultdict(list)
    for r in records:
        groups[(r["holdout"], r["method"])].append(r)

    summary = []
    print("\n================ TRACE SUMMARY ================")
    for h in holdouts:
        print(f"\nHeld-out time {h:g}")
        names = []
        for r in records:
            if np.isclose(r["holdout"], h) and r["method"] not in names:
                names.append(r["method"])
        for name in names:
            rr = groups[(h, name)]
            row = {"holdout": h, "method": name, "n_folds": len(rr)}
            for k in ["W1", "W2", "MMD", "TF_W2", "TF_mean_RMSE"]:
                v = np.asarray([q[k] for q in rr], dtype=float)
                row[k + "_mean"] = float(v.mean())
                row[k + "_sd"] = float(v.std())
            summary.append(row)
            print(
                f"{name:24s} W2 {row['W2_mean']:.3f}±{row['W2_sd']:.3f} | "
                f"MMD {row['MMD_mean']:.4f}±{row['MMD_sd']:.4f} | "
                f"TF_W2 {row['TF_W2_mean']:.3f}±{row['TF_W2_sd']:.3f}"
            )

    summary_path = f"{args.tag}_summary.csv"
    write_rows(summary_path, summary)
    print("===================================================")
    print(f"[saved] {results_path}")
    print(f"[saved] {summary_path}")
    print("[configuration] TRACE default: 4000 steps; rho=1.0; ridge=1.0.")


if __name__ == "__main__":
    main()