"""Simulation and evaluation helpers (NOTEARS tutorial utilities)."""
def is_dag(W: np.ndarray, tol: float = 1e-8) -> bool:
"""Return True if weighted matrix W encodes a DAG."""
d = W.shape[0]
m = (np.abs(W) > tol).astype(float)
acc = m.copy()
for _ in range(d):
if np.trace(acc) > tol:
return False
acc = acc @ m
return True
def simulate_dag(d: int, s0: int, graph_type: str = "ER", seed: int | None = None) -> np.ndarray:
"""Simulate a random binary DAG adjacency (A[i,j]=1 => X_j -> X_i)."""
rng = np.random.default_rng(seed)
if graph_type.upper() == "ER":
prob = min(1.0, s0 / max(d * (d - 1) / 2, 1))
B = np.triu(rng.random((d, d)) < prob, k=1).astype(int)
else: # scale-free-ish upper triangular
B = np.zeros((d, d), dtype=int)
edges = 0
while edges < s0:
i, j = rng.integers(0, d, size=2)
if i < j and B[i, j] == 0:
B[i, j] = 1
edges += 1
return B
def simulate_parameter(B: np.ndarray, low: float = 0.5, high: float = 2.0, seed: int | None = None) -> np.ndarray:
"""Draw edge weights for a binary DAG."""
rng = np.random.default_rng(seed)
W = np.zeros_like(B, dtype=float)
mask = B.astype(bool)
signs = rng.choice([-1.0, 1.0], size=mask.sum())
W[mask] = signs * rng.uniform(low, high, size=mask.sum())
return W
def simulate_linear_sem(W: np.ndarray, n: int, sem_type: str = "gauss", seed: int | None = None) -> np.ndarray:
"""Generate data from a linear SEM in topological order.
With an upper-triangular weight matrix ``W`` (``W[i, j] != 0`` means
``X_i -> X_j``), each node is drawn as ``X_j = sum_i W[i, j] X_i + Z_j``.
Sampling node-by-node in index order is exact here because parents always
have a smaller index than their children.
"""
rng = np.random.default_rng(seed)
d = W.shape[0]
X = np.zeros((n, d))
for j in range(d):
parents = np.where(W[:, j] != 0)[0]
noise = rng.standard_normal(n) if sem_type == "gauss" else rng.exponential(size=n)
if len(parents):
X[:, j] = X[:, parents] @ W[parents, j] + noise
else:
X[:, j] = noise
return X
def count_accuracy(B_true: np.ndarray, B_est: np.ndarray) -> dict:
"""Compare estimated graph to ground truth (FDR, TPR, FPR, SHD, nnz)."""
m = graph_recovery_metrics(B_true, B_est)
tp, fp, fn = m["tp"], m["fp"], m["fn"]
tn = int((~np.eye(B_true.shape[0], dtype=bool)).sum()) - tp - fp - fn
fdr = fp / (tp + fp) if tp + fp else 0.0
tpr = m["recall"]
fpr = fp / (fp + tn) if fp + tn else 0.0
return {"fdr": fdr, "tpr": tpr, "fpr": fpr, "shd": m["shd"], "nnz": int(B_est.sum())}
def simulate_nonlinear_sem_custom(W: np.ndarray, n: int, sem_type: str = "mlp", noise_sd: float = 0.5, seed: int | None = None) -> np.ndarray:
"""Nonlinear SEM with additive noise (A[i,j]: j->i)."""
rng = np.random.default_rng(seed)
d = W.shape[0]
G = nx.DiGraph((j, i) for i in range(d) for j in range(d) if W[i, j] != 0)
order = list(nx.topological_sort(G)) if G.number_of_edges() else list(range(d))
X = np.zeros((n, d))
for i in order:
pa = np.where(W[i, :] != 0)[0]
if len(pa) == 0:
X[:, i] = rng.standard_normal(n)
else:
z = X[:, pa] @ W[i, pa]
if sem_type == "mlp":
X[:, i] = np.tanh(z) + rng.normal(0, noise_sd, n)
elif sem_type == "mim":
X[:, i] = z * np.sin(z) + rng.normal(0, noise_sd, n)
elif sem_type == "gp":
X[:, i] = np.sin(z) + rng.normal(0, noise_sd, n)
else:
X[:, i] = z + rng.standard_normal(n)
return X
