causalscale / scripts /validate_depmap_f1.py

Release 4.0.0: repository is the replication package (protocol, records, audit, conventions, fixed toolkit)

b890fab verified about 1 month ago

10.3 kB

"""
DepMap F1 Validation: causalscale LowRankGNN vs STRING/TRRUST Gold Standard
============================================================================
Verifies the paper claim: "LowRankGNN identifies 28 causal edges (F1>0.97) on DepMap"
using the actual causalscale package (not standalone prototype).

Strategy:
1. Load DepMap CRISPR CERES scores (1208 cell lines x 18531 genes)
2. Run causalscale CausalDiscovery with method="lowrank"
3. Get top edges
4. Compare against STRING+TRRUST curated gold standard
5. GO enrichment analysis
"""
import sys
import os, os, time, json, gzip
import numpy as np
import pandas as pd
from collections import defaultdict

# Add causalscale to path
# causalscale installed via pip install -e .
import torch
print(f"CUDA: {torch.cuda.is_available()}")
print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")

# ───────────────────────────────────────────────────────────────────
# 1. LOAD DEPMAP DATA
# ───────────────────────────────────────────────────────────────────
print("\n[1] Loading DepMap CRISPR data...")
t0 = time.time()

# DATA: Path to DepMap CRISPR data. Download from https://depmap.org/ and set DATA accordingly.
DATA = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "depmap")
# Use CRISPRGeneEffect (CERES scores) - the gold standard for gene dependency
df = pd.read_csv(f"{DATA}/CRISPRGeneEffect.csv", index_col=0)
print(f"  Raw shape: {df.shape}")

# Parse gene names from column headers (format: "GENE (ID)")
gene_names = []
for col in df.columns:
    gene = col.split(" (")[0] if " (" in col else col
    gene_names.append(gene)

# Drop genes with >50% NaN
nan_frac = df.isna().mean()
keep_genes = nan_frac[nan_frac < 0.5].index
df = df[keep_genes]
gene_names_filtered = [gene_names[df.columns.get_loc(c)] for c in df.columns]
print(f"  After NaN filter: {df.shape}")

# Fill remaining NaN with 0
X_raw = df.fillna(0).values.astype(np.float32)
print(f"  Final matrix: {X_raw.shape[0]} cell lines x {X_raw.shape[1]} genes")

# Standardize
from sklearn.preprocessing import StandardScaler
X = StandardScaler().fit_transform(X_raw).astype(np.float32)
d = X.shape[1]
print(f"  d={d}, n={X.shape[0]}, memory={X.nbytes/1e6:.1f} MB")
print(f"  Load time: {time.time()-t0:.1f}s")

# ───────────────────────────────────────────────────────────────────
# 2. RUN CAUSALSCALE LOWRANKGNN
# ───────────────────────────────────────────────────────────────────
print(f"\n[2] Running causalscale LowRankGNN (d={d}, rank=64)...")
t1 = time.time()

from causalscale.core.lowrank import LowRankGNN, train_lowrank_gnn

result = train_lowrank_gnn(
    X,
    rank=64,
    epochs=500,        # more epochs for genome-scale data
    threshold=0.3,
    lr=0.005,
    device="cuda",
    verbose=True
)

train_time = time.time() - t1
W = result["adjacency"]
n_edges = result["gnn_edges"]
print(f"\n  Training complete: {train_time/60:.1f} min")
print(f"  Edges found (|W|>0.3): {n_edges}")
print(f"  Correlation-reconstruction F1: {result['f1']:.4f}")

# Get top edges
edges = []
for i in range(d):
    for j in range(d):
        if i != j and abs(W[i, j]) > 0.3:
            edges.append((i, j, float(W[i, j])))
edges.sort(key=lambda x: -abs(x[2]))
print(f"  Top edges: {len(edges)}")

# ───────────────────────────────────────────────────────────────────
# 3. LOAD STRING + TRRUST GOLD STANDARD
# ───────────────────────────────────────────────────────────────────
print("\n[3] Loading STRING/TRRUST gold standard...")

# Build gene name -> index mapping
gene_to_idx = {g: i for i, g in enumerate(gene_names_filtered)}

# STRING PPI
# STRING_PATH: Download from https://string-db.org/ and set accordingly.
STRING_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "validation", "string_ppi_full.txt.gz")
string_pairs = set()
try:
    with gzip.open(STRING_PATH, "rt", encoding="utf-8", errors="ignore") as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) >= 2:
                string_pairs.add((parts[0], parts[1]))
    print(f"  STRING pairs: {len(string_pairs)}")
except Exception as e:
    print(f"  STRING load error: {e}")
    string_pairs = set()

# TRRUST
# TRRUST_PATH: Download from https://www.grnpedia.org/trrust/ and set accordingly.
TRRUST_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "validation", "trrust_human.tsv")
trrust_pairs = set()
try:
    with open(TRRUST_PATH, encoding="utf-8") as f:
        for line in f:
            parts = line.strip().split("\t")
            if len(parts) >= 2:
                trrust_pairs.add((parts[0].upper(), parts[1].upper()))
    print(f"  TRRUST pairs: {len(trrust_pairs)}")
except Exception as e:
    print(f"  TRRUST load error: {e}")
    trrust_pairs = set()

# Merge gold standard
gold_pairs = string_pairs | trrust_pairs
print(f"  Gold standard (UNION): {len(gold_pairs)} pairs")

# Map to indices
gold_idx_pairs = set()
for src, tgt in gold_pairs:
    if src in gene_to_idx and tgt in gene_to_idx:
        gold_idx_pairs.add((gene_to_idx[src], gene_to_idx[tgt]))
print(f"  Gold pairs mapped to causalscale indices: {len(gold_idx_pairs)}")

# ───────────────────────────────────────────────────────────────────
# 4. COMPUTE F1
# ───────────────────────────────────────────────────────────────────
print("\n[4] Computing F1 vs gold standard...")

# Predicted edges (top N and various thresholds)
def compute_f1(pred_edges, gold, top_n=None):
    """Compute F1 for top-N predicted edges vs gold standard."""
    if top_n:
        pred_set = set((i, j) for i, j, _ in pred_edges[:top_n])
    else:
        pred_set = set((i, j) for i, j, _ in pred_edges)
    
    tp = len(pred_set & gold)
    fp = len(pred_set - gold)
    fn = len(gold - pred_set)
    
    prec = tp / max(tp + fp, 1)
    rec = tp / max(tp + fn, 1)
    f1 = 2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0
    
    return {"tp": tp, "fp": fp, "fn": fn, "prec": prec, "rec": rec, "f1": f1}

# Try multiple thresholds
for thresh in [0.1, 0.2, 0.3, 0.4, 0.5]:
    top_edges = [(i, j, float(W[i, j])) for i in range(d) for j in range(d)
                 if i != j and abs(W[i, j]) > thresh]
    top_edges.sort(key=lambda x: -abs(x[2]))
    
    if len(top_edges) > 0:
        stats = compute_f1(top_edges, gold_idx_pairs)
        print(f"  threshold={thresh}: edges={len(top_edges)}, TP={stats['tp']}, "
              f"FP={stats['fp']}, FN={stats['fn']}, F1={stats['f1']:.4f}")

# Also try top-28 (matching paper claim)
top_28 = edges[:28]
stats_28 = compute_f1(top_28, gold_idx_pairs)
print(f"\n  Top-28 edges (paper claim):")
print(f"    TP={stats_28['tp']}, FP={stats_28['fp']}, FN={stats_28['fn']}")
print(f"    Precision={stats_28['prec']:.4f}, Recall={stats_28['rec']:.4f}, F1={stats_28['f1']:.4f}")

# ───────────────────────────────────────────────────────────────────
# 5. GO ENRICHMENT (simplified)
# ───────────────────────────────────────────────────────────────────
print("\n[5] Edge gene pairs (for GO enrichment)...")
top_edge_genes = []
for i, j, w in edges[:50]:
    gn_i = gene_names_filtered[i]
    gn_j = gene_names_filtered[j]
    top_edge_genes.append((gn_i, gn_j, w))
    if len(top_edge_genes) <= 20:
        print(f"  {gn_i} -> {gn_j}: {w:.4f}")

# ───────────────────────────────────────────────────────────────────
# SAVE RESULTS
# ───────────────────────────────────────────────────────────────────
result_json = {
    "engine": "causalscale v3.0.0 LowRankGNN",
    "data": "DepMap 24Q2 CRISPR CERES",
    "shape": [int(X.shape[0]), int(X.shape[1])],
    "rank": 64,
    "epochs": 500,
    "train_time_s": round(train_time, 1),
    "total_edges_found": n_edges,
    "top28_f1": round(stats_28["f1"], 4),
    "top28_precision": round(stats_28["prec"], 4),
    "top28_recall": round(stats_28["rec"], 4),
    "top28_tp": stats_28["tp"],
    "top28_fp": stats_28["fp"],
    "top28_fn": stats_28["fn"],
    "gold_standard_size": len(gold_idx_pairs),
    "top_edges": [(gene_names_filtered[i], gene_names_filtered[j], round(w, 4))
                  for i, j, w in edges[:100]]
}

OUT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results")
with open(f"{OUT}/exp13_depmap_causalscale_f1.json", "w") as f:
    json.dump(result_json, f, indent=2)
