transfer_base.py

#!/usr/bin/env python3
"""UGRN Transfer Base: Feature extraction + transfer learning for GRN inference."""

import argparse
import json
import os
import random

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.cluster import KMeans
from sklearn.metrics import average_precision_score, roc_auc_score
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm

from ugrn.config import get_genelink_label_dir, select_experiments
from ugrn.dataset import load_expression_matrix, load_genelink_splits
from ugrn.model import scModel
from ugrn.vocab import GeneVocab


def seed_everything(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


class ModelArgs:
    def __init__(self, vocab_len, config_path=None):
        self.__dict__.update(
            {
                "embsize": 256,
                "nlayers": 6,
                "nheads": 8,
                "use_mvc": True,
                "d_hid": 256,
                "dropout": 0.15,
                "pad_token": "<pad>",
                "pad_value": 0,
                "mask_value": -1,
                "cls_value": 0,
                "n_bins": 51,
                "input_emb_style": "continuous",
                "cell_emb_style": "cls",
                "model_structure": "transformer",
                "vocab": None,
                "ntoken": vocab_len,
            }
        )
        if config_path and os.path.exists(config_path):
            with open(config_path) as f:
                self.__dict__.update(
                    {
                        k: v
                        for k, v in json.load(f).items()
                        if k
                        in [
                            "embsize",
                            "nlayers",
                            "nheads",
                            "use_mvc",
                            "d_hid",
                            "dropout",
                        ]
                    }
                )


def load_my_model(ckpt_path, device):
    p = os.path.dirname(ckpt_path)
    v_path = None
    for candidate in [os.path.join(p, "vocab.json"), "vocab.json"]:
        if os.path.exists(candidate):
            v_path = candidate
            break
    if not v_path:
        raise FileNotFoundError("vocab.json not found. Please place it in the model directory or repo root.")
    vocab = GeneVocab.from_file(v_path)
    for s in ["<pad>", "<cls>", "<eoc>"]:
        if s not in vocab:
            vocab.append_token(s)
    model = scModel(vocab, args=ModelArgs(len(vocab), os.path.join(p, "args.json")))
    sd = torch.load(ckpt_path, map_location=device)
    model.load_state_dict(
        {k[7:] if k.startswith("module.") else k: v for k, v in sd.items()},
        strict=False,
    )
    return model.to(device).eval(), vocab


@torch.no_grad()
def my_model_forward_batch(model, vocab, g_ids, vals, device):
    bs = vals.shape[0]
    src_ids = torch.tensor(g_ids, device=device).expand(bs, -1)
    src_val = torch.tensor(vals, device=device).float()
    inp_ids = torch.cat(
        [torch.full((bs, 1), vocab["<cls>"], device=device), src_ids], 1
    )
    inp_vals = torch.cat(
        [torch.full((bs, 1), 0.0, device=device), src_val], 1
    )
    out = model(
        inp_ids,
        inp_vals,
        src_key_padding_mask=torch.zeros_like(inp_ids, dtype=torch.bool),
    )
    return (
        out.get("mvc_output", out.get("mlm_output"))
    )[:, 1:].cpu().numpy()


def _load_id_name_map(path):
    df = pd.read_csv(path)
    id_c = (
        "index"
        if "index" in df
        else next(
            (c for c in df if np.issubdtype(df[c].dtype, np.number)), "_row_id"
        )
    )
    if id_c == "_row_id":
        df["_row_id"] = np.arange(len(df))
    nm_c = next(
        (c for c in ["Gene", "TF", "Target", "gene", "tf", "target"] if c in df),
        None,
    ) or next(
        (c for c in df if not np.issubdtype(df[c].dtype, np.number)), str(id_c)
    )
    return dict(zip(df[id_c].astype(int), df[nm_c].astype(str)))


def extract_dataset_features(net, ds, args, device, model, vocab):
    ld = get_genelink_label_dir(args.genelink_root, net, ds, args.num_tfs)
    if not os.path.exists(ld):
        return None
    splits = load_genelink_splits(
        ld,
        seed=args.seed,
        train_ratio=args.train_ratio,
        val_ratio=args.val_ratio,
        pos_neg_ratio=args.pos_neg_ratio,
    )
    try:
        all_p = np.concatenate([s[0] for s in splits[:3] if s[0] is not None], 0)
        all_l = np.concatenate([s[1] for s in splits[:3] if s[1] is not None], 0)
    except Exception:
        return None

    id2g = {
        **_load_id_name_map(os.path.join(ld, "TF.csv")),
        **_load_id_name_map(os.path.join(ld, "Target.csv")),
    }
    en, em = load_expression_matrix(ld)
    if em.shape[0] != len(en):
        em = em.T
    n2i = {g.upper(): i for i, g in enumerate(en)}
    g2i = vocab.get_stoi()

    v_idx, v_ids, d_pos, temp_indices = [], [], {}, []
    for i in sorted(set(all_p.flatten())):
        g = id2g.get(int(i), "").upper()
        if g in g2i and g in n2i:
            d_pos[int(i)] = len(v_idx)
            v_idx.append(int(i))
            v_ids.append(g2i[g])
            temp_indices.append(n2i[g])
    if not v_idx:
        return None

    mask = np.isin(all_p[:, 0], v_idx) & np.isin(all_p[:, 1], v_idx)
    valid_p = all_p[mask]
    valid_l = all_l[mask]
    v_ids_arr = np.array(v_ids, dtype=np.int64)
    feat_parts = []

    if args.feature_type == "embedding":
        if hasattr(model, "embedding"):
            emb_layer = model.embedding
        elif hasattr(model, "encoder") and hasattr(model.encoder, "embedding"):
            emb_layer = model.encoder.embedding
        else:
            emb_layer = list(model.modules())[1]

        emb_weight = emb_layer.weight.detach().cpu().numpy()
        feat_parts.append(
            np.array(
                [
                    emb_weight[g2i[id2g[s].upper()]] + emb_weight[g2i[id2g[t].upper()]]
                    for s, t in valid_p
                ],
                dtype=np.float32,
            )
        )

    elif args.feature_type == "perturbation":
        X_expr = em[temp_indices, :].T
        base_arrs = (
            KMeans(
                n_clusters=args.num_clusters, n_init=10, random_state=args.seed
            )
            .fit(X_expr)
            .cluster_centers_
            if args.num_clusters > 1
            else X_expr.mean(axis=0, keepdims=True)
        ).astype(np.float32)
        u_genes, all_c_feats = list(set(valid_p.flatten())), []
        for k in range(args.num_clusters):
            curr_base = base_arrs[k]
            y_base = my_model_forward_batch(
                model, vocab, v_ids_arr, curr_base[None, :], device
            )[0]
            p_map = {}
            for i in tqdm(
                range(0, len(u_genes), 32), desc=f"Perturbing (C{k})"
            ):
                bg = u_genes[i : i + 32]
                for f in args.input_factors:
                    binp = np.tile(curr_base, (len(bg), 1))
                    for idx, g in enumerate(bg):
                        binp[idx, d_pos[g]] *= f
                    delta = (
                        my_model_forward_batch(model, vocab, v_ids_arr, binp, device)
                        - y_base
                    )
                    for idx, g in enumerate(bg):
                        p_map.setdefault(g, {})[f] = delta[idx]
            all_c_feats.append(
                np.array(
                    [
                        [p_map[s][f][d_pos[t]] for f in args.input_factors]
                        + [p_map[t][f][d_pos[s]] for f in args.input_factors]
                        + [curr_base[d_pos[s]], curr_base[d_pos[t]]]
                        for s, t in valid_p
                    ],
                    dtype=np.float32,
                )
            )
        feat_parts.append(np.concatenate(all_c_feats, axis=1))

    if not feat_parts:
        return None
    features = np.concatenate(feat_parts, axis=1)
    return (
        pd.DataFrame(features, columns=[f"feat_{i}" for i in range(features.shape[1])])
        .assign(TF=valid_p[:, 0], Target=valid_p[:, 1], Label=valid_l.astype(int))
    )


class SimpleMLP(nn.Module):
    def __init__(self, d_in, dh=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, dh),
            nn.ReLU(),
            nn.Linear(dh, dh // 2),
            nn.ReLU(),
            nn.Linear(dh // 2, 1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        return self.net(x)


def train_model_subset(model, df, dev, c_start, c_end, ep=15, bs=128):
    X = torch.tensor(df.iloc[:, c_start:c_end].values, dtype=torch.float32).to(dev)
    y = (
        torch.tensor(df["Label"].values, dtype=torch.float32)
        .unsqueeze(1)
        .to(dev)
    )
    model.train()
    opt, crit = optim.Adam(model.parameters(), lr=1e-3), nn.BCELoss()
    for _ in range(ep):
        for bx, by in DataLoader(TensorDataset(X, y), bs, shuffle=True):
            opt.zero_grad()
            crit(model(bx), by).backward()
            opt.step()
    return model


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--exp_name", default="ours_transfer_single")
    p.add_argument("--save_root", default="results/transfer_base")
    p.add_argument("--feature_dir", default="save_features")
    p.add_argument("--save_features", action="store_true")
    p.add_argument("--model_ckpt", default="./checkpoints/model.pt", help="Path to model checkpoint.")
    p.add_argument(
        "--genelink_root",
        default="./data/GENELink/Dataset/Benchmark Dataset",
        help="Path to GENELink Benchmark Dataset root.",
    )
    p.add_argument("--benchmark", default="all")
    p.add_argument("--dataset", default="all")
    p.add_argument("--train_source", nargs="+", default=["STRING:hESC"])
    p.add_argument("--num_tfs", type=int, default=500)
    p.add_argument(
        "--feature_type",
        default="perturbation",
        choices=["embedding", "perturbation"],
    )
    p.add_argument("--input_factors", nargs="+", type=float, default=[0.0])
    p.add_argument("--force_extract", action="store_true")
    p.add_argument("--num_clusters", type=int, default=1)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--epochs", type=int, default=50)
    p.add_argument("--runs", type=int, default=3)
    p.add_argument("--downsample_ratio", type=float, default=1.0)
    p.add_argument("--pos_neg_ratio", type=float, default=1.0)
    p.add_argument("--train_ratio", type=float, default=0.8)
    p.add_argument("--val_ratio", type=float, default=0.1)
    p.add_argument("--add_name", type=str, default="")
    p.add_argument("--classifier_model", type=str, default="mlp")
    args = p.parse_args()
    seed_everything(args.seed)
    args.add_name = (
        f"{args.add_name}_K{args.num_clusters}_{args.feature_type}"
        if args.add_name
        else f"K{args.num_clusters}_{args.feature_type}"
    )
    os.makedirs(args.save_root, exist_ok=True)
    os.makedirs(args.feature_dir, exist_ok=True)
    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    exps = select_experiments(args.benchmark, args.dataset)
    cache, model, vocab = {}, None, None

    print(f"\n=== Phase 1: Feature Extraction ({args.feature_type}) ===")
    for net, ds in exps:
        fp = os.path.join(
            args.feature_dir,
            f"{net}_{ds}_{args.feature_type}_K{args.num_clusters}_features.parquet",
        )
        if os.path.exists(fp) and not args.force_extract:
            cache[(net, ds)] = pd.read_parquet(fp)
            print(f"[Loaded Cache] {net}-{ds}")
        else:
            if not model:
                if not args.model_ckpt:
                    raise ValueError("--model_ckpt is required for feature extraction.")
                model, vocab = load_my_model(args.model_ckpt, dev)
            df = extract_dataset_features(net, ds, args, dev, model, vocab)
            if df is not None:
                cache[(net, ds)] = df
                print(f"[Done] {net}-{ds}")
                if args.save_features:
                    df.to_parquet(fp, index=False)

    print(f"\n=== Phase 2: Training (Source: {args.train_source}) ===")
    tr_k = [k for k in cache if f"{k[0]}:{k[1]}" in args.train_source]
    if not tr_k:
        print(f"[Error] No training datasets match {args.train_source}")
        return
    tr_df = pd.concat([cache[k] for k in tr_k], axis=0)
    if args.downsample_ratio < 1:
        tr_df = tr_df.sample(frac=args.downsample_ratio, random_state=args.seed)

    total_dim = len([c for c in tr_df.columns if c.startswith("feat_")])
    res = {k: [] for k in cache}

    for i in range(args.runs):
        seed_everything(args.seed + i)
        print(f"[Run {i + 1}] Training...")

        m = train_model_subset(
            SimpleMLP(total_dim).to(dev), tr_df, dev, 0, total_dim, args.epochs
        )

        for k in cache:
            if k in tr_k:
                res[k].append({"test_auroc": 0.0, "test_aupr": 0.0})
            else:
                m.eval()
                with torch.no_grad():
                    pred = (
                        m(
                            torch.tensor(
                                cache[k].filter(like="feat_").values,
                                dtype=torch.float32,
                            ).to(dev)
                        )
                        .cpu()
                        .numpy()
                        .flatten()
                    )
                try:
                    auc = (
                        roc_auc_score(cache[k]["Label"], pred)
                        if cache[k]["Label"].nunique() > 1
                        else 0.5
                    )
                    aupr = (
                        average_precision_score(cache[k]["Label"], pred)
                        if cache[k]["Label"].nunique() > 1
                        else 0.0
                    )
                except Exception:
                    auc, aupr = 0.5, 0.0
                res[k].append({"test_auroc": auc, "test_aupr": aupr})

    rows = []
    for (n, d), v in res.items():
        aurocs = [x["test_auroc"] for x in v]
        auprs = [x["test_aupr"] for x in v]
        rows.append(
            {
                "Network": n,
                "Dataset": d,
                "Runs": len(v),
                "Mean Test AUROC": np.mean(aurocs),
                "Mean Test AUPR": np.mean(auprs),
                "Std Test AUROC": np.std(aurocs),
                "Std Test AUPR": np.std(auprs),
            }
        )

    odf = pd.DataFrame(rows)
    cols = [
        "Network",
        "Dataset",
        "Runs",
        "Mean Test AUROC",
        "Mean Test AUPR",
        "Std Test AUROC",
        "Std Test AUPR",
    ]
    odf = odf[cols]

    test_metrics = odf[
        ~odf.apply(lambda x: f"{x['Network']}:{x['Dataset']}" in args.train_source, axis=1)
    ]
    avg_row = {
        "Network": "-",
        "Dataset": "Test Avg",
        "Runs": args.runs,
        "Mean Test AUROC": test_metrics["Mean Test AUROC"].mean(),
        "Mean Test AUPR": test_metrics["Mean Test AUPR"].mean(),
        "Std Test AUROC": test_metrics["Std Test AUROC"].mean(),
        "Std Test AUPR": test_metrics["Std Test AUPR"].mean(),
    }

    odf = pd.concat([odf, pd.DataFrame([avg_row])], ignore_index=True)
    out = os.path.join(args.save_root, f"{args.exp_name}_{args.add_name}_summary.xlsx")
    odf.to_excel(out, index=False)
    print(f"[Saved] {out}")


if __name__ == "__main__":
    main()