Dataset Generation Code for Data-Driven Discretizations of Quasi-Lagrangian Turbulence

This code contains:

• SpectralDNS solver in JAX, with example configs:
  • Decaying Taylor-Green vortex in 2D from TVF-SPH (Adami et al., 2013) and in 3D from Brachet et al. (1983).
  • 2D Kolmogorov flow, see JAX-CFD (Kochkov et al., 2021).
  • 3D Homogeneous Isotropic Turbulence (HIT) with forcing (see Lamorgese et al. (2005) and Mortensen and Langtangen (2016)) and without forcing following Tian et al. (2023). Initialization follows Rogallo (1981).
• Utilities to relax Smoothed Particle Hydrodynamics (SPH) particles based on JAX-SPH (Toshev et al., 2024).
• Utilities to evolve SPH particles in parallel to the spectralDNS solver (Mortensen and Langtangen, 2016).
• Utilities to generate datasets compatible with LagrangeBench (Toshev et al., 2023).

Install

poetry install
source .venv/bin/activate
pip install "jax[cuda12]==0.4.29"

Getting Started

• Validation run with reference kinetic energy after 10 steps.

python main.py config=configs/tgv_validate.yaml

• To generate an HIT dataset with $32^3$ particles, we run spectral DNS on a $256^3$ grid and spectrally coarsen/filter to $32^3$.

python main.py config=configs/hit.yaml

Datasets

The two datasets used in the paper can be regenerated with the following scripts:

• sbatch gen_dataset/slurm_kolm2d_64_1.sh - 80 min/traj x 20 trajs
• sbatch gen_dataset/slurm_hit3d_32_1.sh - 60 min/traj x 20 trajs

At the bottom of these scripts are the commands to convert the simulations into a dataset file.