TRACE / biochime_regulon_core.py

Initial TRACE implementation

4c4d79d verified 2 days ago

18.6 kB

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

	#!/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 = ((2s3 - 3s*2 + 1)yi + (s*3 - 2s*2 + s)hb*mi
	+ (-2s3 + 3s*2)yi1 + (s3 - s2)hbmi1)
	mup = ((6s2 - 6s)yi + (3s*2 - 4s + 1)hbmi
	+ (-6s2 + 6s)yi1 + (3s*2 - 2s)hbmi1) / 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()}