LLM-SAGE — multi-equation discovery on a community-structured network

Reference implementation of the Library–Discovery–Review loop of LLM-SAGE, a three-agent framework for discovering the governing equations of a complex network whose nodes follow different dynamical systems.

This repository contains one self-contained workflow: the multi-equation discovery benchmark of Section 2.2 of the paper — a stochastic block model with three communities whose nodes follow Lorenz, Rössler and Chen dynamics, each node carrying an individual-level coefficient perturbation. Running one command recovers the community-level equation of every state variable of every community, together with the node-specific adjustments.

The pipeline runs without an API key. With no language-model endpoint configured, every agent decision is produced by the deterministic policy documented in llmsage/llm.py and the run artifact records that the decision came from that policy, so a reader can always tell which decisions were model-driven.

This is a code release, not a data release. No datasets are bundled; the benchmark is generated from the published parameters by llmsage/data.py, so runs are reproducible from seed.

---

Contents

File	Role
multi_equation_discovery.py	the entry point — one command runs the whole loop
llmsage/data.py	stochastic block model + Lorenz/Rössler/Chen integration
llmsage/library.py	the candidate function library and its families
llmsage/sparse.py	the three sparse-regression backends and the stability evidence
llmsage/agents.py	the three agents, the two-level discovery and the iterative loop
llmsage/llm.py	the JSON decision contracts and the offline deterministic policy
llmsage/metrics.py	R², WMAPE, coefficient error, support F1
llmsage/clustering.py	fuzzy K-means for community inference
configs/multi_dynamics_network.json	every parameter of a run, documented in place
tests/test_pipeline.py	unit and contract tests (python -m unittest)

---

Install

Only two dependencies, both used for numerical work:

pip install -r requirements.txt      # numpy, scipy

The language-model endpoint is reached with urllib from the standard library, so enabling it adds no dependency. PyYAML is optional: it is only needed to read a .yaml configuration, which is why the shipped configuration is .json.

Run

python multi_equation_discovery.py                            # full benchmark, T = 100
python multi_equation_discovery.py --steps 2000               # fast smoke test
python multi_equation_discovery.py --config configs/multi_dynamics_network.json
python multi_equation_discovery.py --llm                      # optional, needs OPENAI_API_KEY
python multi_equation_discovery.py --start-families self_dynamics   # watch REPAIR rebuild the library

Useful flags: --seed, --derivative-method {savgol,analytic}, --infer-communities (use fuzzy K-means instead of the given partition), --output, --model, --base-url, --api-key-env.

--steps 2000 exists to check that the pipeline runs end to end in a few seconds; its numbers are not the results. On a shortened trajectory the individual-level relative error in particular becomes meaningless — that metric is normalised by the norm of the true node deviation, which is near zero over a short window, and it is printed on a scale where a full run reads 81.96% and a 2 000-step run reads five digits. Quote the default run.

Outputs land in results/multi_equation_discovery/:

• equations.json — the full machine-readable result: every discovered equation with its coefficients, the per-term stability evidence, the candidate table the Discovery Agent saw, both agent decisions and the ground-truth comparison.
• report.md — the same result as a human-readable comparison against the ground truth.

Output

The console prints its own progress, one line per discovered equation, then a summary. This is the default run, with the middle eight equations elided:

==============================================================================
LLM-SAGE  |  multi-equation discovery on a community-structured network
==============================================================================
LLM endpoint: disabled - deterministic policy (set OPENAI_API_KEY to enable)
[data] 60 nodes / 3 communities / 10001 samples at dt=0.01 (duration 100), mean degree 4.3
[Library Agent] derivatives: savgol (relative RMSE vs analytic 8.88e-03)
[Library Agent] communities: provided partition, sizes [20, 20, 20]
[split] train/val/test = 6001 / 2000 / 2000 samples (contiguous blocks)

[round 1/5] library: 13 terms from families ['constant', 'self_dynamics', 'polynomial_self', 'network_mean']
    C0 dz/dt: dz/dt = -2.6632*z +0.9814*x*y   [selected=ridge_a1e-08_t11.7153, R2_val=0.9952, R2_test=0.9953]
    ... (eight further equations)
    review: ACCEPT=8, REPAIR=1

------------------------------------------------------------------------------
one-step prediction: mean WMAPE 3.715e-03 | worst 1.110e-02 | worst relative error 1.572e-02
structure: mean support F1 0.863 | mean coefficient relative error 32.93% | min R2_test 0.9701
mean individual-level relative error 81.96%
------------------------------------------------------------------------------

wrote results/multi_equation_discovery/equations.json
wrote results/multi_equation_discovery/report.md

---

Results

One-step prediction is the quantity the method is judged on: each node's next state is predicted from the observed neighbour states through the discovered equation, which isolates the error of the vector field from the rollout drift that accumulates separately.

The nine community equations of the default run — python multi_equation_discovery.py, seed 20260713, one round of the loop. x̄_N denotes the neighbour mean of a state. The table is transcribed from results/multi_equation_discovery/equations.json of that run:

