GitHub – PolymathicAI/The_well: A 15TB Collection of Physics Simulation Datasets

PolymathicAI released The Well, a 15TB collection of 16 physics simulation datasets for machine learning, covering domains like fluid dynamics and supernova explosions. The open-source repository provides tools for downloading, streaming, and benchmarking surrogate models.

Welcome to the Well, a large-scale collection of machine learning datasets containing numerical simulations of a wide variety of spatiotemporal physical systems. The Well draws from domain scientists 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 for accelerating research in machine learning and computational sciences. Once the Well package installed and the data downloaded you can use them in your training pipeline. python from the well.data import WellDataset from torch.utils.data import DataLoader trainset = WellDataset well base path="path/to/base", well dataset name="name of the dataset", well split name="train" train loader = DataLoader trainset for batch in train loader: ... For more information regarding the interface, please refer to the API https://github.com/PolymathicAI/the well/tree/master/docs/api.md and the tutorials https://github.com/PolymathicAI/the well/blob/master/docs/tutorials/dataset.ipynb . If you plan to use The Well datasets to train or evaluate deep learning models, we recommend to use a machine with enough computing resources. We also recommend creating a new Python =3.10 environment to install the Well. For instance, with venv https://docs.python.org/3/library/venv.html : python -m venv path/to/env source path/to/env/activate/bin The Well package can be installed directly from PyPI. pip install the well It can also be installed from source. For this, clone the repository https://github.com/PolymathicAI/the well and install the package with its dependencies. git clone https://github.com/PolymathicAI/the well cd the well pip install . Depending on your acceleration hardware, you can specify --extra-index-url to install the relevant PyTorch version. For example, use pip install . --extra-index-url https://download.pytorch.org/whl/cu121 to install the dependencies built for CUDA 12.1. If you want to run the benchmarks, you should install additional dependencies. pip install the well benchmark The Well datasets range between 6.9GB and 5.1TB of data each, for a total of 15TB for the full collection. Ensure that your system has enough free disk space to accomodate the datasets you wish to download. Once the well is installed, you can use the the-well-download command to download any dataset of The Well. the-well-download --base-path path/to/base --dataset active matter --split train If --dataset and --split are omitted, all datasets and splits will be downloaded. This could take a while Most of the Well datasets are also hosted on Hugging Face https://huggingface.co/collections/polymathic-ai/the-well-67e129f4ca23e0447395d74c . Data can be streamed directly from the hub using the following code. python from the well.data import WellDataset from torch.utils.data import DataLoader The following line may take a couple of minutes to instantiate the datamodule trainset = WellDataset well base path="hf://datasets/polymathic-ai/", access from HF hub well dataset name="active matter", well split name="train", train loader = DataLoader trainset for batch in train loader: ... For better performance in large training, we advise downloading the data locally downloading-the-data instead of streaming it over the network. The repository allows benchmarking surrogate models on the different datasets that compose the Well. Some state-of-the-art models are already implemented in models https://github.com/PolymathicAI/the well/tree/master/the well/benchmark/models , while dataset classes https://github.com/PolymathicAI/the well/tree/master/the well/data handle the raw data of the Well. The benchmark relies on a training script https://github.com/PolymathicAI/the well/blob/master/the well/benchmark/train.py that uses hydra https://hydra.cc/ to instantiate various classes e.g. dataset, model, optimizer from configuration files https://github.com/PolymathicAI/the well/tree/master/the well/benchmark/configs . For instance, to run the training script of default FNO architecture on the active matter dataset, launch the following commands: cd the well/benchmark python train.py experiment=fno server=local data=active matter Each argument corresponds to a specific configuration file. In the command above server=local indicates the training script to use local.yaml https://github.com/PolymathicAI/the well/tree/master/the well/benchmark/configs/server/local.yaml , which just declares the relative path to the data. The configuration can be overridden directly or edited with new YAML files. Please refer to hydra documentation https://hydra.cc/ for editing configuration. You can use this command within a sbatch script to launch the training with Slurm. The model benchmarked in the original paper of the Well have been designed as a a simple baseline. They should not be considered as state-of-the-art. We hope that the community will build upon these results to develop better architectures for PDE surrogate modeling. Most of the checkpoints of the models are available on Hugging Face https://huggingface.co/collections/polymathic-ai/the-well-benchmark-models-67e69bd7cd8e60229b5cd43e . To load a specific checkpoint follow the example below of the FNO model trained on the active matter dataset. python from the well.benchmark.models import FNO model = FNO.from pretrained "polymathic-ai/FNO-active matter" - The dataset viscoelastic instability has been deprecated due to processing errors in the data. It remains available for backwards comparisons. However viscoelastic instability v2 contains the same data without the processing error. This project has been led by the Polymathic AI https://polymathic-ai.org/ organization, in collaboration with researchers from the Flatiron Institute, University of Colorado Boulder, University of Cambridge, New York University, Rutgers University, Cornell University, University of Tokyo, Los Alamos Natioinal Laboratory, University of California, Berkeley, Princeton University, CEA DAM, and University of Liège. If you find this project useful for your research, please consider citing @article{ohana2024well, title={The well: a large-scale collection of diverse physics simulations for machine learning}, author={Ohana, Ruben and McCabe, Michael and Meyer, Lucas and Morel, Rudy and Agocs, Fruzsina and Beneitez, Miguel and Berger, Marsha and Burkhart, Blakesly and Dalziel, Stuart and Fielding, Drummond and others}, journal={Advances in Neural Information Processing Systems}, volume={37}, pages={44989--45037}, year={2024} } For questions regarding this project, please contact Ruben Ohana https://rubenohana.github.io/ and Michael McCabe https://mikemccabe210.github.io/ at {rohana,mmccabe}@flatironinstitute.org. To report a bug in the data or the code , request a feature or simply ask a question, you can open an issue https://github.com/PolymathicAI/the well/issues on the repository https://github.com/PolymathicAI/the well .