JAX-SPH (Toshev et al., 2024) is a modular JAX-based weakly compressible SPH framework, which implements the following SPH routines:
- Standard SPH (Adami et al., 2012)
- Transport velocity SPH (Adami et al., 2013)
- Riemann SPH (Zhang et al., 2017)
Currently, the code can only be installed by cloning this repository. We recommend using a Poetry or python3-venv
environment.
poetry config virtualenvs.in-project true
poetry install
source .venv/bin/activate
Later, you just need to source .venv/bin/activate
to activate the environment.
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt
pip install -e . # to install jax_sph in interactive mode
Later, you just need to source venv/bin/activate
to activate the environment.
If you want to use a CUDA GPU, you first need a running Nvidia driver. And then just follow the instructions here. The whole process could look like this:
source .venv/bin/activate
pip install --upgrade "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
In the following, a quick setup guide for different cases is presented.
- Standard SPH 2D Taylor Green vortex
python main.py --case=TGV --solver=SPH --dim=2 --dx=0.02 --t-end=5 --seed=123 --write-h5 --write-every=25 --data-path="data/tgv2d_notvf/"
- Transport velocity SPH 2D Taylor Green vortex
python main.py --case=TGV --tvf=1.0 --solver=SPH --dim=2 --dx=0.02 --t-end=5 --seed=123 --write-h5 --write-every=25 --data-path="data/tgv2d_notvf/"
- Riemann SPH 2D Taylor Green vortex
python main.py --case=TGV --tvf=1.0 --solver=RIE --dim=2 --dx=0.02 --t-end=5 --seed=123 --write-h5 --write-every=25 --data-path="data/tgv2d_notvf/"
- Thermal diffusion
python main.py --case=HT --solver=SPH --density-evolution --heat-conduction --dim=2 --dx=0.02 --t-end=1.5 --write-h5 --write-vtk --r0-noise-factor=0.05 --outlet-temperature-derivative --data-path="data/therm_diff/"
To train and test our Solver-in-the-Loop model, run the script in ./experiments/sitl.py. This file relies on LagrangeBench, which can be installed by pip install lagrangebench
. For more information on the training and inference setup, visit the LagrangeBench website.
The presented inverse problem of finding the initial state of a 100-step-long SPH simulation can be fully reproduced using the notebook ./experiments/inverse.ipynb.
The presented validation of the gradients through the solver can be fully reproduced using the notebook ./experiments/grads.ipynb
To set up a case, just add it to the cases/
directory. Every case should inherit from SimulationSetup
in jax_sph/case_setup.py
or from another case. Running a case in relaxation mode --mode=rlx
overwrites the selected case, and passed CLI arguments overwrite any argument.
If you wish to contribute, please run
pre-commit install
upon installation to automate the code linting and formatting checks.
The main reference for this code is toshev2024jaxsph
. If you refer to the code used for dataset generation in LagrangeBench, please cite toshev2023lagrangebench
directly.
@inproceedings{toshev2024jaxsph,
title = {JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework},
author = {Artur Toshev and Harish Ramachandran and Jonas A. Erbesdobler and Gianluca Galletti and Johannes Brandstetter and Nikolaus A. Adams},
booktitle = {ICLR 2024 Workshop on AI4DifferentialEquations In Science},
year = {2024},
url = {https://openreview.net/forum?id=8X5PXVmsHW}
}
@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 initial idea for JAX-SPH is due to Fabian Fritz and Ludger Paehler, which has led to the first validated JAX implementation of the 3D Taylor-Green vortex simulated with the transport velocity SPH formulation. Since then, Artur Toshev has maintained the codebase and extended it in various ways. The people who have provided useful feedback include, but are not limited to: Stefan Adami, Xiangyu Hu, Fabian Fritz, Christopher Zöller, Fabian Thiery, Johannes Brandstetter, and Ludger Paehler. Special thanks to Nikolaus Adams, who has supervised the project from the beginning.
- Artur Toshev - developed and maintains the codebase; selected the SPH algorithms and validated most of them; designed the simulation cases and ML experiments.
- Fabian Fritz - provided the first validated transport velocity SPH implementation of the 3D Taylor-Green vortex in JAX.
- Jonas Erbesdobler - implemented Riemann SPH; improved and added solver validation scripts; contributed to refactoring the codebase.
- Harish Ramachandran - implemented thermal diffusion and the inverse problem; helped in the initial phase of Solver-in-the-Loop.
- Gianluca Galletti - validated the gradients through the solver; implemented and tuned Solver-in-the-Loop.