gset 0.1.0__py3-none-any.whl

gset/__init__.py

@@ -0,0 +1,11 @@
+from .graph import build_graph
+from .embeddings import (train_gcn, coexpr_embedding,
+                         train_deepsets, expr_embedding, train_both)
+from .signals import kl_vector, kl_score, ot_residual, ot_score
+from .lags import lag_matrix, symmetric_score
+from .programs import find_programs, program_order, core_genes, programs_from_queries
+from .selection import gset_context, bridge_genes, hvg_ranking, tf_selection, random_selection
+from .metrics import relay_coverage, within_between_gap
+from .pipeline import transition_vectors
+from .reproducible import make_reproducible
+from .api import select_genes, matrices_from_anndata

gset/api.py

@@ -0,0 +1,109 @@
+"""select_genes: GSET in one function.
+
+    from gset import select_genes
+    genes = select_genes(adata, budget=30, n_programs=3)
+
+It follows the steps of the paper:
+  1. train the encoders and compute the KL and OT signals of every transition
+  2. find the gene programs, and their order
+  3. choose the bridge genes between consecutive programs and the core genes of the programs
+"""
+import numpy as np
+
+from .lags import lag_matrix
+from .pipeline import transition_vectors
+from .programs import concat_embedding, find_programs, mean_profiles, program_order, programs_from_queries
+from .selection import gset_context
+
+
+def matrices_from_anndata(adata, time_key="time"):
+    """The count matrices (cells x genes), one per timepoint, in the order of the time.
+    The counts are in adata.X and the time of every cell is in adata.obs[time_key]."""
+    X = adata.X.toarray() if hasattr(adata.X, "toarray") else np.asarray(adata.X)
+    times = np.asarray(adata.obs[time_key])
+    return [X[times == t].astype(float) for t in np.unique(times)]
+
+
+def _runs_from(data, time_key):
+    """Every run is a dict {timepoint: matrix}. Here: one run, or one for every replicate."""
+    if hasattr(data, "obs"):                      # an AnnData object
+        data = matrices_from_anndata(data, time_key)
+    if isinstance(data[0], np.ndarray):           # one run: a list of matrices
+        data = [data]
+    runs = [dict(enumerate(run)) for run in data]
+    return runs, list(range(len(runs[0])))
+
+
+def _split_budget(budget, n_cores, n_bridges, n_kl, n_ot):
+    """How many genes are cores, bridges from the KL signal and bridges from the OT signal."""
+    if n_cores is None and n_bridges is None:
+        n_cores = budget // 3                     # a third of the budget for the cores
+        n_bridges = budget - n_cores              # two thirds for the bridges
+    elif n_cores is None:
+        n_cores = budget - n_bridges
+    elif n_bridges is None:
+        n_bridges = budget - n_cores
+    if n_kl is None and n_ot is None:
+        n_ot = int(round(0.3 * n_bridges))        # 70% of the bridges from KL, 30% from OT
+    if n_kl is None:
+        n_kl = n_bridges - n_ot
+    if n_ot is None:
+        n_ot = n_bridges - n_kl
+    return n_cores + n_bridges, n_kl, n_ot
+
+
+def select_genes(data, budget=30, n_programs=3, n_cores=None, n_bridges=None, n_kl=None, n_ot=None,
+                 query_genes=None, program_size=30, labels=None, order=None,
+                 coexpr_encoder="simulation", time_key="time", seed=42, n_epochs=None):
+    """Select a compact set of genes from time-resolved single-cell data.
+
+    Returns the sorted list of the indices of the selected genes: the bridge genes between consecutive
+    programs and the core genes of the programs.
+
+    Data
+      data           an AnnData object (raw counts in X, the time of every cell in obs[time_key]) or a list of
+                     count matrices (cells x genes), one per timepoint in the order of the time. With several
+                     replicates, a list of such lists. The cells of different timepoints do not have to be matched.
+      time_key       the column of obs with the time of the cells ("time")
+
+    Number of genes
+      budget         how many genes to select (30)
+      n_cores        how many of them are core genes (default: a third of the budget)
+      n_bridges      how many are bridges (default: the rest). If you give both n_cores and n_bridges, the budget
+                     is their sum.
+      n_kl, n_ot     how many of the bridges come from the KL signal without lag and from the OT signal with lag 1
+                     (default: 70% and 30%)
+
+    Programs (the groups of genes that work together). Choose one of the three ways:
+      n_programs     the number of programs, found by clustering the genes (3)
+      query_genes    the indices of one gene for every program, in the order in which the programs become active.
+                     The program of a query gene is the gene and its program_size - 1 nearest genes (program_size
+                     is 30), and the other genes go to the program with the nearest centroid.
+      labels, order  the program of every gene (0, 1, 2, ...) and the order of the programs (by default 0, 1, 2, ...)
+
+    Encoders
+      coexpr_encoder  the encoder of the co-expression embedding (the KL signal). "simulation" (default): the
+                     encoder of the simulations, for panels of tens to a few hundred genes. "real_data": the encoder
+                     of the experiments on the real data, for thousands of genes (rmsexplorer/encoder.py).
+      seed, n_epochs the training. n_epochs: 300 for "simulation", 200 and 150 for "real_data" (default).
+    """
+    runs, times = _runs_from(data, time_key)
+
+    # 1. the signals of every transition
+    vectors = transition_vectors(runs, times, seed=seed, n_epochs=n_epochs, coexpr_encoder=coexpr_encoder)
+    kl_lag0 = lag_matrix(vectors["kl"], 0)
+    ot_lag1 = lag_matrix(vectors["ot"], 1)
+    z_concat = concat_embedding(vectors["gcn"], runs, vectors["active"], embed=vectors["embed"])
+
+    # 2. the programs and their order
+    if query_genes is not None:
+        labels, order = programs_from_queries(z_concat, query_genes, program_size)
+    elif labels is None:
+        labels = find_programs(z_concat, n_programs)
+        order, _ = program_order(labels, mean_profiles(runs, vectors["active"]))
+    elif order is None:
+        order = list(range(int(np.max(labels)) + 1))
+
+    # 3. the bridges and the cores
+    budget, n_kl, n_ot = _split_budget(budget, n_cores, n_bridges, n_kl, n_ot)
+    return gset_context(budget, n_kl, n_ot, np.asarray(labels), order, kl_lag0, ot_lag1, z_concat)

