Skip to content

Getting started: discover a PDE from field data

KD recovers the Burgers equation from a bundled benchmark dataset whose equation is known. This notebook loads the data, inspects the field, fits a symbolic model in one call, compares the result with the known equation, and plots the convergence curve.

import kd

ds = kd.load_burgers()
print(ds.ground_truth)

u_t = -u * u_x + 0.1 * u_xx

Look at the data

The dataset contains a field u sampled on an (x, t) grid. Before fitting, kd.preview prints the grid, the field, and the left-hand-side derivative that KD will use.

import matplotlib.pyplot as plt

kd.preview(ds)

field = ds.get_field(ds.lhs_field).detach().cpu().numpy()
x = ds.get_coords("x")
t = ds.get_coords("t")

fig, ax = plt.subplots(figsize=(7, 4))
mesh = ax.pcolormesh(t, x, field, shading="auto", cmap="viridis")
ax.set_xlabel("t")
ax.set_ylabel("x")
ax.set_title("Burgers field u(x, t)")
fig.colorbar(mesh, ax=ax, label="u")
plt.show()

Dataset: burgers
Axes:
     x | n=256  | range [-8.000, 7.938] | step 0.0625 (uniform)
     t | n=201  | range [0.000, 10.000] | step 0.05 (uniform)
Fields:
     u | dtype=float64 | shape=(256, 201) | min=-1.000 max=1.000 mean=0.000 (NaN=0)
LHS: u_t  (field='u', axis='t')
Status: ready to fit

png

Fit a model

SGA-PDE searches over symbolic expression trees and scores candidate right-hand sides against the target derivative. generations sets the search budget. The value below is small enough to run quickly and large enough to recover a Burgers-like equation.

model = kd.Model(algorithm="sga", generations=5, seed=0, verbose=False).fit(ds)
print(model.best_expr_)
print(model.best_score_)

u_t = -1*mul(u_x, u) + 0.1002*diff2_x(u)
-28.776750720652373

Compare to ground truth

The discovered equation is KD's best symbolic right-hand side for u_t. The exact string can differ from the compact textbook form, so the check is whether it contains the nonlinear advection and diffusion terms.

print(f"Discovered  : {model.best_expr_}")
print(f"Ground truth: {ds.ground_truth}")

Discovered  : u_t = -1*mul(u_x, u) + 0.1002*diff2_x(u)
Ground truth: u_t = -u * u_x + 0.1 * u_xx

Visualize

kd.viz.plots.plot_convergence draws the best-score curve straight from the fitted result. It reads the recorded score series, masks the non-finite warm-up sentinel so the y axis stays readable, labels that axis with the engine's own score kind (AIC for SGA), and adds a subtitle when the curve is flat (best score reached on the first iteration). The return value is a list of warning strings, empty when the curve was drawn.

from kd.viz.plots import plot_convergence

fig, ax = plt.subplots(figsize=(7, 4))
messages = plot_convergence(model.result_, ax)
for message in messages:
    print(message)
plt.show()

png

Next steps

This notebook used SGA-PDE. KD provides seven algorithms behind the same kd.Model interface: