PDE-KD feature maps: pde-kd-frame and pde-kd-3d

Weights of the two learned feature maps of PDE-KD, the distributional evaluation metric of A Distributional Evaluation Metric for Chaotic Physical Systems (Gong, Jiang, He, Li, 2026). PDE-KD is a kernel discrepancy (the squared MMD under KID's cubic polynomial kernel) between a population of simulated states and a population of reference states, both embedded by one of these maps, with independent trajectories (initial conditions) as the unit of evidence.

Code, documentation and every experiment of the paper: https://github.com/HarryGong31/pde-kd.

folder	map	input	output	files
pde-kd-frame/	2D ResNet-18 ("PDE-KD" in the paper)	one frame, 13 canonical channels at 512 x 512	512-d feature (global mean pool); a 46-wide head	multichan-v4-equal-ext36-20260624-000703-step22176.npz (45 MB), its .sha256.json, head_layout.json
pde-kd-3d/	3D ResNet-18 ("PDE-KD (3D)" in the paper)	a window of 20 consecutive frames, 13 channels at 512 x 512	512-d feature	clip3d-t20-noaux-ft96-20260715-step129312.npz (133 MB), its .sha256.json

Use

The pdekd package downloads these files itself, pinned to one commit of this repository, and checks every file's and every tensor's sha256:

pip install torch     # CPU or CUDA build for your platform
pip install "pdekd @ git+https://github.com/HarryGong31/pde-kd"

from pdekd import FeatureMap, KidBucket, two_sample_test

fmap = FeatureMap("rayleigh_benard", "pde-kd-frame")   # or "pde-kd-3d"; device="cuda"
x, x_ics = fmap.embed(reference_trajectories)          # [N, 512], {ic_key: rows}
y, y_ics = fmap.embed(generated_trajectories)
kd, diag = KidBucket(x, x_ics, y, y_ics).estimate(x_ics, y_ics)
p = two_sample_test(x, x_ics, y, y_ics)["p_value"]     # IC-level permutation test

A trajectory is a loader item {"native": [T, nx, ny, C] float32 in physical units, "source_path": str, "traj_index": int} on one of the 20 systems with committed channel statistics: the nine training classes below and eleven further datasets (Turbulent Radiative Layer 2D, HRRR, ERA5 and eight RealPDEBench sets, of which the paper uses the two FSI ones). The repository's README has a runnable example on arrays and one on The Well, the CLIs for whole datasets, and the tests the paper uses (the exact test above for disjoint initial conditions, the shared-IC test for rollouts that start from the reference's own initial conditions).

The files land in the asset root: $PDEKD_ASSETS, by default ~/.cache/pdekd ($XDG_CACHE_HOME/pdekd) for an installed package. To install them by hand (an offline machine), download this repository elsewhere, check it with sha256sum -c SHA256SUMS, put the files of pde-kd-frame/ under $PDEKD_ASSETS/weights/ (and head_layout.json at $PDEKD_ASSETS/) and those of pde-kd-3d/ under $PDEKD_ASSETS/weights3d/, and set PDEKD_OFFLINE=1. Keep the file names: the library reads the training run and step off them, and with them the tag of every feature bank it writes.

Format

Each .npz holds float32 tensors in PyTorch layout under the module's own parameter names (no pickle). pdekd/feature_map/model.py (2D) and model3d.py (3D) load them strictly, and the .sha256.json lists every tensor's shape and sha256, which the loader checks. head_layout.json describes the 2D map's packed 46-wide head: 18 coefficient outputs, 9 class logits and 19 regime logits.

Training

Both maps were trained in JAX/Flax on TPU v5e and exported once to PyTorch; on Rayleigh-Benard validation frames the port matches the JAX forward pass to 6.6e-7 (2D features) and 1.8e-6 (3D features), max absolute difference.

• Objective. Supervised identification of the physics: 9-way class cross-entropy, plus 2.0 x masked MSE on the class's PDE coefficients (e.g. Rayleigh and Prandtl numbers, mapped to [0, 1]), plus 0.5 x cross-entropy on its regime labels.
• Input. Each system's fields are scattered into 13 canonical channel slots, z-scored per (system, channel) with train-split statistics and resized bilinearly to 512 x 512. No augmentation.
• pde-kd-frame. One random frame of one random trajectory per sample, classes sampled uniformly; AdamW, peak learning rate 1e-4, 22,176 steps (global batch 256, then 1,024); 11.1 M frames; 1,298 v5e chip-hours.
• pde-kd-3d. A continuation of the same 2D training (step 46,368), inflated to 3D (every convolution except the 1 x 1 shortcut projections gains a temporal width of three frames, initialized to reproduce the 2D network frame by frame), then fine-tuned on 20-frame clips with the same data and objective for 129,312 steps of 32 clips.

Data. Nine classes, 70,770 / 1,813 / 1,710 train / validation / test trajectories; the splits hold out initial conditions, not coefficient values.

class	source	supervised by
shear flow, Rayleigh-Benard, Euler (open and periodic boundaries), active matter	The Well (CC-BY-4.0)	PDE coefficients
viscoelastic instability, Gray-Scott	The Well (CC-BY-4.0)	regimes / patterns
compressible Euler (CE), Navier-Stokes (NS)	PDEgym of Poseidon (CC-BY-NC-4.0)	regimes

Two overlaps with The Well's evaluation splits are disclosed in the paper's appendix "Evaluation Data": one training directory (Euler, open boundaries, gamma = 1.22) held The Well's test trajectories 0-49 of that setting, so both maps saw those 50 test trajectories in training; and The Well itself duplicates trajectories across its splits (shear flow: 89 of 112 test and 88 of 112 validation trajectories are bitwise copies of training ones; Euler, periodic boundaries: 2 test and 1 validation).

Scope

The maps were trained on the nine classes above, on grids of at most 512 per axis. Other systems need their channel statistics and a channel-slot assignment (pdekd/specs.py); the paper's appendix "Transfer to Out-of-Distribution Data" reports how the metric behaves on weather data (HRRR, ERA5) and on RealPDEBench. PDE-KD compares distributions of states; it is not a paired error on grid points, and two populations drawn from one law give a KD near 0.

License

The weights are released under the MIT license, as is the code. They were trained partly on PDEgym (CC-BY-NC-4.0) and on The Well (CC-BY-4.0); we state the training data so that users can judge those terms for their own use. The repository's THIRD_PARTY.md lists the terms of every third-party component and dataset.

Citation

@article{gong2026distributional,
  title   = {A Distributional Evaluation Metric for Chaotic Physical Systems},
  author  = {Gong, Harry and Jiang, Yehong and He, Kaiming and Li, Zongyi},
  journal = {arXiv preprint},
  year    = {2026}
}