print(f"\nResults saved to exp13_depmap_causalscale_f1.json")
print(f"\nTotal wall time: {(time.time()-t0)/60:.1f} min")

	"""
	DepMap F1 Validation: causalscale LowRankGNN vs STRING/TRRUST Gold Standard
	============================================================================
	Verifies the paper claim: "LowRankGNN identifies 28 causal edges (F1>0.97) on DepMap"
	using the actual causalscale package (not standalone prototype).

	Strategy:
	1. Load DepMap CRISPR CERES scores (1208 cell lines x 18531 genes)
	2. Run causalscale CausalDiscovery with method="lowrank"
	3. Get top edges
	4. Compare against STRING+TRRUST curated gold standard
	5. GO enrichment analysis
	"""
	import sys
	import os, os, time, json, gzip
	import numpy as np
	import pandas as pd
	from collections import defaultdict

	# Add causalscale to path
	# causalscale installed via pip install -e .
	import torch
	print(f"CUDA: {torch.cuda.is_available()}")
	print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")

	# ───────────────────────────────────────────────────────────────────
	# 1. LOAD DEPMAP DATA
	# ───────────────────────────────────────────────────────────────────
	print("\n[1] Loading DepMap CRISPR data...")
	t0 = time.time()

	# DATA: Path to DepMap CRISPR data. Download from https://depmap.org/ and set DATA accordingly.
	DATA = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "depmap")
	# Use CRISPRGeneEffect (CERES scores) - the gold standard for gene dependency
	df = pd.read_csv(f"{DATA}/CRISPRGeneEffect.csv", index_col=0)
	print(f" Raw shape: {df.shape}")

	# Parse gene names from column headers (format: "GENE (ID)")
	gene_names = []
	for col in df.columns:
	gene = col.split(" (")[0] if " (" in col else col
	gene_names.append(gene)

	# Drop genes with >50% NaN
	nan_frac = df.isna().mean()
	keep_genes = nan_frac[nan_frac < 0.5].index
	df = df[keep_genes]
	gene_names_filtered = [gene_names[df.columns.get_loc(c)] for c in df.columns]
	print(f" After NaN filter: {df.shape}")

	# Fill remaining NaN with 0
	X_raw = df.fillna(0).values.astype(np.float32)
	print(f" Final matrix: {X_raw.shape[0]} cell lines x {X_raw.shape[1]} genes")

	# Standardize
	from sklearn.preprocessing import StandardScaler
	X = StandardScaler().fit_transform(X_raw).astype(np.float32)
	d = X.shape[1]
	print(f" d={d}, n={X.shape[0]}, memory={X.nbytes/1e6:.1f} MB")
	print(f" Load time: {time.time()-t0:.1f}s")

	# ───────────────────────────────────────────────────────────────────
	# 2. RUN CAUSALSCALE LOWRANKGNN
	# ───────────────────────────────────────────────────────────────────
	print(f"\n[2] Running causalscale LowRankGNN (d={d}, rank=64)...")
	t1 = time.time()

	from causalscale.core.lowrank import LowRankGNN, train_lowrank_gnn

	result = train_lowrank_gnn(
	X,
	rank=64,
	epochs=500, # more epochs for genome-scale data
	threshold=0.3,
	lr=0.005,
	device="cuda",
	verbose=True
	)

	train_time = time.time() - t1
	W = result["adjacency"]
	n_edges = result["gnn_edges"]
	print(f"\n Training complete: {train_time/60:.1f} min")
	print(f" Edges found (\|W\|>0.3): {n_edges}")
	print(f" Correlation-reconstruction F1: {result['f1']:.4f}")

	# Get top edges
	edges = []
	for i in range(d):
	for j in range(d):
	if i != j and abs(W[i, j]) > 0.3:
	edges.append((i, j, float(W[i, j])))
	edges.sort(key=lambda x: -abs(x[2]))
	print(f" Top edges: {len(edges)}")

	# ───────────────────────────────────────────────────────────────────
	# 3. LOAD STRING + TRRUST GOLD STANDARD
	# ───────────────────────────────────────────────────────────────────
	print("\n[3] Loading STRING/TRRUST gold standard...")

	# Build gene name -> index mapping
	gene_to_idx = {g: i for i, g in enumerate(gene_names_filtered)}

	# STRING PPI
	# STRING_PATH: Download from https://string-db.org/ and set accordingly.
	STRING_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "validation", "string_ppi_full.txt.gz")
	string_pairs = set()
	try:
	with gzip.open(STRING_PATH, "rt", encoding="utf-8", errors="ignore") as f:
	for line in f:
	parts = line.strip().split()
	if len(parts) >= 2:
	string_pairs.add((parts[0], parts[1]))
	print(f" STRING pairs: {len(string_pairs)}")
	except Exception as e:
	print(f" STRING load error: {e}")
	string_pairs = set()

	# TRRUST
	# TRRUST_PATH: Download from https://www.grnpedia.org/trrust/ and set accordingly.
	TRRUST_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "validation", "trrust_human.tsv")
	trrust_pairs = set()
	try:
	with open(TRRUST_PATH, encoding="utf-8") as f:
	for line in f:
	parts = line.strip().split("\t")
	if len(parts) >= 2:
	trrust_pairs.add((parts[0].upper(), parts[1].upper()))
	print(f" TRRUST pairs: {len(trrust_pairs)}")
	except Exception as e:
	print(f" TRRUST load error: {e}")
	trrust_pairs = set()

	# Merge gold standard
	gold_pairs = string_pairs \| trrust_pairs
	print(f" Gold standard (UNION): {len(gold_pairs)} pairs")

	# Map to indices
	gold_idx_pairs = set()
	for src, tgt in gold_pairs:
	if src in gene_to_idx and tgt in gene_to_idx:
	gold_idx_pairs.add((gene_to_idx[src], gene_to_idx[tgt]))
	print(f" Gold pairs mapped to causalscale indices: {len(gold_idx_pairs)}")

	# ───────────────────────────────────────────────────────────────────
	# 4. COMPUTE F1
	# ───────────────────────────────────────────────────────────────────
	print("\n[4] Computing F1 vs gold standard...")

	# Predicted edges (top N and various thresholds)
	def compute_f1(pred_edges, gold, top_n=None):
	"""Compute F1 for top-N predicted edges vs gold standard."""
	if top_n:
	pred_set = set((i, j) for i, j, _ in pred_edges[:top_n])
	else:
	pred_set = set((i, j) for i, j, _ in pred_edges)

	tp = len(pred_set & gold)
	fp = len(pred_set - gold)
	fn = len(gold - pred_set)

	prec = tp / max(tp + fp, 1)
	rec = tp / max(tp + fn, 1)
	f1 = 2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0

	return {"tp": tp, "fp": fp, "fn": fn, "prec": prec, "rec": rec, "f1": f1}

	# Try multiple thresholds
	for thresh in [0.1, 0.2, 0.3, 0.4, 0.5]:
	top_edges = [(i, j, float(W[i, j])) for i in range(d) for j in range(d)
	if i != j and abs(W[i, j]) > thresh]
	top_edges.sort(key=lambda x: -abs(x[2]))

	if len(top_edges) > 0:
	stats = compute_f1(top_edges, gold_idx_pairs)
	print(f" threshold={thresh}: edges={len(top_edges)}, TP={stats['tp']}, "
	f"FP={stats['fp']}, FN={stats['fn']}, F1={stats['f1']:.4f}")

	# Also try top-28 (matching paper claim)
	top_28 = edges[:28]
	stats_28 = compute_f1(top_28, gold_idx_pairs)
	print(f"\n Top-28 edges (paper claim):")
	print(f" TP={stats_28['tp']}, FP={stats_28['fp']}, FN={stats_28['fn']}")
	print(f" Precision={stats_28['prec']:.4f}, Recall={stats_28['rec']:.4f}, F1={stats_28['f1']:.4f}")

	# ───────────────────────────────────────────────────────────────────
	# 5. GO ENRICHMENT (simplified)
	# ───────────────────────────────────────────────────────────────────
	print("\n[5] Edge gene pairs (for GO enrichment)...")
	top_edge_genes = []
	for i, j, w in edges[:50]:
	gn_i = gene_names_filtered[i]
	gn_j = gene_names_filtered[j]
	top_edge_genes.append((gn_i, gn_j, w))
	if len(top_edge_genes) <= 20:
	print(f" {gn_i} -> {gn_j}: {w:.4f}")

	# ───────────────────────────────────────────────────────────────────
	# SAVE RESULTS
	# ───────────────────────────────────────────────────────────────────
	result_json = {
	"engine": "causalscale v3.0.0 LowRankGNN",
	"data": "DepMap 24Q2 CRISPR CERES",
	"shape": [int(X.shape[0]), int(X.shape[1])],
	"rank": 64,
	"epochs": 500,
	"train_time_s": round(train_time, 1),
	"total_edges_found": n_edges,
	"top28_f1": round(stats_28["f1"], 4),
	"top28_precision": round(stats_28["prec"], 4),
	"top28_recall": round(stats_28["rec"], 4),
	"top28_tp": stats_28["tp"],
	"top28_fp": stats_28["fp"],
	"top28_fn": stats_28["fn"],
	"gold_standard_size": len(gold_idx_pairs),
	"top_edges": [(gene_names_filtered[i], gene_names_filtered[j], round(w, 4))
	for i, j, w in edges[:100]]
	}

	OUT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results")
	with open(f"{OUT}/exp13_depmap_causalscale_f1.json", "w") as f:
	json.dump(result_json, f, indent=2)
	print(f"\nResults saved to exp13_depmap_causalscale_f1.json")
	print(f"\nTotal wall time: {(time.time()-t0)/60:.1f} min")