DOCSEARCHbyAlgolia
Submit Docs

JAX-Fluids Documentation

repository·main·Indexed 20 days ago

https://github.com/tumaer/jaxfluids

A fully-differentiable CFD solver for 3D, compressible single-phase and two-phase flows, designed for research at the intersection of Machine Learning and Computational Fluid Dynamics. It supports CPU and GPU execution, featuring explicit time stepping (Euler, RK2, RK3), adaptive high-order reconstruction (WENO, TENO), various Riemann solvers (including NN-based versions), and the Level-set method for multi-phase simulations on Cartesian grids.

Tokens
26.7K
Snippets
67
Records
98
Agent score
70%

What's inside JAX-Fluids

  1. Overview of the jaxfluids.levelset package

    main
    The jaxfluids.levelset package provides tools for level-set methods in fluid dynamics simulations. It is organized into several specialized submodules that handle geometry calculation, interface properties, level-set creation, management, and reinitialization. This package is designed to work within the JAX ecosystem, supporting differentiable fluid dynamics workflows.

  2. Overview of the jaxfluids package structure

    main

    The jaxfluids package is a differentiable fluid dynamics library built with JAX. It is organized into several specialized subpackages and modules that handle different aspects of fluid simulation:

    Core Subpackages

    • jaxfluids.forcing: External forcing terms.
    • jaxfluids.iles: Improved Large Eddy Simulation (ILES) capabilities.
    • jaxfluids.io_utils: Input/Output utilities.
    • jaxfluids.levelset: Level-set methods.
    • jaxfluids.materials: Material property definitions.
    • jaxfluids.post_process: Post-processing tools.
    • jaxfluids.shock_sensor: Detection of shock waves.
    • jaxfluids.solvers: Numerical solvers.
    • jaxfluids.stencils: Spatial discretization stencils.
    • jaxfluids.time_integration: Temporal discretization methods.
    • jaxfluids.turb: Turbulence modeling.

    Key Modules

    • jaxfluids.boundary_condition: Defines boundary conditions for the domain.
    • jaxfluids.domain_information: Manages domain geometry and metadata.
    • jaxfluids.flux_computation: Handles the calculation of numerical fluxes.
    • jaxfluids.initializer: Provides tools to set initial states for simulations.
    • jaxfluids.input_reader: Reads simulation input configurations.
    • jaxfluids.simulation_manager: Orchestrates the simulation lifecycle.
    • jaxfluids.space_solver: Manages spatial discretization and solving.
    • jaxfluids.unit_handler: Manages physical units and conversions.
    • jaxfluids.utilities: General helper functions.

  3. Overview of JAX-Fluids capabilities

    main

    JAX-Fluids is a differentiable computational fluid dynamics (CFD) package designed for compressible two-phase flows. It is built with a modular, object-oriented architecture and supports both CPU and GPU execution.

    Key technical capabilities include:

    • Numerical Methods: Explicit time stepping (Euler, RK2, RK3) and adaptive high-order reconstruction (WENO-3/5/7, WENO-CU6, WENO-3NN, TENO).
    • Physics & Solvers: Riemann solvers (Lax-Friedrichs, Rusanov, HLL, HLLC, Roe) and an implicit turbulence SGS model (ALDM).
    • Multi-phase & Geometry: Two-phase simulations using the Level-set method (supporting arbitrary solid boundaries) on a Cartesian grid.
    • Control & Forcing: Support for temperature, mass flow rate, and turbulence forcings.
    • Data I/O: Configuration via JSON files and output in H5 and XDMF formats.

  4. Use level set stencils in jaxfluids

    main

    The jaxfluids.stencils.levelset package provides various numerical stencil implementations for level set methods. These stencils are used to compute derivatives and handle interface tracking in fluid simulations. The package is organized into submodules based on the specific numerical scheme used:

    • First-order derivatives: deriv_first_order and deriv_first_order_subcell_fix (for subcell fixes).
    • High-order Upwind Central (HOUC) schemes: houc3, houc5, and houc7 modules.
    • Weighted Essentially Non-Oscillatory (WENO) schemes: weno3_deriv and weno5_deriv modules.

  5. Available Riemann Solvers in jaxfluids.solvers.riemann_solvers

    main

    The jaxfluids.solvers.riemann_solvers package provides various Riemann solvers for fluid dynamics simulations. These solvers are used to resolve discontinuities and waves at cell interfaces.

    Available solver modules include:

    • Classical Solvers:

      • AUSMP: AUSM+-type solvers.
      • HLL: Harten-Lax-van Leer solver.
      • HLLC: HLL solver with contact discontinuity resolution.
      • HLLCLM: HLLC-type solver with contact and low-Mach number improvements.
      • LaxFriedrichs: Lax-Friedrichs numerical flux.
      • Rusanov: Rusanov (Local Lax-Friedrichs) solver.

    • Neural Network (NN) Based Solvers:

      • RiemannNN: Neural network-based Riemann solver.
      • RusanovNN: Neural network-based Rusanov solver.

    • Utilities:

      • eigendecomposition: Tools for performing eigen-analysis of the flux Jacobians.
      • signal_speeds: Modules for calculating wave/signal speeds required by the solvers.
      • riemann_solver: The core module defining the Riemann solver interface.

  7. Explore the jaxfluids.stencils package

    main

    The jaxfluids.stencils package provides tools for numerical differentiation and spatial reconstruction, which are essential for implementing fluid dynamics solvers. The package is organized into several specialized subpackages and modules:

    Subpackages

    • jaxfluids.stencils.derivative: Tools for computing spatial derivatives.
    • jaxfluids.stencils.levelset: Stencils specifically designed for level-set methods.
    • jaxfluids.stencils.reconstruction: Tools for reconstructing spatial information from discrete data.

    Key Modules

    • jaxfluids.stencils.spatial_derivative: Contains implementations for computing derivatives in space.
    • jaxfluids.stencils.spatial_reconstruction: Contains implementations for spatial reconstruction tasks.

  8. Explore the jaxfluids.solvers package structure

    main

    The jaxfluids.solvers package provides various numerical methods and schemes for solving fluid dynamics equations. It is organized into subpackages and modules that handle different aspects of the solver logic:

    • Subpackages:
      • jaxfluids.solvers.riemann_solvers: Contains Riemann solvers used for calculating fluxes at cell interfaces.
    • Modules:
      • jaxfluids.solvers.flux_splitting_scheme: Implements flux splitting methods.
      • jaxfluids.solvers.high_order_godunov: Implements high-order Godunov-type schemes.
      • jaxfluids.solvers.source_term_solver: Handles the integration or solving of source terms within the fluid equations.

  9. Available material models in jaxfluids.materials

    main

    The jaxfluids.materials package provides several differentiable material models for fluid dynamics simulations. You can choose from the following submodules depending on the physical properties required for your simulation:

    • ideal_gas: Standard ideal gas model.
    • safe_ideal_gas: A version of the ideal gas model designed for numerical stability (likely preventing negative pressures or temperatures).
    • stiffened_gas: A model used for high-pressure fluids, often used in shock physics.
    • tait: The Tait equation of state, typically used for modeling liquids (e.g., water) where compressibility is low but non-zero.
    • material: The base abstraction/interface for defining material properties.
    • material_manager: A utility for managing or switching between different material models.

  10. Use the jaxfluids.iles package for Incompressible Euler Solvers

    main

    The jaxfluids.iles package provides implementations of Incompressible Euler Solvers (ILES). It is organized into submodules based on the numerical schemes used for spatial discretization, specifically focusing on different orders of the Weighted Essentially Non-Oscillatory (WENO) scheme.

    Available submodules include:

    • jaxfluids.iles.ALDM: The base ALDM (Advection-Lattice-Divergence-Method) module.
    • jaxfluids.iles.ALDM_WENO1: ALDM using a 1st-order WENO scheme.
    • jaxfluids.iles.ALDM_WENO3: ALDM using a 3rd-order WENO scheme.
    • jaxfluids.iles.ALDM_WENO5: ALDM using a 5th-order WENO scheme.

  11. Use the jaxfluids.shock_sensor package for shock detection

    main

    The jaxfluids.shock_sensor package provides tools for detecting shocks in fluid dynamics simulations. It is organized into two primary submodules:

    1. jaxfluids.shock_sensor.ducros: Implements shock detection based on the Ducros criterion.
    2. jaxfluids.shock_sensor.shock_sensor: Provides general shock sensor functionality.

    You can access these capabilities through the main jaxfluids.shock_sensor entry point.

  12. Explore the jaxfluids subpackages

    main

    The jaxfluids package is organized into several specialized subpackages for fluid dynamics simulation. Depending on your task, you may need to interact with these modules:

    • Physics & Models: jaxfluids.forcing, jaxfluids.iles, jaxfluids.materials, jaxfluids.turb (turbulence).
    • Numerical Methods: jaxfluids.solvers, jaxfluids.stencils, jaxfluids.time_integration, jaxfluids.shock_sensor.
    • Geometry & Domain: jaxfluids.levelset.
    • Data & I/O: jaxfluids.io_utils, jaxfluids.post_process.

    Use these subpackages to extend or customize specific parts of the simulation pipeline, such as defining new materials or implementing custom time integration schemes.