Source-linked AI summary

The Well: a Large-Scale Collection of Diverse Physics Simulations for Machine Learning

Ruben Ohana, Michael McCabe, Lucas Meyer, Rudy Morel, Fruzsina J. Agocs, Miguel Beneitez, Marsha Berger, Blakesley Burkhart, Keaton Burns, Stuart B. Dalziel, Drummond B. Fielding, Daniel Fortunato, Jared A. Goldberg, Keiya Hirashima, Yan-Fei Jiang, Rich R. Kerswell, Suryanarayana Maddu, Jonah Miller, Payel Mukhopadhyay, Stefan S. Nixon, Jeff Shen, Romain Watteaux, Bruno Régaldo-Saint Blancard, François Rozet, Liam H. Parker, Miles Cranmer, Shirley Ho

arXiv:2412.00568v2cs.LGphysics.flu-dyn

TL;DR

Existing physical-dynamics datasets often cover limited task classes or provide too few snapshots and overly large samples for broad evaluation. The Well responds with a 15TB, 16-dataset collection and unified interface, whose baselines reveal substantial variation across architectures and difficulty in universal modeling. Its main scope remains challenging, accessible benchmark data, with limitations in resolution, mesh complexity, dimensionality, and analysis breadth.

  • Problem

    Physical-dynamics datasets often cover limited task classes, while high-resolution datasets may have too few snapshots or overly large samples for broad surrogate-model evaluation.

  • Method

    The Well curates 15TB of numerical simulations across 16 physical scenarios with a shared specification, PyTorch interface, and example surrogate-model baselines.

  • Results

    Baseline performance splits across spatial-domain and spectral architectures: 9/17 experiments favor U-net-type models and 8/17 favor spectral models, while longer rollouts reduce some one-step advantages.

  • Takeaways & Limitations

    The Well provides diverse, complex benchmark problems for developing and evaluating physical-dynamics surrogates beyond one-model-fits-all assumptions.

  • Takeaways & Limitations

    The datasets mainly use uniformly sampled, manageable-resolution domains, often 2D, and the release focuses on data rather than comprehensive metrics and analysis.

Abstract

from arXiv · show

Machine learning based surrogate models offer researchers powerful tools for accelerating simulation-based workflows. However, as standard datasets in this space often cover small classes of physical behavior, it can be difficult to evaluate the efficacy of new approaches. To address this gap, we introduce the Well: a large-scale collection of datasets containing numerical simulations of a wide variety of spatiotemporal physical systems. The Well draws from domain experts and numerical software developers to provide 15TB of data across 16 datasets covering diverse domains such as biological systems, fluid dynamics, acoustic scattering, as well as magneto-hydrodynamic simulations of extra-galactic fluids or supernova explosions. These datasets can be used individually or as part of a broader benchmark suite. To facilitate usage of the Well, we provide a unified PyTorch interface for training and evaluating models. We demonstrate the function of this library by introducing example baselines that highlight the new challenges posed by the complex dynamics of the Well. The code and data is available at https://github.com/PolymathicAI/the_well.

1 Introduction

Simulation is central to computational science, but analytical PDE solutions are usually infeasible and numerical simulation is expensive. The Well addresses evaluation and data-scale gaps by providing diverse, large-scale simulation datasets for surrogate modeling.

  • Numerical simulation supports forecasting, optimization, and parameter inference when analytical solutions to most PDEs are infeasible.Numerical methods solve discretized versions of PDEs describing physical behavior under material assumptions.
  • Deep-learning surrogates face adoption challenges because practical problems are more complex than current evaluation datasets, while high accuracy can require substantial data.High-quality simulations also require specialized software, domain expertise, and months of supercomputer time.
  • The Well provides 15 TB of high-quality numerical simulations across 16 datasets spanning application domains, scales, and governing equations.The collection was produced with domain experts and numerical software developers.
  • Each dataset contains temporally coarsened snapshots across multiple initial conditions or physical parameters, with enough snapshots to study simulation stability.

