Skip to content

Design

KD is a platform for discovering partial differential equations from data, organized as a set of algorithm plugins. Each plugin implements the algorithm-specific search procedure and scientific logic, while the platform provides the data processing, execution, visualization, and experiment infrastructure shared across algorithms. Seven algorithms are currently integrated: five published by our group (SGA-PDE, DLGA-PDE, DISCOVER, EqGPT, LLM4ED) and two external baselines (PySR, PySINDy).

Data is prepared once and an algorithm is chosen by name:

model = kd.Model(algorithm="sga").fit(data)

The fitted model carries the equation, the score, and the record of the run. Selecting a different algorithm changes only the algorithm argument; the surrounding code and the way the result is read stay the same.

Division of responsibility

The algorithm decides how to differentiate the data, how to search, how to score a candidate, and which equation to return. The platform implements the components that all seven algorithms need in the same form.

These components are optional. An algorithm can implement any of them itself: EqGPT runs two of its three modes without platform derivatives, and SGA-PDE and LLM4ED fit candidates using code included with the plugin. Every plugin must provide two things: the contract that lets the platform drive it, and a declaration of what the algorithm accepts. The declaration covers which equation forms and data layouts it supports, where its derivatives come from, the name and direction of its score, and which parameters survive a resume. kd.instrument_schemas() returns those declarations at runtime, so a program driving the platform and the parameter tables on this site read the same source rather than a table maintained by hand.

Platform and plugins

Shared machinery

Derivatives are computed by the platform rather than stored in the dataset. kd.PDEDataset holds observations only: fields, axes, and which field-axis-order triple is the left-hand side. Derivatives are computed on demand, either by finite differences on the grid or by training a surrogate of the field and differentiating that network, according to what the algorithm declared it needs. They cannot be stored in the dataset, because the autograd derivatives come from a training run that has not happened when the dataset is written.

The remaining components under a search are shared in the same way: executing a term string into a column of numbers, assembling the design matrix, solving the least-squares problem, computing MSE, NMSE, R² and AIC, and training a surrogate network. Six of the operators route through guarded implementations, so a near-zero denominator or an overflowing exponent produces a finite value rather than an inf or a nan. A new algorithm implements its search and uses these components for the rest.

Data requirements · Derivative sources

Shared notation

The seven algorithms use different internal representations: SGA-PDE searches expression trees, PySR returns SymPy expressions, and PySINDy and DLGA-PDE select columns from a candidate-term library. Before an answer reaches the platform it is converted to the same term strings (mul(u, u_x), diff_x(diff2_x(u))), and the equation is assembled from terms, coefficients and a left-hand side. The internal structures stay inside their own package: results, checkpoints, figures and batch records all carry the shared notation.

The platform keeps no per-algorithm result format. Without a shared notation, two equations could be compared only by the string each algorithm prints, and much of the difference between two such strings is rendering rather than content: reordering a KdV candidate library makes the same law print as add(diff_x(diff2_x(u)), mul(u,u_x)) in one run and add(mul(u,u_x), diff_x(diff2_x(u))) in the next. kd.law_signature therefore reads the set of terms rather than the string, and two equations with the same structure receive the same structure_key.

Equation representation

Score comparability

Each algorithm returns its own objective value: SGA-PDE reports AIC, where lower is better, and DISCOVER reports a reward, where higher is better. The two are not the same quantity, and there is no conversion between them.

The platform does not combine them into a single score. Two scores are comparable when the score_kind, the dataset and the scorer all match. Combining them changes none of those three conditions and only obscures the mismatch. To compare two algorithms, refit the candidate terms against the same data and read the nmse from that fit, which is what kd.evaluate_terms does. For the same reason, best_score_ stays in the algorithm's own units and is never rewritten by the platform.

Score conventions and comparability

Figures from one call

kd.VizEngine renders from the per-iteration record inside the result object, so the call does not change with the algorithm:

kd.VizEngine(output_dir="out/kdv").render_all(
    model.result_, algorithm=model.algorithm_, dataset=data
)

