{ "cells": [ { "cell_type": "markdown", "id": "score_search-0", "metadata": {}, "source": [ "# Score-based search (XGES / hill climbing / Tabu search / exact search / GIES)\n", "\n", "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Abel-ai-lab/andrey/blob/main/examples/score_search.ipynb)\n", "\n", "XGES, hill climbing, and Tabu search look for the graph with the best BIC score. XGES searches\n", "over equivalence classes, as GES does, with deletions first. Hill climbing adds, removes, or\n", "reverses one edge of a DAG at a time and stops when no move improves the score; Tabu search keeps\n", "going past that point and returns the best DAG it visits. All three return a CPDAG and suit\n", "linear-Gaussian data. Exact search finds the best-scoring DAG on a few variables, and GIES also\n", "learns from interventions." ] }, { "cell_type": "markdown", "id": "score_search-1", "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": "score_search-2", "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": "score_search-3", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import pandas as pd\n", "from IPython.display import display\n", "\n", "import andrey\n", "import andrey.metrics as metrics\n", "import andrey.viz\n", "from andrey.data import benchmarks" ] }, { "cell_type": "markdown", "id": "score_search-4", "metadata": {}, "source": [ "## Data\n", "\n", "`linear_gauss_er` draws an Erdos-Renyi graph with linear mechanisms and Gaussian noise: here 10\n", "variables and 2,000 samples." ] }, { "cell_type": "code", "execution_count": 3, "id": "score_search-5", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "2000 observations, 10 variables, seed 7\n" ] }, { "data": { "text/html": [ "0123456789" ], "text/plain": [ "'0123456789'" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "seed = 7\n", "dataset = benchmarks.scm(\"linear_gauss_er\", d=10).sample(n=2000, seed=seed, scale=\"standardize\")\n", "X, truth = dataset.data, dataset.observed_graph\n", "print(f\"{X.shape[0]} observations, {X.shape[1]} variables, seed {seed}\")\n", "display(andrey.viz.draw(truth))" ] }, { "cell_type": "markdown", "id": "score_search-6", "metadata": {}, "source": [ "## Fit XGES, hill climbing, and Tabu search\n", "\n", "`metadata[\"score\"]` is the BIC score of the returned graph, which each search minimizes. A larger\n", "`lambda_value` weighs the BIC penalty more and gives a sparser graph. With\n", "`extended_search=False`, XGES skips its extended search for a faster fit, and `tabu_length` sets\n", "how many recently visited DAGs Tabu search may not return to." ] }, { "cell_type": "code", "execution_count": 4, "id": "score_search-7", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "XGES: cpdag, score -22406.2\n" ] }, { "data": { "text/html": [ "0123456789" ], "text/plain": [ "'0123456789'" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "Hill climbing: cpdag, score -22356.9\n" ] }, { "data": { "text/html": [ "0123456789" ], "text/plain": [ "'0123456789'" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "Tabu search: cpdag, score -22368.4\n" ] }, { "data": { "text/html": [ "0123456789" ], "text/plain": [ "'0123456789'" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "results = {\n", " \"XGES\": andrey.xges(X),\n", " \"Hill climbing\": andrey.hc(X),\n", " \"Tabu search\": andrey.tabu(X),\n", "}\n", "for name, result in results.items():\n", " print(f\"{name}: {result.structure.kind}, score {result.metadata['score']:.1f}\")\n", " display(andrey.viz.draw(result.structure))" ] }, { "cell_type": "markdown", "id": "score_search-8", "metadata": {}, "source": [ "## Compare with GES over ten datasets\n", "\n", "Structural Hamming distance (SHD) counts errors against the truth's CPDAG; lower is better. The\n", "empty graph's SHD is the number of true edges, so an estimate at or above it has not improved on\n", "returning no edges. One dataset can favor any method, so the table covers ten seeds and their\n", "median." ] }, { "cell_type": "code", "execution_count": 5, "id": "score_search-9", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
seed0123456789median
Empty graph1515151515151515151515.0
GES211001017711170010.0
XGES08091016000.5
Hill climbing286082516221141312.0
Tabu search60012091612003.5
\n", "
" ], "text/plain": [ "seed 0 1 2 3 4 5 6 7 8 9 median\n", "Empty graph 15 15 15 15 15 15 15 15 15 15 15.0\n", "GES 21 10 0 10 17 7 11 17 0 0 10.0\n", "XGES 0 8 0 9 1 0 1 6 0 0 0.5\n", "Hill climbing 28 6 0 8 25 16 22 11 4 13 12.0\n", "Tabu search 6 0 0 1 20 9 16 12 0 0 3.5" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "methods = {\n", " \"GES\": lambda X: andrey.ges(X),\n", " \"XGES\": lambda X: andrey.xges(X),\n", " \"Hill climbing\": lambda X: andrey.hc(X),\n", " \"Tabu search\": lambda X: andrey.tabu(X),\n", "}\n", "rows = {}\n", "for s in range(10):\n", " draw = benchmarks.scm(\"linear_gauss_er\", d=10).sample(n=2000, seed=s, scale=\"standardize\")\n", " d = draw.observed_graph.n_nodes\n", " empty = andrey.GraphStructure.from_numpy(np.zeros((d, d), dtype=np.uint8), kind=\"cpdag\")\n", " rows[s] = {\"Empty graph\": metrics.score(empty, draw.observed_graph)[\"shd\"]}\n", " for name, fit in methods.items():\n", " rows[s][name] = metrics.score(fit(draw.data), draw.observed_graph)[\"shd\"]\n", "shd = pd.DataFrame(rows)\n", "shd.columns.name = \"seed\"\n", "shd[\"median\"] = shd.median(axis=1)\n", "shd" ] }, { "cell_type": "markdown", "id": "score_search-10", "metadata": {}, "source": [ "Tabu search starts with hill climbing's moves, so its score is never worse than hill\n", "climbing's on the same data. The [0.3.0 release card](https://andrey.abel.ai/docs/benchmarks/0.3.0.html)\n", "compares XGES with the xges package and the other two with bnlearn, on larger graphs with ten rows\n", "per variable." ] }, { "cell_type": "markdown", "id": "score_search-11", "metadata": {}, "source": [ "## Exact search\n", "\n", "`exact_search` scores every parent set of every variable and returns a DAG with the best BIC\n", "score. Its run time at least doubles with each added variable, so keep it to about 20 variables." ] }, { "cell_type": "code", "execution_count": 6, "id": "score_search-12", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Exact search: dag\n" ] }, { "data": { "text/html": [ "0123456789" ], "text/plain": [ "'0123456789'" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "exact = andrey.exact_search(X)\n", "print(\"Exact search:\", exact.structure.kind)\n", "display(andrey.viz.draw(exact.structure))" ] }, { "cell_type": "markdown", "id": "score_search-13", "metadata": {}, "source": [ "## Interventions: GIES\n", "\n", "`gies` also takes the rows' interventions: `True` where a row's value of a variable was set from\n", "outside. A chain from observational data stays undirected; here the last 1,000 rows set `x1`, which\n", "orients both of its edges. The columns come in an order drawn from the seed, and the interventions\n", "follow them." ] }, { "cell_type": "code", "execution_count": 7, "id": "score_search-14", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[('x1', 'x2', 'directed'), ('x0', 'x1', 'directed')]\n" ] }, { "data": { "text/html": [ "x1x2x0" ], "text/plain": [ "'x1x2x0'" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "CHAIN = [(\"x0\", \"x1\"), (\"x1\", \"x2\")] # the true edges, cause first\n", "\n", "\n", "def sample_intervened(\n", " seed: int, n: int = 3000, shuffle: bool = True\n", ") -> tuple[pd.DataFrame, np.ndarray]:\n", " \"\"\"Draw n rows of the chain, the last 1,000 setting x1 from outside, and the interventions.\"\"\"\n", " rng = np.random.default_rng(seed)\n", " x0 = rng.standard_normal(n)\n", " x1 = 0.8 * x0 + rng.standard_normal(n)\n", " x1[n - 1000 :] = rng.standard_normal(1000) # the last 1,000 rows set x1 from outside\n", " x2 = 0.8 * x1 + rng.standard_normal(n)\n", " frame = pd.DataFrame({\"x0\": x0, \"x1\": x1, \"x2\": x2})\n", " frame = frame[rng.permutation(frame.columns)] if shuffle else frame\n", " intervened = np.zeros((n, 3), dtype=bool)\n", " intervened[n - 1000 :, frame.columns.get_loc(\"x1\")] = True\n", " return frame, intervened\n", "\n", "\n", "gies_out = andrey.gies(*sample_intervened(seed=1))\n", "print(gies_out.structure.oriented_edges())\n", "display(andrey.viz.draw(gies_out.structure))" ] }, { "cell_type": "markdown", "id": "score_search-15", "metadata": {}, "source": [ "## Use another dataset\n", "\n", "Change `seed`, or the `d` and `n` passed to `benchmarks.scm` and `sample`, and rerun." ] } ], "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.14.4" } }, "nbformat": 4, "nbformat_minor": 5 }