2 Related Work

Physical-dynamics datasets have expanded from common reference simulations toward more complex and specialized problems, but breadth, snapshot counts, and sample scale remain limited. The Well is positioned within this shift toward broader and more demanding evaluation data.

  • Large, curated, diverse datasets have supported progress in other machine-learning fields, including internet-scale language data and billion-scale text-image collections.
  • Early physical-dynamics datasets offered varied reference simulations and were rapidly adopted.
  • More recent datasets address complex specialized simulations and open new applications, but typically cover only a small number of tasks.
  • High-resolution simulation datasets pursue more ambitious problems, yet limited snapshots and large individual samples often restrict their use.

3 Diving into the Well

The Well frames surrogate modeling as autoregressive prediction of PDE solution fields and supplies a standardized, extensible dataset interface. Its datasets span diverse physical systems, including acoustic, biological, fluid, astrophysical, and geophysical phenomena.

  • Surrogate modeling: Surrogate models estimate PDE solution functions from inputs such as initial or boundary conditions, often by predicting successive discretized time states.The autoregressive formulation applies model f repeatedly across sampled times.
  • Dataset format: The Well provides 16 uniformly gridded datasets totaling 15TB, sampled at constant time intervals and stored with shared metadata specifications in HDF5 files.A PyTorch interface exposes the datasets for machine-learning workflows.
  • Dataset format: The PyTorch interface accepts files following the shared specification without code changes, enabling third-party datasets to enter the benchmarking pipeline.
  • Physical scenarios: The collection includes acoustic scattering, active matter, and three-dimensional convective-envelope simulations with distinct wave, hydrodynamic, and stellar dynamics.The acoustic and active-matter descriptions emphasize geometry and long-range interactions, while the stellar simulations capture inherently three-dimensional convection.
  • Physical scenarios: Other scenarios cover compressible Euler flows, Gray–Scott reaction-diffusion, periodic acoustic scattering, MHD turbulence, shallow-water dynamics, and neutron-star merger environments.These datasets span biological, fluid, acoustic, geophysical, and astrophysical applications.
  • Physical scenarios: The collection also includes Rayleigh-Bénard convection and Rayleigh-Taylor instability, where buoyancy and density perturbations generate evolving fluid structures.

4 Benchmark

The benchmark evaluates simple surrogate-model baselines with coarse and fine metrics, revealing architecture-dependent performance, rollout difficulty, and opportunities for physics-aware analysis.

  • Baselines: The baselines use 12-hour single-H100 budgets to predict the next snapshot from a four-timestep history.The settings are intended as reasonable off-the-shelf choices rather than peak-performance configurations.
  • Evaluation Metrics: VRMSE is averaged across physical fields and time windows because centered normalization better suits non-negative fields than NRMSE.The benchmark reports one-step and longer-rollout losses, with lower VRMSE indicating better performance.
  • Results: CNextU-net leads 8 of 17 next-step experiments, while 9 favor U-net models and 8 favor spectral models.The split suggests that spatial-versus-spectral preferences vary across problems, with no clear explanation from boundary conditions.
  • Results: Longer autoregressive rollouts perform notably worse, indicating that one-step training alone struggles to maintain prediction quality over time.The performance gap between model classes persists across time-averaged rollout windows.
  • Results: One-step and rollout losses can disagree because rollouts begin at simulation starts, whereas one-step tests use sliding windows sampled from ground truth.Dissipative physics can also make later windows smoother and easier to predict, despite generally large normalized losses.
  • Moving Beyond the Baselines: Per-field and frequency-binned metrics expose behavior hidden by aggregate scores: pressure dominates some losses, low frequencies remain predictable, and high frequencies diverge sooner.The library is intended to support physically meaningful architecture analysis beyond single-valued metrics.
  • Further Challenges: The baselines are powerful but naive because they do not account for dataset-specific physical characteristics.The Well is positioned for follow-up work on conservation laws, boundary conditions, long-term stability, inverse scattering, super-resolution, and generalization.

