Package: lagrangebench
Table of ContentsInstallationUsageDatasetsDirectory StructureContributingCitationInstallationStandalone libraryInstall the corelagrangebenchlibrary from PyPi aspipinstalllagrangebench--extra-index-url=https://download.pytorch.org/whl/cpuNote that by defaultlagrangebenchis installed without JAX GPU support. For that follow the instructions in theGPU supportsection.CloneClone this GitHub repositorygitclonehttps://github.com/tumaer/lagrangebench.gitcdlagrangebenchInstall the dependencies withPoetry (>=1.6.0)poetry install --only mainAlternatively, a requirements file is provided. It directly installs the CUDA version of JAX.pip install -r requirements_cuda.txtFor a CPU version of the requirements file, one could usedocs/requirements.txt.GPU supportTo run JAX on GPU, follow theJax CUDA guide, or in general runpipinstall--upgradejax[cuda11_pip]==0.4.20-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.html# or, for cuda 12pipinstall--upgradejax[cuda12_pip]==0.4.20-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlMacOSCurrently, only the CPU installation works. You will need to change a few small things to get it going:Clone installation: inpyproject.tomlchange the torch version from2.1.0+cputo2.1.0. Then, remove thepoetry.lockfile and runpoetry install --only main.Configs: You will need to setf64: Falseandnum_workers: 0in theconfigs/files.Although the currentjax-metal==0.0.5librarysupports jax in general, there seems to be a missing feature used byjax-mdrelated to padding -> seethis issue.UsageStandalone benchmark libraryA general tutorial is provided in the example notebook "Training GNS on the 2D Taylor Green Vortex" under./notebooks/tutorial.ipynbon theLagrangeBench repository. The notebook covers the basics of LagrangeBench, such as loading a dataset, setting up a case, training a model from scratch and evaluating its performance.Running in a local clone (main.py)Alternatively, experiments can also be set up withmain.py, based on extensive YAML config files and cli arguments (checkconfigs/). By default, the arguments have priority as: 1) passed cli arguments, 2) YAML config and 3)defaults.py(lagrangebenchdefaults).When loading a saved model with--model_dirthe config from the checkpoint is automatically loaded and training is restarted. For more details check theexperiments/directory and therun.pyfile.TrainFor example, to start aGNSrun from scratch on the RPF 2D dataset usepython main.py --config configs/rpf_2d/gns.yamlSome model presets can be found in./configs/.If--mode=all, then training (--mode=train) and subsequent inference (--mode=infer) on the test split will be run in one go.Restart trainingTo restart training from the last checkpoint in--model_dirusepython main.py --model_dir ckp/gns_rpf2d_yyyymmdd-hhmmssInferenceTo evaluate a trained model from--model_diron the test split (--test) usepython main.py --model_dir ckp/gns_rpf2d_yyyymmdd-hhmmss/best --rollout_dir rollout/gns_rpf2d_yyyymmdd-hhmmss/best --mode infer --testIf the default--out_type_infer=pklis active, then the generated trajectories and ametricsYYYY_MM_DD_HH_MM_SS.pklfile will be written to the--rollout_dir. The metrics file contains all--metrics_inferproperties for each generated rollout.DatasetsThe datasets are hosted on Zenodo under the DOI:10.5281/zenodo.10021925. When creating a new dataset instance, the data is automatically downloaded. Alternatively, to manually download them use thedownload_data.shshell script, either with a specific dataset name or "all". NamelyTaylor Green Vortex 2D:bash download_data.sh tgv_2d datasets/Reverse Poiseuille Flow 2D:bash download_data.sh rpf_2d datasets/Lid Driven Cavity 2D:bash download_data.sh ldc_2d datasets/Dam break 2D:bash download_data.sh dam_2d datasets/Taylor Green Vortex 3D:bash download_data.sh tgv_3d datasets/Reverse Poiseuille Flow 3D:bash download_data.sh rpf_3d datasets/Lid Driven Cavity 3D:bash download_data.sh ldc_3d datasets/All:bash download_data.sh all datasets/NotebooksWe provide three notebooks that show LagrangeBench functionalities, namely:tutorial.ipynb, with a general overview of LagrangeBench library, with training and evaluation of a simple GNS model,datasets.ipynb, with more details and visualizations on the datasets, andgns_data.ipynb, showing how to train models within LagrangeBench on the datasets from the paperLearning to Simulate Complex Physics with Graph Networks.Directory structure📦lagrangebench
 ┣ 📂case_setup     # Case setup manager
 ┃ ┣ 📜case.py      # CaseSetupFn class
 ┃ ┣ 📜features.py  # Feature extraction
 ┃ ┗ 📜partition.py # Alternative neighbor list implementations
 ┣ 📂data           # Datasets and dataloading utils
 ┃ ┣ 📜data.py      # H5Dataset class and specific datasets
 ┃ ┗ 📜utils.py
 ┣ 📂evaluate       # Evaluation and rollout generation tools
 ┃ ┣ 📜metrics.py
 ┃ ┗ 📜rollout.py
 ┣ 📂models         # Baseline models
 ┃ ┣ 📜base.py      # BaseModel class
 ┃ ┣ 📜egnn.py
 ┃ ┣ 📜gns.py
 ┃ ┣ 📜linear.py
 ┃ ┣ 📜painn.py
 ┃ ┣ 📜segnn.py
 ┃ ┗ 📜utils.py
 ┣ 📂train          # Trainer method and training tricks
 ┃ ┣ 📜strats.py    # Training tricks
 ┃ ┗ 📜trainer.py   # Trainer method
 ┣ 📜defaults.py    # Default values
 ┗ 📜utils.pyContributingWelcome! We highly appreciateGithub issuesandPRs.You can also chat with us onDiscord.Contributing GuidelineIf you want to contribute to this repository, you will need the dev depencencies, i.e.