equation	true (community level)	discovered	one-step rel. err	one-step WMAPE	review
C0 dx/dt	−10.1x + 10y + 0.1x̄_N	−10.1065x + 9.9835y	3.06e-03	2.66e-03	ACCEPT
C0 dy/dt	28x − 1.1y − xz + 0.1ȳ_N	24.7857x − 0.9051xz − 0.2275y	5.29e-03	4.40e-03	REPAIR
C0 dz/dt	−2.7667z + xy + 0.1z̄_N	−2.6632z + 0.9814xy	2.20e-03	1.60e-03	ACCEPT
C1 dx/dt	−y − z − 0.1x + 0.1x̄_N	−1.0088y − 0.9595z − 0.0699x + 0.0532x̄_N	9.04e-05	6.86e-05	ACCEPT
C1 dy/dt	x + 0.1y + 0.1ȳ_N	0.9852x + 0.1339y + 0.0670ȳ_N	6.30e-05	6.03e-05	ACCEPT
C1 dz/dt	−5.8z + xz + 0.2 + 0.1z̄_N	−5.7051z + 1.0122xz + 0.1279z̄_N	9.26e-04	9.12e-04	ACCEPT
C2 dx/dt	−35.1x + 35y + 0.1x̄_N	−35.0380x + 34.9468y	1.33e-02	9.64e-03	ACCEPT
C2 dy/dt	27.9y − 7x − xz + 0.1ȳ_N	28.2185y − 8.6639x − 0.9552xz	1.57e-02	1.11e-02	ACCEPT
C2 dz/dt	−3.1z + xy + 0.1z̄_N	−2.9675z + 0.9647xy	4.51e-03	2.99e-03	ACCEPT

Summary: one-step prediction mean WMAPE 3.71e-03, worst 1.11e-02; mean relative error 5.02e-03, worst 1.57e-02; mean support F1 0.863; R² on the held-out test block mean 0.9912, worst 0.9701; review gate ACCEPT 8 / REPAIR 1.

Two things this table shows honestly, because a reader should not have to find them:

• The single REPAIR is the Lorenz y-equation. Its split between the collinear x and x*z terms is not identifiable at the community level (measured correlation 0.958); the gate reports REPAIR rather than a confident wrong answer. See Deviations below — the one-step prediction of that same equation is still 5.3e-03, because the trajectory barely explores the ambiguous direction.
• The coupling term is retained in the three Rössler equations and dropped in the other six. Where it is kept it is recovered close to the true 0.1: 0.0532 (C1 dx/dt), 0.0670 (C1 dy/dt) and 0.1279 (C1 dz/dt). The pattern is not arbitrary — the neighbour-mean column is selected exactly where the derivative it enters is of comparable size to the coupling's contribution, and dropped where the local dynamics dominate. The Rössler community is the slow one, so a coupling of 0.1 is a substantial part of its derivative; in the Lorenz and Chen communities the same 0.1 sits next to derivatives one to two orders of magnitude larger and the sparse estimators discard it. In C1 dx/dt the fit also keeps a spurious −0.0699x alongside the coupling term. Recovering the network structure reliably, rather than the local dynamics, is a genuine limitation of fitting each equation independently — the aggregate support F1 of 0.863 is the measure of it, and the eight spurious or missing terms across nine equations are visible in report.md.

The per-term coefficient relative errors are in equations.json and in the generated report.md, and they are not all small, for the reason above.

What the code does, and where each part comes from

The mapping between the paper's components and this implementation:

