JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework

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)

Installation

Currently, the code can only be installed by cloning this repository. We recommend using a Poetry or python3-venv environment.

Using Poetry (recommended)

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.

Using python3-venv

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.

GPU Support

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

Getting Started

In the following, a quick setup guide for different cases is presented.

Running an SPH Simulation

• 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/"

Solver-in-the-Loop

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.

Inverse Problem

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.

Gradient Validation

The presented validation of the gradients through the solver can be fully reproduced using the notebook ./experiments/grads.ipynb

Setting up a case

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.

Development and Contribution

If you wish to contribute, please run

pre-commit install

upon installation to automate the code linting and formatting checks.

Citation

The main reference for this code is the ICLR'24 workshop paper toshev2024jax. If you refer to the code used for dataset generation in LagrangeBench, please cite toshev2024lagrangebench directly.

@article{toshev2024jax,
  title={JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework},
  author={Toshev, Artur P and Ramachandran, Harish and Erbesdobler, Jonas A and Galletti, Gianluca and Brandstetter, Johannes and Adams, Nikolaus A},
  journal={arXiv preprint arXiv:2403.04750},
  year={2024}
}

@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}
}

Acknowledgements

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.

Contributors

• 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.