jaxfluids package

Subpackages

Submodules

jaxfluids.boundary_condition module

class jaxfluids.boundary_condition.BoundaryCondition(domain_information: DomainInformation, material_manager: MaterialManager, unit_handler: UnitHandler, boundary_types: Dict, wall_velocity_functions: Dict, dirichlet_functions: Dict, neumann_functions: Dict, levelset_type: str)[source]

Bases: object

The BoundaryCondition class implements functionality to enforce user- specified boundary conditions. Boundary conditions are enforced on the primitive variables and the level-set field (for two-phase simulations only).

Boundary conditions for the primitive variables: 1) Periodic 2) Symmetric 3) No-slip walls 4) Dirichlet 5) Neumann

Boundary conditions for the level-set field: 1) Periodic 2) Symmetry 3) Zero-Gradient

assign_corners() → Dict[source]

Identifies the boundary type pairs at the corners (2D) / edges (3D) of computational domain. This is necessary to fill the halo cells that are located at the diagonal extension of the domain.

Returns:

Dictionary containing the boundary type pairs at each boundary corner/edge location

Return type:

Dict

dirichlet(cons: Array, primes: Array, location: str, functions: Callable | float, current_time: float, b_range: List) → Tuple[Array, Array][source]

Fills the halo cells of the primitive and conservative variable buffer at the specified location according to the dirichlet boundary condition.

Parameters:

  • cons (jnp.ndarray) – Conservative variable buffer

  • primes (jnp.ndarray) – Primitive variable buffer

  • location (str) – Boundary location

  • functions (Union[Callable, float]) – Dirichlet functions

  • current_time (float) – Current physical simulation time

  • b_range (List) – List containing the spatial range of the boundary at the specified location

Returns:

Primitive and conservative variable buffers with filled halos at specified location

Return type:

Tuple[jnp.ndarray, jnp.ndarray]

fill_boundary_levelset(levelset: Array) → Array[source]

Fills the levelset buffer halo cells.

Parameters:

levelset (jnp.ndarray) – Levelset buffer

Returns:

Levelset buffer with filled halo cells

Return type:

jnp.ndarray

fill_boundary_primes(cons: Array, primes: Array, current_time: float) → Tuple[Array, Array][source]

Fills the halo cells of the primitive and conservative variable buffers.

Parameters:

  • cons (jnp.ndarray) – Buffer of conservative variables

  • primes (jnp.ndarray) – Buffer of primitive variables

  • current_time (float) – Current physical simulation time

Returns:

Primitive and conservative variable buffer with filled halo cells

Return type:

Tuple[jnp.ndarray, jnp.ndarray]

fill_corners_levelset(levelset: Array) → Array[source]

Fills the levelset buffer halo cells that are located at the diagional extension of the domain.

Parameters:

levelset (jnp.ndarray) – Levelset buffer

Returns:

Levelset buffer with filled halo cells at the corners

Return type:

jnp.ndarray

fill_corners_primes(primes: Array) → Array[source]

Fills the prime buffer halo cells that are located at the diagonal extension of the domain.

Parameters:

primes (jnp.ndarray) – Buffer of the primitive variables

Returns:

Primitive variables buffer with filled halo cells at the corners

Return type:

jnp.ndarray

get_slices_fill(location: str, b_range: List) → Tuple[source]

Computes the slice objects to fill the halos depending on the range at the specified boundary location.

Parameters:

  • location (str) – Boundary location

  • b_range (List) – List containing the spatial range of the boundary at the specified location

Returns:

Slice objects

Return type:

Tuple

get_slices_retrieve(location: str, boundary_type: str, b_range: List) → Tuple[source]

Computes the slice objects to retrieve the values depending on the range at the specified boundary location.

Parameters:

  • location (str) – Boundary location

  • boundary_type (str) – Boundary location

  • b_range (List) – List containing the spatial range of the boundary at the specified location

Returns:

Slice objects

Return type:

Tuple

neumann(cons: Array, primes: Array, location: str, functions: Callable | float, current_time: float, b_range: List) → Tuple[Array, Array][source]

Fills the halo cells of the primitive and conservative variable buffer at the specified location according to the neumann boundary condition.

Parameters:

  • cons (jnp.ndarray) – Conservative variable buffer

  • primes (jnp.ndarray) – Primitive variable buffer

  • location (str) – Boundary location

  • functions (Union[Callable, float]) – Neumann functions

  • current_time (float) – Current physical simulation time

  • b_range (List) – List containing the spatial range of the boundary at the specified location

