<div align="center">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://svgshare.com/i/11hG.svg">
<source media="(prefers-color-scheme: light)" srcset="https://svgshare.com/i/11hG.svg">
<img alt="LagrangeBench Logo: Lagrangian Fluid Mechanics Benchmarking Suite" src="https://svgshare.com/i/11hG.svg" width=550pt>
</picture>
<!-- [](https://lagrangebench.readthedocs.io/en/latest/index.html)
[](https://arxiv.org/abs/2309.16342)
[](https://github.com/tumaer/lagrangebench/actions/workflows/ruff.yml)
[](https://github.com/astral-sh/ruff)
[](https://pypi.org/project/lagrangebench/)
[](https://pypi.org/project/lagrangebench/) -->
[](https://arxiv.org/abs/2309.16342)
[](https://lagrangebench.readthedocs.io/en/latest/index.html)
[](https://pypi.org/project/lagrangebench/)
[](https://colab.research.google.com/github/tumaer/lagrangebench/blob/main/notebooks/tutorial.ipynb)
[](https://discord.gg/Ds8jRZ78hU)
[](https://github.com/tumaer/lagrangebench/actions/workflows/tests.yml)
[](https://codecov.io/gh/tumaer/lagrangebench)
[](https://github.com/tumaer/lagrangebench/blob/main/LICENSE)
</div>
NeurIPS page with video and slides [here](https://neurips.cc/virtual/2023/poster/73681).
## Table of Contents
1. [**Installation**](#installation)
1. [**Usage**](#usage)
1. [**Datasets**](#datasets)
1. [**Pretrained Models**](#pretrained-models)
1. [**Directory Structure**](#directory-structure)
1. [**Contributing**](#contributing)
1. [**Citation**](#citation)
## Installation
### Standalone library
Install the core `lagrangebench` library from PyPi as
```bash
python3.10 -m venv venv
source venv/bin/activate
pip install lagrangebench --extra-index-url=https://download.pytorch.org/whl/cpu
```
Note that by default `lagrangebench` is installed without JAX GPU support. For that follow the instructions in the [GPU support](#gpu-support) section.
### Clone
Clone this GitHub repository
```bash
git clone https://github.com/tumaer/lagrangebench.git
cd lagrangebench
```
Install the dependencies with __Poetry (>=1.6.0)__
```
poetry install --only main
```
Alternatively, a requirements file is provided. It directly installs the CUDA version of JAX.
```
pip install -r requirements_cuda.txt
```
For a CPU version of the requirements file, one could use `docs/requirements.txt`.
### GPU support
To run JAX on GPU, follow [Installing JAX](https://jax.readthedocs.io/en/latest/installation.html), or in general run
```bash
pip install -U "jax[cuda12]==0.4.29"
```
> Note: as of 27.06.2024, to make our GNN models **deterministic** on GPUs, you need to set `os.environ["XLA_FLAGS"] = "--xla_gpu_deterministic_ops=true"`. However, all current models rely of `scatter_sum`, and this operation seems to be slower than running a normal for-loop in Python, when executed in deterministic mode, see [#17844](https://github.com/google/jax/issues/17844) and [#10674](https://github.com/google/jax/discussions/10674).
### MacOS
Currently, only the CPU installation works. You will need to change a few small things to get it going:
- Clone installation: in `pyproject.toml` change the torch version from `2.1.0+cpu` to `2.1.0`. Then, remove the `poetry.lock` file and run `poetry install --only main`.
- Configs: You will need to set `dtype=float32` and `train.num_workers=0`.
Although the current [`jax-metal==0.0.5` library](https://pypi.org/project/jax-metal/) supports jax in general, there seems to be a missing feature used by `jax-md` related to padding -> see [this issue](https://github.com/google/jax/issues/16366#issuecomment-1591085071).
## Usage
### Standalone benchmark library
A general tutorial is provided in the example notebook "Training GNS on the 2D Taylor Green Vortex" under `./notebooks/tutorial.ipynb` on the [LagrangeBench repository](https://github.com/tumaer/lagrangebench). 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 with `main.py`, based on extensive YAML config files and cli arguments (check [`configs/`](configs/)). By default, the arguments have priority as 1) passed cli arguments, 2) YAML config and 3) [`defaults.py`](lagrangebench/defaults.py) (`lagrangebench` defaults).
When loading a saved model with `load_ckp` the config from the checkpoint is automatically loaded and training is restarted. For more details check the [`runner.py`](lagrangebench/runner.py) file.
**Train**
For example, to start a _GNS_ run from scratch on the RPF 2D dataset use
```
python main.py config=configs/rpf_2d/gns.yaml
```
Some model presets can be found in `./configs/`.
If `mode=all` is provided, then training (`mode=train`) and subsequent inference (`mode=infer`) on the test split will be run in one go.
**Restart training**
To restart training from the last checkpoint in `load_ckp` use
```
python main.py load_ckp=ckp/gns_rpf2d_yyyymmdd-hhmmss
```
**Inference**
To evaluate a trained model from `load_ckp` on the test split (`test=True`) use
```
python main.py load_ckp=ckp/gns_rpf2d_yyyymmdd-hhmmss/best rollout_dir=rollout/gns_rpf2d_yyyymmdd-hhmmss/best mode=infer test=True
```
If the default `eval.infer.out_type=pkl` is active, then the generated trajectories and a `metricsYYYY_MM_DD_HH_MM_SS.pkl` file will be written to `eval.rollout_dir`. The metrics file contains all `eval.infer.metrics` properties for each generated rollout.
### Notebooks
We provide three notebooks that show LagrangeBench functionalities, namely:
- [`tutorial.ipynb`](notebooks/tutorial.ipynb) [](https://colab.research.google.com/github/tumaer/lagrangebench/blob/main/notebooks/tutorial.ipynb), with a general overview of LagrangeBench library, with training and evaluation of a simple GNS model,
- [`datasets.ipynb`](notebooks/datasets.ipynb) [](https://colab.research.google.com/github/tumaer/lagrangebench/blob/main/notebooks/datasets.ipynb), with more details and visualizations of the datasets, and
- [`gns_data.ipynb`](notebooks/gns_data.ipynb) [](https://colab.research.google.com/github/tumaer/lagrangebench/blob/main/notebooks/gns_data.ipynb), showing how to train models within LagrangeBench on the datasets from the paper [Learning to Simulate Complex Physics with Graph Networks](https://arxiv.org/abs/2002.09405).
## Datasets
The datasets are hosted on Zenodo under the DOI: [10.5281/zenodo.10021925](https://zenodo.org/doi/10.5281/zenodo.10021925). If a dataset is not found in `dataset.src`, the data is automatically downloaded. Alternatively, to manually download the datasets use the `download_data.sh` shell script, either with a specific dataset name or "all". Namely
- __Taylor 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/`
## Pretrained Models
We provide pretrained model weights of our default GNS and SEGNN models on each of the 7 LagrangeBench datasets. You can download and run the checkpoints given below. In the table, we also provide the 20-step error measures on the full test split.
| Dataset | Model | MSE<sub>20</sub> | Sinkhorn | MSE<sub>E<sub>kin</sub></sub> |
| ------- |-------------------------------------------------------------------------------------- | ------ | ------ | ------ |
| 2D TGV | [GNS-10-128](https://drive.google.com/file/d/19TO4PaFGcryXOFFKs93IniuPZKEcaJ37/view) | 5.9e-6 | 3.2e-7 | 4.9e-7 |
| | [SEGNN-10-64](https://drive.google.com/file/d/1llGtakiDmLfarxk6MUAtqj6sLleMQ7RL/view) | 4.4e-6 | 2.1e-7 | 5.0e-7 |
| 2D RPF | [GNS-10-128](https://drive.google.com/file/d/1uYusVlP1ykUNuw58vo7Wss-xyTMmopAn/view) | 4.0e-6 | 2.5e-7 | 2.7e-5 |
| | [SEGNN-10-64](https://drive.google.com/file/d/108dZVWs2qxAvKiboeEBW-nIcv-aslhYP/view) | 3.4e-6 | 2.5e-7 | 1.4e-5 |
| 2D LDC | [GNS-10-128](https://drive.google.com/file/d/1JvdsW0H6XrgC2_cwV3pP66cAm9j1-AXc/view) | 1.5e-5 | 1.1e-6 | 6.1e-7 |
| | [SEGNN-10-64](https://drive.google.com/file/d/1D_wgs2pD9pTXoJK76yi-R0K2tY_T6lPn/view) | 2.1e-5 | 3.7e-6 | 1.6e-5 |
| 2D DAM | [GNS-10-128](https://drive.google.com/file/d/16bJz3VfSMxOG1II8kCg5DlzGhjvdip2p/view) | 3.1e-5 | 1.4e-5 | 1.1e-4 |
| | [SEGNN-10-64](https://drive.google.com/file/d/1_6rHxK81vzrdIMPtJ7rIkeoUgsTeKmSn/view) | 4.1e-5 | 2.3e-5 | 5.2e-4 |
| 3D TGV | [GNS-10-128](https://drive.google.com/file/d/1DEkXxrebS9eyLSMlc_ztHrqlh29NgLXC/view) | 5.8e-3 | 4.7e-6 | 4.8e-2 |
| | [SEGNN-10-64](https://drive.google.com/file/d/1ivJnHTgfbQ0IJujc5O0CUoQNiGU4zi_d/view) | 5.0e-3 | 4.9e-6 | 3.9e-2 |
| 3D RPF | [GNS-10-128](https://drive.google.com/file/d/1yo-qgShLd1sgS1u5zkMXdJvhuPBwEQQE/view) | 2.1e-5 | 3.3e-7 | 1.8e-6 |
| | [SEGNN-10-64](https://drive.google.com/file/d/1Qczh3Z_z0grTuRuPDHyiYLzV1zg7Liz9/view) | 1.7e-5 | 2.7e-7 | 1.7e-6 |
| 3D LDC | [GNS-10-128](https://drive.google.com/file/d/1b3IIkxk5VcWiT8Oyqg1wex8-ZfJv2g_v/view) | 4.1e-5 | 3.2e-7 | 1.9e-8 |
| | [SEGNN-10-64](https://drive.google.com/file/d/1ZIg7FXc1l3C4ekc9WvVvjHEl5KKxOA_U/view) | 4.1e-5 | 2.9e-7 | 2.5e-8 |
To reproduce the numbers in the table, e.g., on 2D TGV with GNS, follow these steps:
```bash
# download the checkpoint (1) through the browser or
# (2) using the file ID from the URL, i.e., for 2D TGV + GNS
gdown 19TO4PaFGcryXOFFKs93IniuPZKEcaJ37
# unzip the downloaded file `gns_tgv2d.zip`
python -c "import shutil; shutil.unpack_archive('gns_tgv2d.zip', 'gns_tgv2d')"
# evaluate the model on the test split
python main.py gpu=$GPU_ID mode=infer eval.test=True load_ckp=gns_tgv2d/best
```
## Directory structure
```
📦lagrangebench
┣ 📂case_setup # Case setup manager
┃ ┣ 📜case.py # CaseSetupFn class
┃ ┗ 📜features.py # Feature extraction
┣ 📂data # Datasets and dataloading utils
┃ ┣ 📜data.py # H5Dataset class and specific datasets
┃ ┗ 📜utils.py
┣ 📂evaluate # Evaluation and rollout generation tools
┃ ┣ 📜metrics.py
┃ ┣ 📜rollout.py
┃ ┗ 📜utils.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
┣ 📜runner.py # Runner wrapping training and inference
┗ 📜utils.py
```
## Contributing
Welcome! We highly appreciate [Github issues](https://github.com/tumaer/lagrangebench/issues) and [PRs](https://github.com/tumaer/lagrangebench/pulls).
You can also chat with us on [**Discord**](https://discord.gg/Ds8jRZ78hU).
### Contributing Guideline
If you want to contribute to this repository, you will need the dev dependencies, i.e.
install the environment with `poetry install` without the ` --only main` flag.
Then, we also recommend you install the pre-commit hooks
if you don't want to manually run `pre-commit run` before each commit. To sum up:
```bash
git clone https://github.com/tumaer/lagrangebench.git
cd lagrangebench
poetry install
source $PATH_TO_LAGRANGEBENCH_VENV/bin/activate
# install pre-commit hooks defined in .pre-commit-config.yaml
# ruff is configured in pyproject.toml
pre-commit install
# if you want to bump the version in both pyproject.toml and __init__.py, do
poetry self add poetry-bumpversion
poetry version patch # or minor/major
```
After you have run `git add <FILE>` and try to `git commit`, the pre-commit hook will
fix the linting and formatting of `<FILE>` before you are allowed to commit.
You should also run the tests locally before creating a PR. Do this simply by:
```bash
# pytest is configured in pyproject.toml
pytest
```
### Clone vs Library
LagrangeBench can be installed by cloning the repository or as a standalone library. This offers more flexibility, but it also comes with its disadvantages: the necessity to implement some things twice. If you change any of the following things, make sure to update its counterpart as well:
- General setup in `lagrangebench/runner.py` and `notebooks/tutorial.ipynb`
- Configs in `configs/` and `lagrangebench/defaults.py`
- Zenodo URLs in `download_data.sh` and `lagrangebench/data/data.py`
- Dependencies in `pyproject.toml`, `requirements_cuda.txt`, and `docs/requirements.txt`
- Library version in `pyproject.toml` and `lagrangebench/__init__.py`
## Citation
The paper (at NeurIPS 2023 Datasets and Benchmarks) can be cited as:
```bibtex
@article{toshev2024lagrangebench,
title={Lagrangebench: A lagrangian fluid mechanics benchmarking suite},
author={Toshev, Artur and Galletti, Gianluca and Fritz, Fabian and Adami, Stefan and Adams, Nikolaus},
journal={Advances in Neural Information Processing Systems},
volume={36},
year={2024}
}
```
The associated datasets can be cited as:
```bibtex
@dataset{toshev_2024_10491868,
author = {Toshev, Artur P. and Adams, Nikolaus A.},
title = {LagrangeBench Datasets},
month = jan,
year = 2024,
publisher = {Zenodo},
doi = {10.5281/zenodo.10491868},
url = {https://doi.org/10.5281/zenodo.10491868}
}
```
### Publications
The following further publications are based on the LagrangeBench codebase:
1. [Learning Lagrangian Fluid Mechanics with E(3)-Equivariant Graph Neural Networks (GSI 2023)](https://arxiv.org/abs/2305.15603), A. P. Toshev, G. Galletti, J. Brandstetter, S. Adami, N. A. Adams
2. [Neural SPH: Improved Neural Modeling of Lagrangian Fluid Dynamics (ICML 2024)](https://arxiv.org/abs/2402.06275), A. P. Toshev, J. A. Erbesdobler, N. A. Adams, J. Brandstetter
Raw data
{
"_id": null,
"home_page": "https://lagrangebench.readthedocs.io/",
"name": "lagrangebench",
"maintainer": null,
"docs_url": null,
"requires_python": "<=3.11,>=3.9",
"maintainer_email": null,
"keywords": "smoothed-particle-hydrodynamics, benchmark-suite, lagrangian-dynamics, graph-neural-networks, lagrangian-particles",
"author": "Artur Toshev, Gianluca Galletti",
"author_email": "artur.toshev@tum.de",
"download_url": "https://files.pythonhosted.org/packages/49/72/6ec85cab8bfb69e123a11012873edbdbdf0513e09049892a3fa79a8f385c/lagrangebench-0.2.0.tar.gz",
"platform": null,
"description": "<div align=\"center\">\n\n<picture>\n <source media=\"(prefers-color-scheme: dark)\" srcset=\"https://svgshare.com/i/11hG.svg\">\n <source media=\"(prefers-color-scheme: light)\" srcset=\"https://svgshare.com/i/11hG.svg\">\n<img alt=\"LagrangeBench Logo: Lagrangian Fluid Mechanics Benchmarking Suite\" src=\"https://svgshare.com/i/11hG.svg\" width=550pt>\n</picture>\n<!-- [](https://lagrangebench.readthedocs.io/en/latest/index.html)\n[](https://arxiv.org/abs/2309.16342)\n[](https://github.com/tumaer/lagrangebench/actions/workflows/ruff.yml)\n[](https://github.com/astral-sh/ruff)\n[](https://pypi.org/project/lagrangebench/)\n[](https://pypi.org/project/lagrangebench/) -->\n\n[](https://arxiv.org/abs/2309.16342)\n[](https://lagrangebench.readthedocs.io/en/latest/index.html)\n[](https://pypi.org/project/lagrangebench/)\n[](https://colab.research.google.com/github/tumaer/lagrangebench/blob/main/notebooks/tutorial.ipynb)\n[](https://discord.gg/Ds8jRZ78hU)\n\n[](https://github.com/tumaer/lagrangebench/actions/workflows/tests.yml)\n[](https://codecov.io/gh/tumaer/lagrangebench)\n[](https://github.com/tumaer/lagrangebench/blob/main/LICENSE)\n\n</div>\n\nNeurIPS page with video and slides [here](https://neurips.cc/virtual/2023/poster/73681).\n\n## Table of Contents\n\n1. [**Installation**](#installation)\n1. [**Usage**](#usage)\n1. [**Datasets**](#datasets)\n1. [**Pretrained Models**](#pretrained-models)\n1. [**Directory Structure**](#directory-structure)\n1. [**Contributing**](#contributing)\n1. [**Citation**](#citation)\n\n## Installation\n### Standalone library\nInstall the core `lagrangebench` library from PyPi as\n```bash\npython3.10 -m venv venv\nsource venv/bin/activate\npip install lagrangebench --extra-index-url=https://download.pytorch.org/whl/cpu\n```\n\nNote that by default `lagrangebench` is installed without JAX GPU support. For that follow the instructions in the [GPU support](#gpu-support) section.\n\n### Clone\nClone this GitHub repository\n```bash\ngit clone https://github.com/tumaer/lagrangebench.git\ncd lagrangebench\n```\n\nInstall the dependencies with __Poetry (>=1.6.0)__\n```\npoetry install --only main\n```\n\nAlternatively, a requirements file is provided. It directly installs the CUDA version of JAX.\n```\npip install -r requirements_cuda.txt\n```\nFor a CPU version of the requirements file, one could use `docs/requirements.txt`.\n\n### GPU support\nTo run JAX on GPU, follow [Installing JAX](https://jax.readthedocs.io/en/latest/installation.html), or in general run\n```bash\npip install -U \"jax[cuda12]==0.4.29\"\n```\n\n> Note: as of 27.06.2024, to make our GNN models **deterministic** on GPUs, you need to set `os.environ[\"XLA_FLAGS\"] = \"--xla_gpu_deterministic_ops=true\"`. However, all current models rely of `scatter_sum`, and this operation seems to be slower than running a normal for-loop in Python, when executed in deterministic mode, see [#17844](https://github.com/google/jax/issues/17844) and [#10674](https://github.com/google/jax/discussions/10674).\n\n\n### MacOS\nCurrently, only the CPU installation works. You will need to change a few small things to get it going:\n- Clone installation: in `pyproject.toml` change the torch version from `2.1.0+cpu` to `2.1.0`. Then, remove the `poetry.lock` file and run `poetry install --only main`.\n- Configs: You will need to set `dtype=float32` and `train.num_workers=0`.\n\nAlthough the current [`jax-metal==0.0.5` library](https://pypi.org/project/jax-metal/) supports jax in general, there seems to be a missing feature used by `jax-md` related to padding -> see [this issue](https://github.com/google/jax/issues/16366#issuecomment-1591085071).\n\n## Usage\n### Standalone benchmark library\nA general tutorial is provided in the example notebook \"Training GNS on the 2D Taylor Green Vortex\" under `./notebooks/tutorial.ipynb` on the [LagrangeBench repository](https://github.com/tumaer/lagrangebench). 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.\n\n### Running in a local clone (`main.py`)\nAlternatively, experiments can also be set up with `main.py`, based on extensive YAML config files and cli arguments (check [`configs/`](configs/)). By default, the arguments have priority as 1) passed cli arguments, 2) YAML config and 3) [`defaults.py`](lagrangebench/defaults.py) (`lagrangebench` defaults).\n\nWhen loading a saved model with `load_ckp` the config from the checkpoint is automatically loaded and training is restarted. For more details check the [`runner.py`](lagrangebench/runner.py) file.\n\n**Train**\n\nFor example, to start a _GNS_ run from scratch on the RPF 2D dataset use\n```\npython main.py config=configs/rpf_2d/gns.yaml\n```\nSome model presets can be found in `./configs/`.\n\nIf `mode=all` is provided, then training (`mode=train`) and subsequent inference (`mode=infer`) on the test split will be run in one go.\n\n\n**Restart training**\n\nTo restart training from the last checkpoint in `load_ckp` use\n```\npython main.py load_ckp=ckp/gns_rpf2d_yyyymmdd-hhmmss\n```\n\n**Inference**\n\nTo evaluate a trained model from `load_ckp` on the test split (`test=True`) use\n```\npython main.py load_ckp=ckp/gns_rpf2d_yyyymmdd-hhmmss/best rollout_dir=rollout/gns_rpf2d_yyyymmdd-hhmmss/best mode=infer test=True\n```\n\nIf the default `eval.infer.out_type=pkl` is active, then the generated trajectories and a `metricsYYYY_MM_DD_HH_MM_SS.pkl` file will be written to `eval.rollout_dir`. The metrics file contains all `eval.infer.metrics` properties for each generated rollout.\n\n### Notebooks\nWe provide three notebooks that show LagrangeBench functionalities, namely:\n- [`tutorial.ipynb`](notebooks/tutorial.ipynb) [](https://colab.research.google.com/github/tumaer/lagrangebench/blob/main/notebooks/tutorial.ipynb), with a general overview of LagrangeBench library, with training and evaluation of a simple GNS model,\n- [`datasets.ipynb`](notebooks/datasets.ipynb) [](https://colab.research.google.com/github/tumaer/lagrangebench/blob/main/notebooks/datasets.ipynb), with more details and visualizations of the datasets, and\n- [`gns_data.ipynb`](notebooks/gns_data.ipynb) [](https://colab.research.google.com/github/tumaer/lagrangebench/blob/main/notebooks/gns_data.ipynb), showing how to train models within LagrangeBench on the datasets from the paper [Learning to Simulate Complex Physics with Graph Networks](https://arxiv.org/abs/2002.09405).\n\n## Datasets\nThe datasets are hosted on Zenodo under the DOI: [10.5281/zenodo.10021925](https://zenodo.org/doi/10.5281/zenodo.10021925). If a dataset is not found in `dataset.src`, the data is automatically downloaded. Alternatively, to manually download the datasets use the `download_data.sh` shell script, either with a specific dataset name or \"all\". Namely\n- __Taylor Green Vortex 2D__: `bash download_data.sh tgv_2d datasets/`\n- __Reverse Poiseuille Flow 2D__: `bash download_data.sh rpf_2d datasets/`\n- __Lid Driven Cavity 2D__: `bash download_data.sh ldc_2d datasets/`\n- __Dam break 2D__: `bash download_data.sh dam_2d datasets/`\n- __Taylor Green Vortex 3D__: `bash download_data.sh tgv_3d datasets/`\n- __Reverse Poiseuille Flow 3D__: `bash download_data.sh rpf_3d datasets/`\n- __Lid Driven Cavity 3D__: `bash download_data.sh ldc_3d datasets/`\n- __All__: `bash download_data.sh all datasets/`\n\n## Pretrained Models\nWe provide pretrained model weights of our default GNS and SEGNN models on each of the 7 LagrangeBench datasets. You can download and run the checkpoints given below. In the table, we also provide the 20-step error measures on the full test split.\n\n| Dataset | Model | MSE<sub>20</sub> | Sinkhorn | MSE<sub>E<sub>kin</sub></sub> |\n| ------- |-------------------------------------------------------------------------------------- | ------ | ------ | ------ |\n| 2D TGV | [GNS-10-128](https://drive.google.com/file/d/19TO4PaFGcryXOFFKs93IniuPZKEcaJ37/view) | 5.9e-6 | 3.2e-7 | 4.9e-7 |\n| | [SEGNN-10-64](https://drive.google.com/file/d/1llGtakiDmLfarxk6MUAtqj6sLleMQ7RL/view) | 4.4e-6 | 2.1e-7 | 5.0e-7 |\n| 2D RPF | [GNS-10-128](https://drive.google.com/file/d/1uYusVlP1ykUNuw58vo7Wss-xyTMmopAn/view) | 4.0e-6 | 2.5e-7 | 2.7e-5 |\n| | [SEGNN-10-64](https://drive.google.com/file/d/108dZVWs2qxAvKiboeEBW-nIcv-aslhYP/view) | 3.4e-6 | 2.5e-7 | 1.4e-5 |\n| 2D LDC | [GNS-10-128](https://drive.google.com/file/d/1JvdsW0H6XrgC2_cwV3pP66cAm9j1-AXc/view) | 1.5e-5 | 1.1e-6 | 6.1e-7 |\n| | [SEGNN-10-64](https://drive.google.com/file/d/1D_wgs2pD9pTXoJK76yi-R0K2tY_T6lPn/view) | 2.1e-5 | 3.7e-6 | 1.6e-5 |\n| 2D DAM | [GNS-10-128](https://drive.google.com/file/d/16bJz3VfSMxOG1II8kCg5DlzGhjvdip2p/view) | 3.1e-5 | 1.4e-5 | 1.1e-4 |\n| | [SEGNN-10-64](https://drive.google.com/file/d/1_6rHxK81vzrdIMPtJ7rIkeoUgsTeKmSn/view) | 4.1e-5 | 2.3e-5 | 5.2e-4 |\n| 3D TGV | [GNS-10-128](https://drive.google.com/file/d/1DEkXxrebS9eyLSMlc_ztHrqlh29NgLXC/view) | 5.8e-3 | 4.7e-6 | 4.8e-2 |\n| | [SEGNN-10-64](https://drive.google.com/file/d/1ivJnHTgfbQ0IJujc5O0CUoQNiGU4zi_d/view) | 5.0e-3 | 4.9e-6 | 3.9e-2 |\n| 3D RPF | [GNS-10-128](https://drive.google.com/file/d/1yo-qgShLd1sgS1u5zkMXdJvhuPBwEQQE/view) | 2.1e-5 | 3.3e-7 | 1.8e-6 |\n| | [SEGNN-10-64](https://drive.google.com/file/d/1Qczh3Z_z0grTuRuPDHyiYLzV1zg7Liz9/view) | 1.7e-5 | 2.7e-7 | 1.7e-6 |\n| 3D LDC | [GNS-10-128](https://drive.google.com/file/d/1b3IIkxk5VcWiT8Oyqg1wex8-ZfJv2g_v/view) | 4.1e-5 | 3.2e-7 | 1.9e-8 |\n| | [SEGNN-10-64](https://drive.google.com/file/d/1ZIg7FXc1l3C4ekc9WvVvjHEl5KKxOA_U/view) | 4.1e-5 | 2.9e-7 | 2.5e-8 |\n\nTo reproduce the numbers in the table, e.g., on 2D TGV with GNS, follow these steps:\n```bash\n# download the checkpoint (1) through the browser or \n# (2) using the file ID from the URL, i.e., for 2D TGV + GNS\ngdown 19TO4PaFGcryXOFFKs93IniuPZKEcaJ37\n# unzip the downloaded file `gns_tgv2d.zip`\npython -c \"import shutil; shutil.unpack_archive('gns_tgv2d.zip', 'gns_tgv2d')\"\n# evaluate the model on the test split\npython main.py gpu=$GPU_ID mode=infer eval.test=True load_ckp=gns_tgv2d/best\n```\n\n## Directory structure\n```\n\ud83d\udce6lagrangebench\n \u2523 \ud83d\udcc2case_setup # Case setup manager\n \u2503 \u2523 \ud83d\udcdccase.py # CaseSetupFn class\n \u2503 \u2517 \ud83d\udcdcfeatures.py # Feature extraction\n \u2523 \ud83d\udcc2data # Datasets and dataloading utils\n \u2503 \u2523 \ud83d\udcdcdata.py # H5Dataset class and specific datasets\n \u2503 \u2517 \ud83d\udcdcutils.py\n \u2523 \ud83d\udcc2evaluate # Evaluation and rollout generation tools\n \u2503 \u2523 \ud83d\udcdcmetrics.py\n \u2503 \u2523 \ud83d\udcdcrollout.py\n \u2503 \u2517 \ud83d\udcdcutils.py\n \u2523 \ud83d\udcc2models # Baseline models\n \u2503 \u2523 \ud83d\udcdcbase.py # BaseModel class\n \u2503 \u2523 \ud83d\udcdcegnn.py\n \u2503 \u2523 \ud83d\udcdcgns.py\n \u2503 \u2523 \ud83d\udcdclinear.py\n \u2503 \u2523 \ud83d\udcdcpainn.py\n \u2503 \u2523 \ud83d\udcdcsegnn.py\n \u2503 \u2517 \ud83d\udcdcutils.py\n \u2523 \ud83d\udcc2train # Trainer method and training tricks\n \u2503 \u2523 \ud83d\udcdcstrats.py # Training tricks\n \u2503 \u2517 \ud83d\udcdctrainer.py # Trainer method\n \u2523 \ud83d\udcdcdefaults.py # Default values\n \u2523 \ud83d\udcdcrunner.py # Runner wrapping training and inference\n \u2517 \ud83d\udcdcutils.py\n```\n\n\n## Contributing\nWelcome! We highly appreciate [Github issues](https://github.com/tumaer/lagrangebench/issues) and [PRs](https://github.com/tumaer/lagrangebench/pulls).\n\nYou can also chat with us on [**Discord**](https://discord.gg/Ds8jRZ78hU).\n\n### Contributing Guideline\nIf you want to contribute to this repository, you will need the dev dependencies, i.e.\ninstall the environment with `poetry install` without the ` --only main` flag.\nThen, we also recommend you install the pre-commit hooks\nif you don't want to manually run `pre-commit run` before each commit. To sum up:\n\n```bash\ngit clone https://github.com/tumaer/lagrangebench.git\ncd lagrangebench\npoetry install\nsource $PATH_TO_LAGRANGEBENCH_VENV/bin/activate\n\n# install pre-commit hooks defined in .pre-commit-config.yaml\n# ruff is configured in pyproject.toml\npre-commit install\n\n# if you want to bump the version in both pyproject.toml and __init__.py, do\npoetry self add poetry-bumpversion\npoetry version patch # or minor/major\n```\n\nAfter you have run `git add <FILE>` and try to `git commit`, the pre-commit hook will\nfix the linting and formatting of `<FILE>` before you are allowed to commit.\n\nYou should also run the tests locally before creating a PR. Do this simply by:\n\n```bash\n# pytest is configured in pyproject.toml\npytest\n```\n\n### Clone vs Library\nLagrangeBench can be installed by cloning the repository or as a standalone library. This offers more flexibility, but it also comes with its disadvantages: the necessity to implement some things twice. If you change any of the following things, make sure to update its counterpart as well:\n- General setup in `lagrangebench/runner.py` and `notebooks/tutorial.ipynb`\n- Configs in `configs/` and `lagrangebench/defaults.py`\n- Zenodo URLs in `download_data.sh` and `lagrangebench/data/data.py`\n- Dependencies in `pyproject.toml`, `requirements_cuda.txt`, and `docs/requirements.txt`\n- Library version in `pyproject.toml` and `lagrangebench/__init__.py`\n\n\n## Citation\nThe paper (at NeurIPS 2023 Datasets and Benchmarks) can be cited as:\n```bibtex\n@article{toshev2024lagrangebench,\n title={Lagrangebench: A lagrangian fluid mechanics benchmarking suite},\n author={Toshev, Artur and Galletti, Gianluca and Fritz, Fabian and Adami, Stefan and Adams, Nikolaus},\n journal={Advances in Neural Information Processing Systems},\n volume={36},\n year={2024}\n}\n```\n\nThe associated datasets can be cited as:\n```bibtex\n@dataset{toshev_2024_10491868,\n author = {Toshev, Artur P. and Adams, Nikolaus A.},\n title = {LagrangeBench Datasets},\n month = jan,\n year = 2024,\n publisher = {Zenodo},\n doi = {10.5281/zenodo.10491868},\n url = {https://doi.org/10.5281/zenodo.10491868}\n}\n```\n\n\n### Publications\nThe following further publications are based on the LagrangeBench codebase:\n\n1. [Learning Lagrangian Fluid Mechanics with E(3)-Equivariant Graph Neural Networks (GSI 2023)](https://arxiv.org/abs/2305.15603), A. P. Toshev, G. Galletti, J. Brandstetter, S. Adami, N. A. Adams\n2. [Neural SPH: Improved Neural Modeling of Lagrangian Fluid Dynamics (ICML 2024)](https://arxiv.org/abs/2402.06275), A. P. Toshev, J. A. Erbesdobler, N. A. Adams, J. Brandstetter\n",
"bugtrack_url": null,
"license": "MIT",
"summary": "LagrangeBench: A Lagrangian Fluid Mechanics Benchmarking Suite",
"version": "0.2.0",
"project_urls": {
"Documentation": "https://lagrangebench.readthedocs.io/",
"Homepage": "https://lagrangebench.readthedocs.io/",
"Repository": "https://github.com/tumaer/lagrangebench"
},
"split_keywords": [
"smoothed-particle-hydrodynamics",
" benchmark-suite",
" lagrangian-dynamics",
" graph-neural-networks",
" lagrangian-particles"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "3703fba7c58df1f061c9e224ffddc50cb6d8166206c5d561ff49444889a1b1c3",
"md5": "f2f935077e2a59862bafa50259276a30",
"sha256": "1336e1ee709aaf61c5fa7c03d840381474e0192c2022338a1db2c6ea83776e5b"
},
"downloads": -1,
"filename": "lagrangebench-0.2.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "f2f935077e2a59862bafa50259276a30",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": "<=3.11,>=3.9",
"size": 61721,
"upload_time": "2024-07-07T23:55:23",
"upload_time_iso_8601": "2024-07-07T23:55:23.385624Z",
"url": "https://files.pythonhosted.org/packages/37/03/fba7c58df1f061c9e224ffddc50cb6d8166206c5d561ff49444889a1b1c3/lagrangebench-0.2.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "49726ec85cab8bfb69e123a11012873edbdbdf0513e09049892a3fa79a8f385c",
"md5": "eaa9cd35c02b35caca4737c874840b8f",
"sha256": "137f716e697599ba886350962c450321ef91765183e82d675d2317c7ce59ac85"
},
"downloads": -1,
"filename": "lagrangebench-0.2.0.tar.gz",
"has_sig": false,
"md5_digest": "eaa9cd35c02b35caca4737c874840b8f",
"packagetype": "sdist",
"python_version": "source",
"requires_python": "<=3.11,>=3.9",
"size": 56207,
"upload_time": "2024-07-07T23:55:24",
"upload_time_iso_8601": "2024-07-07T23:55:24.944747Z",
"url": "https://files.pythonhosted.org/packages/49/72/6ec85cab8bfb69e123a11012873edbdbdf0513e09049892a3fa79a8f385c/lagrangebench-0.2.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-07-07 23:55:24",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "tumaer",
"github_project": "lagrangebench",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "lagrangebench"
}
JSON
