{
"cells": [
{
"cell_type": "markdown",
"id": "calm-0",
"metadata": {},
"source": [
"# Continuous optimization (CALM)\n",
"\n",
"[](https://colab.research.google.com/github/Abel-ai-lab/andrey/blob/main/examples/calm.ipynb)\n",
"\n",
"CALM learns a weighted DAG by gradient descent instead of searching over graphs. It first estimates\n",
"the moral graph, which leaves only the pairs of variables that can share an edge. It then\n",
"minimizes the Gaussian likelihood plus a penalty per edge, under a constraint that the graph is\n",
"acyclic. The fit runs on torch, so it needs the `torch` extra."
]
},
{
"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. CALM needs torch, which Colab provides; the next cell\n",
"installs it with uv in a local kernel that lacks it."
]
},
{
"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": "calm-install",
"metadata": {},
"outputs": [],
"source": [
"from importlib.util import find_spec\n",
"\n",
"if not IN_COLAB and find_spec(\"torch\") is None:\n",
" !uv pip install --python {sys.executable} --torch-backend=auto \"torch>=2.2\""
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "calm-1",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\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": "calm-2",
"metadata": {},
"source": [
"## Data\n",
"\n",
"`linear_gauss_er` draws an Erdos-Renyi graph with linear mechanisms and Gaussian noise: here 10\n",
"variables and 100 samples, ten per variable, the size the 0.7.0 benchmark measured."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "calm-3",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"100 observations, 10 variables, seed 7\n"
]
},
{
"data": {
"text/html": [
""
],
"text/plain": [
"''"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"seed = 7\n",
"dataset = benchmarks.scm(\"linear_gauss_er\", d=10).sample(n=100, 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": "calm-4",
"metadata": {},
"source": [
"## Fit CALM\n",
"\n",
"`andrey.calm` returns a weighted DAG. Its optimizer draws random numbers and is seeded on every\n",
"call, so a run repeats on one machine. The fit takes 10 to 30 seconds on a CPU."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "calm-5",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"dag, 10 edges, seed 0\n"
]
},
{
"data": {
"text/html": [
""
],
"text/plain": [
"''"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"out = andrey.calm(X)\n",
"print(f\"{out.structure.kind}, {out.metadata['n_edges']} edges, seed {out.metadata['seed']}\")\n",
"display(andrey.viz.draw(out.structure))"
]
},
{
"cell_type": "markdown",
"id": "calm-6",
"metadata": {},
"source": [
"## Check accuracy\n",
"\n",
"SHD counts errors against the true graph; lower is better. Gaussian data identify the graph only up\n",
"to its equivalence class, so `mec_shd` compares the two CPDAGs. The empty graph's SHD is the number\n",
"of true edges, so an estimate at or above it has not improved on returning no edges.\n",
"One dataset can land on either side: on the 0.7.0 benchmark's five datasets of this size,\n",
"CALM's median was 6 errors against the empty graph's 10."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "calm-7",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CALM SHD: 14\n",
"Empty graph SHD: 15\n"
]
}
],
"source": [
"d = truth.n_nodes\n",
"empty = andrey.GraphStructure.from_numpy(np.zeros((d, d), dtype=np.uint8), kind=\"dag\")\n",
"print(\"CALM SHD:\", metrics.score(out, truth)[\"mec_shd\"])\n",
"print(\"Empty graph SHD:\", metrics.score(empty, truth)[\"mec_shd\"])"
]
},
{
"cell_type": "markdown",
"id": "calm-8",
"metadata": {},
"source": [
"## Cost and limits\n",
"\n",
"Every optimizer step computes a determinant and a matrix power of the full weight matrix, so a\n",
"step's cost grows with the cube of the variable count. The median fit in the `andrey.calm`\n",
"docstring took 18 s at 10 variables and 59 s at 50, on 4 cores. The fit is nonconvex and can stop\n",
"at a local optimum, and one dataset can favor any method: change `seed`, `d`, or `n` 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
}