Kemp-Tenenbaum form discovery: MATLAB to Python

A test-driven translation of Charles Kemp's formdiscovery1.0 MATLAB code, which implements the structural form discovery model from:

Kemp, C., & Tenenbaum, J. B. (2008). The discovery of structural form. PNAS, 105(31), 10687–10692.

The model fits graph structures (partitions, chains, orders, rings, hierarchies, trees, grids, cylinders) to feature, similarity, or relational data and returns the best structure of each form together with its posterior probability.

The model is fully translated and verified against the original code running in GNU Octave. The Python port runs the original masterrun demo end to end:

• With Octave's random permutations replayed, all nine demo runs (chain, ring, tree × three feature data sets) reproduce the Octave baseline: final scores to ten significant digits, and identical final graphs and growth histories.
• Without replay (numpy permutations, scipy optimiser, three seeds), every run finds the same clustering as Octave (adjusted Rand index 1, same cluster count) with scores within about 1e-5 relative.
• Deterministic code (utilities, priors, graph construction, likelihoods, relational model) matches Octave to near machine precision. The one irreducible gap is the branch-length optimiser: Octave's fminunc stops early, so log-evidence values agree to about 2e-4 relative, and Python's optimum is never worse than Octave's.

Anomalies found along the way (for example, on the animals data the code prefers a hierarchy where the paper reports a tree) are logged in ANOMALIES.md.

• Paper-level checks: every synthetic data set recovers its true form and colors gives a ring, in Octave and in Python.

• The display code (draw_dot, graph_to_dot, dot_to_graph, graph_draw) is ported with pygraphviz and matplotlib, and gives the same DOT text and node positions as the original. networkx, plotly and pyvis backends and GraphML/DOT exports are also available.

• Speed: per call, fast-mode scoring takes about half of Octave's time and relational scoring a fifth. Slow mode, with scipy's trust-exact optimiser, takes 1.4-3.6 times Octave's fminunc. Whole runs take 0.8-1.4 times Octave's, 0.25 on relational data. The model keeps BLAS to one thread, because OpenBLAS threads made it up to 20 times slower on these small matrices; the results do not change. Run python tools/bench_perf.py to measure it on your machine.

Progress is tracked in RalphLoops/loop0001/PROGRESS.md and iterations.md.

conda env create -f environment.yml      # env "fd": Python, numpy/scipy, pygraphviz, PySide6 (no Octave)
conda activate fd
pip install -e .

formdiscovery run                                            # masterrun demo: chain,ring,tree × datasets 1-3
formdiscovery run --structures dirringnoself --datasets 4 --seed 3 --out results/
formdiscovery run --help

From Python:

from formdiscovery.run import masterrun
res = masterrun()            # MasterResults: names, ll, cluster counts, z, graphs, ps per run

environment.yml is Python-only. The Octave oracle that produced the test fixtures needs conda env update -f legacy/environment-octave.yml as well (see "Legacy: the Octave oracle").

Data sets and structure names follow the original setps.m (20 data sets, 24 forms); both can be given by name or 1-based index.

The notebook examples/formdiscovery_demo.ipynb runs the whole masterrun demo, compares it with Octave and draws the results with every backend. python tools/gen_demo_notebook.py rebuilds and runs it.

from formdiscovery.run import masterrun, masterrun_ps, runmodel, save_results

ps = masterrun_ps()                       # defaultps(setps()) with reloutsideinit = 'overd'
res = masterrun(ps, thisstruct=[1, 3, 5], thisdata=[0, 1, 2], seed=1, log=print)
res.modellike[:, :, 0]                    # log posterior per (structure, data set)
graph = res.structure[5, 2, 0]            # the tree fitted to demo_tree_feat
names = res.names[0, 2]
save_results(res, "results/resultsdemo", ps)   # .npz + .json, read back with load_results

ll, graph, names, bestglls, bestgraph = runmodel(ps, 5, 2, 1, rng=1)   # a single run

In Python, the structure and data set indices are 0-based (ps.structures[5] == 'tree', ps.data[2] == 'demo_tree_feat'). The CLI takes names or MATLAB's 1-based indices. rng is a seed or a permutation provider (formdiscovery.rng). runmodel and structurefit limit BLAS to one thread while they run (formdiscovery.threads, needs threadpoolctl). Wrap your own loops over graph_like in with formdiscovery.threads.blas_threads(): for the same speed-up. On the command line the option is formdiscovery run --blas-threads N (0 keeps the library's setting).

draw_dot(adj, labels, backend=...) is draw_dot.m. It returns MATLAB's (xret, yret, labels), and return_figure=True also returns the figure.

from formdiscovery.viz.draw import draw_dot, draw_graph, draw_results, ProgressFigures
from formdiscovery.viz.networkx_backend import to_networkx, to_graphml, to_dot

draw_results(res, path="results.png", backend="networkx", layout="kamada_kawai")
draw_results(res, path="results.html", backend="plotly")    # all runs, interactive

fig = draw_graph(graph, names, backend="plotly")            # hover: object -> cluster,
fig.show()                                                  #        cluster -> members
draw_graph(graph, names, backend="pyvis").write_html("tree.html")

