shodhbench 0.1.1

Physics-first evaluation framework for fluid and molecular AI models.

ShodhBench

ShodhBench is a physics-first benchmark for grading physical AI models against canonical fluid and molecular trajectories. The benchmark stores ground-truth answer keys in tokenizer-ready Zarr layouts, then scores model predictions by physical violations instead of pixel MSE.

ShodhBench v0.1: Provisional Smoke Leaderboard

The benchmark generation pipeline and mathematical grader are the primary contributions of this v0.1 release. Metrics below are evaluated on n=1 validation smoke samples from shard_0, sample_0 for Taylor-Green 3D and Cylinder Flow 2D. Full dataset validation sweeps are currently computing.

Lower is better for every numeric ShodhBench metric. No undisputed scalar rank is implied yet: the metrics expose different failure modes.

Trade-Off Table

Metric Family	Current Smoke Winner	UNIPHY-1B / Legacy EXP-401	Vanilla PINN	Interpretation
Mass divergence, max	UNIPHY-1B / Legacy EXP-401	0.046534	2.293518	UNIPHY-1B has much stronger hard-constraint behavior on divergence.
Mass divergence, mean	UNIPHY-1B / Legacy EXP-401	0.000905	0.042048	UNIPHY-1B reduces average divergence by a large margin.
Energy drift	Vanilla PINN	1.000000	0.012380	The current UNIPHY-1B state-rollout adapter loses kinetic energy over long rollout.
Spectral L2 error	Vanilla PINN	0.087823	0.056413	The PINN is lower on this n=1 smoke average.

Model	Params	Evaluation	Max Divergence	Mean Divergence	Energy Drift	Spectral L2 Error	Status
UNIPHY-1B / Legacy EXP-401 state-rollout adapter	1,086,137,368	Taylor + Cylinder average, shard_0 sample_0	0.046534	0.000905	1.000000	0.087823	Strong divergence control; severe energy rollout failure
Vanilla PINN	50,692	Taylor + Cylinder average, shard_0 sample_0	2.293518	0.042048	0.012380	0.056413	Small open-source baseline; lower energy drift in this smoke test

UNIPHY-1B (Legacy EXP-401) utilizes an autoregressive state-rollout adapter which suffers from severe numerical diffusion, resulting in near-total kinetic energy loss. This benchmark validates the necessity of Symplectic ODE architectures for long-horizon rollouts.

The EXP-401 row uses exported predictions from autoregressive_state_rollout_from_pred_fluid, with native 128^3 feedback for all 50 steps, one final downsample to ShodhBench shape, pressure_proxy_from_density = (rho - rho_ref) / 3, and rho_ref = 1.0. This is a ShodhBench adapter smoke, not a derivative/ODE rollout.

Detailed Breakdown

Model	Params	Dataset	Sample	Steps	Max Divergence	Mean Divergence	Energy Drift	Spectral L2 Error
UNIPHY-1B / Legacy EXP-401 state-rollout adapter	1,086,137,368	Taylor-Green 3D	shard_0 sample_0	50-frame rollout	0.053350	0.000813	1.000000	0.133621
UNIPHY-1B / Legacy EXP-401 state-rollout adapter	1,086,137,368	Cylinder Flow 2D	shard_0 sample_0	50-frame rollout	0.039718	0.000996	1.000000	0.042025
Vanilla PINN	50,692	Taylor-Green 3D	shard_0 sample_0	3000	4.501651	0.075477	0.004256	0.001231
Vanilla PINN	50,692	Cylinder Flow 2D	shard_0 sample_0	3000	0.085386	0.008620	0.020504	0.111596

Baseline artifacts were produced on the locked production smoke samples. Public dataset mirrors and validation-sweep artifacts will be linked here after release. The PINN parameter count can be reproduced with python3 scripts/count_pinn_parameters.py. Native EXP-401 energy diagnostics show the energy loss is already present before the final 64^3 downsample.

