the-well

name

the-well

description

Download, load, explore, and benchmark The Well — Polymathic AI's 15TB collection of 16 diverse physics-simulation datasets (fluid dynamics, MHD, acoustic scattering, active matter, supernovae, reaction-diffusion) for training and evaluating spatiotemporal ML surrogate / physics-foundation models. Use when the user mentions "The Well", "polymathic", "the_well", "WellDataset", "물리 시뮬레이션 데이터셋", "PDE surrogate", "물리 서로게이트 모델", "물리 파운데이션 모델", "spatiotemporal simulation dataset", "neural PDE benchmark", "학습용 물리 시뮬레이션 데이터", or wants to pick/stream/download a Well dataset, wire it into a PyTorch DataLoader, visualize a simulation rollout, or benchmark a surrogate model. Do NOT use for generic HuggingFace dataset loading (use hf-datasets), non-physics ML data, or CFD solver authoring.

The Well — Physics Simulation Datasets for ML

Polymathic AI's The Well: 15TB, 16 spatiotemporal physics-simulation datasets for training/evaluating ML surrogate models. This skill routes a request to the cheapest path that answers it — stream a few frames for exploration, or download a split for training — and never assumes 15TB is on disk.

• Package: the_well (PyPI), Python ≥ 3.10, PyTorch interface via WellDataset.
• Paper: arXiv:2412.00568 (NeurIPS 2024). Repo/docs: github.com/PolymathicAI/the_well.
• Data mirror: HuggingFace polymathic-ai/* (streamable, no full download needed).

Full dataset list → references/dataset-catalog.md. API/CLI details → references/api-reference.md. Model benchmarking → references/benchmark-guide.md.

---

Step 1: Detect Runtime (never assume)

Run these and build the decision tree from the results. Prefer the repo canonical interpreter (.venv/bin/python, 3.12) per this project's python-runtime rule.

# interpreter
PY=".venv/bin/python"; [ -x "$PY" ] || PY="python3"; $PY --version
# package present?
$PY -c "import the_well, importlib.metadata as m; print('the_well', m.version('the_well'))" 2>/dev/null || echo "NO the_well"
# torch + device
$PY -c "import torch;print('torch',torch.__version__,'mps',torch.backends.mps.is_available(),'cuda',torch.cuda.is_available())" 2>/dev/null || echo "NO torch"
# CLI shim
command -v the-well-download >/dev/null && echo "CLI ok" || echo "NO cli"
# HF auth (needed only for gated/large private pulls; public streaming works without)
$PY -c "import huggingface_hub as h;print('hf', h.whoami()['name'])" 2>/dev/null || echo "NO hf-auth (public streaming still fine)"
# disk headroom at target base path
df -h . | tail -1

Decision tree

Detection result	Path
NO the_well	Install: VIRTUAL_ENV="$PWD/.venv" uv pip install the_well (add [benchmark] only if benchmarking). Then re-run Step 1.
NO torch	the_well pulls torch as a dep; the install above fixes it. On Apple Silicon, MPS is used automatically.
Goal = explore / few frames / visualize	Stream from HF (Step 3, Method B) — zero full download.
Goal = train / repeated epochs / benchmark	Download the split (Step 3, Method A) — check df -h first; splits range 6.9GB–5.1TB.
df headroom < requested dataset	Fall back to streaming or a smaller dataset/split; warn the user before any multi-GB pull.

Exit gate: you know (a) interpreter, (b) whether the_well importable, (c) stream-vs-download, (d) disk headroom.

---

Step 2: Choose the Dataset & Split

Ask the user for a physical domain if not given; otherwise map their intent to a dataset using references/dataset-catalog.md (16 datasets across fluids, MHD, astro, acoustics, biology, reaction-diffusion). Defaults below keep the skill from stalling.

Parameter	Default	Notes
dataset	active_matter	Smallest well-rounded 2D set — good first pull / demo.
split	train	Also valid, test.
base_path	hf://datasets/polymathic-ai/	HF streaming root (no download). For local: data/the_well.
mode	explore	explore → stream+visualize; train → download+DataLoader; benchmark → hydra.
n_preview_frames	4	Frames to pull for a quick look.

If the user names a phenomenon ("turbulence", "supernova", "Rayleigh–Bénard", "MHD"), resolve it via the catalog rather than guessing the exact dataset string.

---

Step 3: Acquire the Data

Pick ONE method from Step 1's decision tree. Do not download unless the goal needs repeated access AND disk headroom is confirmed.

Method A — Download a split (training / benchmarking)

the-well-download --base-path data/the_well --dataset active_matter --split train
# omit --split to fetch all splits; see references/api-reference.md for every flag

Gate: confirm df -h headroom ≥ split size before running; echo the expected size to the user.

Method B — Stream from HuggingFace (exploration / low disk)

No download. Point well_base_path at the HF root; frames are fetched lazily.

from the_well.data import WellDataset
ds = WellDataset(
    well_base_path="hf://datasets/polymathic-ai/",  # streaming root
    well_dataset_name="active_matter",
    well_split_name="train",
)

---

Step 4: Load into a PyTorch Pipeline

from the_well.data import WellDataset
from torch.utils.data import DataLoader

ds = WellDataset(
    well_base_path="data/the_well",        # or the hf:// root from Method B
    well_dataset_name="active_matter",
    well_split_name="train",
)
loader = DataLoader(ds, batch_size=1, num_workers=2)
sample = next(iter(loader))               # dict — verified keys below

A sample is a dict (verified on active_matter, the_well==1.2.0): input_fields / output_fields (T, H, W[, D], C), constant_scalars, boundary_conditions, space_grid, input_time_grid, output_time_grid. Control the window with constructor kwargs n_steps_input / n_steps_output (default 1/1), use_normalization=True, return_grid, flatten_tensors, full_trajectory_mode. Field semantics per dataset → references/dataset-catalog.md; full kwarg list → references/api-reference.md.

Gate: print sample tensor shapes and dtypes; confirm they match the expected dataset dims before any heavy loop.

---

Step 5: Explore / Visualize (mode=explore)

Write a short script (native Write, then run with $PY) that pulls n_preview_frames, prints field names + shapes, and renders a rollout to PNG/GIF. Keep to compositor-cheap matplotlib. Save under outputs/the-well/{dataset}/. On macOS interactive sessions, open the final image (per this repo's display-generated-images rule); skip open in headless/cron runs.

Gate: image exists and is > 5KB before reporting success.

---

Step 6: Benchmark a Surrogate (mode=benchmark, optional)

Only when the user wants model evaluation. Requires the_well[benchmark] (hydra, extra deps). The repo ships SOTA baselines (e.g. FNO, U-Net variants) + a hydra train.py. Route to references/benchmark-guide.md for config layout, model list, and metrics (VRMSE etc.). Heavy runs → dispatch to a subagent / GPU host, not the main context ([[loop-monitor-cost-guard]]).

---

Step 7: Respond to the User

Output these numbered sections:

1. Selection — dataset, split, domain (1 line), why it matched the request.
2. Access method — stream vs download, and disk cost (GB) if downloaded.
3. What was produced — tensor shapes of a sample, and file path(s) of any visualization.
4. Next step — one concrete follow-up (train loop, another dataset, benchmark).
5. Provenance & license — cite The Well (arXiv:2412.00568) and note dataset licensing must be checked in the repo/HF card before redistribution (do not assert a license from memory).

---

Reference Files

• references/dataset-catalog.md — all 16 datasets: domain, dimensionality, fields, size tier, resolver keywords.
• references/api-reference.md — WellDataset params, the-well-download flags, streaming, normalization, metadata schema.
• references/benchmark-guide.md — hydra config layout, bundled models, metrics, GPU guidance.