"""Synthetic model generation, following Appendix D.1 of arXiv:2606.00278. Two procedures are specified in the paper and both are implemented verbatim: :func:`sample_linear_model` (used for Figures 2, 3, 8, 9) and :func:`sample_graphical_model` (used for Figure 5). """ from __future__ import annotations import numpy as np from .graphical import ( StatementGraph, has_confounding_path, transitive_closure, ) from .linear import A_from_gamma __all__ = [ "sample_linear_model", "perturb_statements", "sample_graphical_model", "inject_graph_errors", "marginal_admg", ] # -------------------------------------------------------------------------- # D.1, linear procedure (steps 1-6) # -------------------------------------------------------------------------- def sample_linear_model(n: int, m: int, p: float, rng: np.random.Generator): """Sample a ground-truth linear Gaussian SEM and its true statements. Returns ``(Sigma_obs, A_true, Gamma_obs)`` where ``Sigma_obs`` is the correlation matrix of the ``n`` observed variables, ``Gamma_obs`` the marginalised causal coefficient matrix and ``A_true = (I - Gamma_obs)^{-1}`` the matrix of correct bivariate causal statements. Steps, quoting Appendix D.1: 1. draw causal coefficients for ``n + m`` variables from a standard normal and error variances from an exponential distribution with variance 1; 2. zero each coefficient independently with probability ``1 - p``; 3. compute the covariance matrix assuming independent errors; 4. rescale error variances and coefficients so every variable has unit variance; 5. marginalise to ``n`` uniformly chosen observed variables (Lemma 2.1); 6. read off ``A = (I - Gamma)^{-1}``. """ N = n + m Gamma = np.tril(rng.normal(size=(N, N)), -1) Gamma *= (rng.random((N, N)) < p) Gamma = np.tril(Gamma, -1) # exponential with variance 1 => scale 1 (mean 1, variance 1) err_var = rng.exponential(scale=1.0, size=N) A_full = A_from_gamma(Gamma) Sigma = A_full @ np.diag(err_var) @ A_full.T # Step 4: rescale so that every variable has unit variance. Writing # S = diag(sd), the standardised model has Gamma' = S^{-1} Gamma S and # error variances err_var / sd^2. sd = np.sqrt(np.diag(Sigma)) Gamma = Gamma * np.outer(1.0 / sd, sd) Sigma = Sigma / np.outer(sd, sd) # Step 5: marginalise to n observed variables. hidden = rng.choice(N, size=m, replace=False) obs = np.array(sorted(set(range(N)) - set(hidden.tolist()))) Z = np.array(sorted(hidden.tolist())) if m > 0: G_YY = Gamma[np.ix_(obs, obs)] G_YZ = Gamma[np.ix_(obs, Z)] G_ZZ = Gamma[np.ix_(Z, Z)] G_ZY = Gamma[np.ix_(Z, obs)] Gamma_obs = G_YY + G_YZ @ np.linalg.solve(np.eye(m) - G_ZZ, G_ZY) else: Gamma_obs = Gamma.copy() Gamma_obs = np.tril(Gamma_obs, -1) Sigma_obs = Sigma[np.ix_(obs, obs)] A_true = A_from_gamma(Gamma_obs) return Sigma_obs, A_true, Gamma_obs def perturb_statements(A_true: np.ndarray, sigma: float, rng: np.random.Generator) -> np.ndarray: """D.1 step 7: add centred Gaussian noise of variance ``sigma`` to each bivariate statement. ``sigma`` is a *variance*, as written in the paper, so the added noise has standard deviation ``sqrt(sigma)``.""" n = A_true.shape[0] A = A_true.copy() noise = rng.normal(scale=np.sqrt(sigma), size=(n, n)) idx = np.tril_indices(n, -1) A[idx] = A_true[idx] + noise[idx] np.fill_diagonal(A, 1.0) return np.tril(A) # -------------------------------------------------------------------------- # D.1, graphical procedure # -------------------------------------------------------------------------- def marginal_admg(D: np.ndarray, Bd: np.ndarray, keep: np.ndarray) -> StatementGraph: """Definition 3.2: marginalise a mixed graph onto the vertices ``keep``.""" N = D.shape[0] full = StatementGraph(N, D, Bd) hidden = sorted(set(range(N)) - set(keep.tolist())) k = len(keep) # 1. directed edge v -> w if a directed path exists with all intermediate # vertices hidden. reach = D.copy() for h in hidden: # allow routing only through hidden vertices reach = reach | (np.outer(reach[:, h], reach[h, :]) & ~np.eye(N, dtype=bool)) Dm = reach[np.ix_(keep, keep)].copy() np.fill_diagonal(Dm, False) # 2. bidirected edge v <-> w if a confounding path exists with all # intermediate vertices hidden. We test on the subgraph induced by # {v, w} together with the hidden vertices, which is exactly the # "all intermediates hidden" restriction. Bm = np.zeros((k, k), bool) for a in range(k): for b in range(a + 1, k): v, w = int(keep[a]), int(keep[b]) sub = sorted(set(hidden) | {v, w}) idx = np.array(sub) sg = StatementGraph(len(sub), D[np.ix_(idx, idx)], Bd[np.ix_(idx, idx)]) if has_confounding_path(sg, sub.index(v), sub.index(w)): Bm[a, b] = Bm[b, a] = True return StatementGraph(k, Dm, Bm) def sample_graphical_model(n: int, m: int, p: float, rng: np.random.Generator): """Sample ground truth and the correct graphical bivariate statements. Steps 1-4 of the graphical procedure in Appendix D.1. The returned statement graph is the union of the pairwise marginal ADMGs, and is by construction graphically compatible (Lemma 3.5). """ N = n + m perm = rng.permutation(N) D = np.zeros((N, N), bool) for a in range(N): for b in range(a + 1, N): if rng.random() < p: D[perm[a], perm[b]] = True Bd = np.zeros((N, N), bool) hidden = rng.choice(N, size=m, replace=False) keep = np.array(sorted(set(range(N)) - set(hidden.tolist()))) marg = marginal_admg(D, Bd, keep) # Step 4: marginalise further onto each pair and take the union. Under a # pairwise marginalisation every other vertex is hidden, so the directed # part of the union is exactly the transitive closure of the observed # graph, and a bidirected edge appears wherever a confounding path exists. Bs = np.zeros((n, n), bool) Ds = transitive_closure(marg.D) for a in range(n): for b in range(a + 1, n): if has_confounding_path(marg, a, b): Bs[a, b] = Bs[b, a] = True return StatementGraph(n, Ds, Bs) def inject_graph_errors(g: StatementGraph, k: int, rng: np.random.Generator) -> StatementGraph: """D.1 step 5: toggle ``k`` uniformly chosen edge slots. A slot is either one of the ``n(n-1)`` ordered pairs (a directed edge) or one of the ``C(n,2)`` unordered pairs (a bidirected edge); the chosen slot is deleted if present and inserted if absent. Slots are drawn without replacement so that exactly ``k`` statements are altered. """ n = g.n slots = [("d", u, v) for u in range(n) for v in range(n) if u != v] slots += [("b", u, v) for u in range(n) for v in range(u + 1, n)] pick = rng.choice(len(slots), size=k, replace=False) D, Bd = g.D.copy(), g.Bd.copy() for idx in pick: kind, u, v = slots[int(idx)] if kind == "d": D[u, v] = not D[u, v] else: Bd[u, v] = Bd[v, u] = not Bd[u, v] return StatementGraph(n, D, Bd)