5 Conclusion

The Well contributes a diverse, large-scale benchmark for physical dynamics prediction while remaining bounded by resolution, dimensionality, and future hardware constraints.

  • Limitations: The datasets largely use uniformly sampled domains at manageable resolutions, whereas many engineering applications require higher resolutions, complex meshes, and three-dimensional data.The release is primarily data-focused, with more extensive metrics and analysis explored by other work.
  • Conclusion: The Well contains 15 TB across 16 physical scenarios developed with domain experts to provide complexity, volume, and diversity.The authors present it as a resource intended to stimulate surrogate-modeling developments and new workflows.

A Datasheet for The Well

The Well addresses limitations in existing physics datasets by combining diverse, complex simulations under a common structure. It provides simulation trajectories, metadata, and task-oriented splits for machine-learning use.

  • Motivation: The Well was created to bridge the gap between diverse low-complexity benchmarks and complex single-domain simulation datasets.Its construction involved domain experts and numerical software developers to identify simulations that are distinct, non-trivial for learning, and manageable in resolution.
  • Instances and format: The collection contains simulation snapshots of physical state variables on discrete spatial and temporal grids, with metadata describing fields, dimensions, and simulation parameters.Datasets follow a common schema and are stored in self-documenting HDF5 files accessible through a unified interface.
  • Tasks: The Well supports temporal rollout tasks and additional challenges, but it does not assign fixed labels to instances.Users predict future values from historical values, while other supported tasks include challenges documented per dataset.
  • Splits: Within each dataset and simulation-parameter set, trajectories generally use an 80/10/10 train-validation-test split along initial conditions.Datasets with too few simulations use temporally blocked splitting instead.
  • Data characteristics: The dataset is self-contained, contains no human data, and has no cross-linkage between different datasets.Temporal and spatial relationships are inherent in the storage format.

Q36 What (other) tasks could the dataset be used for?

Beyond temporal prediction, the Well supports several AI-for-Science tasks and is distributed openly so users can select individual datasets or use the collection as a benchmark suite.

  • Alternative tasks: Several datasets support inverse acoustic scattering, superresolution, and stability challenges in addition to conventional temporal prediction.These tasks are documented for individual datasets in Section C.
  • Distribution and contributions: The dataset is intended for open distribution with contribution guidelines for new contributors.Contributions require coordination with Polymathic AI for validation, formatting, manageable size, and preliminary benchmarking.
  • Distribution and contributions: The data is hosted by the Flatiron Institute and can be downloaded directly with provided code or accessed through a Globus endpoint.Because the collection is modular, users can obtain individual datasets without downloading unused data.

Q42 When will the dataset be distributed?

The Well is maintained as an extensible collection hosted by the Flatiron Institute, with open access routes, contributor procedures, and ongoing support for existing datasets.

  • Access: Users can download the data through provided code or access it through a Globus endpoint hosted by the Flatiron Institute.Individual datasets can be used separately or together, avoiding the need to download the full collection.
  • Access: The data is distributed openly, with download information provided through the project’s GitHub repository.The Flatiron Institute hosts and maintains the data and Globus endpoint, with planned availability of subsets through Hugging Face.
  • Licensing: Available generating code is released under CC-BY-4.0, although some datasets were generated with proprietary code that cannot be released.This creates a boundary between the licensing of available code and the generation software used for some datasets.
  • Maintenance and updates: The collection is intended to remain extensible through contributor guidelines and collaboration with the Polymathic AI team.New datasets must be validated, formatted, size-checked, and preliminarily benchmarked before integration.
  • Maintenance and updates: Existing datasets will be maintained, while future collection changes are planned as additions rather than replacements.The authors state that this approach prevents existing data from becoming out-of-date.

B.3 Data Specification