G = to_networkx(graph, names)          # DiGraph; node kind/label/cluster_id, edge W and 1/W
to_graphml(graph, "tree.graphml", names)                    # for Gephi / Cytoscape

# MATLAB's progress figures 1-3 (pre-clean, post-clean, best split / result)
show = ProgressFigures(outdir="figures", backend="networkx")
runmodel(ProgressFigures.enable(ps), 5, 2, 1, show=show)

From the shell:

formdiscovery run --structures chain,ring,tree --datasets 1,2,3 --out results/ --figures figs/
formdiscovery draw results/resultsdemo.npz --out results.png                 # neato + matplotlib
formdiscovery draw results/resultsdemo.npz --out results.png --backend networkx --layout kamada_kawai
formdiscovery draw results/resultsdemo.npz --out results.html --backend plotly
formdiscovery draw results/resultsdemo.npz --out tree.html --backend pyvis --runs 8
formdiscovery draw results/resultsdemo.npz --out tree.svg --graphviz --runs 8  # Graphviz renders (pygraphviz)

Display quirks of the original are kept by default. Each undirected edge is drawn as two arrows (KI-37); use undirected='lines' for plain lines. draw_dot's neato flags are malformed (KI-32/35); use flags='intended' for the documented ones. See KNOWN_ISSUES.md.

A PySide6 desktop app (pip install -e .[gui]; in the fd env it is already installed):

formdiscovery gui                                                  # pick a file with "Open data file…"
formdiscovery gui data/demo_chain_feat.mat                         # start with a data set loaded
formdiscovery gui --demo                                           # open demo_chain_feat and run chain

Open any .mat data file (the shipped ones or your own: data as features, a similarity matrix or a relational struct, optional names). The left column describes the data set and has the forms (chain, ring, tree preselected), seed, speed, drawing backend and "Draw best splits". Run fits every selected form, each in its own worker thread ("Parallel runs" at a time; one by default), so the window stays responsive and Stop works at any time. The middle canvas redraws the graph each time the model transforms it (pre-clean / post-clean per depth, best splits if chosen, then the inferred graph), with node positions kept between frames. When the run ends, the Statistics column shows the final score and its parts (log prior from graph_prior, log likelihood from graph_like), the clusters and their members, the score per depth (bestglls) as a chart, wall time and frame count. "Export results…" writes the .npz + .json pair that formdiscovery run writes (read back with load_results), and "Save figure…" saves the graph as PNG or SVG.