Quick Start

After the PyPI release:

python -m pip install shodhbench

Run a local smoke evaluation:

import numpy as np
import shodhbench

dataset = shodhbench.load("taylor_green_3d")
truth = dataset.trajectory
prediction = truth + np.random.default_rng(7).normal(0.0, 0.01, truth.shape).astype(truth.dtype)
scores = shodhbench.evaluate_physics(prediction, dataset)
print(scores)

Locked V1 Schema

Fluid samples are channel-first and sample-local:

trajectory: [num_samples, 50, 4, Z, Y, X]
sdf_mask:   [num_samples, 1, 1, Z, Y, X]
params:     [num_samples, P]
channels:   [vx, vy, vz, pressure]

Taylor-Green uses [num_samples, 50, 4, 64, 64, 64] in production. Its sdf_mask is repeated per sample and filled with 1.0 because the whole domain is fluid.

Cylinder flow is physically 2D but padded for 3D tokenizer compatibility:

trajectory: [num_samples, 50, 4, 1, 64, 64]
sdf_mask:   [num_samples, 1, 1, 1, 64, 64]

Molecular samples use padded graph tensors:

trajectory: [num_samples, 50, 1024, 9]
atom_mask:  [num_samples, 1024]
features:   [x, y, z, vx, vy, vz, atom_type_id, charge, mass]

Editable Local Smoke Test

The local defaults are intentionally tiny so you can verify the package from a laptop:

python3 -m pip install -e .
python3 scripts/generate_datasets.py --dataset all --samples 1 --fluid-resolution 16
python3 test.py

test.py regenerates a small Taylor-Green store, adds Gaussian noise as a dummy model prediction, and prints the physics scores.

Production Generation Target

The locked production sweep is:

20,000 Taylor-Green samples, Reynolds 100-5000
20,000 cylinder samples, varying inlet velocity and cylinder radius
10,000 water samples, temperature 273K-600K
T = 50 timesteps per sample
Fluid channels = [vx, vy, vz, pressure]
Water features = [x, y, z, vx, vy, vz, atom_type_id, charge, mass]

The small local CLI exposes the production shape contract:

python3 scripts/generate_datasets.py --production-spec

Do not run that command on a laptop. The actual cluster writer is:

python3 scripts/generate_cluster.py \
  --dataset taylor_green_3d \
  --output gs://<your-bucket>/shodhbench_taylor_green.zarr \
  --num-samples 20000 \
  --steps 50 \
  --resolution 64 \
  --chunk-size 2048 \
  --zarr-sample-chunk 1 \
  --backend jax \
  --overwrite

python3 scripts/generate_cluster.py \
  --dataset cylinder_flow_2d \
  --output gs://<your-bucket>/shodhbench_cylinder.zarr \
  --num-samples 20000 \
  --steps 50 \
  --resolution 64 \
  --chunk-size 2048 \
  --zarr-sample-chunk 8 \
  --backend jax \
  --overwrite

python3 scripts/generate_cluster.py \
  --dataset water_box_h2o \
  --output gs://<your-bucket>/shodhbench_water.zarr \
  --num-samples 10000 \
  --steps 50 \
  --chunk-size 2048 \
  --zarr-sample-chunk 8 \
  --backend jax \
  --overwrite

On the cluster:

git clone <repo-url>
cd shodhbench
python3 -m pip install -e ".[cluster]"
gcloud auth application-default login

Then run the three commands above. Use --dry-run first to print the final shapes and metadata without writing.

For the current 8-node v6e-8 TPU layout, use explicit shard IDs instead of TPU WORKER_ID, because each separate node reports WORKER_ID=0. The Class-A optimized production layout is:

Taylor-Green: 8 shards x 2,500 samples, trajectory chunks [4, 50, 4, 64, 64, 64]
Cylinder:     8 shards x 2,500 samples, trajectory chunks [8, 50, 4, 1, 64, 64]
Water H2O:    8 shards x 1,250 samples, trajectory chunks [8, 50, 1024, 9]