gset/cardamom_eval.py

@@ -0,0 +1,56 @@
+"""Scoring of an inferred network (from CardamomOT) against the true network.
+
+The inferred matrix belongs to a context of genes. Its first row and column are
+the stimulus, the other rows and columns follow the order of the genes.
+"""
+import numpy as np
+from sklearn.metrics import average_precision_score, roc_auc_score
+
+
+def _auroc(y, score):
+    return float(roc_auc_score(y, score)) if 0 < y.sum() < len(y) else None
+
+
+def _auprc(y, score):
+    return float(average_precision_score(y, score)) if y.sum() > 0 else None
+
+
+def restrict(network, genes):
+    """The part of a (genes + 1) x (genes + 1) matrix that belongs to the given genes."""
+    index = [0] + [g + 1 for g in genes]
+    return network[np.ix_(index, index)]
+
+
+def evaluate_network(inferred, true_network, genes, labels):
+    """Scores for all edges inside the context and for the relay edges.
+
+    all_edge: AUROC and AUPRC over all pairs of genes of the context.
+    relay_ab: relay edges from program a to program b. A relay edge with a gene
+              outside the context cannot be found and gets the score 0
+              ("uncovered edges count as missed").
+    """
+    genes = list(genes)
+    position = {g: k for k, g in enumerate(genes)}
+    scores = inferred[1:, 1:]
+    n = len(genes)
+
+    truth = restrict(true_network, genes)[1:, 1:]
+    off_diagonal = ~np.eye(n, dtype=bool)
+    y = (np.abs(truth) > 1e-6)[off_diagonal].astype(int)
+    s = np.abs(scores)[off_diagonal]
+    result = {"all_edge": {"auroc": _auroc(y, s), "auprc": _auprc(y, s),
+                           "n_true": int(y.sum()), "base_rate": float(y.mean())}}
+
+    for a in range(int(labels.max())):
+        b = a + 1
+        pairs = [(int(i), int(j)) for i in np.where(labels == a)[0]
+                 for j in np.where(labels == b)[0] if i != j]
+        is_edge = np.array([int(abs(true_network[i + 1, j + 1]) > 1e-6) for i, j in pairs])
+        kept = np.array([(i in position and j in position) for i, j in pairs])
+        score = np.array([abs(scores[position[i], position[j]]) if (i in position and j in position) else 0.0
+                          for i, j in pairs])
+        result[f"relay_{a}{b}"] = {
+            "n_true": int(is_edge.sum()), "n_covered": int((is_edge * kept).sum()),
+            "auroc": _auroc(is_edge, score), "auprc": _auprc(is_edge, score),
+            "base_rate": float(is_edge.mean())}
+    return result

gset/clustering.py

@@ -0,0 +1,62 @@
+"""Scores for how well a distance matrix between genes recovers the programs."""
+import numpy as np
+from scipy.sparse import csr_matrix
+from scipy.sparse.csgraph import shortest_path
+from sklearn.cluster import AgglomerativeClustering
+from sklearn.manifold import MDS
+from sklearn.metrics import (adjusted_rand_score, normalized_mutual_info_score,
+                             roc_auc_score)
+
+
+def cosine_distance(V):
+    norms = np.linalg.norm(V, axis=1, keepdims=True)
+    norms = np.where(norms == 0, 1e-12, norms)
+    Vn = V / norms
+    return (1.0 - np.clip(Vn @ Vn.T, -1.0, 1.0)).astype(np.float32)
+
+
+def max_rank_distance(W1, W2):
+    """Larger of the two normalised ranks of every entry."""
+    n = W1.shape[0]
+
+    def ranks(W):
+        return np.argsort(np.argsort(W.flatten())).reshape(W.shape) / (n * n)
+
+    return np.maximum(ranks(W1), ranks(W2)).astype(np.float32)
+
+
+def score_distance_matrix(W, labels, n_programs, complete_with_paths=False):
+    """Cluster the genes and compare with the true programs.
+
+    The distances are turned into 10 coordinates by multidimensional scaling
+    and the genes are clustered with Ward linkage. ARI and NMI compare the
+    clusters with the true programs. AUROC asks if two genes of the same
+    program are closer than two genes of different programs.
+
+    complete_with_paths: for sparse distance matrices (zero means "no edge")
+    the missing distances are first filled with shortest paths.
+    """
+    if complete_with_paths:
+        W_full = shortest_path(csr_matrix(W), method="D", directed=False)
+        finite = W_full[np.isfinite(W_full)]
+        W_full[~np.isfinite(W_full)] = (finite.max() if len(finite) else 1.0) * 1.5
+    else:
+        W_full = W.astype(float)
+        W_full = (W_full + W_full.T) / 2
+    np.fill_diagonal(W_full, 0.0)
+
+    n_genes = len(labels)
+    n_mds = min(10, n_genes - 1)
+    mds = MDS(n_components=n_mds, dissimilarity="precomputed",
+              random_state=42, normalized_stress="auto")
+    coords = mds.fit_transform(W_full)
+    clusters = AgglomerativeClustering(n_clusters=min(n_programs, n_genes - 1),
+                                       linkage="ward").fit_predict(coords)
+
+    rows, cols = np.triu_indices(n_genes, k=1)
+    same = (labels[rows] == labels[cols]).astype(int)
+    return {
+        "ARI": float(adjusted_rand_score(labels, clusters)),
+        "NMI": float(normalized_mutual_info_score(labels, clusters)),
+        "AUROC": float(roc_auc_score(same, -W[rows, cols])),
+    }

gset/data.py

