Examples

GES quickstart#

Open in Colab Download this notebook

Greedy Equivalence Search (GES) recovers the CPDAG of the Markov-equivalence class from observational linear-Gaussian data. We generate a three-node chain and recover its structure.

Start#

Choose a CPU runtime and run this cell first. In Colab it installs andrey-core==0.9.0 from PyPI; local kernels use the installed environment.

import os
import sys

IN_COLAB = "google.colab" in sys.modules

if IN_COLAB:
    os.environ["ANDREY_DEVICE"] = "cpu"
    !pip install --quiet "andrey-core[viz]==0.9.0" pandas

import numpy as np
import pandas as pd

import andrey

A linear-Gaussian structural equation model – a chain x0 -> x1 -> x2. The columns come in an order drawn from the seed, so their order says nothing about the causal order; shuffle=False keeps x0, x1, x2.

CHAIN = [("x0", "x1"), ("x1", "x2")]  # the true edges, cause first


def sample(seed: int, n: int = 1500, shuffle: bool = True) -> pd.DataFrame:
    """Draw n rows of the chain, its columns in an order drawn from the seed."""
    rng = np.random.default_rng(seed)
    x0 = rng.standard_normal(n)
    x1 = 0.9 * x0 + rng.standard_normal(n)
    x2 = 0.9 * x1 + rng.standard_normal(n)
    frame = pd.DataFrame({"x0": x0, "x1": x1, "x2": x2})
    return frame[rng.permutation(frame.columns)] if shuffle else frame


data = sample(seed=0)
print("columns:", list(data.columns))

columns: ['x0', 'x2', 'x1']

Fit GES and inspect the recovered CPDAG (endpoint-mark adjacency) and its BIC score. The chain is Markov-equivalent to its undirected skeleton, so the CPDAG leaves the two edges unoriented.

out = andrey.ges(data)
print("kind:", out.structure.kind)
print("CPDAG adjacency, in column order:")
print(out.structure.to_numpy())
print("edges:", out.structure.oriented_edges())
print("BIC score:", round(out.metadata["score"], 4))

kind: cpdag
CPDAG adjacency, in column order:
[[0 0 1]
 [0 0 1]
 [1 1 0]]
edges: [('x0', 'x1', 'undirected'), ('x2', 'x1', 'undirected')]
BIC score: 1.0951

import andrey.viz

andrey.viz.draw(out.structure)
x0x2x1