"""Diagnostic probe for DrugRank-Flow weak retrieval. Dumps hard statistics to results/diagnostics/alignment_diag.json so we can reason about WHY gap_emb <-> drug_proj alignment is only ~2-3x chance: 1. Transcriptional signal strength by dose: ||target_mean - source_mean|| (is 10nM basically vehicle-noise?) 2. gap_emb collapse: pairwise cosine spread + per-dim std over many conditions 3. drug_proj separation: pairwise cosine among the 189 gallery drugs 4. Same-drug cross-dose consistency of gap_emb (tests dose-confounding) 5. Retrieval Hit@10 / median-rank stratified by dose 6. gap_emb vs drug_proj: for a batch, cosine(true pair) vs cosine(best wrong) Run: python scripts/diagnose_alignment.py \ --checkpoint outputs/drug_rank/phase2_best.pt \ --config configs/drug_rank_phase2.yaml """ from __future__ import annotations import argparse import json import logging import os import sys from pathlib import Path import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) from gidflow.data.sciplex_dataset import Sciplex3Dataset from gidflow.models.population_encoder import PopulationEncoder from gidflow.models.gap_encoder import GapEncoder from gidflow.models.drug_encoder import DrugEncoder from gidflow.models.drug_gene_bridge import DrugGeneBridge logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(message)s", datefmt="%H:%M:%S") log = logging.getLogger(__name__) ANNOTATION_DIR = Path("/data/boom/ICLR/data/annotation") CELL_LINE_MAP = {"A549": 1, "K562": 2, "MCF7": 3} def build_and_load(cfg, ckpt_path, device): m = cfg["model"] num_proteins = len(json.load(open(ANNOTATION_DIR / "protein_target_vocab.json"))) source_enc = PopulationEncoder(num_genes=m["num_genes"], hidden_dim=m["encoder_hidden"], output_dim=m["encoder_output"]).to(device) target_enc = PopulationEncoder(num_genes=m["num_genes"], hidden_dim=m["encoder_hidden"], output_dim=m["encoder_output"]).to(device) gap_enc = GapEncoder(input_dim=m["encoder_output"], hidden_dim=m["gap_hidden"], output_dim=m["gap_output"], proj_dim=m["gap_proj_dim"], num_cell_lines=m["num_cell_lines"], num_genes=m["num_genes"]).to(device) drug_enc = DrugEncoder(encoding=m.get("drug_encoder", "morgan"), emb_dim=m["drug_emb_dim"], freeze=True).to(device) bridge = DrugGeneBridge(num_proteins=num_proteins, drug_emb_dim=m["drug_emb_dim"], hidden_dim=m["bridge_hidden_dim"], proj_dim=m["bridge_proj_dim"], protein_emb_dim=m["bridge_protein_emb_dim"]).to(device) ck = torch.load(ckpt_path, map_location=device, weights_only=False) source_enc.load_state_dict(ck["source_enc"]); target_enc.load_state_dict(ck["target_enc"]) gap_enc.load_state_dict(ck["gap_enc"]); drug_enc.load_state_dict(ck["drug_enc"]); bridge.load_state_dict(ck["bridge"]) for mod in (source_enc, target_enc, gap_enc, drug_enc, bridge): mod.eval() return source_enc, target_enc, gap_enc, drug_enc, bridge @torch.no_grad() def encode_conditions(dataset, conds, models, device, pass_cell_line): source_enc, target_enc, gap_enc, drug_enc, bridge = models X = dataset._X gaps, dnames, doses, cls, deltas = [], [], [], [], [] for cond in conds: veh = np.asarray(cond["vehicle_cell_idx"]); drg = np.asarray(cond["drug_cell_idx"]) if len(veh) == 0 or len(drg) == 0: continue ns = min(64, len(veh)); nt = min(64, len(drg)) rng = np.random.default_rng(0) s = rng.choice(veh, ns, replace=False); t = rng.choice(drg, nt, replace=False) src = torch.from_numpy(np.asarray(X[s], np.float32))[None].to(device) tgt = torch.from_numpy(np.asarray(X[t], np.float32))[None].to(device) sm = torch.ones(1, ns, dtype=torch.bool, device=device) tm = torch.ones(1, nt, dtype=torch.bool, device=device) z_s = source_enc(src, sm); z_t = target_enc(tgt, tm) if pass_cell_line: cl = torch.tensor([CELL_LINE_MAP.get(cond["cell_line"], 0)], device=device) g = gap_enc(z_s, z_t, cell_line_ids=cl)["gap_emb"] else: g = gap_enc(z_s, z_t)["gap_emb"] gaps.append(g[0].cpu().numpy()) dnames.append(cond["drug_name"]); doses.append(float(cond["dose"])); cls.append(cond["cell_line"]) deltas.append(float(np.linalg.norm(np.asarray(X[t]).mean(0) - np.asarray(X[s]).mean(0)))) return np.array(gaps), dnames, np.array(doses), cls, np.array(deltas) @torch.no_grad() def build_gallery(drug_order, smiles_map, drug_enc, bridge, device): projs = [] for start in range(0, len(drug_order), 32): names = drug_order[start:start+32] smis = [smiles_map.get(n, "C") or "C" for n in names] emb = drug_enc(smis); projs.append(bridge(emb)["drug_proj"].cpu().numpy()) return np.concatenate(projs, 0) def main(): ap = argparse.ArgumentParser() ap.add_argument("--checkpoint", default="outputs/drug_rank/phase2_best.pt") ap.add_argument("--config", default="configs/drug_rank_phase2.yaml") ap.add_argument("--n-conditions", type=int, default=400) args = ap.parse_args() import yaml cfg = yaml.safe_load(open(args.config)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") models = build_and_load(cfg, args.checkpoint, device) source_enc, target_enc, gap_enc, drug_enc, bridge = models dc = cfg["data"]; mc = cfg["model"] smiles_csv = os.path.join(dc["annotation_dir"], "drug_annotation_master.csv") dataset = Sciplex3Dataset(h5ad_path=dc["sciplex3_h5ad"], n_hvg=mc["num_genes"], max_source_cells=64, max_target_cells=64, seed=42, drug_emb_dim=mc["drug_emb_dim"], preprocessed_path=None, drug_smiles_csv=smiles_csv) drug_order = json.load(open(ANNOTATION_DIR / "drug_order.json")) drug_to_idx = {n: i for i, n in enumerate(drug_order)} import csv as _csv smiles_map = {} with open(os.path.join(dc["annotation_dir"], "drug_annotation_master.csv")) as f: for row in _csv.DictReader(f): if row.get("drug_name") and row.get("smiles"): smiles_map[row["drug_name"].strip()] = row["smiles"].strip() out = {} # ---- Signal strength by dose (all conditions) ---- # Cache vehicle means per cell line ONCE (vehicle_cell_idx is shared per line). veh_mean_cache = {} def _veh_mean(cell_line, veh_idx): if cell_line not in veh_mean_cache: veh_mean_cache[cell_line] = dataset._X[np.asarray(veh_idx)].mean(0) return veh_mean_cache[cell_line] by_dose = {} for cond in dataset._conditions: veh = np.asarray(cond["vehicle_cell_idx"]); drg = np.asarray(cond["drug_cell_idx"]) if len(veh) == 0 or len(drg) == 0: continue d = float(np.linalg.norm(dataset._X[drg].mean(0) - _veh_mean(cond["cell_line"], veh))) by_dose.setdefault(float(cond["dose"]), []).append(d) out["signal_by_dose"] = {str(k): {"mean_delta_norm": round(float(np.mean(v)), 4), "std": round(float(np.std(v)), 4), "n": len(v)} for k, v in sorted(by_dose.items())} # vehicle-vehicle noise floor: split vehicle cells of one cell line in half veh_all = dataset._conditions[0]["vehicle_cell_idx"] rng = np.random.default_rng(1) noise = [] for _ in range(20): perm = rng.permutation(np.asarray(veh_all)); h = len(perm)//2 noise.append(float(np.linalg.norm(dataset._X[perm[:h]].mean(0) - dataset._X[perm[h:2*h]].mean(0)))) out["vehicle_noise_floor"] = round(float(np.mean(noise)), 4) # ---- Encode a sample of conditions (BOTH with and without cell_line) ---- rng2 = np.random.default_rng(7) sample = list(rng2.choice(len(dataset._conditions), min(args.n_conditions, len(dataset._conditions)), replace=False)) conds = [dataset._conditions[i] for i in sample] for tag, pass_cl in [("no_cellline_TRAINMODE", False), ("with_cellline_EVALMODE", True)]: gaps, dnames, doses, cls, deltas = encode_conditions(dataset, conds, models, device, pass_cl) gaps_n = gaps / (np.linalg.norm(gaps, axis=1, keepdims=True) + 1e-8) # gap collapse sim = gaps_n @ gaps_n.T off = sim[~np.eye(len(sim), dtype=bool)] # gallery gallery = build_gallery(drug_order, smiles_map, drug_enc, bridge, device) gallery_n = gallery / (np.linalg.norm(gallery, axis=1, keepdims=True) + 1e-8) scores = gaps_n @ gallery_n.T # [N, 189] true_idx = np.array([drug_to_idx.get(n, -1) for n in dnames]) valid = true_idx >= 0 ranks = [] for i in np.where(valid)[0]: order = np.argsort(-scores[i]) ranks.append(int(np.where(order == true_idx[i])[0][0]) + 1) ranks = np.array(ranks) # by dose retrieval vdoses = doses[valid] dose_ret = {} for dv in sorted(set(vdoses.tolist())): rr = ranks[vdoses == dv] dose_ret[str(dv)] = {"hit@10": round(float((rr <= 10).mean()), 4), "median_rank": float(np.median(rr)), "n": int(len(rr))} out[tag] = { "gap_offdiag_cosine_mean": round(float(off.mean()), 4), "gap_offdiag_cosine_std": round(float(off.std()), 4), "gap_perdim_std_mean": round(float(gaps.std(0).mean()), 4), "true_pair_cosine_mean": round(float(np.mean([scores[i, true_idx[i]] for i in np.where(valid)[0]])), 4), "overall_hit@10": round(float((ranks <= 10).mean()), 4), "overall_median_rank": float(np.median(ranks)), "retrieval_by_dose": dose_ret, } # ---- drug_proj separation (shared) ---- gallery = build_gallery(drug_order, smiles_map, drug_enc, bridge, device) gallery_n = gallery / (np.linalg.norm(gallery, axis=1, keepdims=True) + 1e-8) gsim = gallery_n @ gallery_n.T goff = gsim[~np.eye(len(gsim), dtype=bool)] out["drug_proj_separation"] = { "pairwise_cosine_mean": round(float(goff.mean()), 4), "pairwise_cosine_std": round(float(goff.std()), 4), "pairwise_cosine_max": round(float(goff.max()), 4), "n_drugs": len(gallery), } # ---- Same-drug cross-dose gap consistency (train mode encoding) ---- gaps, dnames, doses, cls, deltas = encode_conditions(dataset, conds, models, device, pass_cell_line=False) gaps_n = gaps / (np.linalg.norm(gaps, axis=1, keepdims=True) + 1e-8) from collections import defaultdict drug_groups = defaultdict(list) for i, n in enumerate(dnames): drug_groups[n].append(i) within, across = [], [] for n, idxs in drug_groups.items(): if len(idxs) >= 2: for a in range(len(idxs)): for b in range(a+1, len(idxs)): within.append(float(gaps_n[idxs[a]] @ gaps_n[idxs[b]])) # random across-drug pairs rng3 = np.random.default_rng(3) for _ in range(2000): i, j = rng3.integers(0, len(gaps_n), 2) if dnames[i] != dnames[j]: across.append(float(gaps_n[i] @ gaps_n[j])) out["same_drug_gap_consistency"] = { "within_drug_cosine_mean": round(float(np.mean(within)), 4) if within else None, "across_drug_cosine_mean": round(float(np.mean(across)), 4) if across else None, "n_within_pairs": len(within), "note": "within should be >> across if gap_emb is drug-specific; if similar, dose/noise dominates", } os.makedirs("results/diagnostics", exist_ok=True) outpath = "results/diagnostics/alignment_diag.json" json.dump(out, open(outpath, "w"), indent=2) log.info("Saved %s", outpath) print(json.dumps(out, indent=2)) if __name__ == "__main__": main()