@@ -0,0 +1,64 @@
+import json
+import numpy as np
+from pathlib import Path
+
+
+def load_run(path, keep_stimulus=False):
+    """Read one HARISSA output file.
+
+    The file has one row of times, one row with the stimulus and then one row
+    per gene. The first column only holds labels, so it is skipped.
+    Returns the matrix (cells, genes) and the time of every cell.
+    """
+    raw = np.loadtxt(path, delimiter="\t")
+    times = raw[0, 1:].astype(int)
+    first_row = 1 if keep_stimulus else 2
+    X = raw[first_row:, 1:].T
+    return X, times
+
+
+def load_simulation(sim_dir, keep_stimulus=False):
+    """Load all runs of a simulation folder.
+
+    Returns a list of runs (each a dict: time -> cells x genes matrix),
+    the program label of every gene and the sorted list of timepoints.
+    """
+    sim_dir = Path(sim_dir)
+    labels = np.load(sim_dir / "cluster_ids.npy")
+    runs = []
+    all_times = set()
+    for path in sorted((sim_dir / "data").glob("data_*.txt")):
+        X, times = load_run(path, keep_stimulus)
+        run = {}
+        for t in np.unique(times):
+            run[int(t)] = X[times == t]
+            all_times.add(int(t))
+        runs.append(run)
+    return runs, labels, sorted(all_times)
+
+
+def load_all_cells(sim_dir):
+    """All cells of all runs in the order of the files (run after run).
+
+    Returns the matrix (cells, genes) and the time of every cell.
+    """
+    parts, times = [], []
+    for path in sorted((Path(sim_dir) / "data").glob("data_*.txt")):
+        X, t = load_run(path)
+        parts.append(X)
+        times.append(t)
+    return np.vstack(parts), np.concatenate(times)
+
+
+def active_timepoints(runs, timepoints):
+    """Timepoints at which the genes are switched on (mean count above 1)."""
+    keep = []
+    for t in timepoints:
+        mean = np.concatenate([run[t] for run in runs]).mean()
+        if mean > 1.0:
+            keep.append(t)
+    return keep if len(keep) >= 2 else timepoints
+
+
+def load_metadata(sim_dir):
+    return json.loads((Path(sim_dir) / "metadata.json").read_text())

gset/embeddings.py

