Initial TRACE implementation

.gitignore

@@ -0,0 +1,13 @@
+__pycache__/
+*.py[cod]
+.venv/
+*.log
+*.tmp
+.DS_Store
+.RData
+.Rhistory
+outputs/
+results/
+checkpoints/
+inputs/*.npz
+inputs/*.csv

README.md

@@ -1,15 +1,77 @@
 ---
 license: cc-by-nc-nd-4.0
+tags:
+- single-cell
+- temporal-dynamics
+- flow-matching
+- transcriptomics
 ---
 
-# TRACE
+# TRACE — Transcriptional Regulator-Modulated Reference Geometry
 
-**Transcriptional Regulator-Modulated Reference Geometry for Temporal Single-Cell Flow Matching**
+**Research code accompanying the TRACE ICLR 2027 submission.**
 
-TRACE incorporates transcription-factor regulon-derived directional information into cubic-Hermite reference paths for conditional flow matching.
+TRACE (formerly BioCHIME in the development scripts) modulates the reference path of temporal single-cell conditional flow matching using transcription-factor regulon derivatives. It first couples retained population snapshots with balanced optimal transport, constructs PCHIP/Hermite reference paths, maps regulon activity derivatives into PCA coordinates with a ridge fit \(Z pprox AW+b\), calibrates their norms, and blends the mapped tangents with geometric PCHIP tangents before shape-preserving projection.
 
-The method is evaluated on temporal single-cell RNA-seq datasets using held-out population reconstruction.
+This is an **initial code-only release**. Input matrices, frozen TF-target prior, pretrained weights, and third-party benchmark implementations are not included. The code has been distilled from the original research scripts; the reorganized runner has **not yet been checked for numerical equivalence against archived paper runs**.
 
-## Implementation
+## Submitted TRACE configuration
 
-Code and reproducibility instructions will be added incrementally.
+| Parameter | Value |
+| --- | --- |
+| Optimization steps | **4,000** |
+| TF steering coefficient \(
+ho\) | **1.0** |
+| Ridge regularization \(\lambda\) | 1.0 |
+| Latent PCA dimensions | 30 |
+| Hidden layers / width | 3 / 256 |
+| Activation / time encoding | SiLU / 8 Fourier frequencies |
+| Learning rate / optimizer | 0.001 / Adam |
+| Batch size | 128 |
+| Held-out repeats | 5 (seeds 0–4) |
+| ODE integration | 50 Euler steps |
+| OT | Balanced EMD, squared Euclidean cost | 
+
+## Files
+
+- `biochime_regulon_core.py` — method core implementing TRACE
+- `benchmark_table.py` — original neural field, OT, reference path and Euler integration helpers.
+- `run_trace_benchmark.py` — simplified runner for linear CFM, PCHIP cubic CFM, TRACE, and paired regulatory permutations.
+- `tests/test_trace_core.py` — numerical unit tests for spline, shape-projection and ridge helpers.
+
+## Install and test
+
+```bash
+python -m pip install -r requirements.txt
+python -m unittest discover -s tests -v
+```
+
+## Run a Gupta day-7 holdout
+
+Provide the study-specific input files (not distributed with this initial release):
+
+- `inputs/gupta_hvg_matrix.npz` containing `X` (cells × genes), `t` (biological time), `var_names` (gene names), and optionally `used_scale`.
+- `inputs/tcell_regulatory_prior_frozen29.csv` with signed TF-target edges; see `RegulonScorer` for required columns and coverage rules.
+
+```bash
+python run_trace_benchmark.py \
+  --mat inputs/gupta_hvg_matrix.npz \
+  --prior inputs/tcell_regulatory_prior_frozen29.csv \
+  --holdouts 7 \
+  --steps 4000 --rho 1.0 \
+  --folds 5 --ncomp 30 --batch 128 --ridge 1.0 \
+  --candidate_source all --edge_weight_mode sign \
+  --tag results/gupta_d7
+```
+
+The runner produces fold-level results, mean ± population-SD summaries, and TF coverage tables. For external experiments, keep `--rho 1.0` fixed and provide Giles/DSM matrices with holdouts `8,15` and `8,28`, respectively.
+
+## Reproducibility limitations
+
+- The original uploaded `src/` archive did not include the data preparation / NPZ construction scripts or frozen signed prior, so the exact 29-TF panel and initial gene selection cannot yet be independently recreated from this package.
+- `benchmark_table.ot_plan()` attempts exact EMD but has an approximate fallback if EMD fails. Confirm whether the fallback was triggered in archived paper runs before claiming exact reproduction.
+- This release does not yet reproduce all manuscript components (Slingshot curvature, external method baselines, interpolation-family ablations, or saved checkpoints).
+
+## Review and licensing
+
+Keep this lab-owned repository **private during double-blind review**, and do not link it as anonymous code. The present Hub metadata retains the repository's CC BY-NC-ND 4.0 designation, which restricts derivative works and may not be the desired license for software; confirm the final code license with the lab before public release.

benchmark_table.py

@@ -0,0 +1,1487 @@
+#!/usr/bin/env python3
+"""
+benchmark_table.py
+
+Core baseline utilities for held-out interior-day population imputation.
+
+The functions here can be imported by 02_bracket_heldout.py.
+"""
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import ot
+
+
+# ============================================================
+# Neural velocity field
+# ============================================================
+
+def fourier(x, n=8):
+    f = (
+        2.0 ** torch.arange(n, device=x.device, dtype=x.dtype)
+    ) * math.pi
+
+    a = x * f
+
+    return torch.cat(
+        [
+            torch.sin(a),
+            torch.cos(a),
+        ],
+        dim=-1,
+    )
+
+
+class Field(nn.Module):
+    def __init__(self, d, h=256, nt=8):
+        super().__init__()
+
+        self.nt = nt
+
+        self.net = nn.Sequential(
+            nn.Linear(d + 2 * nt, h),
+            nn.SiLU(),
+
+            nn.Linear(h, h),
+            nn.SiLU(),
+
+            nn.Linear(h, h),
+            nn.SiLU(),
+
+            nn.Linear(h, d),
+        )
+
+    def forward(self, z, tau):
+
+        if tau.ndim == 1:
+            tau = tau[:, None]
+
+        return self.net(
+            torch.cat(
+                [
+                    z,
+                    fourier(tau, self.nt),
+                ],
+                dim=-1,
+            )
+        )
+
+
+# ============================================================
+# Optimal transport
+# ============================================================
+
+def ot_plan(X0, X1):
+    """
+    Balanced quadratic-cost OT plan.
+
+    Exact EMD is attempted first.
+    If it fails, fall back to a simple Sinkhorn-like scaling procedure.
+    """
+
+    X0 = np.asarray(X0, dtype=np.float64)
+    X1 = np.asarray(X1, dtype=np.float64)
+
+    C = (
+        (
+            X0[:, None, :]
+            -
+            X1[None, :, :]
+        ) ** 2
+    ).sum(-1)
+
+    a = np.ones(len(X0), dtype=np.float64) / len(X0)
+    b = np.ones(len(X1), dtype=np.float64) / len(X1)
+
+    try:
+        return ot.emd(
+            a,
+            b,
+            C,
+            numItermax=2_000_000,
+        )
+
+    except Exception:
+
+        Cn = C / (C.max() + 1e-12)
+
+        K = np.exp(
+            -Cn / 0.05
+        )
+
+        u = np.ones_like(a)
+        v = np.ones_like(b)
+
+        for _ in range(300):
+
+            u = a / (
+                K @ v + 1e-12
+            )
+
+            v = b / (
+                K.T @ u + 1e-12
+            )
+
+        return (
+            u[:, None]
+            *
+            K
+            *
+            v[None, :]
+        )
+
+
+def align_all(clouds, rng, batch=128):
+    """
+    Sample approximately OT-aligned tuples across all retained clouds.
+
+    Returns
+    -------
+    Z : ndarray
+        Shape (K, B, d)
+    """
+
+    B = min(
+        batch,
+        min(
+            len(c)
+            for c in clouds
+        ),
+    )
+
+    first_idx = rng.choice(
+        len(clouds[0]),
+        B,
+        replace=False,
+    )
+
+    z = [
+        clouds[0][first_idx]
+    ]
+
+    for k in range(
+        1,
+        len(clouds),
+    ):
+
+        nxt_idx = rng.choice(
+            len(clouds[k]),
+            B,
+            replace=False,
+        )
+
+        nxt = clouds[k][nxt_idx]
+
+        P = ot_plan(
+            z[-1],
+            nxt,
+        )
+
+        Pn = P / (
+            P.sum(
+                axis=1,
+                keepdims=True,
+            )
+            +
+            1e-12
+        )
+
+        selected = np.array(
+            [
+                rng.choice(
+                    B,
+                    p=Pn[i],
+                )
+                for i in range(B)
+            ]
+        )
+
+        z.append(
+            nxt[selected]
+        )
+
+    return np.stack(
+        z,
+        axis=0,
+    )
+
+
+# ============================================================
+# Spline interpolation
+# ============================================================
+
+def pchip(x, Y, t):
+    """
+    Vectorized shape-preserving cubic Hermite interpolation.
+
+    Parameters
+    ----------
+    x : (K,)
+        Knot times.
+
+    Y : (K, B, d)
+        Aligned tuples.
+
+    t : (B,)
+        Evaluation times.
+
+    Returns
+    -------
+    mu : (B, d)
+        Interpolated value.
+
+    mup : (B, d)
+        Time derivative.
+    """
+
+    K = len(x)
+
+    h = np.diff(x)
+
+    d = (
+        Y[1:]
+        -
+        Y[:-1]
+    ) / h[:, None, None]
+
+    m = np.zeros_like(Y)
+
+    # Interior PCHIP slopes
+    for k in range(
+        1,
+        K - 1,
+    ):
+
+        d0 = d[k - 1]
+        d1 = d[k]
+
+        h0 = h[k - 1]
+        h1 = h[k]
+
+        w0 = 2 * h1 + h0
+        w1 = h1 + 2 * h0
+
+        mono = (
+            d0 * d1
+        ) > 0
+
+        with np.errstate(
+            divide="ignore",
+            invalid="ignore",
+        ):
+
+            mk = (
+                w0 + w1
+            ) / (
+                w0 / d0
+                +
+                w1 / d1
+            )
+
+        m[k] = np.where(
+            mono,
+            mk,
+            0.0,
+        )
+
+    def endpoint_slope(
+        h0,
+        h1,
+        d0,
+        d1,
+    ):
+
+        me = (
+            (
+                2 * h0 + h1
+            )
+            *
+            d0
+            -
+            h0 * d1
+        ) / (
+            h0 + h1
+        )
+
+        me = np.where(
+            np.sign(me)
+            !=
+            np.sign(d0),
+            0.0,
+            me,
+        )
+
+        me = np.where(
+            (
+                np.sign(d0)
+                !=
+                np.sign(d1)
+            )
+            &
+            (
+                np.abs(me)
+                >
+                3 * np.abs(d0)
+            ),
+            3 * d0,
+            me,
+        )
+
+        return me
+
+    m[0] = endpoint_slope(
+        h[0],
+        h[1] if K > 2 else h[0],
+        d[0],
+        d[1] if K > 2 else d[0],
+    )
+
+    m[-1] = endpoint_slope(
+        h[-1],
+        h[-2] if K > 2 else h[-1],
+        d[-1],
+        d[-2] if K > 2 else d[-1],
+    )
+
+    B = Y.shape[1]
+
+    idx = np.clip(
+        np.searchsorted(
+            x,
+            t,
+            side="right",
+        )
+        -
+        1,
+        0,
+        K - 2,
+    )
+
+    xi = x[idx]
+    xi1 = x[idx + 1]
+
+    hi = xi1 - xi
+
+    s = (
+        (
+            t - xi
+        )
+        /
+        hi
+    )[:, None]
+
+    yi = Y[
+        idx,
+        np.arange(B),
+    ]
+
+    yi1 = Y[
+        idx + 1,
+        np.arange(B),
+    ]
+
+    mi = m[
+        idx,
+        np.arange(B),
+    ]
+
+    mi1 = m[
+        idx + 1,
+        np.arange(B),
+    ]
+
+    hb = hi[:, None]
+
+    mu = (
+        (
+            2 * s**3
+            -
+            3 * s**2
+            +
+            1
+        )
+        *
+        yi
+        +
+        (
+            s**3
+            -
+            2 * s**2
+            +
+            s
+        )
+        *
+        hb
+        *
+        mi
+        +
+        (
+            -2 * s**3
+            +
+            3 * s**2
+        )
+        *
+        yi1
+        +
+        (
+            s**3
+            -
+            s**2
+        )
+        *
+        hb
+        *
+        mi1
+    )
+
+    mup = (
+        (
+            6 * s**2
+            -
+            6 * s
+        )
+        *
+        yi
+        +
+        (
+            3 * s**2
+            -
+            4 * s
+            +
+            1
+        )
+        *
+        hb
+        *
+        mi
+        +
+        (
+            -6 * s**2
+            +
+            6 * s
+        )
+        *
+        yi1
+        +
+        (
+            3 * s**2
+            -
+            2 * s
+        )
+        *
+        hb
+        *
+        mi1
+    ) / hb
+
+    return mu, mup
+
+
+def natural(x, Y, t):
+    """
+    Natural cubic spline interpolation.
+    """
+
+    K = len(x)
+
+    h = np.diff(x)
+
+    n = K - 2
+
+    if n <= 0:
+
+        M = np.zeros_like(Y)
+
+    else:
+
+        A = np.zeros(
+            (n, n),
+            dtype=np.float64,
+        )
+
+        rhs = np.zeros(
+            (n,) + Y.shape[1:],
+            dtype=np.float64,
+        )
+
+        for i in range(
+            1,
+            K - 1,
+        ):
+
+            r = i - 1
+
+            A[r, r] = 2 * (
+                h[i - 1]
+                +
+                h[i]
+            )
+
+            if r > 0:
+                A[
+                    r,
+                    r - 1,
+                ] = h[i - 1]
+
+            if r < n - 1:
+                A[
+                    r,
+                    r + 1,
+                ] = h[i]
+
+            rhs[r] = 6 * (
+                (
+                    Y[i + 1]
+                    -
+                    Y[i]
+                )
+                /
+                h[i]
+                -
+                (
+                    Y[i]
+                    -
+                    Y[i - 1]
+                )
+                /
+                h[i - 1]
+            )
+
+        M = np.zeros_like(Y)
+
+        # Solve rather than explicitly invert A.
+        flat_rhs = rhs.reshape(
+            n,
+            -1,
+        )
+
+        flat_solution = np.linalg.solve(
+            A,
+            flat_rhs,
+        )
+
+        M[
+            1:K - 1
+        ] = flat_solution.reshape(
+            rhs.shape
+        )
+
+    B = Y.shape[1]
+
+    idx = np.clip(
+        np.searchsorted(
+            x,
+            t,
+            side="right",
+        )
+        -
+        1,
+        0,
+        K - 2,
+    )
+
+    xi = x[idx]
+    xi1 = x[idx + 1]
+
+    hh = xi1 - xi
+
+    a = (
+        (
+            xi1 - t
+        )
+        /
+        hh
+    )[:, None]
+
+    b = (
+        (
+            t - xi
+        )
+        /
+        hh
+    )[:, None]
+
+    hb = hh[:, None]
+
+    yi = Y[
+        idx,
+        np.arange(B),
+    ]
+
+    yi1 = Y[
+        idx + 1,
+        np.arange(B),
+    ]
+
+    mi = M[
+        idx,
+        np.arange(B),
+    ]
+
+    mi1 = M[
+        idx + 1,
+        np.arange(B),
+    ]
+
+    mu = (
+        a * yi
+        +
+        b * yi1
+        +
+        (
+            (
+                a**3 - a
+            )
+            *
+            mi
+            +
+            (
+                b**3 - b
+            )
+            *
+            mi1
+        )
+        *
+        (
+            hb**2
+        )
+        /
+        6
+    )
+
+    mup = (
+        (
+            yi1 - yi
+        )
+        /
+        hb
+        +
+        hb
+        /
+        6
+        *
+        (
+            (
+                1
+                -
+                3 * a**2
+            )
+            *
+            mi
+            +
+            (
+                3 * b**2
+                -
+                1
+            )
+            *
+            mi1
+        )
+    )
+
+    return mu, mup
+
+
+def linear(x, Y, t):
+    """
+    Piecewise-linear interpolation.
+    """
+
+    B = Y.shape[1]
+
+    idx = np.clip(
+        np.searchsorted(
+            x,
+            t,
+            side="right",
+        )
+        -
+        1,
+        0,
+        len(x) - 2,
+    )
+
+    xi = x[idx]
+    xi1 = x[idx + 1]
+
+    hb = (
+        xi1 - xi
+    )[:, None]
+
+    s = (
+        (
+            t - xi
+        )
+        /
+        (
+            xi1 - xi
+        )
+    )[:, None]
+
+    yi = Y[
+        idx,
+        np.arange(B),
+    ]
+
+    yi1 = Y[
+        idx + 1,
+        np.arange(B),
+    ]
+
+    mu = (
+        1 - s
+    ) * yi + s * yi1
+
+    mup = (
+        yi1 - yi
+    ) / hb
+
+    return mu, mup
+
+
+SPL = {
+    "linear": linear,
+    "natural": natural,
+    "pchip": pchip,
+}
+
+
+# ============================================================
+# Conditional flow matching
+# ============================================================
+
+def fit_flow(
+    clouds_by_day,
+    dn,
+    spline,
+    dev,
+    steps,
+    seed,
+    batch=128,
+    lr=1e-3,
+):
+    """
+    Fit a conditional velocity field to one of the reference interpolants.
+    """
+
+    torch.manual_seed(seed)
+
+    rng = np.random.default_rng(
+        seed
+    )
+
+    d = next(
+        iter(
+            clouds_by_day.values()
+        )
+    ).shape[1]
+
+    field = Field(
+        d
+    ).to(dev)
+
+    opt = torch.optim.Adam(
+        field.parameters(),
+        lr=lr,
+    )
+
+    x = np.array(
+        sorted(dn),
+        dtype=float,
+    )
+
+    clouds = [
+        clouds_by_day[t]
+        for t in x
+    ]
+
+    ni = len(x) - 1
+
+    for step in range(steps):
+
+        Z = align_all(
+            clouds,
+            rng,
+            batch,
+        )
+
+        B = Z.shape[1]
+
+        per = B // ni
+
+        ts = []
+
+        for i in range(ni):
+
+            m = per + (
+                1
+                if i < B - per * ni
+                else 0
+            )
+
+            ts.append(
+                rng.uniform(
+                    x[i],
+                    x[i + 1],
+                    size=m,
+                )
+            )
+
+        t = np.concatenate(
+            ts
+        )[:B]
+
+        rng.shuffle(
+            t
+        )
+
+        mu, mup = SPL[
+            spline
+        ](
+            x,
+            Z,
+            t,
+        )
+
+        tt = torch.tensor(
+            t[:, None],
+            dtype=torch.float32,
+            device=dev,
+        )
+
+        mu_t = torch.tensor(
+            mu,
+            dtype=torch.float32,
+            device=dev,
+        )
+
+        mup_t = torch.tensor(
+            mup,
+            dtype=torch.float32,
+            device=dev,
+        )
+
+        pred = field(
+            mu_t,
+            tt,
+        )
+
+        loss = (
+            (
+                pred
+                -
+                mup_t
+            ) ** 2
+        ).mean()
+
+        opt.zero_grad()
+
+        loss.backward()
+
+        nn.utils.clip_grad_norm_(
+            field.parameters(),
+            1.0,
+        )
+
+        opt.step()
+
+    field.eval()
+
+    return field
+
+
+def ode(
+    field,
+    z0,
+    ts,
+    te,
+    dev,
+    N=50,
+):
+    """
+    Forward Euler integration of the learned field.
+    """
+
+    z = z0.clone()
+
+    dt = (
+        te - ts
+    ) / N
+
+    with torch.no_grad():
+
+        for k in range(N):
+
+            tcur = (
+                ts
+                +
+                (
+                    te - ts
+                )
+                *
+                k
+                /
+                N
+            )
+
+            tau = torch.full(
+                (
+                    len(z),
+                    1,
+                ),
+                float(tcur),
+                device=dev,
+            )
+
+            z = (
+                z
+                +
+                dt
+                *
+                field(
+                    z,
+                    tau,
+                )
+            )
+
+    return z.cpu().numpy()
+
+
+# ============================================================
+# Simple held-out baselines
+# ============================================================
+
+def bracketing_days(
+    retained_days,
+    heldout_day,
+):
+    """
+    Return nearest observed retained day before and after heldout_day.
+
+    Example
+    -------
+    Gupta d7:
+        retained = [2.5, 14, 21]
+        -> left=2.5, right=14
+
+    Gupta d14:
+        retained = [2.5, 7, 21]
+        -> left=7, right=21
+    """
+
+    retained = np.array(
+        sorted(
+            float(x)
+            for x in retained_days
+        ),
+        dtype=float,
+    )
+
+    left = retained[
+        retained < heldout_day
+    ]
+
+    right = retained[
+        retained > heldout_day
+    ]
+
+    if len(left) == 0:
+        raise ValueError(
+            f"No retained day before heldout={heldout_day}"
+        )
+
+    if len(right) == 0:
+        raise ValueError(
+            f"No retained day after heldout={heldout_day}"
+        )
+
+    return (
+        float(left.max()),
+        float(right.min()),
+    )
+
+
+def no_move(
+    X_source,
+    rng=None,
+    n_out=None,
+):
+    """
+    No-move baseline.
+
+    Predict that the held-out population is simply the previous observed
+    source population.
+
+    If n_out is supplied, resample source cells to n_out.
+    """
+
+    X_source = np.asarray(
+        X_source
+    )
+
+    if n_out is None:
+        return X_source.copy()
+
+    if rng is None:
+        rng = np.random.default_rng(
+            0
+        )
+
+    idx = rng.choice(
+        len(X_source),
+        size=n_out,
+        replace=(
+            n_out
+            >
+            len(X_source)
+        ),
+    )
+
+    return X_source[
+        idx
+    ].copy()
+
+
+def mccann(
+    Xa,
+    Xb,
+    alpha,
+    rng,
+    n_out=None,
+):
+    """
+    McCann / displacement interpolation under an OT coupling.
+
+    Parameters
+    ----------
+    Xa, Xb
+        Bracketing populations.
+
+    alpha
+        Fractional position in time:
+            alpha = (t_h - t_left) / (t_right - t_left)
+
+    n_out
+        Number of samples drawn from the interpolated distribution.
+        Defaults to len(Xa).
+    """
+
+    Xa = np.asarray(
+        Xa
+    )
+
+    Xb = np.asarray(
+        Xb
+    )
+
+    P = ot_plan(
+        Xa,
+        Xb,
+    )
+
+    total = P.sum()
+
+    if total <= 0:
+        raise RuntimeError(
+            "OT plan has zero total mass."
+        )
+
+    if n_out is None:
+        n_out = len(Xa)
+
+    idx = rng.choice(
+        P.size,
+        size=n_out,
+        replace=True,
+        p=(
+            P / total
+        ).ravel(),
+    )
+
+    pi, pj = np.unravel_index(
+        idx,
+        P.shape,
+    )
+
+    return (
+        (
+            1.0 - alpha
+        )
+        *
+        Xa[pi]
+        +
+        alpha
+        *
+        Xb[pj]
+    )
+
+
+def plain_ot_impute(
+    clouds_by_day,
+    retained_days,
+    heldout_day,
+    seed=0,
+    n_fit=1500,
+    n_out=1500,
+):
+    """
+    Plain OT / McCann interpolation baseline.
+
+    Uses only the two retained observed populations immediately bracketing
+    the held-out time.
+
+    To keep exact EMD computationally manageable, each bracketing cloud is
+    subsampled to at most n_fit cells before constructing the OT plan.
+
+    The same n_fit should be used for every method/dataset in the benchmark.
+    """
+
+    rng = np.random.default_rng(
+        seed
+    )
+
+    t0, t1 = bracketing_days(
+        retained_days,
+        heldout_day,
+    )
+
+    Xa = np.asarray(
+        clouds_by_day[t0]
+    )
+
+    Xb = np.asarray(
+        clouds_by_day[t1]
+    )
+
+    if len(Xa) > n_fit:
+
+        ia = rng.choice(
+            len(Xa),
+            n_fit,
+            replace=False,
+        )
+
+        Xa_fit = Xa[
+            ia
+        ]
+
+    else:
+
+        Xa_fit = Xa
+
+    if len(Xb) > n_fit:
+
+        ib = rng.choice(
+            len(Xb),
+            n_fit,
+            replace=False,
+        )
+
+        Xb_fit = Xb[
+            ib
+        ]
+
+    else:
+
+        Xb_fit = Xb
+
+    alpha = (
+        heldout_day - t0
+    ) / (
+        t1 - t0
+    )
+
+    pred = mccann(
+        Xa_fit,
+        Xb_fit,
+        alpha,
+        rng,
+        n_out=n_out,
+    )
+
+    return pred
+
+
+def no_move_impute(
+    clouds_by_day,
+    retained_days,
+    heldout_day,
+    seed=0,
+    n_out=1500,
+):
+    """
+    Held-out no-move baseline using the nearest earlier retained population.
+    """
+
+    rng = np.random.default_rng(
+        seed
+    )
+
+    t0, _ = bracketing_days(
+        retained_days,
+        heldout_day,
+    )
+
+    return no_move(
+        clouds_by_day[t0],
+        rng=rng,
+        n_out=n_out,
+    )
+
+
+# ============================================================
+# FSI / spline image helper
+# ============================================================
+
+def fsi(
+    clouds_by_day,
+    dn,
+    t_h,
+    rng,
+):
+    """
+    Push the first observed population through a sequence of barycentric
+    OT maps, then fit a natural spline through the resulting trajectories.
+
+    Retained here for compatibility with the existing benchmark code.
+    """
+
+    x = np.array(
+        sorted(dn),
+        dtype=float,
+    )
+
+    clouds = [
+        clouds_by_day[t]
+        for t in x
+    ]
+
+    cur = clouds[0]
+
+    imgs = [
+        cur
+    ]
+
+    for k in range(
+        1,
+        len(clouds),
+    ):
+
+        P = ot_plan(
+            cur,
+            clouds[k],
+        )
+
+        Pn = P / (
+            P.sum(
+                axis=1,
+                keepdims=True,
+            )
+            +
+            1e-12
+        )
+
+        cur = (
+            Pn
+            @
+            clouds[k]
+        )
+
+        imgs.append(
+            cur
+        )
+
+    Y = np.stack(
+        imgs,
+        axis=0,
+    )
+
+    mu, _ = natural(
+        x,
+        Y,
+        np.full(
+            Y.shape[1],
+            t_h,
+        ),
+    )
+
+    return mu
+
+
+# ============================================================
+# Evaluation
+# ============================================================
+
+def w2(
+    A,
+    B,
+    mx=1500,
+    seed=0,
+):
+    """
+    Empirical W2 using exact balanced OT after paired-size subsampling.
+    """
+
+    r = np.random.default_rng(
+        seed
+    )
+
+    A = np.asarray(
+        A
+    )
+
+    B = np.asarray(
+        B
+    )
+
+    if len(A) > mx:
+
+        A = A[
+            r.choice(
+                len(A),
+                mx,
+                replace=False,
+            )
+        ]
+
+    if len(B) > mx:
+
+        B = B[
+            r.choice(
+                len(B),
+                mx,
+                replace=False,
+            )
+        ]
+
+    C = (
+        (
+            A[:, None, :]
+            -
+            B[None, :, :]
+        ) ** 2
+    ).sum(-1)
+
+    cost = ot.emd2(
+        ot.unif(
+            len(A)
+        ),
+        ot.unif(
+            len(B)
+        ),
+        C,
+        numItermax=2_000_000,
+    )
+
+    return float(
+        np.sqrt(
+            max(
+                cost,
+                0.0,
+            )
+        )
+    )
+
+
+def mse_means(
+    A,
+    B,
+):
+    """
+    Squared error between population means.
+    """
+
+    A = np.asarray(
+        A
+    )
+
+    B = np.asarray(
+        B
+    )
+
+    return float(
+        (
+            (
+                A.mean(0)
+                -
+                B.mean(0)
+            ) ** 2
+        ).mean()
+    )
+
+
+# ============================================================
+# Convenience function for held-out baselines
+# ============================================================
+
+def basic_baseline_predictions(
+    clouds_by_day,
+    retained_days,
+    heldout_day,
+    seed=0,
+    n_fit=1500,
+    n_out=1500,
+):
+    """
+    Return the two simple reference baselines needed for the first TRACE
+    comparison.
+
+    Returns
+    -------
+    dict
+        {
+            "No_move": ...,
+            "OT_McCann": ...
+        }
+    """
+
+    pred_no_move = no_move_impute(
+        clouds_by_day,
+        retained_days,
+        heldout_day,
+        seed=seed,
+        n_out=n_out,
+    )
+
+    pred_ot = plain_ot_impute(
+        clouds_by_day,
+        retained_days,
+        heldout_day,
+        seed=seed,
+        n_fit=n_fit,
+        n_out=n_out,
+    )
+
+    return {
+        "No_move": pred_no_move,
+        "OT_McCann": pred_ot,
+    }
+
+
+# ============================================================
+# Import check
+# ============================================================
+
+if __name__ == "__main__":
+
+    print(
+        "benchmark_table.py loaded successfully."
+    )
+
+    print(
+        "This file is intended to be imported by "
+        "02_bracket_heldout.py."
+    )
+
+    print(
+        "Available simple baselines: "
+        "No_move, OT_McCann."
+    )
+
+    print(
+        "Available CFM splines: "
+        "linear, natural, pchip."
+    )
+

biochime_regulon_core.py

@@ -0,0 +1,481 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import csv
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Sequence, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+# -----------------------------------------------------------------------------
+# Regulatory prior
+# -----------------------------------------------------------------------------
+@dataclass
+class TFRegulon:
+    tf: str
+    targets: np.ndarray          # column indices in current expression matrix
+    weights: np.ndarray          # signed CollecTRI edge weights
+    n_total_targets: int
+    n_present_targets: int
+    coverage: float
+    programs: Tuple[str, ...]
+    tf_weight: float
+    reproducibility: float
+    perturbed_cl13: bool
+
+
+class RegulonScorer:
+    """Simple, dependency-light signed TF activity scorer.
+
+    Activity_r(cell) = sum_g w_rg * z_g(cell) / sqrt(sum_g w_rg^2)
+
+    where z_g is standardized using retained TRAINING cells only.
+
+    This transparent scoring method is the signed target-gene prior
+    used in the TRACE implementation; it is not TF transcript abundance.
+    """
+
+    def __init__(
+        self,
+        prior_csv: str,
+        var_names: Sequence[str],
+        min_targets: int = 5,
+        min_coverage: float = 0.15,
+        candidate_source: str = "taiji",
+        max_tfs: int = 0,
+        edge_weight_mode: str = "weight",
+    ):
+        self.prior_csv = prior_csv
+        self.var_names = np.asarray([str(x) for x in var_names])
+        self.lookup = {g.upper(): i for i, g in enumerate(self.var_names)}
+        self.min_targets = int(min_targets)
+        self.min_coverage = float(min_coverage)
+        self.candidate_source = candidate_source
+        self.max_tfs = int(max_tfs)
+        self.edge_weight_mode = str(edge_weight_mode).lower()
+        if self.edge_weight_mode not in {"weight", "sign", "sign_resources"}:
+            raise ValueError("edge_weight_mode must be one of: weight, sign, sign_resources")
+
+        rows = []
+        with open(prior_csv, newline="") as f:
+            reader = csv.DictReader(f)
+            required = {"program_id", "tf", "target_gene", "weight", "provenance"}
+            if not required.issubset(set(reader.fieldnames or [])):
+                raise ValueError(
+                    f"{prior_csv} must contain at least {sorted(required)}; "
+                    f"found {reader.fieldnames}"
+                )
+            for r in reader:
+                rows.append(r)
+        if not rows:
+            raise ValueError(f"No rows found in {prior_csv}")
+
+        # Candidate TF metadata comes from Taiji state-association rows.  This
+        # selects TFs biologically linked to the T-cell programs without imposing
+        # a state order on the trajectory.
+        meta: Dict[str, dict] = {}
+        for r in rows:
+            if (r.get("provenance") or "").strip().lower() != "taiji_pagerank":
+                continue
+            tf = (r.get("tf") or "").strip()
+            if not tf:
+                continue
+            m = meta.setdefault(tf, {
+                "programs": set(), "tf_weight": [], "reproducibility": [], "perturbed": False
+            })
+            m["programs"].add((r.get("program_id") or "").strip())
+            try:
+                m["tf_weight"].append(abs(float(r.get("tf_weight") or 1.0)))
+            except Exception:
+                pass
+            try:
+                val = float(r.get("reproducibility") or "nan")
+                if np.isfinite(val):
+                    m["reproducibility"].append(val)
+            except Exception:
+                pass
+            m["perturbed"] = m["perturbed"] or str(r.get("perturbed_cl13", "")).lower() in {"true", "1", "yes"}
+
+        # Collect signed TF-target edges.  The source file repeats some edges in
+        # multiple program blocks; aggregate duplicates by mean edge weight so a
+        # TF-target relation is counted once.
+        edge_vals: Dict[Tuple[str, str], List[float]] = {}
+        total_targets: Dict[str, set] = {}
+        for r in rows:
+            if (r.get("provenance") or "").strip().lower() != "collectri":
+                continue
+            tf = (r.get("tf") or "").strip()
+            tg = (r.get("target_gene") or "").strip()
+            if not tf or not tg:
+                continue
+            if candidate_source == "taiji" and tf not in meta:
+                continue
+            try:
+                if self.edge_weight_mode == "weight":
+                    # Legacy BioCHIME behavior.  In the assembled T-cell prior this
+                    # column may include state-specific scaling, so the paper
+                    # configuration uses state-independent sign mode.
+                    w = float(r.get("weight") or 0.0)
+                else:
+                    sg = float(r.get("sign") or 0.0)
+                    if sg == 0.0:
+                        raw = float(r.get("weight") or 0.0)
+                        sg = float(np.sign(raw))
+                    if self.edge_weight_mode == "sign":
+                        w = sg
+                    else:
+                        nr = max(float(r.get("n_resources") or 1.0), 1.0)
+                        w = sg * np.sqrt(nr)
+            except Exception:
+                continue
+            if not np.isfinite(w) or abs(w) < 1e-15:
+                continue
+            edge_vals.setdefault((tf, tg), []).append(w)
+            total_targets.setdefault(tf, set()).add(tg)
+
+        # Rank candidates only if requested.  We use reproducibility first, then
+        # Taiji weight.  This ranking is based on external prior metadata, not on
+        # the held-out expression data.
+        all_tfs = sorted(total_targets)
+        if max_tfs > 0:
+            def rank_key(tf: str):
+                m = meta.get(tf, {})
+                rep = np.nanmean(m.get("reproducibility", [np.nan])) if m.get("reproducibility") else -np.inf
+                tw = np.nanmax(m.get("tf_weight", [0.0])) if m.get("tf_weight") else 0.0
+                if not np.isfinite(rep):
+                    rep = -1e9
+                return (rep, tw)
+            all_tfs = sorted(all_tfs, key=rank_key, reverse=True)[:max_tfs]
+
+        regulons: List[TFRegulon] = []
+        for tf in all_tfs:
+            tg_all = sorted(total_targets[tf])
+            present_idx, present_w = [], []
+            for tg in tg_all:
+                j = self.lookup.get(tg.upper())
+                if j is None:
+                    continue
+                vals = edge_vals.get((tf, tg), [])
+                if not vals:
+                    continue
+                present_idx.append(j)
+                present_w.append(float(np.mean(vals)))
+            nt = len(tg_all)
+            npres = len(present_idx)
+            cov = npres / max(nt, 1)
+            if npres < min_targets or cov < min_coverage:
+                continue
+            m = meta.get(tf, {})
+            tw = float(np.nanmax(m.get("tf_weight", [1.0]))) if m.get("tf_weight") else 1.0
+            reps = m.get("reproducibility", [])
+            rep = float(np.nanmax(reps)) if reps else float("nan")
+            regulons.append(TFRegulon(
+                tf=tf,
+                targets=np.asarray(present_idx, dtype=int),
+                weights=np.asarray(present_w, dtype=np.float64),
+                n_total_targets=nt,
+                n_present_targets=npres,
+                coverage=cov,
+                programs=tuple(sorted(x for x in m.get("programs", set()) if x)),
+                tf_weight=tw,
+                reproducibility=rep,
+                perturbed_cl13=bool(m.get("perturbed", False)),
+            ))
+
+        if not regulons:
+            raise ValueError(
+                "No TF regulons passed coverage filters. Try --min_targets 3 "
+                "and/or --min_coverage 0.05, or rebuild the expression matrix "
+                "with more genes / forced inclusion of prior target genes."
+            )
+        self.regulons = regulons
+        self.tf_names = [r.tf for r in regulons]
+        self.gene_mu = None
+        self.gene_sd = None
+
+    def fit_gene_scaler(self, X_train: np.ndarray):
+        X_train = np.asarray(X_train, dtype=np.float64)
+        self.gene_mu = X_train.mean(0)
+        self.gene_sd = X_train.std(0)
+        self.gene_sd[self.gene_sd < 1e-8] = 1.0
+        return self
+
+    def score(self, X: np.ndarray) -> np.ndarray:
+        if self.gene_mu is None or self.gene_sd is None:
+            raise RuntimeError("Call fit_gene_scaler(X_train) before score(X)")
+        X = np.asarray(X, dtype=np.float64)
+        Zg = (X - self.gene_mu) / self.gene_sd
+        out = np.empty((len(X), len(self.regulons)), dtype=np.float32)
+        for j, r in enumerate(self.regulons):
+            w = r.weights
+            denom = np.sqrt(np.sum(w * w)) + 1e-12
+            out[:, j] = (Zg[:, r.targets] @ w / denom).astype(np.float32)
+        return out
+
+    def coverage_rows(self) -> List[dict]:
+        rows = []
+        for r in self.regulons:
+            rows.append({
+                "tf": r.tf,
+                "n_targets_total": r.n_total_targets,
+                "n_targets_present": r.n_present_targets,
+                "coverage": r.coverage,
+                "programs": ";".join(r.programs),
+                "tf_weight": r.tf_weight,
+                "reproducibility": r.reproducibility,
+                "perturbed_cl13": r.perturbed_cl13,
+            })
+        return rows
+
+
+def standardize_activity(A_train: np.ndarray, A: np.ndarray | None = None):
+    A_train = np.asarray(A_train, dtype=np.float64)
+    mu = A_train.mean(0)
+    sd = A_train.std(0)
+    sd[sd < 1e-8] = 1.0
+    if A is None:
+        return ((A_train - mu) / sd).astype(np.float32), mu.astype(np.float32), sd.astype(np.float32)
+    return ((np.asarray(A) - mu) / sd).astype(np.float32), mu.astype(np.float32), sd.astype(np.float32)
+
+
+def fit_ridge_map(A: np.ndarray, Z: np.ndarray, lam: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
+    """Fit Z ~= A @ W + b; derivatives map as dZ/dt ~= dA/dt @ W."""
+    A = np.asarray(A, dtype=np.float64)
+    Z = np.asarray(Z, dtype=np.float64)
+    am, zm = A.mean(0), Z.mean(0)
+    Ac, Zc = A - am, Z - zm
+    G = Ac.T @ Ac + float(lam) * np.eye(A.shape[1])
+    W = np.linalg.solve(G, Ac.T @ Zc)
+    b = zm - am @ W
+    return W.astype(np.float32), b.astype(np.float32)
+
+
+# -----------------------------------------------------------------------------
+#Hermite helpers
+# -----------------------------------------------------------------------------
+def pchip_slopes(x: np.ndarray, Y: np.ndarray) -> np.ndarray:
+    x = np.asarray(x, dtype=float)
+    K = len(x)
+    h = np.diff(x)
+    d = (Y[1:] - Y[:-1]) / h[:, None, None]
+    m = np.zeros_like(Y)
+    if K == 2:
+        m[0] = d[0]
+        m[1] = d[0]
+        return m
+    for k in range(1, K - 1):
+        d0, d1 = d[k - 1], d[k]
+        h0, h1 = h[k - 1], h[k]
+        w0, w1 = 2 * h1 + h0, h1 + 2 * h0
+        mono = (d0 * d1) > 0
+        with np.errstate(divide="ignore", invalid="ignore"):
+            mk = (w0 + w1) / (w0 / d0 + w1 / d1)
+        m[k] = np.where(mono, mk, 0.0)
+
+    def endpoint(h0, h1, d0, d1):
+        me = ((2 * h0 + h1) * d0 - h0 * d1) / (h0 + h1)
+        me = np.where(np.sign(me) != np.sign(d0), 0.0, me)
+        me = np.where((np.sign(d0) != np.sign(d1)) & (np.abs(me) > 3 * np.abs(d0)), 3 * d0, me)
+        return me
+
+    m[0] = endpoint(h[0], h[1], d[0], d[1])
+    m[-1] = endpoint(h[-1], h[-2], d[-1], d[-2])
+    return m
+
+
+def project_shape_safe_slopes(x: np.ndarray, Y: np.ndarray, M: np.ndarray) -> np.ndarray:
+    """Conservative coordinate-wise projection that prevents PCHIP-like overshoot."""
+    x = np.asarray(x, dtype=float)
+    h = np.diff(x)
+    d = (Y[1:] - Y[:-1]) / h[:, None, None]
+    out = np.asarray(M).copy()
+    K = len(x)
+    if K == 2:
+        sec = d[0]
+        for k in [0, 1]:
+            ok = np.sign(out[k]) == np.sign(sec)
+            out[k] = np.where(ok, out[k], 0.0)
+            out[k] = np.sign(sec) * np.minimum(np.abs(out[k]), 3.0 * np.abs(sec))
+        return out
+
+    for k, sec in [(0, d[0]), (K - 1, d[-1])]:
+        ok = np.sign(out[k]) == np.sign(sec)
+        out[k] = np.where(ok, out[k], 0.0)
+        out[k] = np.sign(sec) * np.minimum(np.abs(out[k]), 3.0 * np.abs(sec))
+
+    for k in range(1, K - 1):
+        dl, dr = d[k - 1], d[k]
+        mono = (dl * dr) > 0
+        sgn = np.sign(dl)
+        cap = 3.0 * np.minimum(np.abs(dl), np.abs(dr))
+        cand = np.where(np.sign(out[k]) == sgn, out[k], 0.0)
+        cand = sgn * np.minimum(np.abs(cand), cap)
+        out[k] = np.where(mono, cand, 0.0)
+    return out
+
+
+def hermite_eval(x: np.ndarray, Y: np.ndarray, M: np.ndarray, t: np.ndarray):
+    x = np.asarray(x, dtype=float)
+    t = np.asarray(t, dtype=float)
+    K, B, _ = Y.shape
+    idx = np.clip(np.searchsorted(x, t, side="right") - 1, 0, K - 2)
+    xi, xi1 = x[idx], x[idx + 1]
+    hi = xi1 - xi
+    s = ((t - xi) / hi)[:, None]
+    hb = hi[:, None]
+    ar = np.arange(B)
+    yi, yi1 = Y[idx, ar], Y[idx + 1, ar]
+    mi, mi1 = M[idx, ar], M[idx + 1, ar]
+    mu = ((2*s**3 - 3*s**2 + 1)*yi + (s**3 - 2*s**2 + s)*hb*mi
+          + (-2*s**3 + 3*s**2)*yi1 + (s**3 - s**2)*hb*mi1)
+    mup = ((6*s**2 - 6*s)*yi + (3*s**2 - 4*s + 1)*hb*mi
+           + (-6*s**2 + 6*s)*yi1 + (3*s**2 - 2*s)*hb*mi1) / hb
+    return mu, mup
+
+
+def align_joint(
+    clouds_z: Sequence[np.ndarray],
+    clouds_a: Sequence[np.ndarray],
+    rng: np.random.Generator,
+    batch: int = 128,
+):
+    """Use the existing Euclidean OT plan, carrying TF activities with cells."""
+    B = min(batch, min(len(c) for c in clouds_z))
+    i0 = rng.choice(len(clouds_z[0]), B, replace=False)
+    z = [clouds_z[0][i0]]
+    a = [clouds_a[0][i0]]
+    for k in range(1, len(clouds_z)):
+        cand = rng.choice(len(clouds_z[k]), B, replace=False)
+        zc, ac = clouds_z[k][cand], clouds_a[k][cand]
+        import benchmark_table as BT
+        P = BT.ot_plan(z[-1], zc)
+        Pn = P / (P.sum(1, keepdims=True) + 1e-12)
+        jj = np.array([rng.choice(B, p=Pn[i]) for i in range(B)])
+        z.append(zc[jj])
+        a.append(ac[jj])
+    return np.stack(z, 0), np.stack(a, 0)
+
+
+def _match_bio_vector_norm(m_geom: np.ndarray, m_bio: np.ndarray, clip=(0.25, 4.0)) -> np.ndarray:
+    """Make biology contribute direction rather than arbitrary numerical scale."""
+    ng = np.linalg.norm(m_geom, axis=-1, keepdims=True)
+    nb = np.linalg.norm(m_bio, axis=-1, keepdims=True)
+    scale = ng / (nb + 1e-12)
+    scale = np.clip(scale, clip[0], clip[1])
+    scale = np.where((ng > 1e-8) & (nb > 1e-8), scale, 0.0)
+    return m_bio * scale
+
+
+@dataclass
+class BioFitInfo:
+    rho: float
+    ridge_W: np.ndarray
+    ridge_b: np.ndarray
+
+
+def fit_regulon_bio_flow(
+    clouds_z_by_day: Dict[float, np.ndarray],
+    clouds_a_by_day: Dict[float, np.ndarray],
+    dn: Sequence[float],
+    dev: str,
+    steps: int,
+    seed: int,
+    rho: float = 1.0,
+    batch: int = 128,
+    lr: float = 1e-3,
+    ridge: float = 1.0,
+    ridge_map: tuple[np.ndarray, np.ndarray] | None = None,
+    shape_project: bool = True,
+    log_every: int = 0,
+):
+    """Train TRACE with TF-regulon-guided Hermite knot tangents."""
+    if not (0.0 <= rho <= 1.0):
+        raise ValueError("rho must be in [0,1]")
+    torch.manual_seed(seed)
+    rng = np.random.default_rng(seed)
+    x = np.asarray(sorted(dn), dtype=float)
+    zclouds = [clouds_z_by_day[t] for t in x]
+    aclouds = [clouds_a_by_day[t] for t in x]
+    d = zclouds[0].shape[1]
+
+    Apool = np.vstack(aclouds)
+    Zpool = np.vstack(zclouds)
+    if ridge_map is None:
+        W, b = fit_ridge_map(Apool, Zpool, lam=ridge)
+    else:
+        # Negative controls such as TF-channel permutation must use the
+        # ridge map fitted on the TRUE TF identities.  Re-fitting after a
+        # global column permutation would simply permute the rows of W and
+        # make the control algebraically equivalent to the unpermuted model.
+        W, b = ridge_map
+        W = np.asarray(W, dtype=np.float32)
+        b = np.asarray(b, dtype=np.float32)
+        if W.shape != (Apool.shape[1], Zpool.shape[1]):
+            raise ValueError(
+                f"ridge_map W has shape {W.shape}, expected "
+                f"{(Apool.shape[1], Zpool.shape[1])}"
+            )
+        if b.shape != (Zpool.shape[1],):
+            raise ValueError(
+                f"ridge_map b has shape {b.shape}, expected {(Zpool.shape[1],)}"
+            )
+
+    import benchmark_table as BT
+    field = BT.Field(d).to(dev)
+    opt = torch.optim.Adam(field.parameters(), lr=lr)
+    ni = len(x) - 1
+
+    for step in range(steps):
+        Z, A = align_joint(zclouds, aclouds, rng, batch=batch)
+        B = Z.shape[1]
+        per = B // ni
+        ts = []
+        for i in range(ni):
+            n = per + (1 if i < B - per * ni else 0)
+            ts.append(rng.uniform(x[i], x[i + 1], size=n))
+        tq = np.concatenate(ts)[:B]
+        rng.shuffle(tq)
+
+        m_geom = pchip_slopes(x, Z)
+        if rho > 0:
+            m_tf = pchip_slopes(x, A)
+            m_bio = np.einsum("kbr,rd->kbd", m_tf, W)
+            m_bio = _match_bio_vector_norm(m_geom, m_bio)
+            M = (1.0 - rho) * m_geom + rho * m_bio
+            if shape_project:
+                M = project_shape_safe_slopes(x, Z, M)
+        else:
+            M = m_geom
+
+        pos, vel = hermite_eval(x, Z, M, tq)
+        tt = torch.tensor(tq[:, None], dtype=torch.float32, device=dev)
+        pred = field(torch.tensor(pos, dtype=torch.float32, device=dev), tt)
+        target = torch.tensor(vel, dtype=torch.float32, device=dev)
+        loss = ((pred - target) ** 2).mean()
+        opt.zero_grad()
+        loss.backward()
+        nn.utils.clip_grad_norm_(field.parameters(), 1.0)
+        opt.step()
+        if log_every and (step + 1) % log_every == 0:
+            print(f"[BioCHIME rho={rho:g}] {step+1}/{steps} loss={float(loss):.6f}", flush=True)
+
+    field.eval()
+    return field, BioFitInfo(rho=rho, ridge_W=W, ridge_b=b)
+
+
+def permute_within_day(clouds_a_by_day: Dict[float, np.ndarray], seed: int):
+    """Negative control preserving per-day TF distributions but breaking cell correspondence."""
+    rng = np.random.default_rng(seed)
+    return {t: A[rng.permutation(len(A))].copy() for t, A in clouds_a_by_day.items()}
+
+
+def permute_tf_channels(clouds_a_by_day: Dict[float, np.ndarray], seed: int):
+    """Wrong-biology control: keep values but scramble TF identities globally."""
+    rng = np.random.default_rng(seed)
+    r = next(iter(clouds_a_by_day.values())).shape[1]
+    p = rng.permutation(r)
+    return {t: A[:, p].copy() for t, A in clouds_a_by_day.items()}

requirements.txt

@@ -0,0 +1,5 @@
+numpy>=1.24,<3
+scipy>=1.11
+scikit-learn>=1.3
+POT>=0.9
+torch>=2.0

run_trace_benchmark.py

@@ -0,0 +1,308 @@
+#!/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()

tests/test_trace_core.py

@@ -0,0 +1,65 @@
+"""Numerical tests for the TRACE method-specific components (no training data needed)."""
+import sys
+import unittest
+from pathlib import Path
+
+import numpy as np
+from scipy.interpolate import PchipInterpolator
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from biochime_regulon_core import (  # noqa: E402
+    pchip_slopes, hermite_eval, project_shape_safe_slopes,
+    _match_bio_vector_norm, fit_ridge_map, permute_tf_channels,
+)
+
+
+class TestTRACECore(unittest.TestCase):
+    def test_pchip_matches_scipy(self):
+        rng = np.random.default_rng(17)
+        for k in (2, 3, 4):
+            x = np.cumsum(np.r_[0.0, rng.uniform(0.1, 0.7, k-1)])
+            y = rng.normal(size=(k, 15, 6))
+            m = pchip_slopes(x, y)
+            ref = PchipInterpolator(x, y, axis=0)
+            np.testing.assert_allclose(m, ref.derivative()(x), atol=1e-10)
+            t = rng.uniform(x[0], x[-1], len(y[0]))
+            pos, vel = hermite_eval(x, y, m, t)
+            np.testing.assert_allclose(pos, np.stack([ref(tt)[i] for i, tt in enumerate(t)]), atol=1e-10)
+            np.testing.assert_allclose(vel, np.stack([ref.derivative()(tt)[i] for i, tt in enumerate(t)]), atol=1e-10)
+
+    def test_projected_slope_preserves_monotone_intervals(self):
+        x = np.array([0.0, 0.4, 1.0])
+        z = np.array([[[0., 0.]], [[1., 2.]], [[2., 3.]]])
+        raw = np.array([[[20., 20.]], [[-50., 10.]], [[20., -15.]]])
+        m = project_shape_safe_slopes(x, z, raw)
+        for lo, hi in [(0., 0.4), (0.4, 1.0)]:
+            t = np.linspace(lo, hi, 29)
+            # evaluate one query time for one trajectory at a time
+            for tt in t:
+                pos, _ = hermite_eval(x, z, m, np.array([tt]))
+                self.assertTrue(np.all(pos[0] >= z[0, 0] - 1e-9))
+                self.assertTrue(np.all(pos[0] <= z[-1, 0] + 1e-9))
+
+    def test_ridge_activity_to_state_direction(self):
+        rng = np.random.default_rng(23)
+        a = rng.normal(size=(200, 7))
+        w = rng.normal(size=(7, 4))
+        b = rng.normal(size=4)
+        z = a @ w + b
+        fitted, intercept = fit_ridge_map(a, z, lam=0.0)
+        np.testing.assert_allclose(fitted, w, atol=1e-5)
+        np.testing.assert_allclose(intercept, b, atol=1e-5)
+
+    def test_norm_clipping_and_permutation(self):
+        m_geom = np.ones((3, 12, 5), dtype=float)
+        m_bio = np.ones_like(m_geom) * 100.0
+        matched = _match_bio_vector_norm(m_geom, m_bio)
+        np.testing.assert_allclose(matched, m_bio * 0.25)
+        a = {0.0: np.arange(18).reshape(6, 3), 1.0: np.arange(18).reshape(6, 3) + 1}
+        perm = permute_tf_channels(a, seed=4)
+        self.assertEqual(perm[0.0].shape, a[0.0].shape)
+        self.assertEqual(sorted(perm[0.0].T.tolist()), sorted(a[0.0].T.tolist()))
+
+
+if __name__ == '__main__':
+    unittest.main()