#!/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()