install the environment withpoetry installwithout the--only mainflag.
Then, we also recommend you to install the pre-commit hooks
if you don't want to manually runpre-commit runbefore each commit. To sum up:gitclonehttps://github.com/tumaer/lagrangebench.gitcdlagrangebench
poetryinstallsource$PATH_TO_LAGRANGEBENCH_VENV/bin/activate# install pre-commit hooks defined in .pre-commit-config.yaml# ruff is configured in pyproject.tomlpre-commitinstallAfter you have rungit add <FILE>and try togit commit, the pre-commit hook will
fix the linting and formatting of<FILE>before you are allowed to commit.You should also run the unit tests locally before creating a PR. Do this simply by:# pytest is configured in pyproject.tomlpytestCitationThe paper (at NeurIPS 2023 Datasets and Benchmarks) can be cited as:@inproceedings{toshev2023lagrangebench,title={LagrangeBench: A Lagrangian Fluid Mechanics Benchmarking Suite},author={Artur P. Toshev and Gianluca Galletti and Fabian Fritz and Stefan Adami and Nikolaus A. Adams},year={2023},url={https://arxiv.org/abs/2309.16342},booktitle={37th Conference on Neural Information Processing Systems (NeurIPS 2023) Track on Datasets and Benchmarks},}The associated datasets can be cited as:@dataset{toshev_2023_10021926,author={Toshev, Artur P. and Adams, Nikolaus A.},title={LagrangeBench Datasets},month=oct,year=2023,publisher={Zenodo},version={0.0.1},url={https://zenodo.org/doi/10.5281/zenodo.10021925},doi={10.5281/zenodo.10021925},}PublicationsThe following further publcations are based on the LagrangeBench codebase:Learning Lagrangian Fluid Mechanics with E(3)-Equivariant Graph Neural Networks (GSI 2023), A. P. Toshev, G. Galletti, J. Brandstetter, S. Adami, N. A. Adams