Returns:

Primitive and conservative variable buffers with filled halos at specified location

Return type:

Tuple[jnp.ndarray, jnp.ndarray]

periodic(cons: Array, primes: Array, location: str, b_range: List) → Tuple[Array, Array][source]

Fills the halos of the conservative and primitive variable buffer at the specified location according to the periodic boundary condition.

Parameters:

  • cons (jnp.ndarray) – Conservative variable buffer

  • primes (jnp.ndarray) – Primitive variable buffer

  • location (str) – Boundary location

  • b_range (List) – List containing the spatial range of the boundary at the specified location

Returns:

Primitive and conservative variable buffers with filled halos at specified location

Return type:

Tuple[jnp.ndarray, jnp.ndarray]

periodic_levelset(levelset: Array, location: str)[source]
symmetry(cons: Array, primes: Array, location: str, b_range: List) → Tuple[Array, Array][source]

Fills the halo cells of the primitive and conservative variable buffer at the specified location according to the symmetric boundary condition.

Parameters:

  • cons (jnp.ndarray) – Conservative variable buffer

  • primes (jnp.ndarray) – Primitive variable buffer

  • location (str) – Boundary location

  • b_range (List) – List containing the spatial range of the boundary at the specified location

Returns:

Primitive and conservative variable buffers with filled halos at specified location

Return type:

Tuple[jnp.ndarray, jnp.ndarray]

symmetry_levelset(levelset: Array, location: str)[source]
wall(cons: Array, primes: Array, location: str, functions: Dict, current_time: float, b_range: List) → Tuple[Array, Array][source]

Fills the halo cells of the primitive and conservative variable buffer at the specified location according to the no-slip wall boundary condition.

Parameters:

  • cons (jnp.ndarray) – Conservative variable buffer

  • primes (jnp.ndarray) – Primitive variable buffer

  • location (str) – Boundary location

  • functions (Dict) – Wall velocity functions

  • current_time (float) – Current physical simulation time

  • b_range (List) – List containing the spatial range of the boundary at the specified location

Returns:

Primitive and conservative variable buffers with filled halos at specified location

Return type:

Tuple[jnp.ndarray, jnp.ndarray]

jaxfluids.domain_information module

class jaxfluids.domain_information.DomainInformation(dim: int, nx: int, ny: int, nz: int, nh_conservatives: int, nh_geometry: int, domain_size: Dict)[source]

Bases: object

This DomainInformation class holds information about the computational domain, e.g., mesh, number of cells, extension in each spatial direction, active axis, domain slice objects etc..

compute_mesh_xi(nxi: int, domain_size_xi: List) → Tuple[Array, Array, Array][source]

Computes cell center coordinates, cell face coordinates and cell sizes in the specified direction

Parameters:

  • nxi (int) – Number of cells in xi direction

  • domain_size_xi (List) – Domain size in xi direction

Returns:

Cell center coordinates, cell face coordinates and cell sizes in xi direction

Return type:

Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

jaxfluids.flux_computation module

class jaxfluids.flux_computation.FluxComputer(numerical_setup: Dict, material_manager: MaterialManager, domain_information: DomainInformation)[source]

Bases: object

The Flux Computer sets up the user-specified flux function for the calculation of the convective terms. The flux calculation is called in the space solver by compute_convective_flux_xi().

There are three general options for the convective flux function. 1) High-order Godunov Scheme 2) Flux-splitting Scheme 3) ALDM Scheme

compute_convective_flux_xi(primes: Array, cons: Array, axis: int, ml_parameters_dict: Dict | None = None, ml_networks_dict: Dict | None = None) → Array[source]

Computes the convective fluxes.

Parameters:

  • primes (jnp.ndarray) – Primitive variable buffer

  • cons (jnp.ndarray) – Conservative variable buffer

  • axis (int) – Spatial direction

  • ml_parameters_dict (Union[Dict, None], optional) – Dictionary of neural network weights, defaults to None

  • ml_networks_dict (Union[Dict, None], optional) – Dictionary of neural network architectures, defaults to None

Returns:

Convective fluxes in axis direction

Return type:

jnp.ndarray

jaxfluids.initializer module

class jaxfluids.initializer.Initializer(input_reader: InputReader)[source]

Bases: object