The "Results" table under the graph ranks the forms by ll (the winner in bold, as masterrun's modellike would pick it); click a row to see that form's final graph and statistics. The slider under the graph scrubs back through the frames the shown form sent (the last 500 per form are kept); at its right end the canvas is live again. Parallel runs share one Python process, so they overlap little (the search holds the GIL; ANOMALIES A27).

Keyboard shortcuts (also in the File and Run menus): Ctrl+O open, Ctrl+R or F5 run, Esc or Ctrl+. stop, Ctrl+E export results, Ctrl+Shift+S save figure, Ctrl+Q quit. A file that cannot be read, a failed run or a failed export opens a warning box (a failed run's traceback is under "Show Details…"). The file dialog starts in the last directory a file was loaded from (~/.config/formdiscovery/gui.ini). The window title names the file. The GUI needs no Octave (only the tests compare against it).

The GUI is a thin shell over the port (src/formdiscovery/gui/); its only hook into the model is search.run_hooks (per-thread cancel and per-depth callbacks, off by default, so results are unchanged). Screenshots, with a description of each: examples/gui/ (QT_QPA_PLATFORM=offscreen python tools/gui_screenshots.py).

The Octave oracle has been frozen and removed from main. Everything below that mentions legacy/ (the MATLAB sources, the fixture scripts, the Octave tools) now lives only at the git tag octave-oracle-final; legacy/FREEZE.md there records the environment and the final verification. The committed fixtures remain the reference.

• Every MATLAB function has an Octave fixture script in legacy/tests_octave/ that runs the original code on real inputs and saves inputs and outputs to tests/fixtures/*.mat. The pytest suite compares the Python translation against those fixtures, so the tests run without Octave. Tests marked octave additionally drive Octave live through oct2py. They are skipped when it is not installed, except in strict mode (the fd env's Python, or RALPH_REQUIRE_OCTAVE=1), where a skipped Octave test fails the run; tests/test_gate_env.py checks the toolchain itself. Octave 10.3.0 and oct2py 6.1.1 come from legacy/environment-octave.yml, added to the Python-only environment.yml env; every committed fixture was produced with those versions.
• Structural outputs are compared exactly; deterministic floats at rtol=1e-10; optimiser-dependent values by optimality (objective and gradient norm no worse than Octave's) plus a documented tolerance.
• Randomness enters only through randperm. An Octave shim (legacy/matlab/octave_shims/) records or replays permutations, and formdiscovery.rng provides matching Python providers, so the search heuristics are compared with identical random choices.
• The Octave baseline runs (legacy/matlab/run_baseline.m) for the feature and relational demo grids are committed under tests/fixtures/baseline/.
• Every process runs one BLAS/OpenMP thread (OPENBLAS_NUM_THREADS, OMP_NUM_THREADS, MKL_NUM_THREADS = 1, formdiscovery.threads.PIN_ENV). The test session, the Oct2Py Octave and every octave-cli the tools start are pinned; speed comes from running processes side by side. Pinned and unpinned runs give identical .mat contents (legacy/tools/mat_compare.py DIR_A DIR_B compares two output directories), except for gibbs.mat. Its grid x synthgrid run depends on the BLAS thread count, so it is regenerated unpinned (ANOMALIES A19, BLAS_DEFAULT_FIXTURES in tests/conftest.py).
• Fixtures regenerate in parallel: legacy/tools/gen_fixtures.py --jobs N (default cores/2) runs one octave-cli per fixture script (paperlevel split into its 45 runs), so a script that reads another fixture waits for it. It keeps a log per task and prints a timing table. All 31 fixtures take 7 min 56 s with --jobs 16 against 68 min 49 s with --jobs 1. The two runs and the committed fixtures have the same content (--compare DIR: loaded arrays to all digits, timing fields skipped, temp-directory names masked).
• The baselines regenerate in parallel too: legacy/tools/gen_baselines.py --kind feat|rel --jobs N runs one octave-cli per (structure, dataset) pair, each into its own directory. A merge (legacy/matlab/baseline_merge.m, plus a copy of the results/ trees) then writes what the serial run_baseline writes. The merged output has the same file set and content as the committed tests/fixtures/baseline/{feat,rel} (--compare DIR; only timings.seconds is skipped). With --jobs 16 feat takes 5.8 s and rel 14.6 s, against 33–46 s and 140–176 s for serial run_baseline.
• The gate runs in parallel with pytest-xdist (-n 16). Each worker has its own Octave process (its own session-scoped octave fixture), and tests write only to tmp_path. A run that leaves a new file under tests/ fails. Under xdist the tests that take 30 s or more (LONG_TESTS in tests/conftest.py) start first, one per worker (--maxschedchunk 1). The gate then takes about 6 min (the longest test, gibbs test_live_fresh_seeds, alone takes about 6 min) against about 40 min serial. Three runs at -n 16 and one at -n 8 gave the same pass set.
• legacy/tools/compare_live.py --pairs S:D[:SEED] ... --jobs N runs Octave and Python side by side. For each (structure, dataset, seed) a worker runs run_baseline in its own octave-cli with the randperm shim logging every draw. It then runs Python's runmodel on the same pair, replaying those draws without oracles: Python uses its own optimizer and tie breaking, and once a draw no longer fits it goes on with numpy. The output is one table: Octave and Python ll, relative difference, ARI, cluster counts, how far the replay went, and each side's wall time. The default set (the 9 feature and 54 relational baseline pairs, seed 1) passes PLAN §7.1 on all 63 rows. Octave reproduces the committed baseline exactly. The relational scores agree to 1e-15, and the feature scores to 7.4e-6 relative, with ARI 1 everywhere. This takes 26 s with --jobs 16 against 3 min 35 s with --jobs 1, and the rows are identical apart from the times.

Before and after the parallel harness (RalphLoops/loop0002, 32-core machine; every parallel run gives the same contents as the serial one):

Before this work the gate ran under a Python without oct2py, so the live Octave tests were skipped; it now runs them all and fails if one is skipped.

make test                             # default gate: fixture-only, no Octave (~1 min, 16 workers)
python -m pytest -q -m "not slow and not octave" -n 16   # the same without make
make fixture-sums                     # tests/fixtures against tests/fixtures/SHA256SUMS (also checked by the gate)
make slow                             # long runs (Python only)
git checkout octave-oracle-final      # the frozen oracle: `make legacy-check`, legacy/tools/gen_fixtures.py --compare, ...

The original MATLAB code, run in GNU Octave 10.3.0, is the reference the port was checked against. On 2026-10-01 it was frozen: the strict live gate passed (3866 tests), all 31 fixtures and both baseline grids regenerated identical to the committed files, and the tree was tagged octave-oracle-final. The next commit removed legacy/ from main. To reconstruct it:

git checkout octave-oracle-final
conda env create -f environment.yml && conda env update -f legacy/environment-octave.yml
make legacy-check                     # strict live-Octave gate

Policy since the freeze: a change that alters any fixture value must first be checked against the oracle at the tag, and the reason recorded in ANOMALIES.md (see CLAUDE.md).

Python function names match the MATLAB file names (graph_like_conn, gibbs_clean, ...), and each docstring cites the source lines it translates. Indices are 0-based in Python and converted only at the I/O boundary. Quirks of the original are replicated by default; the only deliberate deviations are that keyboard debug stops become exceptions and runmodel writes to an explicit output directory instead of changing the working directory.

Original MATLAB code: Charles Kemp (2008), with contributions from Kevin Murphy (BNT), Thomas Minka (lightspeed), Leon Peshkin (Graphviz interface), Carl Rasmussen, John Burkardt, and Michael Kay, as listed in the original README.txt.

Translation by Alex Linhares using Claude Code.