{
"cells": [
{
"cell_type": "markdown",
"id": "ges_quickstart-0",
"metadata": {
"cell_marker": "\"\"\""
},
"source": [
"# GES quickstart\n",
"\n",
"[](https://colab.research.google.com/github/Abel-ai-lab/andrey/blob/main/examples/ges_quickstart.ipynb)\n",
"\n",
"Greedy Equivalence Search (GES) recovers the CPDAG of the Markov-equivalence class from\n",
"observational linear-Gaussian data. We generate a three-node chain and recover its structure."
]
},
{
"cell_type": "markdown",
"id": "7fb27b941602401d91542211134fc71a",
"metadata": {},
"source": [
"## Start\n",
"\n",
"Choose a CPU runtime and run this cell first. In Colab it installs `andrey-core==0.9.0` from PyPI;\n",
"local kernels use the installed environment."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "acae54e37e7d407bbb7b55eff062a284",
"metadata": {
"tags": [
"colab-setup"
]
},
"outputs": [],
"source": [
"import os\n",
"import sys\n",
"\n",
"IN_COLAB = \"google.colab\" in sys.modules\n",
"\n",
"if IN_COLAB:\n",
" os.environ[\"ANDREY_DEVICE\"] = \"cpu\"\n",
" !pip install --quiet \"andrey-core[viz]==0.9.0\" pandas"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "ges_quickstart-1",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import pandas as pd\n",
"\n",
"import andrey"
]
},
{
"cell_type": "markdown",
"id": "ges_quickstart-2",
"metadata": {},
"source": [
"A linear-Gaussian structural equation model -- a chain `x0 -> x1 -> x2`. The columns come in an\n",
"order drawn from the seed, so their order says nothing about the causal order; `shuffle=False`\n",
"keeps `x0, x1, x2`."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "ges_quickstart-3",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"columns: ['x0', 'x2', 'x1']\n"
]
}
],
"source": [
"CHAIN = [(\"x0\", \"x1\"), (\"x1\", \"x2\")] # the true edges, cause first\n",
"\n",
"\n",
"def sample(seed: int, n: int = 1500, shuffle: bool = True) -> pd.DataFrame:\n",
" \"\"\"Draw n rows of the chain, its columns in an order drawn from the seed.\"\"\"\n",
" rng = np.random.default_rng(seed)\n",
" x0 = rng.standard_normal(n)\n",
" x1 = 0.9 * x0 + rng.standard_normal(n)\n",
" x2 = 0.9 * x1 + rng.standard_normal(n)\n",
" frame = pd.DataFrame({\"x0\": x0, \"x1\": x1, \"x2\": x2})\n",
" return frame[rng.permutation(frame.columns)] if shuffle else frame\n",
"\n",
"\n",
"data = sample(seed=0)\n",
"print(\"columns:\", list(data.columns))"
]
},
{
"cell_type": "markdown",
"id": "ges_quickstart-4",
"metadata": {},
"source": [
"Fit GES and inspect the recovered CPDAG (endpoint-mark adjacency) and its BIC score. The chain is\n",
"Markov-equivalent to its undirected skeleton, so the CPDAG leaves the two edges unoriented."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "ges_quickstart-5",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"kind: cpdag\n",
"CPDAG adjacency, in column order:\n",
"[[0 0 1]\n",
" [0 0 1]\n",
" [1 1 0]]\n",
"edges: [('x0', 'x1', 'undirected'), ('x2', 'x1', 'undirected')]\n",
"BIC score: 1.0951\n"
]
}
],
"source": [
"out = andrey.ges(data)\n",
"print(\"kind:\", out.structure.kind)\n",
"print(\"CPDAG adjacency, in column order:\")\n",
"print(out.structure.to_numpy())\n",
"print(\"edges:\", out.structure.oriented_edges())\n",
"print(\"BIC score:\", round(out.metadata[\"score\"], 4))"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "ges_quickstart-6",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
""
],
"text/plain": [
"''"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import andrey.viz\n",
"\n",
"andrey.viz.draw(out.structure)"
]
},
{
"cell_type": "markdown",
"id": "ges_quickstart-7",
"metadata": {},
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}