The Well uses a shared, self-documenting HDF5 specification and PyTorch interface to represent multidimensional simulation data consistently across datasets. Metadata records dataset identity, spatial structure, trajectories, fields, parameters, and variation across samples and time.

  • Shared specification: The Well provides a unified data specification and PyTorch-based interface for HDF5 archives with a shared format.The specification is designed to support consistent processing across the collection.
  • Dataset metadata: Metadata defines simulation parameters, dataset name, grid type, spatial dimensions, trajectory count, coordinates, and boundary-condition groups.The example specification describes a 2D Cartesian simulation with dimensions B x T x W x H.
  • Fields and variation: Fields are stored as datasets whose shapes encode batch, time, spatial, and optional tensor dimensions.Field metadata indicates whether dimensions, samples, and time vary for each field.
  • Specification scope: The specification is self-documenting and contains metadata sufficient for machine processing, but it does not provide Croissant descriptions for the HDF5 data.The authors avoided converting the 15TB collection to another format because the current Croissant specification does not support HDF5.
  • Interface behavior: The GenericWellDataset returns all fields as a dictionary, allowing users to choose input-output arrangements for their objectives.Default processors include time-invariant fields as model inputs but exclude them from targets.

C Dataset Details

The Well spans diverse physical systems and simulation methods, with datasets designed to expose challenges from discontinuities, multiscale dynamics, expensive computation, and parameterized behavior.

  • All numerical simulations use uniform grids, uniform time steps, and single-precision fp32.
  • Acoustic scattering: Acoustic-scattering datasets vary material-density configurations to test learned models on sharp discontinuities and irregular structures.The configurations include single discontinuities, Gaussian bumps, gradients, constant fields, smoothed noise, and inclusions.
  • Acoustic scattering: The acoustic simulations use TVD Clawpack solvers with CFL-controlled time steps and reflective or open boundary conditions.Each simulation took approximately 15 minutes on 64 Icelake CPU cores.
  • MHD turbulence: MHD simulations vary sonic and Alfvénic Mach numbers across sub-Alfvénic and super-Alfvénic regimes while remaining non-self-gravitating and scale-free.On 64 CPU cores, each simulation takes 48 hours; rescaling requires preserving both Mach numbers.
  • Post-neutron-star merger: Post-neutron-star-merger simulations must resolve magnetorotational instability and long-term outflows while tracking electron fraction Ye relevant to kilonova signals.The authors identify these simulations as computationally expensive and challenging.
  • Turbulent radiative layers: Radiative-layer simulations show that cooling and mass-transfer rates scale with relative phase velocity and cooling time according to an analytic theory.The reported relation is ṠEcool∝ṠM∝v^3/4 and cooling time^-1/4.

D Additional Tasks of Interest

The Well supports challenges beyond forecasting, including resolution, dimensionality, time-step, parameter, long-horizon, inverse, and simulation-acceleration tasks.

  • Super-resolution: The Well supports super-resolution by providing MHD and supernova_explosion data at multiple resolutions.These data enable inference of unresolved scales and generalization from lower-resolution training to higher-resolution targets.
  • Generalization: Models can be tested for generalization across time steps and unseen physical-parameter ranges with different behavior.Rayleigh–Taylor data vary time steps, while several datasets provide filterable parameter ranges.
  • Long-horizon prediction: Steady-state prediction and stable long-term forecasting are proposed for datasets that reach equilibrium or provide multi-year trajectories.The planetswe trajectories extend for three model years, and convective_envelope_rsg and gray_scott_reaction_diffusion eventually reach steady states.
  • Inverse problems: Inverse scattering reverses the forward task by predicting material densities from pressure-field evolution.The relevant datasets are acoustic_scattering and helmholtz_staircase.
  • Simulation acceleration: Simulation acceleration targets post_neutron_star_merger and turbulence_gravity_cooling, whose simulations take months to generate.Accurate predictions could substantially reduce the cost relative to the generating process.

E.1 Standard Methodology

