Metadata-Version: 2.4
Name: shodhbench
Version: 0.1.0
Summary: Physics-first evaluation framework for fluid and molecular AI models.
Author: Adya
License-Expression: MIT
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.24
Requires-Dist: zarr>=2.16
Provides-Extra: dev
Requires-Dist: pytest>=8; extra == "dev"
Provides-Extra: sim
Requires-Dist: jax>=0.4; extra == "sim"
Requires-Dist: jax-md>=0.2; extra == "sim"
Provides-Extra: cluster
Requires-Dist: jax>=0.4; extra == "cluster"
Requires-Dist: jax-md>=0.2; extra == "cluster"
Requires-Dist: zarr>=2.16; extra == "cluster"
Requires-Dist: gcsfs>=2024.0; extra == "cluster"
Requires-Dist: fsspec>=2024.0; extra == "cluster"
Provides-Extra: baseline
Requires-Dist: torch>=2; extra == "baseline"
Requires-Dist: zarr>=2.16; extra == "baseline"
Requires-Dist: gcsfs>=2024.0; extra == "baseline"
Provides-Extra: tpu
Requires-Dist: jax[tpu]>=0.4; extra == "tpu"
Requires-Dist: zarr>=2.16; extra == "tpu"
Requires-Dist: gcsfs>=2024.0; extra == "tpu"
Requires-Dist: fsspec>=2024.0; extra == "tpu"

# 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 production dataset at `gs://shodhbench-tokyo-prod-stunning-grin-493612-t4/production_20260604_classa_opt/`.
The PINN parameter count can be reproduced with `python3 scripts/count_pinn_parameters.py`.
EXP-401 score artifacts are under `/sharedstorage/innmivshodhaislurmh1/home/shodh_harshita/model_r-D/runs/exp401_shodhbench_predictions_130808/`. Native energy diagnostics show the energy loss is already present before the final `64^3` downsample.

## Quick Start

After the PyPI release:

```bash
python -m pip install shodhbench
```

Run a local smoke evaluation:

```python
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:

```text
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:

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

Molecular samples use padded graph tensors:

```text
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:

```bash
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:

```text
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:

```bash
python3 scripts/generate_datasets.py --production-spec
```

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

```bash
python3 scripts/generate_cluster.py \
  --dataset taylor_green_3d \
  --output gs://shodhbench-public-v1/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://shodhbench-public-v1/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://shodhbench-public-v1/shodhbench_water.zarr \
  --num-samples 10000 \
  --steps 50 \
  --chunk-size 2048 \
  --zarr-sample-chunk 8 \
  --backend jax \
  --overwrite
```

On the cluster:

```bash
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:

```text
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:

```bash
scripts/setup_tpu_8node_env.sh
```

Launch the 50k Class-A optimized generation:

```bash
SHODHBENCH_OUTPUT_ROOT=gs://shodhbench-tokyo-prod-stunning-grin-493612-t4/production_20260608_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

```python
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:

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

python3 scripts/vanilla_pinn_baseline.py \
  --zarr gs://shodhbench-tokyo-prod-stunning-grin-493612-t4/production_20260604_classa_opt/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:

```bash
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`:

```bash
python scripts/vanilla_pinn_baseline_jax.py \
  --zarr gs://shodhbench-tokyo-prod-stunning-grin-493612-t4/production_20260604_classa_opt/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`:

```bash
python scripts/vanilla_pinn_baseline_jax.py \
  --zarr gs://shodhbench-tokyo-prod-stunning-grin-493612-t4/production_20260604_classa_opt/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.