Paper / Supplementary Information	Implementation
Library Agent, Discovery Agent, Review Agent	ask_library, ask_discovery, ask_review in llmsage/llm.py
Library–Discovery–Review iterative loop; ACCEPT / REPAIR / REJECT	run_pipeline in llmsage/agents.py
Fuzzy K-means for community clustering	fuzzy_kmeans in llmsage/clustering.py (Bezdek's FCM, NumPy)
Savitzky–Golay numerical differentiation	savitzky_golay_derivatives, scipy.signal.savgol_filter
Three sparse-regression backends: bootstrap STLSQ, cross-validated Lasso, Ridge threshold scan	sparse.stlsq, sparse.lasso_cv, sparse.ridge_threshold_scan
Three separate stability quantities: cross-backend retention, bootstrap selection frequency, sign consistency	sparse.stability_evidence — returned side by side and never averaged
Pareto front over (validation error, structural complexity)	sparse.pareto_front, then one selection by the Discovery Agent
Two-level discovery: community level, then individual level	discover_equation, then discover_individual_level on the residual
Review gate: derivative fit, coefficient stability, one-step state prediction	review_checks + heuristic_review_decision
WMAPE	metrics.wmape
The three JSON schemas of SI Tables 4–6	SCHEMAS in llmsage/llm.py, validated in validate_decision
Benchmark of SI Section 2.2	simulate_multi_dynamics_network in llmsage/data.py

Two design rules the code enforces

The language model never produces a number that enters an equation. It selects among candidates the sparse-regression backends have already fitted, and among families the library can actually evaluate; coefficients always come from the numerical refit. Every returned identifier is validated against the whitelist of the current run and anything unknown is discarded rather than trusted (validate_decision, filter_families, ask_discovery).

The deterministic gate is authoritative. A language model may refine the reason for a review outcome and the repair proposal, but it cannot upgrade a REJECT into an ACCEPT: when the two disagree about acceptance, the deterministic outcome wins (ask_review).

How the individual level is scored

Only the deviation of a node's coefficients from its community mean is identifiable — a common offset over a community is indistinguishable from the community-level term itself. The estimate is therefore centred over the community before comparison, which removes exactly that unidentifiable component and nothing else (individual_level_accuracy, against the centred ground truth from data.centered_individual_coefficients).

---

Deviations from the paper

These are the places where the shipped code does not use the number the paper quotes. Each is a deliberate, measured choice rather than an oversight, and each is stated here so that a reader does not have to reconstruct it from the source.

individual_sigma = 0.05, not 0.1. Section 2.2 quotes an individual-coefficient perturbation of N(0, 0.1). The shipped benchmark uses 0.05 because 0.1 is not survivable for the Rössler community: the perturbation is applied to a node's own dynamics, and a perturbation of the same size as the -z coefficient that bounds the attractor sends the trajectory to infinity. Measured over 20 seeds, 0.1 escaped for 2 of 20 seeds and 0.05 for 0 of 20, with a maximum state magnitude of 115 and 56 respectively. Both numbers are recorded in simulate_multi_dynamics_network, and the original value is one configuration key away.

Coupling coefficients are not perturbed. The individual-level perturbation is applied only to a node's own dynamics, never to its coupling coefficients. Perturbing the two coefficients of a diffusive edge independently can make that edge anti-diffusive, which is not a perturbation of the system but a different system.

Selection uses a parsimony rule, not BIC. The offline fallback among Pareto-optimal candidates is "the simplest model whose validation error is within 2% of the best on the front" (select_parsimonious). An information criterion such as BIC degenerates here: the pooled validation block has tens of thousands of rows, so the k·log(n) complexity penalty is negligible against the fit term and the criterion collapses into "pick the best fit", which retains every spurious term a dense estimator produced. The relative margin is scale-free and expresses the actual requirement. This is the rule a language-model endpoint overrides.

Thresholds are fractions of the target's RMS, not of the largest coefficient. The usual STLSQ convention thresholds relative to the largest coefficient, which is unusable on these equations: on the Lorenz dy/dt = ρx − xz − y it would discard the y term (coefficient 1) while keeping ρx (coefficient 28), although both are genuine mechanisms. Every threshold here is instead a fraction of the RMS of the quantity being fitted (stlsq_threshold_ratio · RMS(target)), so one setting applies to all nine equations.

That convention is approximately, not exactly, "a coefficient is the term's RMS contribution": column_scale divides each column by its standard deviation and does not centre it, so a scaled coefficient is the term's RMS contribution divided by a per-column factor. Measured on the shipped training design, RMS/σ ranges from 1.16 (y) to 3.28 (z), which means the effective cutoff on a term's actual contribution varies by that factor across columns. The 1 column has zero standard deviation and falls back to unit scale, so it is the one column that is not standardised at all. This is a known imprecision in the threshold's units, and lowering the threshold to compensate is not the fix: the Results section above records what the present setting costs — the true coupling term falls below it in the six fast-dynamics equations — and lowering it would admit the spurious terms the parsimony rule exists to remove.

One coefficient split is not identifiable, and this does not affect prediction. On the Lorenz attractor the correlation between the x and x*z columns is 0.958 (measured, n = 24020 rows), so for C0 dy/dt = ρx − xz − y the community-level split between those two terms is ill-conditioned, and this is also the equation where the individual-level perturbation is strongest — the x*z term has an RMS of 254, so a perturbation of 0.05 contributes ≈ 12.7 to a derivative whose own RMS is ≈ 43. Per-node regressions recover the true coefficients accurately (27.9, 27.0, 27.3, 28.9 against a true 28), but the pooled community-level fit splits the two collinear terms differently. What matters is that the trajectory barely explores the direction in which the split is ambiguous, so the one-step prediction error of that same equation is 5.3e-03 (0.53%) — the inaccuracy is confined to a direction the dynamics do not use. The run artifact reports every metric per equation rather than a single aggregate, so a reader can see exactly where the pipeline is and is not confident.

The reported structure is reconciled with the bootstrap evidence. A term that the resampled ensemble never selects is not a mechanism the data supports, so it is dropped from the selected support and the rest is refitted (min_bootstrap_frequency). Without this step the reported equations carry terms like -0.0001*x^2 that no replicate confirms, which has two visible consequences: the equation is wrong on its face, and the Review Agent can then only answer REPAIR, because a term with no dispersion estimate makes the coefficient-stability check fail by construction. Both symptoms have the same cause and the same fix.

---

Testing

python -m unittest discover -s tests -v

Standard-library unittest only, so the tests run with just numpy and scipy installed.

Citation

If you use this code, please cite the LLM-SAGE paper. The bibliographic entry will be added here once the paper is published.

License

MIT — see LICENSE.