@@ -0,0 +1,167 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import GCNConv
+
+from .graph import build_graph, log_normalize
+
+BETA = 6.0                                 # power of the correlation in the co-expression graph (as in WGCNA)
+THRESH = 0.01                              # edges with a weight below this are dropped
+QUANTILE_LEVELS = np.linspace(1, 99, 49)   # the 49 quantiles of a gene that the DeepSets encoder has to reconstruct
+
+
+# ---------------------------------------------------------------
+# co-expression embedding (graph auto-encoder with a GCN)
+# ---------------------------------------------------------------
+
+class GCNEncoder(nn.Module):
+    """Graph auto-encoder: two graph convolutions turn every gene into a vector of 16 numbers."""
+
+    def __init__(self, n_genes, d_hid=64, d_out=16):
+        super().__init__()
+        self.conv1 = GCNConv(n_genes, d_hid, add_self_loops=False)
+        self.conv2 = GCNConv(d_hid, d_out, add_self_loops=False)
+
+    def forward(self, data):
+        x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))
+        return self.conv2(x, data.edge_index, data.edge_attr)
+
+    def loss(self, data):
+        # the inner product of two embeddings should give the edge weight
+        z = self(data)
+        src, dst = data.edge_index
+        pred = (z[src] * z[dst]).sum(dim=1)
+        return F.mse_loss(pred, data.edge_attr.reshape(-1))
+
+
+def train_gcn(matrices, n_genes, n_epochs=300, seed=42):
+    """Train one GCN on the graphs of all matrices (weights are shared)."""
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    model = GCNEncoder(n_genes)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
+
+    graphs = []
+    for X in matrices:
+        data, n_edges = build_graph(X, BETA, THRESH)
+        if n_edges > 0:
+            graphs.append(data)
+
+    order = np.arange(len(graphs))
+    model.train()
+    for epoch in range(n_epochs):
+        np.random.shuffle(order)
+        for i in order:
+            optimizer.zero_grad()
+            loss = model.loss(graphs[i])
+            loss.backward()
+            optimizer.step()
+    model.eval()
+    return model
+
+
+def coexpr_embedding(model, X):
+    """z_coexpr for one cell x gene matrix."""
+    data, _ = build_graph(X.astype(float), BETA, THRESH)
+    with torch.no_grad():
+        z = model(data).numpy()
+    return np.nan_to_num(z, nan=0.0)
+
+
+# ---------------------------------------------------------------
+# marginal expression embedding (DeepSets on the cells of each gene)
+# ---------------------------------------------------------------
+
+class DeepSetsEncoder(nn.Module):
+    """Turns the expression of one gene in all cells into a vector. The order of the cells does not matter:
+    every cell goes through phi, the results are pooled (mean and std over the cells) and go through rho.
+    The decoder rebuilds the 49 quantiles of the gene from the vector."""
+
+    def __init__(self, d_hid=32, d_out=16, d_target=49):
+        super().__init__()
+        self.phi = nn.Sequential(nn.Linear(1, d_hid), nn.ReLU(), nn.Linear(d_hid, d_hid))
+        self.rho = nn.Sequential(nn.Linear(2 * d_hid, d_hid), nn.ReLU(), nn.Linear(d_hid, d_out))
+        self.decoder = nn.Sequential(nn.Linear(d_out, d_hid), nn.ReLU(), nn.Linear(d_hid, d_target))
+
+    def encode(self, x):
+        # x has shape (cells, genes). Every gene is encoded on its own.
+        h = self.phi(x.unsqueeze(-1))
+        pooled = torch.cat([h.mean(0), h.std(0)], dim=-1)
+        return self.rho(pooled)
+
+    def forward(self, x):
+        z = self.encode(x)
+        return z, self.decoder(z)
+
+
+def quantile_target(X):
+    # 49 quantiles of each gene, used as the reconstruction target
+    Xn = log_normalize(X)
+    q = np.percentile(Xn, QUANTILE_LEVELS, axis=0)
+    return q.T.astype(np.float32)
+
+
+def train_deepsets(matrices, n_epochs=300, seed=42):
+    """Train the DeepSets encoder on the cell x gene matrices of all timepoints.
+
+    The loss has two parts: the vector of a gene must rebuild its quantiles (loss_rec), and the vectors of two
+    random halves of the cells of the gene must be the same (loss_con)."""
+    torch.manual_seed(seed)
+    rng = np.random.default_rng(seed)
+    model = DeepSetsEncoder()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
+
+    # 1. log-normalise and standardise every gene with the mean and std of all timepoints together
+    logs = [log_normalize(X) for X in matrices]
+    pooled = np.vstack(logs)
+    ref_mean = pooled.mean(0, keepdims=True)
+    ref_std = pooled.std(0, keepdims=True) + 1e-8
+
+    data = []
+    for X, Xn in zip(matrices, logs):
+        x_std = (Xn - ref_mean) / ref_std
+        target = quantile_target(X)
+        target = (target - target.mean(0, keepdims=True)) / (target.std(0, keepdims=True) + 1e-8)
+        data.append((x_std, target))
+
+    # 2. training: one step for every matrix in every epoch
+    model.train()
+    for epoch in range(n_epochs):
+        for x_std, target in data:
+            n = x_std.shape[0]
+            perm = rng.permutation(n)
+            half1, half2 = perm[: n // 2], perm[n // 2:]
+            x_all = torch.tensor(x_std, dtype=torch.float32)
+            x1 = torch.tensor(x_std[half1], dtype=torch.float32)
+            x2 = torch.tensor(x_std[half2], dtype=torch.float32)
+            target = torch.tensor(target, dtype=torch.float32)
+
+            optimizer.zero_grad()
+            _, recon = model(x_all)
+            loss_rec = ((recon - target) ** 2).mean()
+            # two halves of the cells of one gene should give the same embedding
+            loss_con = ((model.encode(x1) - model.encode(x2)) ** 2).mean()
+            loss = loss_rec + loss_con
+            loss.backward()
+            optimizer.step()
+    model.eval()
+    model.ref_mean = ref_mean
+    model.ref_std = ref_std
+    return model
+
+
+def expr_embedding(model, X):
+    """z_expr for one cell x gene matrix."""
+    Xn = log_normalize(X.astype(float))
+    x_std = (Xn - model.ref_mean) / model.ref_std
+    with torch.no_grad():
+        z = model.encode(torch.tensor(x_std, dtype=torch.float32)).numpy()
+    return np.nan_to_num(z, nan=0.0).astype(np.float32)
+
+
+def train_both(matrices, n_genes, n_epochs=300, seed=42):
+    """Train the GCN and the DeepSets encoder on the same matrices."""
+    gcn = train_gcn(matrices, n_genes, n_epochs, seed)
+    deepsets = train_deepsets([X.astype(float) for X in matrices], n_epochs, seed)
+    return gcn, deepsets

gset/graph.py

@@ -0,0 +1,50 @@
+import numpy as np
+import torch
+from torch_geometric.data import Data
+
+
+def log_normalize(X):
+    # library size normalisation followed by log1p
+    return np.log1p(X / (X.sum(1, keepdims=True) + 1e-8) * 1e4)
+
+
+def build_graph(X, beta=6.0, thresh=0.01):
+    """Co-expression graph of the genes (like WGCNA).
+
+    X has shape (cells, genes). The edge weight of two genes is
+    |correlation| ** beta. Weights below thresh are dropped.
+    """
+    Xn = log_normalize(X)
+    mean = Xn.mean(0)
+    std = Xn.std(0)
+    Xs = (Xn - mean) / (std + 1e-8)
+    Xs[:, std < 1e-6] = 0.0
+
+    n_genes = Xs.shape[1]
+    corr = (Xs.T @ Xs) / Xs.shape[0]
+    np.fill_diagonal(corr, 0)
+    corr = np.nan_to_num(corr, nan=0.0)
+    A = np.clip(np.abs(corr) ** beta, 0.0, 1.0)
+
+    rows, cols = np.where(A >= thresh)
+    keep = rows < cols
+    rows, cols = rows[keep], cols[keep]
+    weights = A[rows, cols].astype(np.float32)
+
+    if len(rows) == 0:
+        # nothing passed the threshold, so keep the strongest 1% of the pairs
+        iu = np.triu_indices(n_genes, k=1)
+        vals = A[iu]
+        k = max(1, int(round(0.01 * len(vals))))
+        top = np.argpartition(-vals, k - 1)[:k]
+        rows, cols = iu[0][top], iu[1][top]
+        weights = vals[top].astype(np.float32)
+
+    edge_index = torch.tensor(
+        np.stack([np.concatenate([rows, cols]), np.concatenate([cols, rows])]),
+        dtype=torch.long)
+    edge_weight = torch.tensor(np.concatenate([weights, weights]), dtype=torch.float32)
+    data = Data(x=torch.tensor(corr, dtype=torch.float32),
+                edge_index=edge_index,
+                edge_attr=edge_weight.unsqueeze(1))
+    return data, len(rows)

gset/lags.py

@@ -0,0 +1,25 @@
+import numpy as np
+from sklearn.preprocessing import normalize
+
+
+def lag_matrix(vectors, lag):
+    """Lagged cosine similarity between the genes.
+
+    vectors is a list with one (genes x dim) array per transition.
+    Entry [i, j] compares gene i in transition t with gene j in transition
+    t + lag, averaged over t. With lag 0 the matrix is symmetric.
+    """
+    vecs = [normalize(v, norm="l2", axis=1) for v in vectors]
+    if lag == 0:
+        return sum(v @ v.T for v in vecs) / len(vecs)
+    n = len(vecs) - lag
+    return sum(vecs[t] @ vecs[t + lag].T for t in range(n)) / n
+
+
+def symmetric_score(M, labels):
+    """Best partner of every gene in any other program (both directions)."""
+    score = np.zeros(len(labels))
+    for g in range(len(labels)):
+        other = labels != labels[g]
+        score[g] = max(M[g, other].max(), M[other, g].max())
+    return score

gset/metrics.py

@@ -0,0 +1,87 @@
+import numpy as np
+from sklearn.metrics import roc_auc_score, average_precision_score
+from sklearn.preprocessing import normalize
+
+
+def relay_edges(labels, true_network, a, b):
+    """True edges from program a to program b.
+
+    true_network is the interaction matrix of the simulation. Its first row
+    and column belong to the stimulus, so gene i has index i + 1.
+    """
+    edges = []
+    for i in np.where(labels == a)[0]:
+        for j in np.where(labels == b)[0]:
+            if i != j and true_network[i + 1, j + 1] != 0:
+                edges.append((int(i), int(j)))
+    return edges
+
+
+def relay_coverage(genes, labels, true_network):
+    """For every pair of consecutive programs: how many relay edges have both
+    genes in the selected set."""
+    selected = set(genes)
+    result = {}
+    for a in range(int(labels.max())):
+        edges = relay_edges(labels, true_network, a, a + 1)
+        covered = sum(1 for i, j in edges if i in selected and j in selected)
+        result[f"{a}{a + 1}"] = (covered, len(edges))
+    return result
+
+
+def within_between_gap(V, labels):
+    """Mean cosine similarity of gene pairs in the same program minus the mean
+    of pairs in different programs."""
+    Vn = normalize(V, norm="l2", axis=1)
+    S = Vn @ Vn.T
+    same = labels[:, None] == labels[None, :]
+    upper = np.triu(np.ones_like(same, dtype=bool), k=1)
+    return S[same & upper].mean() - S[~same & upper].mean()
+
+
+def load_relay_edges(sim_dir):
+    """Relay edges of a simulation as {(a, b): set of (i, j)}."""
+    import json
+    from pathlib import Path
+    raw = json.loads((Path(sim_dir) / "relay_edges.json").read_text())
+    edges = {}
+    for key, pairs in raw.items():
+        a, b = (int(x) for x in key.split("_"))
+        edges[(a, b)] = {(int(i), int(j)) for i, j in pairs}
+    return edges
+
+
+def candidate_pairs(labels, tf_mask, a, b, targets="any"):
+    """Pairs (i, j) with i a TF of program a and j a gene of program b.
+    targets="tf" keeps only TFs as j."""
+    sources = [i for i in np.where(labels == a)[0] if tf_mask[i]]
+    if targets == "tf":
+        sinks = [j for j in np.where(labels == b)[0] if tf_mask[j]]
+    else:
+        sinks = list(np.where(labels == b)[0])
+    return [(int(i), int(j)) for i in sources for j in sinks]
+
+
+def forward_discrimination(C, edges, pairs):
+    """AUROC of true edges against the other candidate pairs, using the
+    forward reading C[i, j] (i in the earlier program)."""
+    y = [int((i, j) in edges) for i, j in pairs]
+    return roc_auc_score(y, [C[i, j] for i, j in pairs])
+
+
+def direction_statistics(C, edges, pairs):
+    """How often does the forward reading beat the reversed reading?
+
+    Returns the fraction for true edges, the fraction for the other pairs and
+    the AUROC of the asymmetry C[i, j] - C[j, i] (true edges against the other
+    pairs). A symmetric matrix has no direction and gives None.
+    """
+    if np.abs(C - C.T).max() < 1e-9:
+        return None
+    edge_list = sorted(edges)
+    others = [p for p in pairs if p not in edges]
+    asym = lambda ps: [C[i, j] - C[j, i] for i, j in ps]
+    y = [1] * len(edge_list) + [0] * len(others)
+    return {"edges": float(np.mean([x > 0 for x in asym(edge_list)])),
+            "others": float(np.mean([x > 0 for x in asym(others)])),
+            "auroc": float(roc_auc_score(y, asym(edge_list) + asym(others)))}