The Initializer class implements functionality to create a dictionary of initial buffers that is passed to the simulate() method of the SimulationManager class. The initialization() method returns this dictionary. The initial buffers are created in one of the following ways: 1) From a restart file that is specified in the case setup. 2) From turbulent initial condition parameters that are specified in the case setup 3) From the initial primitive buffer that is passed to the initialization() method 4) From the initial conditions for primitive variables specified in case setup Note that if multiple of the above are provided, the priority is 1) - 4).

initialization(user_prime_init: ndarray | Array | None = None) → Dict[str, Dict[str, Array | float]][source]

Creates a dictionary of initial buffers. The buffers are either created from the case setup file or from the user_prime_init argument.

Parameters:

user_prime_init (Union[np.ndarray, jnp.ndarray], optional) – User specified initial buffer for the primitive variables, defaults to None

Returns:

Dictionary of initial buffers

Return type:

Dict[str, Dict[str, Union[jnp.ndarray, float]]]

jaxfluids.input_reader module

class jaxfluids.input_reader.InputReader(case_setup: str | Dict, numerical_setup: str | Dict)[source]

Bases: object

The InputReader class to reads a case setup and a numerical setup for setting up a JAX-FLUIDS simulation. Case setup and numerical setup can be provided as either a path to a json file or as a preloaded dictionary.

info() → Tuple[Dict, Dict][source]
sanity_check() → None[source]

Checks if the case setup and numerical setup is consistent.

jaxfluids.simulation_manager module

class jaxfluids.simulation_manager.SimulationManager(input_reader: InputReader)[source]

Bases: object

The SimulationManager is the top-level class in JAX-FLUIDS. It provides functionality to perform conventional CFD simulations as well as end-to-end optimization of ML models.

The most important methods of the SimulationManager are: 1) simulate() - Performs conventional CFD simulation. 2) feedforward() - Feedforward of a batch of data, i.e.,

advances a batch of initial conditions in time for a fixed amount of steps

  1. do_integration_step() - Performs a single integration step

advance(buffer_dictionary: Dict[str, Dict[str, Array | float]]) → None[source]

Advances the initial buffers in time.

Parameters:

buffer_dictionary – ictionary containing the material field buffers,

levelset quantitiy buffers, time control and mass flow forcing parameters :type buffer_dictionary: Dict

compute_timestep(primes: Array, levelset: Array, volume_fraction: Array, **kwargs) → float[source]

Computes the physical time step size depending on the active physics.

Parameters:

  • primes (jnp.ndarray) – Buffer of primitive variables

  • levelset (jnp.ndarray) – Levelset buffer

  • volume_fraction (jnp.ndarray) – Volume fraction buffer

Returns:

Time step size

Return type:

float

do_integration_step(cons: Array, primes: Array, timestep_size: float, current_time: float, levelset: Array | None = None, volume_fraction: Array | None = None, apertures: List | None = None, reinitialize: bool = False, forcings_dictionary: Dict | None = None, ml_parameters_dict: Dict | None = None, ml_networks_dict: Dict | None = None, **kwargs) → Tuple[Dict, Dict, Dict][source]

Performs an integration step using the specified integration scheme. For twophase simulations a single RK stage consists of the following: 1) Compute right-hand-side of Navier-Stokes and levelset advection equation 2) Transform volume-averaged conservatives to actual conservatives

that can be integrated according to volume fraction

  1. Prepare the conservative and levelset buffer for integration according

    to the present integration scheme

  2. Integrate conservatives

  3. Integrate levelset + reinitialize levelset + fill levelset boundaries

  4. Compute volume fraction and apertures from integrated levelset quantities

  5. Apply the mixing procedure to the integrated conservative variables

  6. Transform mixed conservative variables to volume-averaged conservative variables and compute

    corresponding primitive variables

  7. Extend primitive variables into ghost cells and compute conservative variables in ghost cells from extended primitive

    variables

  8. Fill material boundaries

Parameters:

  • cons (jnp.ndarray) – Buffer of conservative variables

  • primes (jnp.ndarray) – Buffer of primitive variables

  • timestep_size (float) – Current physical time step size

  • current_time (float) – Current physical simulation time

  • levelset (Union[jnp.ndarray, None], optional) – Levelset buffer, defaults to None

  • volume_fraction (Union[jnp.ndarray, None], optional) – Volume fraction buffer, defaults to None

  • apertures (Union[List, None], optional) – Aperture buffers, defaults to None

  • reinitialize (bool, optional) – Flag indicating whether to reinitialize levelset in the present time step, defaults to False

  • forcings_dictionary (Union[Dict, None], optional) – Dictionary containing forcing buffers, defaults to None

  • ml_parameters_dict (Union[Dict, None], optional) – Dictionary containing NN weights, defaults to None

  • ml_networks_dict (Union[Dict, None], optional) – Dictionary containing NN architectures, defaults to None