The engine renders the figures common to all algorithms (convergence, parity, residual distribution and field, the equation and the field it reproduces), the diagnostic panels that the algorithm declares as its own, and a single-file HTML report. The engine does not branch on the algorithm's name; each algorithm declares which panels are its own. A figure that cannot be drawn is listed in the report's warnings rather than omitted silently. An overlay of runs whose scores are of different kinds does not average them, and the figure states that it did not.

Visualization

Checkpoints and resuming

Point checkpoint_dir at a directory and the run saves its state as it goes, plus a final checkpoint when it ends. kd.load_checkpoint_manifest reads the ledger that directory keeps, so a resume point comes from recorded entries rather than from a glob over filenames.

A resume restores the search state, while the configuration comes from the new kd.Model. Each parameter carries a tier declared by the algorithm. Changing a resume-safe parameter takes effect on the resumed run; changing an init-only parameter raises an error that names the parameter along with its stored and live values. A parameter with no registered tier is treated as init-only, so an unregistered parameter causes the resume to be refused rather than silently accepted.

Parameter tiers and resuming

Batch studies

A single fit covers one configuration. Comparing results across seeds, algorithms or datasets requires many runs. kd.harness takes an ordered matrix of entries, each an algorithm with a dataset, a seed and its settings, and runs it into a directory that is sealed on write and re-verified whenever it is reopened. An entry that raises is recorded as a failed outcome, and the remaining entries still run.

The aggregate is computed from those records. Runs are grouped by the canonical set of terms they selected, and each group reports how far the agreement inside it extends.

Batch experiments and evidence

Architecture

Batch
kd.harness
Entry points
kd.Modelkd.evaluate_termskd.preview
Algorithms
SGA-PDEDLGA-PDEDISCOVEREqGPTLLM4EDPySRPySINDy
Search
the iteration loopthe plugin contractrecordscheckpoints
Figures
kd.VizEngineplotsreport
Core
term executionleast squaresmetricsguarded operators
Data
kd.PDEDatasetderivativesloaders
Networks
field modeltrainerlanguage-model clients

Each band uses only the bands below it, and no algorithm's search depends on another's. Figures sit below the algorithms rather than above them because an algorithm declares its own diagnostic panels out of the drawing layer's descriptors.

The layering is enforced rather than only documented. An import graph is built on every push and matched against this list, and the list is exhaustive in both directions: a new module must be placed in a band, and a new algorithm must be independent of the other six. An undeclared module or algorithm fails the build.

The path of one fit

The bands above describe which layers may use which. In execution order, a single kd.Model.fit passes through five stages, one of which is the search:

Before
The caller builds the dataset. It holds observations and nothing derived from them: fields, axes, and which field-axis-order triple is the left-hand side.
kd.PDEDataseta loader, a file, the Hub or a generatoradded noise, if wanted
Setup
The fit resolves and checks the left-hand side, reads what the algorithm declared it needs, and assembles that much and no more. An algorithm that declares it needs no derivatives from the platform receives neither those nor an evaluator, and keeps the term executor alone.
finite differences, or a surrogate network trained and differentiatedterm executorevaluator
Search
One call to prepare, then the same three steps every iteration. This is the whole contract: the platform drives the loop and does not decide what happens inside a step.
propose → evaluate → updatecallbacksone record per iterationa checkpoint per iteration, if a directory was given
Answer
The algorithm returns its final candidate, and the platform assembles the equation and the record of the run into one object. The fit returns here.
model.best_expr_model.best_score_model.result_
After
Figures and an independent re-check are separate calls, made when they are wanted. Neither is a tail stage of the fit.
kd.VizEnginekd.verify_equation

Algorithms use different amounts of the setup stage. DISCOVER uses the shared chain throughout: the platform executes and caches the term columns, solves for the coefficients and computes the metrics. PySINDy, PySR and EqGPT are mixed, with the platform building the design matrix and re-scoring the answer while the algorithm's own backend performs the coefficient regression. SGA-PDE runs its own pipeline end to end, including its own sparse regression, and uses only the loop and the records.