The Well

The Well

The Well

The Well 是由 Polymathic AI（隶属于 Simons Foundation）联合 Flatiron Institute、剑桥大学、纽约大学等多机构合作开源的物理模拟数据集项目。项目总计 15TB，包含 16 个涵盖生物系统、流体力学、声波散射、超新星爆炸等硬核物理现象的数值模拟数据集，旨在为机器学习研究者提供训练跨学科泛化的偏微分方程（PDE）代理模型的数据基础设施。该项目已正式发表于 NeurIPS 2024 Datasets & Benchmarks Track。

核心规格

  • 总规模：15TB
  • 数据集数量：16 个
  • 数据格式：统一 HDF5 格式，均匀网格、恒定时间间隔采样，fp32 精度，附带元数据 YAML 文件
  • 数据形状：(n_traj, n_steps, coord1, coord2, [coord3])
  • 数据划分：80% 训练、10% 验证、10% 测试
  • 生成工具：Dedalus、Athena++、TurMix3D 等最先进数值求解器
  • 计算成本：耗费百万 CPU/GPU 小时

16 个数据集概览

  1. acoustic_scattering（Cartesian 2D, 256×256, 100 steps, 8000 traj）：声波散射模拟
  2. active_matter（Cartesian 2D, 256×256, 81 steps, 360 traj, 51.3 GB）：生物活性物质
  3. convective_envelope_rsg（Spherical, 256×128×256, 100 steps, 29 traj, 570 GB）：红巨星对流包 3D 辐射流体力学
  4. euler_multi_quadrants（Cartesian 2D, 512×512, 100 steps, 10000 traj）：欧拉方程多象限激波
  5. gray_scott_reaction_diffusion（Cartesian 2D, 128×128, 1001 steps, 1200 traj, 154 GB）：Gray-Scott 反应扩散系统
  6. helmholtz_staircase（Cartesian 2D, 1024×256, 50 steps, 512 traj, 52 GB）：Helmholtz 方程阶梯问题
  7. MHD（Cartesian 3D, 64³/256³, 100 steps, 100 traj）：磁流体力学模拟
  8. planetswe（Angular, 256×512, 1008 steps, 120 traj, 186 GB）：行星尺度浅水方程
  9. post_neutron_star_merger（Log-Spherical, 192×128×66, 181 steps, 8 traj, 110 GB）：中子星并合后演化
  10. rayleigh_benard（Cartesian 2D, 512×128, 200 steps, 1750 traj, 358 GB）：Rayleigh-Bénard 对流
  11. rayleigh_taylor_instability（Cartesian 3D, 128³, 120 steps, 45 traj, 256 GB）：Rayleigh-Taylor 不稳定性
  12. shear_flow（Cartesian 2D, 128×256, 200 steps, 1120 traj, 115 GB）：周期剪切流（Navier-Stokes）
  13. supernova_explosion（Cartesian 3D, 64³/128³, 59 steps, 1000 traj）：超新星爆炸
  14. turbulence_gravity_cooling（Cartesian 3D, 64³, 50 steps, 2700 traj, 829 GB）：湍流-重力-冷却
  15. turbulent_radiative_layer_2D（Cartesian 2D, 128×384, 101 steps, 90 traj, 6.9 GB）：湍流辐射层（最小数据集）
  16. turbulent_radiative_layer_3D（Cartesian 3D, 128×128×256, 101 steps, 90 traj, 745 GB）：三维湍流辐射层

使用方法

安装

1
pip install the_well

下载数据

1
the-well-download --base-path /your/data/path --dataset active_matter --split train

加载数据

1
2
3
4
5
6
7
8
9
10
11
from the_well.data import WellDataset
from torch.utils.data import DataLoader

dataset = WellDataset(
    well_base_path="/your/data/path",
    well_dataset_name="active_matter",
    well_split_name="train",
    n_steps_input=1,
    n_steps_output=1,
)
loader = DataLoader(dataset, batch_size=4, shuffle=True)

Hugging Face 流式加载

1
2
3
4
5
trainset = WellDataset(
    well_base_path="hf://datasets/polymathic-ai/",
    well_dataset_name="active_matter",
    well_split_name="train"
)

运行基准模型训练

1
2
cd the_well/benchmark
python train.py experiment=fno server=local data=active_matter

基准模型

项目内置 FNO、TFNO、U-Net 等模型训练脚本，支持 Hydra 配置系统。早期基准测试显示，U-Net 在 active_matter 等数据集上 VRMSE 表现优于 FNO，但整体仍有显著提升空间。

资源链接

意义与影响

The Well 被视为生成式 AI 向科学 AI 转型的里程碑项目。它有望像 Protein Data Bank 催生 AlphaFold 一样，催生下一代科学 AI 基础模型。项目强调"互补而非替代"数值模拟，机器学习模型可快速预测稳态、增长率，释放算力用于更精细的物理模拟，最终加速科学发现。

分享到