Set up all eight TPU nodes:

scripts/setup_tpu_8node_env.sh

Launch the 50k Class-A optimized generation into your own cloud bucket:

SHODHBENCH_OUTPUT_ROOT=gs://<your-bucket>/shodhbench_v1_classa_opt_8node \
  scripts/launch_tpu_8node_classa_opt.sh

The launcher starts one process per TPU node and passes --jax-distributed --jax-num-processes 8 --jax-process-id <0..7> so the Megascale TPU slice initializes correctly. Each node writes a separate shard.

Evaluation

import shodhbench

dataset = shodhbench.load("taylor_green_3d")
predictions = my_ai_model.predict(dataset)
scores = shodhbench.evaluate_physics(predictions, dataset)
print(scores)

Fluid scores include:

• max_divergence: max absolute div(V) in the predicted final field
• mean_divergence: mean absolute div(V) in the predicted final field
• energy_drift: fractional final kinetic-energy error against the ground-truth final frame
• spectral_l2_error: L2 distance between predicted and ground-truth FFT spectra; cylinder samples with Z=1 use a true 2D FFT over [y, x]

Molecular scores currently include energy_drift over the padded atom trajectory using atom_mask.

The public grader hard-fails shape mismatches. Predictions must match the ground-truth tensor shape exactly.

Vanilla PINN Baseline

The repo includes a simple PyTorch PINN baseline for one Taylor-Green or cylinder sample:

python3 -m pip install -e ".[baseline]"

python3 scripts/vanilla_pinn_baseline.py \
  --zarr <PUBLIC_DATA_ROOT>/taylor_green_3d/shard_0.zarr \
  --dataset taylor_green_3d \
  --sample-index 0 \
  --steps 3000 \
  --data-batch 4096 \
  --pde-batch 4096 \
  --output runs/vanilla_pinn_taylor_sample0.npz

This baseline is intentionally plain: an MLP maps (x, y, z, t) to [vx, vy, vz, pressure], using data MSE plus a soft Navier-Stokes residual. It is meant to establish a weak open-source leaderboard baseline, not to be a competitive model.

For TPU-only environments, use the pure-JAX baseline instead. On a fresh TPU VM:

python3 -m venv ~/shodhbench-venv
source ~/shodhbench-venv/bin/activate
python -m pip install -U pip setuptools wheel
python -m pip install -e ".[tpu]" \
  -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

python -c "import jax; print(jax.devices()); print(jax.device_count())"

Run Taylor-Green on production shard_0:

python scripts/vanilla_pinn_baseline_jax.py \
  --zarr <PUBLIC_DATA_ROOT>/taylor_green_3d/shard_0.zarr \
  --dataset taylor_green_3d \
  --sample-index 0 \
  --steps 3000 \
  --data-batch 4096 \
  --pde-batch 1024 \
  --output runs/vanilla_pinn_taylor_sample0_jax.npz \
  --scores-json runs/vanilla_pinn_taylor_sample0_jax_scores.json

Run cylinder on production shard_0:

python scripts/vanilla_pinn_baseline_jax.py \
  --zarr <PUBLIC_DATA_ROOT>/cylinder_flow_2d/shard_0.zarr \
  --dataset cylinder_flow_2d \
  --sample-index 0 \
  --steps 3000 \
  --data-batch 4096 \
  --pde-batch 1024 \
  --output runs/vanilla_pinn_cylinder_sample0_jax.npz \
  --scores-json runs/vanilla_pinn_cylinder_sample0_jax_scores.json

The JAX script refuses to run on CPU unless --allow-cpu is passed, so official TPU baseline results cannot accidentally come from a CPU fallback.

Data Policy

This repo uses canonical public/synthetic decoys only. It does not use proprietary Shodh AI data from Aarti, Jubilant, A123, or any customer/internal source.