simulation

Functions for differentiable fluid simulation.

All units are given in grid squares per unit time.

High-level API

The high-level simulation API consists of two functions:

• simulation.simulate() – run a multi-timestep simulation starting from some initial conditions.

• simulation.step() – advance a simulation by a single step, given the state at the last step.

Both are differentiable with respect to their vector field inputs, letting you backpropagate through the whole simulation process.

The high-level API is available through deltaflow.* as well as deltaflow.simulation.*.

simulation.simulate(timesteps: int, color: jax.numpy.ndarray, velocity: Optional[jax.numpy.ndarray] = None, force: Optional[jax.numpy.ndarray] = None, config: simulation.SimulationConfig = SimulationConfig(delta_t=0.05, density_coeff=1.0, diffusion_coeff=0.001, pressure_iterations=16), return_frames: bool = True, disable_progress_bar: bool = False) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

Run a multi-timestep fluid simulation.

The value of return_frames determines what this function does:

• If True (default), every frame’s color and velocity are returned. This enables animation.

• If False, only the last frame’s color and velocity are returned. This is faster and supports computing gradients through the whole simulation.

All fields must match in their spatial dimensions.

Parameters

• timesteps – The number of timesteps to run the simulation for.

• color – The RGB color field of the fluid at the first frame. Shape: [y, x, 3].

• velocity – The velocity field of the fluid at the first frame. Shape: [y, x, y/x]. If None (default), initialize a zero velocity field.

• force – A static force field to apply at each frame. Shape: [y, x, y/x]. If None (default), initialize a zero force field.

• config – The simulation’s physical configuration. Defaults to a reasonable default configuration.

• return_frames – If True (default), return all color and velocity values as numpy arrays, with timesteps on axis 0. If False, return only the last color and velocity fields.

• disable_progress_bar – If True, disable the printed progress bar.

Returns

• color (jnp.ndarray) – If return_frames is False: the final color field, with shape [y, x, 3]. Otherwise: each frame’s color field, with shape [timesteps, y, x, 3].

• velocity (jnp.ndarray) – If return_frames is False: the final velocity field, with shape [y, x, 2]. Otherwise: each frame’s velocity field, with shape [timesteps, y, x, 2].

simulation.step(color: jax.numpy.ndarray, velocity: jax.numpy.ndarray, force: jax.numpy.ndarray, pressure: jax.numpy.ndarray, config: simulation.SimulationConfig = SimulationConfig(delta_t=0.05, density_coeff=1.0, diffusion_coeff=0.001, pressure_iterations=16)) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]

Advance the simulation by a single timestep.

Supports differentiation through all arguments except config. All fields must match in their spatial dimensions.

This function is JITted, since compiling multiple steps together doesn’t yield much speedup. Using a different spatial resolution or config will trigger recompilation; doing this often will slow down simulation dramatically.

Parameters

• color – The RGB color field of the fluid at the last frame. Shape: [y, x, 3].

• velocity – The velocity field of the fluid at the last frame. Shape: [y, x, y/x].

• force – A static force field to apply at this frame. Shape: [y, x, y/x].

• pressure – The pressure field of the fluid at the last frame. Shape: [y, x]. If this was the first frame, zeroes may be passed instead.

• config – The simulation’s physical configuration. Statically traced, so changing this argument triggers recompilation. Defaults to a reasonable default configuration.

Returns

• color (jnp.ndarray) – The RGB color field of the fluid at the next frame. Shape: [y, x, 3].

• velocity (jnp.ndarray) – The velocity field of the fluid at the next frame. Shape: [y, x, y/x].

• pressure (jnp.ndarray) – The pressure field of the fluid at the next frame. Shape: [y, x].

class simulation.SimulationConfig(delta_t: float = 0.05, density_coeff: float = 1.0, diffusion_coeff: float = 0.001, pressure_iterations: int = 16)

A configuration object determining a simulation’s physical properties.

Parameters

• delta_t (float) – The time elapsed in each timestep. Smaller values produce more accurate results but advance the simulation slower.

• density_coeff (float) – A coefficient on fluid density. Higher values cause fluids to respond slower to pressure gradients. Values too far from 1.0 may produce unrealistic effects.

• diffusion_coeff (float) – A coefficient on diffusion rate and viscosity. Higher values cause faster diffusion and greater viscosity; values should be very small. A value of 0.0 simulates an inviscid flow and simulates faster.

• pressure_iterations (int) – The number of iterations used to compute pressure. Larger values produce more accurate results but simulate slower.

Low-level simulation functions

This API exposes the internals used for individual simulation components. Use at your own risk!

simulation._get_predecessor_coordinates(velocity: jax.numpy.ndarray, delta_t: float) → jax.numpy.ndarray

For each point on the grid, get the fractional coordinates of a particle that would move to the center of that gridsquare at the next timestep.

Parameters

• velocity – The fluid velocity field. Shape: [y, x, y/x].

• delta_t – The time elapsed in each timestep.

Returns

For each grid square, the (fractional) coordinates of the grid square that would move to its center at the next timestep. Shape: [y/x, y, x].

Return type

jnp.ndarray

simulation._advect(field: jax.numpy.ndarray, predecessor_coords: jax.numpy.ndarray) → jax.numpy.ndarray

Transport a vector field by reading values moving into the center of each square.

Parameters

• field – The vector field to transport. Shape: [y, x, any].

• predecessor_coords – The predecessor coordinates computed from the velocity. Shape: [y/x, y, x].

Returns

The field, advected by one timestep.

Return type

jnp.ndarray

simulation._divergence_2d(field: jax.numpy.ndarray) → jax.numpy.ndarray

Compute the divergence of a 2D vector field.

Parameters

field – Any 2D vector field. Shape: [y, x, 2].

Returns

The divergence of field as a scalar field. Shape: [y, x].

Return type

jnp.ndarray

simulation._compute_pressure(advected_velocity: jax.numpy.ndarray, pressure: jax.numpy.ndarray, pressure_iterations: int) → jax.numpy.ndarray

Compute the (unitless) pressure of a fluid.

Uses Jacobi iteration, initialized with the last frame’s pressure.

Parameters

• advected_velocity – The fluid’s velocity field for this frame, advected but without pressure correction. Shape: [y, x, 2].

• pressure – Last frame’s pressure field, used as an initial guess for the pressure solver. If no such estimate is available (e.g. for the first frame), a zero field may be passed. Shape: [y, x].

• pressure_iterations – The number of iterations used to compute pressure. Must be static during JIT tracing.

Returns

The estimated pressure field for the fluid at this frame. Shape: [y, x].

Return type

jnp.ndarray

simulation._diffuse(field: jax.numpy.ndarray, diffusion_coeff: float, delta_t: float) → jax.numpy.ndarray

Average each value in a vector field closer to its neighbors to simulate diffusion and viscosity.

Parameters

• field – The vector field to diffuse. Shape: [y, x, any].

• diffusion_coeff – A coefficient determining the amount of diffusion at each frame. Must be static during JIT tracing.

• delta_t – The time elapsed in each timestep. Must be static during JIT tracing.

Returns

field, with diffusion applied for this frame.

Return type

jnp.ndarray