Examples

Continuous optimization (CALM)#

Open in Colab Download this notebook

CALM learns a weighted DAG by gradient descent instead of searching over graphs. It first estimates the moral graph, which leaves only the pairs of variables that can share an edge. It then minimizes the Gaussian likelihood plus a penalty per edge, under a constraint that the graph is acyclic. The fit runs on torch, so it needs the torch extra.

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. CALM needs torch, which Colab provides; the next cell installs it with uv in a local kernel that lacks it.

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

from importlib.util import find_spec

if not IN_COLAB and find_spec("torch") is None:
    !uv pip install --python {sys.executable} --torch-backend=auto "torch>=2.2"

import numpy as np
from IPython.display import display

import andrey
import andrey.metrics as metrics
import andrey.viz
from andrey.data import benchmarks

Data#

linear_gauss_er draws an Erdos-Renyi graph with linear mechanisms and Gaussian noise: here 10 variables and 100 samples, ten per variable, the size the 0.7.0 benchmark measured.

seed = 7
dataset = benchmarks.scm("linear_gauss_er", d=10).sample(n=100, seed=seed, scale="standardize")
X, truth = dataset.data, dataset.observed_graph
print(f"{X.shape[0]} observations, {X.shape[1]} variables, seed {seed}")
display(andrey.viz.draw(truth))

100 observations, 10 variables, seed 7
0123456789

Fit CALM#

andrey.calm returns a weighted DAG. Its optimizer draws random numbers and is seeded on every call, so a run repeats on one machine. The fit takes 10 to 30 seconds on a CPU.

out = andrey.calm(X)
print(f"{out.structure.kind}, {out.metadata['n_edges']} edges, seed {out.metadata['seed']}")
display(andrey.viz.draw(out.structure))

dag, 10 edges, seed 0
0123456789

Check accuracy#

SHD counts errors against the true graph; lower is better. Gaussian data identify the graph only up to its equivalence class, so mec_shd compares the two CPDAGs. The empty graph’s SHD is the number of true edges, so an estimate at or above it has not improved on returning no edges. One dataset can land on either side: on the 0.7.0 benchmark’s five datasets of this size, CALM’s median was 6 errors against the empty graph’s 10.

d = truth.n_nodes
empty = andrey.GraphStructure.from_numpy(np.zeros((d, d), dtype=np.uint8), kind="dag")
print("CALM SHD:", metrics.score(out, truth)["mec_shd"])
print("Empty graph SHD:", metrics.score(empty, truth)["mec_shd"])

CALM SHD: 14
Empty graph SHD: 15

Cost and limits#

Every optimizer step computes a determinant and a matrix power of the full weight matrix, so a step’s cost grows with the cube of the variable count. The median fit in the andrey.calm docstring took 18 s at 10 variables and 59 s at 50, on 4 cores. The fit is nonconvex and can stop at a local optimum, and one dataset can favor any method: change seed, d, or n and rerun.