The preliminary benchmarks use standardized, resource-constrained training procedures and established surrogate architectures to provide a representative starting point for evaluating the Well.

  • Benchmark procedure: Benchmarks use generic standard-practice choices, including approximately 15–20 million parameters, memory-maximizing batch sizes, AdamW, and validation-selected learning rates.Learning rates are searched over {1×10−4, 5×10−4, 1×10−3, 5×10−3, 1×10−2}.
  • Benchmark procedure: All models train with mean squared error averaged over fields and space, while boundary conditions follow the model architecture through periodic convolutions or zero padding.Fourier-domain convolutions implicitly use periodic boundaries; spatial convolutions use standard zero padding.
  • Benchmark procedure: Each run is limited to 12 hours on one NVIDIA H100, favoring faster models and optimized libraries over architectures without optimized GPU kernels.Single precision is used because several datasets encounter stability issues with mixed or low precision training.
  • Baseline models: The baseline suite uses established models that natively extend to 3D: FNO, TFNO, U-net, and a ConvNext-based modernization of U-net.TFNO is included for improved scalability, while ConvNext blocks replace older U-net convolutional blocks for fairer evaluation.
  • Model scaling: For 2D models, architectures are adjusted primarily through depth or width to reach approximately 15–20 million parameters.The reported configurations specify spectral or spatial filter sizes, hidden or initial dimensions, and block counts.

E.3 Metrics

The Well evaluates surrogate models with spatial and spectral error metrics, emphasizing variance-scaled errors whose interpretation compares predictions against the spatial mean.

  • Spatial metrics: The evaluation includes spatial metrics such as MSE, RMSE, NMSE, NRMSE, VMSE, VRMSE, and bandpass-filtered spectral errors.These metrics compare spatial fields and include normalized and frequency-resolved variants.
  • Normalized metrics: VRMSE is the square-root variant of variance-scaled mean squared error, which normalizes MSE by the variance of the truth.The variance-scaled formulation supports comparison across fields with different variability.
  • Normalized metrics: A VRMSE above 1 indicates worse performance than estimating the target field with its accurate spatial mean.This comparison is against the spatial mean of the target, not the population mean.
  • Spectral metrics: Spectral metrics partition frequencies into disjoint low-, intermediate-, and high-frequency bands based on wavenumber magnitude on a logarithmic scale.The spectral errors are computed after Fourier filtering or bandpass filtering over the selected frequency bands.
  • Spectral metrics: Binned spectral normalized mean square error is normalized to the target’s bin energy, so values at least 1 indicate zero coefficients would outperform the model at that scale.This interpretation is used to make rollout quality more visually interpretable.

E.4 Results

The benchmarks show that generic training can perform poorly on the Well, with accuracy varying across fields and spatial frequencies, while several large or complex datasets remain compute-limited.

  • Overall results: One-step VRMSE and time-averaged rollout losses are reported for models selected by validation performance, with scores above 1 indicating the spatial-mean baseline would be better.The metric is chosen to compare predictions with the easier task of estimating the target’s non-spatially varying mean.
  • Field-level behavior: Performance is not uniform across fields: individual fields can be accurate even when overall dataset performance is poor.The authors suggest unnormalized training losses may partly explain this pattern and motivate normalized losses for general surrogate modeling.
  • Frequency-dependent behavior: Low-frequency modes are tracked better over time, whereas high-frequency modes diverge relatively quickly.The Well’s metrics are intended to expose such frequency-dependent behavior during training and architecture development.
  • Limitations: Four datasets allow fewer than 5 training epochs within the 12-hour limit because of computational limitations or inherent dynamical complexity.These are convective_envelope_rsg, euler_multi_quadrants, turbulence_gravity_cooling, and turbulent_radiative_layer_3D; unrestricted training could improve results.
  • Benchmark reporting: The validation comparison reports dataset–model VRMSE values, with the best result bolded and the mean-value prediction calibrated to a score of 1.Table 6 gives the learning rate and training epochs associated with the validation losses in Table 5.
Loading 2412.00568v2…