Release 4.0.0: repository is the replication package (protocol, records, audit, conventions, fixed toolkit)
b890fab verified | # -*- coding: utf-8 -*- | |
| """The frozen benchmark protocol. | |
| Every number in the paper comes from this file. It is deliberately small: one | |
| generator family, one scoring rule, one tuning budget, and the two solvers that | |
| are compared. Nothing here depends on the `causalscale` package -- that package | |
| is the *object* being measured, not part of the measurement. | |
| Conventions | |
| ----------- | |
| * A DAG is a matrix ``S`` with ``S[i, j] != 0`` meaning i -> j. | |
| * Data are generated as ``X = X @ S + eps``, so ``X[:, j]`` is a function of its | |
| parents plus noise, and columns are standardised. | |
| * Structure F1 compares ``|W| > tau`` against ``|S| > 0`` with the diagonal | |
| masked out, over the full d x d matrix (no triangular shortcut), so an | |
| implementation that stores its coefficients transposed scores worse, not | |
| differently. | |
| Imports are lazy: ``numpy`` is enough to read the records and rebuild every | |
| table and figure. ``torch`` and ``scipy`` are only needed to *re-run* the | |
| solvers. | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| TAU = 0.3 | |
| # --------------------------------------------------------------------- data ---- | |
| def make_er_dag(d: int, seed: int, p: float | None = None) -> np.ndarray: | |
| """Erdos-Renyi DAG with expected degree 2, oriented by a random order. | |
| ``p = 2/(d-1)`` reproduces the density used by the submission under review. | |
| """ | |
| rng = np.random.default_rng(seed) | |
| if p is None: | |
| p = 2.0 / (d - 1) | |
| mask = np.triu(rng.random((d, d)) < p, k=1) | |
| mag = rng.uniform(0.5, 1.0, size=(d, d)) | |
| sign = np.where(rng.random((d, d)) < 0.5, -1.0, 1.0) | |
| return (mask * mag * sign).astype(np.float64) | |
| def make_ba_dag(d: int, seed: int, m: int = 2) -> np.ndarray: | |
| """Scale-free (Barabasi-Albert) DAG. | |
| Each arriving node receives ``m`` incoming edges drawn with probability | |
| proportional to the current degree of existing nodes, so early nodes become | |
| hubs and the in-degree distribution is heavy-tailed. Edges are oriented | |
| along the arrival order, which makes the graph acyclic by construction. | |
| """ | |
| rng = np.random.default_rng(seed + 4242) | |
| order = rng.permutation(d) | |
| W = np.zeros((d, d)) | |
| deg = np.zeros(d) | |
| for t in range(d): | |
| node = int(order[t]) | |
| if t == 0: | |
| deg[node] = 1.0 | |
| continue | |
| cand = order[:t] | |
| w = deg[cand] + 1.0 | |
| w = w / w.sum() | |
| k = min(m, t) | |
| for c in rng.choice(cand, size=k, replace=False, p=w): | |
| W[int(c), node] = rng.choice([-1.0, 1.0]) * rng.uniform(0.5, 1.0) | |
| deg[int(c)] += 1.0 | |
| deg[node] = 1.0 | |
| return W | |
| def sample_sem(S: np.ndarray, n: int, seed: int, sem: str = "linear-gaussian") -> np.ndarray: | |
| """Draw n samples from the structural equation implied by S.""" | |
| rng = np.random.default_rng(seed + 777_000) | |
| d = S.shape[0] | |
| X = np.zeros((n, d)) | |
| for j in range(d): | |
| pa = np.nonzero(S[:, j])[0] | |
| if sem in ("linear-gaussian", "linear-exponential"): | |
| X[:, j] = X[:, pa] @ S[pa, j] | |
| else: | |
| a = 2.0 if sem == "nonlinear-tanh" else 1.0 | |
| for i in pa: | |
| X[:, j] += S[i, j] * np.tanh(a * X[:, i]) | |
| eps = (rng.standard_normal(n) if sem != "linear-exponential" | |
| else rng.exponential(1.0, n) - 1.0) | |
| X[:, j] += eps | |
| return ((X - X.mean(0)) / X.std(0).clip(1e-8)).astype(np.float64) | |
| def inject_confounder(S, X, seed, q: float = 0.0, strength: float = 1.0): | |
| """Add one latent common cause affecting a fraction q of the edges' children.""" | |
| if q <= 0: | |
| return X, 0 | |
| rng = np.random.default_rng(seed + 555_000) | |
| d = S.shape[0] | |
| src, dst = np.nonzero(S) | |
| k = int(round(q * len(src))) | |
| idx = rng.choice(len(src), size=k, replace=False) | |
| coef = np.zeros(d) | |
| for t in idx: | |
| coef[dst[t]] += strength * rng.uniform(0.5, 1.0) | |
| U = rng.standard_normal(X.shape[0]) | |
| X = X + np.outer(U, coef) | |
| return (X - X.mean(0)) / X.std(0).clip(1e-8), k | |
| # ------------------------------------------------------------------ metrics ---- | |
| def score(W_est, S, tau: float = TAU) -> dict: | |
| """Structure F1 of an estimate against a ground-truth DAG.""" | |
| P = np.abs(W_est) > tau | |
| T = np.abs(S) > 0 | |
| np.fill_diagonal(P, False) | |
| np.fill_diagonal(T, False) | |
| tp = int((P & T).sum()) | |
| fp = int((P & ~T).sum()) | |
| fn = int((~P & T).sum()) | |
| prec = tp / max(tp + fp, 1) | |
| rec = tp / max(tp + fn, 1) | |
| f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0 | |
| return dict(f1=f1, prec=prec, rec=rec, shd=fp + fn, | |
| edges=tp + fp, true_edges=tp + fn) | |
| def score_best(W_est, S, tau_grid=None): | |
| """Best achievable F1 over a threshold sweep (reported alongside tau=0.3).""" | |
| if tau_grid is None: | |
| tau_grid = np.arange(0.05, 0.65, 0.05) | |
| best, bt = 0.0, None | |
| for tau in tau_grid: | |
| v = score(W_est, S, tau)["f1"] | |
| if v > best: | |
| best, bt = v, round(float(tau), 2) | |
| return best, bt | |
| # ------------------------------------------------------------------ solvers ---- | |
| def notears_official(X, lam1: float = 0.1, max_iter: int = 100, h_tol: float = 1e-8, | |
| rho_max: float = 1e16, w_threshold: float = TAU, | |
| lr: float = 1e-4): | |
| """Faithful re-implementation of Zheng et al. (2018) `notears_linear`. | |
| scipy L-BFGS-B on the augmented objective, rho0 = 1.0, x10 growth, early stop | |
| at h(W) < 1e-8. Requires scipy only. | |
| """ | |
| import time | |
| from scipy.linalg import expm | |
| from scipy.optimize import minimize | |
| n, d = X.shape | |
| state = {"alpha": 0.0, "rho": 1.0} | |
| def _h(W): | |
| return np.trace(expm(W * W)) - d | |
| def _h_grad(W): | |
| E = expm(W * W) | |
| return W * (E + E.T) | |
| def _func(w): | |
| W = w.reshape(d, d) | |
| R = X - X @ W | |
| loss = 0.5 / n * (R ** 2).sum() + lam1 * np.abs(W).sum() | |
| G = -1.0 / n * (X.T @ R) + lam1 * np.sign(W) | |
| h = _h(W) | |
| a, r = state["alpha"], state["rho"] | |
| return loss + a * h + 0.5 * r * h * h, (G + (a + r * h) * _h_grad(W)).flatten() | |
| rho, alpha, h, W = 1.0, 0.0, np.inf, np.zeros((d, d)) | |
| t0 = time.time() | |
| while rho < rho_max: | |
| state["alpha"], state["rho"] = alpha, rho | |
| sol = minimize(_func, W.flatten(), method="L-BFGS-B", jac=True, | |
| options=dict(maxiter=max_iter, maxfun=max_iter)) | |
| W = sol.x.reshape(d, d) | |
| h = _h(W) | |
| if h <= h_tol or rho >= rho_max: | |
| break | |
| alpha += rho * h | |
| rho *= 10 | |
| if time.time() - t0 > 900: | |
| break | |
| return W, time.time() - t0, float(h) | |
| def notears_adam(X, lam1: float = 0.1, rho0: float = 1.0, factor: float = 10.0, | |
| outer: int = 100, inner: int = 250, lr: float = 2e-3, | |
| seed: int = 0, device: str | None = None): | |
| """The objective of Eq. (1), optimised with Adam. | |
| ``0.5/n ||X - XW||^2 + lam1*||W||_1`` subject to | |
| ``h(W) = tr(exp(W*W)) - d = 0``, with the augmented-Lagrangian schedule | |
| exposed as ``rho0``/``factor`` and the stopping rule as ``outer``/``inner``. | |
| Those four numbers are the ones a reported "collapse" conflates. | |
| Note: the released `notears` package additionally symmetrises W and forces | |
| non-negative weights. That is a quirk of that release, not of the published | |
| method, so it is not reproduced here. | |
| """ | |
| import time | |
| import torch | |
| if device is None: | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| torch.manual_seed(seed) | |
| Xt = torch.tensor(np.asarray(X, dtype=np.float32), device=device) | |
| n, d = Xt.shape | |
| W = torch.zeros(d, d, requires_grad=True, device=device) | |
| alpha, rho, h_prev = 0.0, rho0, float("inf") | |
| opt = torch.optim.Adam([W], lr=lr) | |
| t0 = time.time() | |
| h_c = float("nan") | |
| for _ in range(outer): | |
| for _ in range(inner): | |
| opt.zero_grad() | |
| R = Xt - Xt @ W | |
| h = torch.trace(torch.linalg.matrix_exp(W * W)) - d | |
| loss = (0.5 / n * (R ** 2).sum() + lam1 * W.abs().sum() | |
| + alpha * h + 0.5 * rho * h * h) | |
| loss.backward() | |
| opt.step() | |
| with torch.no_grad(): | |
| h_c = (torch.trace(torch.linalg.matrix_exp(W * W)) - d).item() | |
| if h_c > 0.25 * h_prev: | |
| rho = min(rho * factor, 1e16) | |
| else: | |
| alpha += rho * h_c | |
| h_prev = h_c | |
| if h_c < 1e-8: | |
| break | |
| return W.detach().cpu().numpy(), time.time() - t0, float(h_c) | |
| def run_dagma(X, seed: int = 0, lam1: float = 0.02, **kw): | |
| """Official DAGMA release, used unchanged.""" | |
| import time | |
| from dagma.linear import DagmaLinear | |
| t0 = time.time() | |
| mdl = DagmaLinear(loss_type="l2") | |
| try: | |
| W = mdl.fit(X, lambda1=lam1, s=1.0, verbose=False) | |
| except TypeError: | |
| W = mdl.fit(X, lambda1=lam1, s=1.0) | |
| return np.asarray(W), time.time() - t0, None | |
| # --------------------------------------------------------------- baselines ---- | |
| _LAM_CACHE: dict[int, float] = {} | |
| LAM_GRID = (0.003, 0.01, 0.03, 0.1, 0.3) | |
| def tune_lam(d: int, name: str = "notears") -> float: | |
| """Pick the L1 penalty on held-out DAGs (seeds 1000-1002), never on the test DAG. | |
| This is the tuning budget every method -- baseline included -- receives. It | |
| is cached per (name, d) because the selection is deterministic. | |
| """ | |
| key = (name, d) | |
| if key in _LAM_CACHE: | |
| return _LAM_CACHE[key] | |
| acc = {lam: [] for lam in LAM_GRID} | |
| for vseed in (1000, 1001, 1002): | |
| Sv = make_er_dag(d, vseed) | |
| Xv = sample_sem(Sv, 2 * d, vseed) | |
| for lam in LAM_GRID: | |
| Wv, _, _ = notears_adam(Xv, lam1=lam, outer=60, seed=0) | |
| acc[lam].append(score(Wv, Sv)["f1"]) | |
| best = max(LAM_GRID, key=lambda l: float(np.mean(acc[l]))) | |
| _LAM_CACHE[key] = best | |
| print(f" [tune] d={d} -> lambda1={best} " | |
| f"(held-out F1 {[round(float(np.mean(acc[l])), 3) for l in LAM_GRID]})", | |
| flush=True) | |
| return best | |
| def run_nt_official(X, seed: int = 0, **kw): | |
| """NOTEARS under the schedule its own paper publishes: lam1=0.1, x10, 100 outer.""" | |
| return notears_adam(X, lam1=0.1, rho0=1.0, factor=10.0, outer=100, seed=seed) | |
| def run_nt_tuned(X, seed: int = 0, **kw): | |
| """NOTEARS with its penalty chosen on held-out DAGs instead of inherited.""" | |
| lam = tune_lam(X.shape[1]) | |
| return notears_adam(X, lam1=lam, rho0=1.0, factor=10.0, outer=100, seed=seed) | |
| def run_nt_original_cfg(X, seed: int = 0, **kw): | |
| """The configuration the submission under review describes for its baseline: | |
| rho0 = 0.01, factor = 2, outer = 10. Recorded so the claim can be checked.""" | |
| return notears_adam(X, lam1=0.1, rho0=0.01, factor=2.0, outer=10, seed=seed) | |
| def run_cs_shipped(X, seed: int = 0, **kw): | |
| """The engine shipped by the toolkit under audit, called through its own API. | |
| Imported lazily so that reading the records never requires the package. | |
| """ | |
| import time | |
| import causalscale as cs | |
| t0 = time.time() | |
| m = cs.CausalDiscovery(X, method="cluster_aware", device=kw.get("device", "cuda"), | |
| verbose=False) | |
| m.fit(verbose=False) | |
| return m._network.adjacency, time.time() - t0, None | |
| METHODS = { | |
| "nt_official": run_nt_official, | |
| "nt_tuned": run_nt_tuned, | |
| "nt_original_cfg": run_nt_original_cfg, | |
| "dagma": run_dagma, | |
| "cs_shipped": run_cs_shipped, | |
| } | |
| # What each method needs installed to run. Record-reading needs nothing. | |
| REQUIRES = {"nt_official": "torch", "nt_tuned": "torch", | |
| "nt_original_cfg": "torch", "dagma": "dagma", "cs_shipped": "causalscale"} | |
| def one_run(d, seed, sem="linear-gaussian", method="nt_official", q=0.0, | |
| n_mult=2.0, tau=TAU, topo="er"): | |
| """One (method, d, seed, condition) cell. Returns a JSON-serialisable row.""" | |
| import time | |
| gen = make_er_dag if topo == "er" else make_ba_dag | |
| S = gen(d, seed) | |
| n = int(round(n_mult * d)) | |
| X = sample_sem(S, n, seed, sem=sem) | |
| n_conf = 0 | |
| if q > 0: | |
| X, n_conf = inject_confounder(S, X, seed, q=q) | |
| t0 = time.time() | |
| try: | |
| W_est, t_fit, extra = METHODS[method](X, seed, S_true=S) | |
| except Exception as e: # noqa: BLE001 | |
| return dict(topo=topo, d=d, seed=seed, sem=sem, method=method, q=q, n=n, | |
| error=f"{type(e).__name__}: {e}", status="error", | |
| wall_s=round(time.time() - t0, 2)) | |
| a = score(W_est, S, tau) | |
| b = score(W_est.T, S, tau) | |
| o, otau = score_best(W_est, S) | |
| return dict(topo=topo, d=d, seed=seed, sem=sem, method=method, q=q, n=n, tau=tau, | |
| f1=a["f1"], prec=a["prec"], rec=a["rec"], shd=a["shd"], | |
| edges=a["edges"], true_edges=a["true_edges"], | |
| f1_flipped=b["f1"], f1_oracle=o, tau_oracle=otau, | |
| fit_time_s=round(float(t_fit), 2), wall_s=round(time.time() - t0, 2), | |
| n_conf_edges=n_conf, | |
| h_final=(float(extra) if isinstance(extra, (int, float)) else None), | |
| status="ok") | |