KD discovers the governing partial differential equation from data: give it a
field sampled on a spatiotemporal grid, get back a symbolic PDE. Five
in-house discovery engines (SGA, DLGA, DISCOVER, EqGPT,
LLM4ED) run behind one kd.Model API, sharing a single dataset
interface, term evaluator, and HTML-report visualization.
2D Burgers (
u_t = -u·u_x - u·u_y + 0.01·∇²u): the field's true evolution beside the PDE integrated forward through the platform. Regenerate with examples/14_field_animation_2d.py.
Requires Python >= 3.11 and PyTorch >= 2.0.
git clone -b trunk https://github.com/Scientific-Artificial-Intelligence-Lab/kd.git
cd kd
uv syncimport kd
# Generate a synthetic Burgers dataset.
dataset = kd.generate_burgers_data(nx=64, nt=32, nu=0.1, seed=0)
model = kd.Model(algorithm="sga", generations=30, population=15, seed=0)
model.fit(dataset)
print(model.best_expr_) # u_t = -0.9863*mul(u_x, u) + 0.09955*diff_x(u_x)
print(model.best_score_) # -11.38 (AIC, lower is better)The printed expression is KD's canonical function-call notation (funcall IR): mul(u_x, u) is
u*u_x and diff_x(u_x) is u_xx, so the discovered equation reads
u_t = -0.99*u*u_x + 0.0996*u_xx — the Burgers equation recovered from data
with fitted coefficients (ground truth: u_t = -1*u*u_x + 0.1*u_xx).
The HTML report renders it in standard notation:
The discovered equation as rendered in the report (longer run on the bundled 256×201 Burgers benchmark).
True vs predicted solution (Burgers equation).
See examples/ for runnable scripts covering every engine,
including 09_compare_algorithms.py,
which runs the engines on the same dataset and ranks the discovered
equations on one unified NMSE ruler.
All five engines are re-implementations of algorithms developed in this
lab, refactored onto KD's shared platform. Swap the algorithm= string to
switch:
| Algorithm | algorithm= |
Origin | Approach |
|---|---|---|---|
| SGA | "sga" |
Chen et al. 2022 (SGA-PDE) | Genetic algorithm over symbolic expression trees |
| DLGA | "dlga" |
Xu et al. 2020 | Neural-network surrogate + genetic algorithm |
| DISCOVER | "discover" |
Du et al. 2024 | LSTM controller + policy gradient |
| EqGPT | "eqgpt" |
Xu et al. 2025 (EqGPT) | Pretrained generative GPT proposes candidate PDEs, then reward-guided fine-tuning |
| LLM4ED | "llm4ed" |
Du et al. 2024 (LLM4ED) | An LLM proposes candidate equations as text, scored by a sparse-regression reward |
EqGPT needs its pretrained GPT weights, which are not vendored; see
examples/16_eqgpt.py for where to place them.
LLM4ED runs fully offline with an injected provider, or against any
OpenAI-compatible API (see examples/17_llm4ed.py).
Two external baselines can also be driven through the same facade for
cross-checking: the PySR symbolic regressor (algorithm="pysr", needs
uv sync --extra pysr) and PySINDy's native STLSQ sparse-regression optimizer
over the KD term library (algorithm="pysindy", needs
uv sync --extra pysindy).
Each engine's own settings go through the same call: any field of its config
carries as a keyword argument, so kd.Model(algorithm="pysindy", threshold=0.2, normalize_columns=True) and kd.Model(algorithm="dlga", pop_size=200, epsilon=1e-4) need no per-engine call form. An unknown name is rejected with
the accepted ones listed. kd.instrument_schemas() returns one row per engine —
its config fields with types and defaults, plus the facade parameters
(generations, population, seed, …) that are not config fields — so a
caller holding only JSON can configure any engine without hardcoding names.
The simulated datasets come from this lab's PDE-discovery papers — SGA-PDE (Chen et al., Phys. Rev. Research 4, 023174, 2022), EqGPT (Xu et al., Nat Commun 16, 10255, 2025) and LLM4ED (Du et al., Phys. Fluids 36, 097121, 2024):
| Dataset | Governing PDE | Grid | What it models |
|---|---|---|---|
allen-cahn |
u_t = 0.003·u_xx + u - u³ |
(256, 201) |
phase separation (reaction–diffusion) |
burgers |
u_t = -u·u_x + 0.1·u_xx |
(256, 201) |
shock waves in fluids |
burgers-2d |
u_t = -u·u_x - u·u_y + 0.01·∇²u |
(101, 51, 100) |
2D Burgers flow |
chafee-infante |
u_t = u_xx - u + u³ |
(301, 200) |
reaction–diffusion |
convection-diffusion |
u_t = -u_x + 0.25·u_xx |
(256, 100) |
advection plus diffusion |
eq-6-2-12 |
u_t = -0.1·u_x_t - 0.1·u_x |
(501, 501) |
handbook equation with a mixed space–time derivative |
kdv |
u_t = -u·u_x - 0.0025·u_xxx |
(256, 201) |
shallow-water solitons |
klein-gordon |
u_tt = 0.5·u_xx - 5·u |
(201, 201) |
relativistic wave equation |
llm4ed-fisher |
u_t = 0.02·u_xx + 10·u·(1-u) |
(x, t) |
population growth with spatial spread |
llm4ed-fisher-nonlinear |
u_t = 0.02·(u·u_xx + u_x²) + 10·u·(1-u) |
(x, t) |
Fisher growth with nonlinear diffusion |
llm4ed-heat |
u_t = 0.05·u_xx |
(x, t) |
heat conduction |
pde-compound |
u_t = u·u_xx + u_x² |
(100, 251) |
constructed compound-structure case |
pde-divide |
u_t = -u_x/x + 0.25·u_xx |
(100, 251) |
constructed case with a division term |
wave |
u_tt = u_xx |
(161, 321) |
vibrating string |
Load any bundled dataset with kd.load_burgers(), or browse the catalog
programmatically with kd.list_datasets() / kd.get_dataset(id); each entry
carries its .source and .license (see NOTICE).
kd.generate_burgers_data(), kd.generate_diffusion_data(), … build
synthetic datasets on demand. Remote (HF) entries
are fetched with kd.load_from_hub(id) after installing the hub extra
(uv sync --extra hub) and are cached locally, checksum-verified, and
revision-pinned; browse them with kd.list_remote_datasets().
KD also bundles real-world experimental data — measured, not simulated:
| Dataset | Type | Measured quantity | Size | Reference |
|---|---|---|---|---|
wave-breaking |
wave-tank experiment (Imperial College London) | surface elevation η(t, x) of wave groups approaching breaking |
314,478 points (one of the paper's 12 experiments) | Xu et al., Nat Commun 16, 10255 (2025) |
tlc-cc |
automated chromatography experiment | column retention volumes V_S, V_E vs (R_F, r) |
2 tables × 74 conditions | Xu et al., Nat Commun 16, 832 (2025) |
Wave breaking — surface elevation of focused wave groups approaching
breaking, reconstructed frame by frame from camera images in the wave-tank
experiments of the EqGPT paper. KD bundles one of the paper's 12 experiments
(case N_G2Tp12A100_broad) as scattered (t, x, η) points — a table rather
than a gridded PDEDataset.
TLC-CC — column-chromatography retention volumes measured on an
automated platform (192 compounds, 4 g silica columns), aggregated to mean
start/end retention volumes over 74 (R_F, r) conditions — ready for KD's
scalar symbolic-regression entries
(examples/12,
examples/13).
wb = kd.load_wave_breaking() # η(t, x): scattered wave-tank points
cc = kd.load_tlc_cc(target="start") # X = (R_F, r), y = V_SExperimental background, protocols, and references for both datasets are in the papers and their Supplementary Information (wave breaking, TLC-CC).
Wrap your own arrays (any field on a regular grid) into a PDEDataset:
import torch
import kd
x = torch.linspace(0.0, 1.0, 64) # spatial grid
t = torch.linspace(0.0, 1.0, 32) # time grid
u = torch.rand(64, 32) # your measured field on the (x, t) grid
dataset = kd.PDEDataset.from_arrays(
coords={"x": x, "t": t}, # one 1D array per axis; insertion order sets the axis order
fields={"u": u}, # field shaped (len(x), len(t))
lhs="u_t", # left-hand side of the equation to discover
periodic={"x"}, # optional: periodic axes improve fits
name="my_pde",
ground_truth="u_t = 0.1 * u_xx",
)The same call handles 2D spatial fields: add a y axis and pass an nD field,
e.g. coords={"x": x, "y": y, "t": t} with u shaped (len(x), len(y), len(t)).
You don't have to run a search to use KD's evaluator. evaluate_terms fits
a candidate term set directly; validate_terms classifies terms without
fitting. Both fail loud with a complete per-term rejection report (reason +
hint), so a caller (human or LLM agent) can repair and resubmit:
import kd
result = kd.evaluate_terms(dataset, ["diff2_x(u)", "mul(u, diff_x(u))"])
print(result.coefficients, result.nmse)
report = kd.validate_terms(dataset, ["u_xx", "u + u_x"]) # no fit performed
for v in report.rejected:
print(v.term, "->", v.reason) # "u + u_x" is not canonical funcall IRA blind search starts from "any equation could be here". When part of the law is
already settled physics, fit(dataset, sketch=...) states that part and searches
only the rest:
import kd
from kd.core.equation import (
LhsSpec,
PinnedTerm,
Sketch,
SketchMatchPolicy,
TermConstraint,
TermHole,
)
from kd.core.expr import TermVocabulary
sketch = Sketch(
lhs_spec=LhsSpec("u", "t", 1),
vocabulary=TermVocabulary(
fields=frozenset({"u"}), coordinates=frozenset({"x", "t"})
),
pinned=(PinnedTerm("mul(u,u_x)", -1.0),), # known term WITH its coefficient
anchored=(), # known structure, coefficient free
holes=( # the unknown remainder
TermHole(
id="diffusion",
min_count=1,
max_count=2,
constraint=TermConstraint(max_deriv_order=2),
),
),
match_policy=SketchMatchPolicy(
coeff_atol=1e-9, coeff_rtol=1e-9, support_threshold=0.0
),
)
model = kd.Model(algorithm="sga", generations=30, population=15, seed=0)
model.fit(dataset, sketch=sketch)
outcome = model.result_.sketch_outcome
print(outcome.verdict.overall) # True when every clause of the sketch holds
print(outcome.solution) # published only when overall is TrueA pinned term is subtracted from the regression target before the search and
restored exactly in the solution, so the search cannot spend budget
rediscovering it. A hole declares how many terms may fill it and what shapes
they may take (derivative-order cap, allowed operators, fields, axes). The exit
is certified: outcome.solution is published only when the discovered law
satisfies every clause, and otherwise the run reports outcome.best_candidate
plus the clause that failed. On the bundled 64×51 Burgers field the run above
restores the pinned coefficient at exactly -1.0 and fills the hole with a
second-order term whose coefficient comes out 0.1000 against a ground truth of
0.1.
"sga" and "pysindy" accept sketches today. SGA compiles the sketch natively,
narrowing the search's variable and operator pools at the source; an engine that
cannot honor a clause refuses the fit with a ValueError naming that clause
instead of searching wider than declared.
The full walkthrough, including how the two backends differ, is
examples/21_sketch_discovery.py
(about 30 seconds).
model = kd.Model(algorithm="discover", generations=500, checkpoint_dir="ckpts")
model.fit(dataset) # writes ckpts/checkpoint_*.pt as it goes
# Later (or after a crash), continue from the saved search state:
model = kd.Model(algorithm="discover", generations=200)
model.fit(dataset, resume_from="ckpts/checkpoint_final.pt")The checkpoint restores search state (population / controller weights /
best); generations and other settings come from the new Model.
The checkpoint directory also carries a manifest.json ledger: one entry per
checkpoint recording filename, iteration, best score and expression, algorithm,
seed, config hash and write time. Read it instead of globbing filenames:
for entry in kd.load_checkpoint_manifest("ckpts"):
print(entry.filename, entry.iteration, entry.best_score)The reader is read-only and fail-loud, and verifies the whole directory contract. Call it on a directory whose run has terminated: during a live run a checkpoint file can briefly exist before its ledger entry is appended.
kd.harness runs a declarative experiment matrix and stores sealed evidence:
from pathlib import Path
from kd.harness import ExperimentPlan, PlanEntry, run_plan
plan = ExperimentPlan(
name="burgers-sweep",
entries=tuple(
PlanEntry(
instrument=engine,
dataset_ref="burgers",
seed=seed,
model_kwargs={"generations": 50},
)
for engine in ("sga", "dlga")
for seed in (0, 1, 2)
),
)
result = run_plan(
plan, datasets={"burgers": dataset}, store_root=Path("evidence")
)Entry order is part of the plan's identity, so a plan hash pins the exact
execution matrix. Each run's evidence is stored with an environment
fingerprint; build_consensus aggregates a sealed store across runs, and both
dispatch and consensus results render to markdown or to versioned JSON
artifacts. Execution is serial: the package composes the existing Model
surface and adds no routing or budget logic.
Passing recording=RecordingOptions(...) runs each entry inside a standard run
directory (runs/entry-NNNN/, numbered to match its record) carrying that
entry's iteration events, search phases and optional checkpoints; adding
catalog_path= appends one row per run to a cross-run ledger whose rows point
at those run directories. run_plan also renders a human-readable report.md
into the store root; it is derived and regenerable, while the sealed index and
records stay the authority.
The full chain (plan, run, re-open the sealed store, consensus, Markdown +
JSON) runs in about five seconds in
examples/19_batch_harness.py.
After a fit, render a full HTML report (universal figures plus the fitted engine's own search diagnostics):
import kd
viz = kd.VizEngine(output_dir="out/my_run")
report = viz.render_all(model.result_, algorithm=model.algorithm_, dataset=dataset)
print(report.report) # path to report.html
print(len(report.figures)) # number of figure filesThe report bundles universal figures — the discovered equation rendered in LaTeX, search convergence, a parity plot, residual maps, True/Predicted field comparisons — plus the fitted engine's own diagnostics:
Parity plot from the report: predicted vs actual
u_t for the
discovered Burgers equation (R² = 1.0000).
True / Predicted / Residual panels from a Chafee-Infante fit (
u_t = u_xx - u + u^3 recovered by SGA).
The report also renders the discovered equation as a structure-only expression tree, and (for the SGA engine) the raw genome tree of the best evolved individual, so you can see what the search actually produced versus the sparse equation it was distilled into:
Example: SGA on the built-in Chafee-Infante dataset (recovers
u_t = u_xx - u + u^3). Left: the raw GP genome of the best
individual, still carrying evolved bloat (redundant / zeroed terms). Right: the
discovered equation after sparse selection, operators and derivatives only,
coefficients dropped (they stay in the LaTeX equation figure).
Every result also carries a manifest (dataset fingerprint, seed, KD
version) so a run can be identified and reproduced later.
src/kd/
├── api.py # Model facade: one-line fit() for every engine
├── evaluate.py # evaluate_terms / validate_terms: score terms directly
├── data/ # PDEDataset, synthetic generators, dataset loaders
├── search/ # sga / dlga / discover / eqgpt / llm4ed / pysr / pysindy
├── harness/ # batch experiment plans, evidence store, consensus reports
├── viz/ # VizEngine: HTML reports & figures
└── inspect.py # preview() dataset sanity checks
The SGA, DLGA, DISCOVER, EqGPT, and LLM4ED engines are refactored re-implementations of algorithms developed in this lab; credit for the methods belongs to the original works:
- SGA-PDE: Chen et al., SGA-PDE;
also the source of several bundled datasets (see
NOTICE) - DLGA: Xu et al. 2020
- DISCOVER: Du et al., DISCOVER
- EqGPT: Xu et al., EqGPT,
Nat Commun 16, 10255 (2025); also the source of several bundled
datasets, including the wave-breaking experiments (see
NOTICE) - LLM4ED: Du et al., LLM4ED, Phys. Fluids 36, 097121 (2024)
KD also builds on PySR and PySINDy (optional external baselines), SymPy, and PyTorch.
Apache-2.0. Copyright 2026 Mao, Hao and the Scientific Artificial Intelligence Lab.