Returns:

Tuple of material fields dictionary, levelset quantities dictionary and residual dictionary

Return type:

Tuple[Dict, Dict, Dict]

feed_forward(batch_primes_init: Array, batch_levelset_init: Array, n_steps: int, timestep_size: float, t_start: float, output_freq: int = 1, ml_parameters_dict: Dict | None = None, ml_networks_dict: Dict | None = None) → Tuple[Array, Array][source]

Vectorized version of the _feed_forward() method.

Parameters:

  • batch_primes_init (jnp.ndarray) – batch of initial primitive variable buffers

  • batch_levelset_init (jnp.ndarray) – batch of initial levelset buffers

  • n_steps (int) – Number of integration steps

  • timestep (float) – Physical time step size

  • t_start (float) – Physical start time

  • output_freq (int, optional) – Frequency in time steps for output, defaults to 1

  • ml_parameters_dict (Union[Dict, None], optional) – NN weights, defaults to None

  • ml_networks_dict (Union[Dict, None], optional) – NN architectures, defaults to None

Returns:

_description_

Return type:

Tuple[jnp.ndarray, jnp.ndarray]

initialize(buffer_dictionary: Dict[str, Dict[str, Array | float]]) → None[source]

Initializes the simulation, i.e., creates the output directory, logs the numerical and case setup, and writes the initial output.

Parameters:

buffer_dictionary – Dictionary containing the material field buffers,

levelset quantitiy buffers, time control and mass flow forcing parameters :type buffer_dictionary: Dict

simulate(buffer_dictionary: Dict[str, Dict[str, Array | float]]) → None[source]

Performs a conventional CFD simulation.

Parameters:

buffer_dictionary – Dictionary containing the material field buffers,

levelset quantitiy buffers, time control and mass flow forcing parameters :type buffer_dictionary: Dict[str, Dict[str, Union[jnp.ndarray, float]]]

jaxfluids.space_solver module

class jaxfluids.space_solver.SpaceSolver(domain_information: DomainInformation, material_manager: MaterialManager, numerical_setup: Dict, gravity: Array, levelset_type: Tuple[None, str], levelset_handler: LevelsetHandler | None)[source]

Bases: object

The Space Solver class manages the calculation of the righ-hand-side (i.e., fluxes) of the NSE and, for two-phase simulations, manages the calculation of the rhs of the level-set advection.

Depending on the numerical setup, the calculation of the fluxes for the NSE has contributions from: 1) Convective flux 2) Viscous flux 3) Heat flux 4) Interface exchange flux 5) Volume force flux 6) External forcing

compute_rhs(cons: Array, primes: Array, current_time: float, levelset: Array | None = None, volume_fraction: Array | None = None, apertures: List | None = None, forcings_dictionary: Dict | None = None, ml_parameters_dict: Dict | None = None, ml_networks_dict: Dict | None = None) → Tuple[Array, Array | None, float | None][source]

Computes the right-hand-side of the Navier-Stokes equations depending on active physics and active axis. For levelset simulations with FLUID-FLUID or FLUID-SOLID-DYNAMIC interface interactions, also computes the right-hand-side of the levelset advection.

Parameters:

  • cons (jnp.ndarray) – Buffer of conservative variables

  • primes (jnp.ndarray) – Buffer of primitive variables

  • current_time (float) – Current physical simulation time

  • levelset (jnp.ndarray, optional) – Levelset buffer, defaults to None

  • volume_fraction (jnp.ndarray, optional) – Volume fraction buffer, defaults to None

  • apertures (List, optional) – Aperture buffers, defaults to None

  • forcings_dictionary (Union[Dict, None], optional) – Forcings dictionary, defaults to None

  • ml_parameters_dict (Union[Dict, None], optional) – Dictionary containing NN weights, defaults to None

  • ml_networks_dict (Union[Dict, None], optional) – Dictionary containing NN architectures, defaults to None

Returns:

Tuple containing the right-hand-side buffer of the Navier-Stokes equations, the

right-hand-side buffer of the levelset advection equation and the maximum extension residual of the interface quantities :rtype: Tuple[jnp.ndarray, Union[jnp.ndarray, None], Union[float, None]]

jaxfluids.unit_handler module

