Bayesian Network Structure Learning with K2 Search

Given a dataset and no prior knowledge of how its variables relate, find the directed acyclic graph that best explains the observed data — that's Bayesian network structure learning, and this is an implementation of the K2 search heuristic to do it, scored against three datasets of very different shapes: 8 variables and a few hundred rows (a Titanic-style survival dataset), 13 variables (a wine quality dataset), and 50 variables with 10,000 rows (an anonymized dataset built to stress-test the algorithm at scale, not to be interpreted).

Every number below comes from actually running the current code against the datasets in this repo, not from a lab notebook — python3 wsmall_dataset.py, python3 medium_dataset.py, and python3 wlarge_dataset.py reproduce all three.

The Bayesian score

The core scoring function (bayesian_score in each dataset script) is a BDeu-style score: for each variable, it builds a contingency table of counts against every possible configuration of its parents (statistics), combines that with a Dirichlet prior weighted by an equivalent sample size — ESS — (prior), and evaluates the whole thing in closed form with the log-gamma function rather than needing to integrate anything numerically. Higher (less negative) is better. It decomposes cleanly across variables, which is what makes it possible to search structure greedily in the first place — the score of the whole graph is just the sum of each node's own local score given its current parent set.

K2 search

k2_search walks the given node order once, and for each node in turn, scans the entire node list as candidate parents — adding an edge whenever it doesn't create a cycle and the node hasn't already hit max_parents. That makes node_order a processing priority rather than a hard topological constraint: whichever node is handled first gets first claim on any other node as a parent, since the graph is still empty and almost nothing creates a cycle yet, and nodes handled later in the sequence have fewer cycle-free options left by the time it's their turn. It's a single greedy pass, not an exhaustive search over parent sets, which is what makes it tractable on a 50-variable dataset where an exhaustive search over parent sets would never finish.

That processing-order effect is exactly why, in the wine dataset below, quality ends up as the parent of several chemical properties rather than their child — it happens to get processed early enough in that run's order to claim them before they can claim it.

The real work in this repo is the parameter sweep around that core search: experiment_with_parameters tries multiple node orderings, multiple max_parents ceilings, and (for the medium and large datasets) multiple ESS values, and keeps whichever combination scores best. Since node order changes which claims happen first, and therefore which edges are even possible once cycle-prevention kicks in, trying the natural order, the reversed order, and a random permutation is doing real work here, not just covering bases.

Results

Dataset	Variables	Search space explored	Best configuration	Best Bayesian score
Small (small.csv)	8	max_parents in {2,4,5} x 3 node orders	max_parents=4, reversed order	-3,932.77
Medium (medium.csv)	13	single fixed configuration	max_parents=4, ESS=1.0	-97,011.28
Large (large.csv)	50	ESS in {0.5,1,2,5} x max_parents in {2,3,4} x 3 node orders (36 combinations)	max_parents=2, ESS=5.0, random order	-478,354.55

The small dataset's variable names — age, sex, passengerclass, survived, numsiblings — give away what it actually is: a cut of the Titanic dataset. survived ends up with exactly four parents in the learned graph — sex, passengerclass, numsiblings, and numparentschildren — which is the maximum max_parents allows, and is a genuinely reasonable answer: sex and class are the two variables every analysis of this dataset converges on, and the algorithm found both of them without being told to look for them.

The medium dataset is a wine quality dataset — alcohol, pH, sulphates, volatileacidity, and so on, with alcohol and fixedacidity binned into five categories as a preprocessing step before search. The direction is worth reading carefully here: quality comes out as the parent of alcohol, sulphates, pH, and density, not their child — a reminder that a learned Bayesian network encodes a factorization that best compresses the data under this search procedure, not a verified causal diagram. The variables it groups together (quality with exactly the properties wine chemistry research treats as most predictive of quality) are the genuinely interesting part; which way each arrow points is a function of the node order it was searched under, not a claim about what causes what.

The large dataset's 50 variables are anonymized two-letter codes on purpose — this one isn't about the domain, it's about whether the search still works at a scale where node order and parameter choice actually start to matter for runtime, not just score. It does: the full 36-combination sweep runs in about seven minutes.

A real performance bug, found and fixed by actually running this

Both k2_search and the outer experiment_with_parameters were computing the Bayesian score after every single node in the search — once for logging inside k2_search, and once more, on the finished graph, in the calling function. The first one was pure waste: its result was assigned to a local variable and never read, never returned, never printed. Nothing depended on it. For the small and medium datasets that redundancy was cheap enough not to notice. For the large dataset — 50 variables, 10,000 rows, computed 50 times per search, 36 searches in the parameter sweep — it meant the full sweep hadn't finished after fifteen minutes, and it's the reason the score above wasn't verified locally until that dead computation was removed. After removing it, the identical sweep — same code path, same data, same result, -478,354.55364..., matching to five significant figures — runs in under seven minutes. Nothing about what the algorithm searches or returns changed; the only difference is that it no longer computes something it was about to throw away.

Running this

pip install numpy pandas networkx scipy matplotlib

python3 wsmall_dataset.py    # small.csv  — seconds
python3 medium_dataset.py    # medium.csv — seconds
python3 wlarge_dataset.py    # large.csv  — about 7 minutes, 36-combination sweep

Each script prints the score for every configuration it tries, then writes the best graph it found into results/ as a .gph edge list, a .gml file (medium and large only), and a rendered PNG.

Repository structure

wsmall_dataset.py     8-variable Titanic-style dataset — full max_parents x node-order sweep
medium_dataset.py     13-variable wine quality dataset — single fixed configuration
wlarge_dataset.py     50-variable anonymized dataset — full ESS x max_parents x node-order sweep
small.csv / medium.csv / large.csv    the three datasets
results/              generated graphs (.gph, .gml, .png) — regenerated by running the scripts above