File size: 8,625 Bytes
b890fab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | # Full experiment matrix for KDD 2027 - causalscale paper
# Checkpointing: never re-run completed experiments
import sys, os, json, time, warnings, numpy as np, torch, pandas as pd
warnings.filterwarnings('ignore')
# causalscale installed via pip install -e .
from causalscale.core._notears import run_notears
DEVICE = 'cuda'
OUT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'results')
os.makedirs(OUT, exist_ok=True)
def load(name):
p = os.path.join(OUT, name)
return json.load(open(p, encoding='utf-8')) if os.path.exists(p) else {}
def save(name, data):
with open(os.path.join(OUT, name), 'w', encoding='utf-8') as f:
json.dump(data, f, indent=2)
def er_dag(d, seed=42):
rng = np.random.RandomState(seed)
W = np.zeros((d,d))
for i in range(d):
for j in range(i):
if rng.rand() < 2/(d-1):
W[i,j] = rng.choice([-1,1]) * rng.uniform(0.5, 1.0)
return W
def sf_dag(d, seed=42):
rng = np.random.RandomState(seed)
W, deg = np.zeros((d,d)), np.ones(d)
for i in range(1,d):
probs = deg[:i]/deg[:i].sum()
n_e = max(1, rng.poisson(2))
for p in rng.choice(i, size=min(n_e,i), replace=False, p=probs):
W[i,p] = rng.choice([-1,1]) * rng.uniform(0.5, 1.0)
return W
def gen_data(W, n):
X = np.linalg.inv(np.eye(W.shape[0]) - W) @ np.random.randn(W.shape[0], n)
return X.T.astype(np.float32)
def metrics(W_true, W_est, tau=0.3):
mt = np.abs(W_true) > 0; me = np.abs(W_est) > tau
tp = np.sum(mt & me); fp = np.sum(~mt & me); fn = np.sum(mt & ~me)
p = tp/(tp+fp) if (tp+fp)>0 else 0; r = tp/(tp+fn) if (tp+fn)>0 else 0
f = 2*p*r/(p+r) if (p+r)>0 else 0
return {'f1': round(f,4), 'shd': int(fp+fn), 'prec': round(p,4), 'rec': round(r,4), 'tp': int(tp), 'fp': int(fp), 'fn': int(fn)}
# ============================================================
# EXP1: causalscale on Synthetic ER (6 dims x 5 seeds)
# ============================================================
print('EXP1: causalscale Synthetic ER Benchmark')
ckpt = load('exp1_causalscale_er.json')
for d in [30,50,80,100,150,200]:
for seed in range(5):
key = f'd{d}_s{seed}'
if key in ckpt: print(f' [skip] {key}'); continue
W_true = er_dag(d, seed); X = gen_data(W_true, 2*d)
t0 = time.time()
W, ec, h, _ = run_notears(X, device=DEVICE, outer=30, inner=200, seed=seed)
elapsed = time.time() - t0
ckpt[key] = {'d':d, 'seed':seed, 'true_edges': int(np.sum(np.abs(W_true)>0)), **metrics(W_true, W), 'h': round(h,2), 'cs_edges': ec, 'time': round(elapsed,1)}
save('exp1_causalscale_er.json', ckpt)
print(f' [done] {key}: f1={ckpt[key]["f1"]}, {elapsed:.0f}s')
# ============================================================
# EXP2: causalscale on Scale-Free (2 dims x 3 seeds)
# ============================================================
print('EXP2: causalscale Scale-Free Topology')
ckpt = load('exp2_causalscale_sf.json')
for d in [50, 100]:
for seed in range(3):
key = f'd{d}_s{seed}'
if key in ckpt: print(f' [skip] {key}'); continue
W_true = sf_dag(d, seed); X = gen_data(W_true, 2*d)
t0 = time.time()
W, ec, h, _ = run_notears(X, device=DEVICE, outer=30, inner=200, seed=seed)
ckpt[key] = {'d':d, 'seed':seed, 'true_edges': int(np.sum(np.abs(W_true)>0)), **metrics(W_true, W), 'time': round(time.time()-t0,1)}
save('exp2_causalscale_sf.json', ckpt)
print(f' [done] {key}: f1={ckpt[key]["f1"]}')
# ============================================================
# EXP3: Tau Sensitivity (d=50,100 x 3 seeds x 8 taus)
# ============================================================
print('EXP3: Tau Sensitivity')
ckpt = load('exp3_tau_sensitivity.json')
for d in [50, 100]:
for seed in range(3):
base_key = f'd{d}_s{seed}'
# Check if all taus are done
all_done = all(f'{base_key}_t{tau:.2f}' in ckpt for tau in [0.1,0.15,0.2,0.25,0.3,0.35,0.4,0.5])
if all_done: print(f' [skip] {base_key} (all taus)'); continue
W_true = er_dag(d, seed); X = gen_data(W_true, 2*d)
W, ec, h, _ = run_notears(X, device=DEVICE, outer=30, inner=200, seed=seed)
for tau in [0.1,0.15,0.2,0.25,0.3,0.35,0.4,0.5]:
tk = f'{base_key}_t{tau:.2f}'
if tk in ckpt: continue
ckpt[tk] = {'d':d, 'seed':seed, 'tau':tau, **metrics(W_true, W, tau=tau)}
save('exp3_tau_sensitivity.json', ckpt)
print(f' [done] {base_key}: all 8 taus')
# ============================================================
# EXP4: Multi-Seed Pan-Cancer (33 cancers x 3 seeds)
# ============================================================
print('EXP4: Multi-Seed Pan-Cancer')
ckpt = load('exp4_pancancer_3seed.json')
# TSD: Path to TCGA data. Download from https://xenabrowser.net/ and set TSD accordingly.
TSD = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'data')
for fname in sorted(os.listdir(TSD)):
if not fname.startswith('TCGA_') or not fname.endswith('.tsv') or 'HiSeqV2' not in fname: continue
cancer = fname.replace('TCGA_','').replace('_HiSeqV2.tsv','')
# Check if already have 3 seeds
if cancer in ckpt and len(ckpt[cancer].get('seeds',{})) >= 3:
print(f' [skip] {cancer}'); continue
df = pd.read_csv(os.path.join(TSD, fname), sep='\t', index_col=0, nrows=None)
# Fix orientation: ensure rows=samples, columns=genes
if df.index.name != 'sample' or not str(df.index[0]).startswith('TCGA-'):
df = df.T # genes were rows, flip to samples as rows
# Drop any non-gene columns that may have been read
numeric_cols = df.select_dtypes(include=['float64', 'float32', 'int64']).columns
df = df[numeric_cols]
targets = ['ARID1A','MTOR']
present = [g for g in targets if g in df.columns]
var_genes = df.var().nlargest(100).index.tolist()
selected = list(dict.fromkeys(present + [g for g in var_genes if g not in targets][:100]))
X = df[selected].values.astype(np.float32)
seeds_data = ckpt.get(cancer, {}).get('seeds', {})
for seed in range(3):
sk = str(seed)
if sk in seeds_data: continue
t0 = time.time()
W, ec, h, _ = run_notears(torch.tensor(X, device=DEVICE), device=DEVICE, outer=30, inner=200, seed=seed)
ai = selected.index('ARID1A') if 'ARID1A' in selected else -1
mi = selected.index('MTOR') if 'MTOR' in selected else -1
seeds_data[sk] = {'edges': ec, 'a2m': round(float(W[ai,mi]),4) if ai>=0 and mi>=0 else 0,
'm2a': round(float(W[mi,ai]),4) if ai>=0 and mi>=0 else 0,
'time': round(time.time()-t0,1), 'h': round(h,2)}
ckpt[cancer] = {'n': X.shape[0], 'd': X.shape[1], 'seeds': seeds_data}
save('exp4_pancancer_3seed.json', ckpt)
edges = [v['edges'] for v in seeds_data.values()]
print(f' [done] {cancer} (n={X.shape[0]}): edges={edges}')
# ============================================================
# EXP5: CPU vs GPU Speedup (d=50,80,100)
# ============================================================
print('EXP5: CPU vs GPU Speedup')
ckpt = load('exp5_cpu_vs_gpu.json')
for d in [50, 80, 100]:
key = f'd{d}'
if key in ckpt: print(f' [skip] {key}'); continue
W_true = er_dag(d); X = gen_data(W_true, 2*d)
X_gpu = torch.tensor(X, device='cuda')
# GPU (warmup first)
t0 = time.time(); run_notears(X_gpu, device='cuda', outer=30, inner=200, seed=0)
gpu_t = time.time() - t0
# CPU
t0 = time.time(); run_notears(X_gpu.cpu(), device='cpu', outer=30, inner=200, seed=0)
cpu_t = time.time() - t0
ckpt[key] = {'d': d, 'gpu_s': round(gpu_t,1), 'cpu_s': round(cpu_t,1), 'speedup': round(cpu_t/gpu_t,1)}
save('exp5_cpu_vs_gpu.json', ckpt)
print(f' [done] d={d}: GPU={gpu_t:.1f}s, CPU={cpu_t:.1f}s, {cpu_t/gpu_t:.1f}x')
# ============================================================
# EXP6: Full Metrics Report
# ============================================================
print('\nEXP6: Summary Report')
print('='*60)
for name in sorted(os.listdir(OUT)):
if not name.endswith('.json'): continue
d = load(name)
kb = os.path.getsize(os.path.join(OUT, name))/1024
n = len(d) if isinstance(d, dict) else 'list'
print(f' {name}: {kb:.0f}KB, {n} entries')
print('='*60)
print('DONE. All results in', OUT)
|