class jaxfluids.unit_handler.UnitHandler(density_reference: float, length_reference: float, velocity_reference: float, temperature_reference: float)[source]

Bases: object

The UnitHandler class implements functionaly to solve the NSE in non-dimensional form.

dimensionalize(value: Array | float, quantity: str) → Array | float[source]

Dimensionalizes the given quantity buffer w.r.t. the specified quanty.

Parameters:

  • value (Union[jnp.ndarray, float]) – Non-dimensional quantity buffer

  • quantity (str) – Quantity name

Returns:

Dimensional quantity buffer

Return type:

Union[jnp.ndarray, float]

non_dimensionalize(value: Array | float, quantity: str) → Array | float[source]

Non-dimensionalizes the given buffer w.r.t. the specified quantity.

Parameters:

  • value (Union[jnp.ndarray, float]) – Dimensional quantity buffer

  • quantity (str) – Quantity name

Returns:

Non-dimensional quantity buffer

Return type:

Union[jnp.ndarray, float]

non_dimensionalize_domain_size(domain_size: Dict) → Dict[source]

jaxfluids.utilities module

jaxfluids.utilities.get_conservatives_from_primitives(primes: Array, material_manager: MaterialManager) → Array[source]

Converts primitive variables to conservative variables.

Parameters:

  • primes (jnp.ndarray) – Buffer of primitive variables

  • material_manager (MaterialManager) – Class that calculats material quantities

Returns:

Buffer of conservative variables

Return type:

jnp.ndarray

jaxfluids.utilities.get_fluxes_xi(primes: Array, cons: Array, axis: int) → Array[source]

Computes the physical flux in a specified spatial direction. Cf. Eq. (3.65) in Toro.

Parameters:

  • primes (jnp.ndarray) – Buffer of primitive variables

  • cons (jnp.ndarray) – Buffer of conservative variables

  • axis (int) – Spatial direction along which fluxes are calculated

Returns:

Physical fluxes in axis direction

Return type:

jnp.ndarray

jaxfluids.utilities.get_primitives_from_conservatives(cons: Array, material_manager: MaterialManager) → Array[source]

Converts conservative variables to primitive variables.

Parameters:

  • cons (jnp.ndarray) – Buffer of conservative variables

  • material_manager (MaterialManager) – Class that calculats material quantities

Returns:

Buffer of primitive variables

Return type:

jnp.ndarray

Module contents

JAX-FLUIDS

A fully-differentiable computational fluid dynamics solver for compressible two-phase flows. Solves the compressible single or two-phase Navier-Stokes equations in three spatial dimensions.

Let’s make Machine Learning + Computational Fluid Dynamics great!

Contact: deniz.bezgin@tum.de / aaron.buhendwa@tum.de

Documentation

Documentation available at: https://github.com/tumaer/JAXFLUIDS

Examples

The examples folder includes example CFD simulations in one, two, and three dimensions.

For example in examples/examples_1D/02_sod

>>> python run_sod.py

Available subpackages

forcing

Implements external forces for the Navier-Stokes equations

iles

Implicit Large-Eddy Simulation models

io_utils

Tools for logging and writing output

levelset

Comprehensive level-set implementation for two-phase flows

materials

Materials and material managers

post_process

Tools for post-processing CFD simulations

shock_sensor

Shock-sensors

stencils

Spatial reconstruction and derivative stencils

time_integration

Explicit time integration schemes up to third-order

turb

Tools for turbulent initial conditions and statistics

class jaxfluids.Initializer(input_reader: InputReader)[source]

Bases: object

The Initializer class implements functionality to create a dictionary of initial buffers that is passed to the simulate() method of the SimulationManager class. The initialization() method returns this dictionary. The initial buffers are created in one of the following ways: 1) From a restart file that is specified in the case setup. 2) From turbulent initial condition parameters that are specified in the case setup 3) From the initial primitive buffer that is passed to the initialization() method 4) From the initial conditions for primitive variables specified in case setup Note that if multiple of the above are provided, the priority is 1) - 4).

initialization(user_prime_init: ndarray | Array | None = None) → Dict[str, Dict[str, Array | float]][source]

Creates a dictionary of initial buffers. The buffers are either created from the case setup file or from the user_prime_init argument.

Parameters:

user_prime_init (Union[np.ndarray, jnp.ndarray], optional) – User specified initial buffer for the primitive variables, defaults to None

Returns:

Dictionary of initial buffers

Return type:

Dict[str, Dict[str, Union[jnp.ndarray, float]]]