Writing Fluid Simulations in Φ_Flow

There are two main viewpoints for simulating fluids:

• Eulerian simulations use grids, tracking fluid distribution and fluid velocity at fixed sample points
• Lagrangian simulations track particles that move with the fluid.

Φ_Flow supports both methods to some extent but mainly focuses on Eulerian simulations.

Before we discuss the various operations required for fluid simulations, let's define our variables and initial state. In this case, we will create a 64×96 grid, sampling velocity vectors in staggered form and marker values at the centroids.

from tqdm.notebook import trange
from phi.jax.flow import *  # imports sub-modules + core classes

velocity = StaggeredGrid(Noise(), 'periodic', x=64, y=96)
plot({"velocity": velocity, "vorticity": field.curl(velocity)})

/opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phi/vis/_matplotlib/_matplotlib_plots.py:167: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
  plt.tight_layout()  # because subplot titles can be added after figure creation

Operator Splitting

The Navier-Stokes equations for fluids, $\frac{\partial u}{\partial t} = - (u \cdot \nabla) u - \nu \nabla^2 u - \frac 1 \rho \nabla p + g$, comprise multiple terms.

Operator splitting enables writing fast and stable fluid simulations by sequentially evaluating the different terms. For each of the terms, Φ_Flow provides functions to compute them:

• Advection: advect.semi_lagrangian [Stam 1999], advect.mac_cormack [MacCormack 2002]
• Diffusion: diffuse.explicit, diffuse.implicit
• Pressure projection: fluid.make_incompressible [Chorin and Temam 1968]

All of these functions take in a state variable and return the new state after a certain time dt has passed. In the following example, the velocity is self-advected and made incompressible, while the marker is passively advected.

@jit_compile
def operator_split_step(v, p, dt, viscosity=0.1):
    v = advect.semi_lagrangian(v, v, dt)  # velocity self-advection
    v = diffuse.explicit(v, viscosity, dt)
    v, p = fluid.make_incompressible(v, (), Solve(x0=p))
    return v, p

velocity0, pressure0 = fluid.make_incompressible(velocity)
velocity1, pressure1 = operator_split_step(velocity0, None, dt=1.)
plot({'initial vorticity': field.curl(velocity0), 'after step': field.curl(velocity1)})

/opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).
  warnings.warn("Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).")
/opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).
  warnings.warn("Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).")
/opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phi/vis/_matplotlib/_matplotlib_plots.py:167: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
  plt.tight_layout()  # because subplot titles can be added after figure creation

@jit_compile
def operator_split_step(v, p, dt):
    v = advect.semi_lagrangian(v, v, dt)  # velocity self-advection
    v = diffuse.explicit(v, 0.1, dt)
    v, p = fluid.make_incompressible(v, (), Solve(x0=p))
    return v, p

velocity0, pressure0 = fluid.make_incompressible(velocity)
velocity1, pressure1 = operator_split_step(velocity0, None, dt=1.)
plot({'initial vorticity': field.curl(velocity0), 'after step': field.curl(velocity1)})

/opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).
  warnings.warn("Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).")
/opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phi/vis/_matplotlib/_matplotlib_plots.py:167: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
  plt.tight_layout()  # because subplot titles can be added after figure creation

We can use iterate to compute a trajectory by repeatedly calling operator_split_step. All intermediate states are stacked along the specified dimension which we call time.

velocity_trj, pressure_trj = iterate(operator_split_step, batch(time=100), velocity0, pressure0, dt=1., range=trange)

/opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).
  warnings.warn("Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).")

Alternatively, we could have written a for loop, added all intermediate states to a list, and stacked the results afterward. Now, let's plot this trajectory by animating the time dimension.

plot(field.curl(velocity_trj), animate='time', same_scale=False)

<Figure size 640x480 with 0 Axes>

Higher-Order Simulations

The operator splitting approach is not compatible with more accurate numerical schemes. For more accurate simulations, we can use higher-order spatial schemes as well as time integration. In that case, we define a momentum equation which computes the PDE terms directly, without integrating them in time. The following example computes explicit fourth-order accurate advection and diffusion.

def momentum_equation(v, viscosity=0.1):
    advection = advect.finite_difference(v, v, order=4, implicit=None)
    diffusion = diffuse.finite_difference(v, viscosity, order=4, implicit=None)
    return advection + diffusion

Next, we perform time integration with the incompressibility constraint. This is considerably more expensive than the previous approach but yields much more accurate results.

@jit_compile
def rk4_step(v, p, dt):
    return fluid.incompressible_rk4(momentum_equation, v, p, dt, pressure_order=4)

velocity_trj, pressure_trj = iterate(rk4_step, batch(time=200), velocity0, pressure0, dt=.5, range=trange)
plot(field.curl(velocity_trj.time[::2]), animate='time', same_scale=False)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[7], line 5
      1 @jit_compile
      2 def rk4_step(v, p, dt):
      3     return fluid.incompressible_rk4(momentum_equation, v, p, dt, pressure_order=4)
----> 5 velocity_trj, pressure_trj = iterate(rk4_step, batch(time=200), velocity0, pressure0, dt=.5, range=trange)
      6 plot(field.curl(velocity_trj.time[::2]), animate='time', same_scale=False)

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/math/_functional.py:1300, in iterate(map_function, iterations, f_kwargs, range, measure, substeps, *x0, **f_kwargs_)
   1298 for _i in range(iterations.size):
   1299     for _sub_i in builtin_range(substeps):
-> 1300         x = map_function(*x[:len(x0)], **f_kwargs)
   1301         x = x if isinstance(x, tuple) else (x,)
   1302         if len(x) < len(x0):

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/math/_functional.py:308, in JitFunction.__call__(self, *args, **kwargs)
    303             if len(self.traces) >= 10:
    304                 warnings.warn(f"""Φ-ML: The jit-compiled function '{f_name(self.f)}' was traced {len(self.traces)} times.
    305 Performing many traces may be slow and cause memory leaks.
    306 Re-tracing occurs when the number or types of arguments vary, tensor shapes vary between calls or different auxiliary arguments are given (compared by reference).
    307 Set forget_traces=True to avoid memory leaks when many traces are required. Tracing reason: {trace_check(self, *args, **kwargs)[1]}""", RuntimeWarning)
--> 308             all_natives = native_jit_function(*natives)  # this appends out_key to recorded_mappings
    309             out_key = self.recorded_mappings[key][-1]
    310             buffer_config = out_key.buffer_config

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/backend/jax/_jax_backend.py:252, in JaxBackend.jit_compile.<locals>.run_jit_f(*args)
    249 def run_jit_f(*args):
    250     # print(jax.make_jaxpr(f)(*args))
    251     ML_LOGGER.debug(f"JaxBackend: running jit-compiled '{f.__name__}' with shapes {[self.shape(arg) for arg in args]} and dtypes {[self.dtype(arg) for arg in args]}")
--> 252     return self.as_registered.call(jit_f, *args, name=f"run jit-compiled '{f.__name__}'")

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/backend/_backend.py:415, in Backend.call(self, f, name, *args)
    406 def call(self, f: Callable, *args, name=None):
    407     """
    408     Calls `f(*args)` and returns the result.
    409     This method may be used to register internal calls with the profiler.
   (...)
    413         choose_backend(key).call(custom_function, *args)
    414     """
--> 415     return f(*args)

    [... skipping hidden 11 frame]

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phiml/math/_functional.py:254, in JitFunction._jit_compile.<locals>.jit_f_native(*natives)
    252     aux_tensors = assemble_tensors(aux_natives, aux_specs)
    253     attached_aux_args_ = assemble_tree(aux_tree, aux_tensors)
--> 254 f_output = self.f(**kwargs, **attached_aux_args_)  # Tensor or tuple/list of Tensors
    255 tree, out_tensors = disassemble_tree((f_output, self._extract_tensors), cache=True)
    256 result_natives, result_shapes, specs = disassemble_tensors(out_tensors, expand=True)

Cell In[7], line 3, in rk4_step(v, p, dt)
      1 @jit_compile
      2 def rk4_step(v, p, dt):
----> 3     return fluid.incompressible_rk4(momentum_equation, v, p, dt, pressure_order=4)

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phi/physics/fluid.py:335, in incompressible_rk4(pde, velocity, pressure, dt, pressure_order, pressure_solve, **pde_aux_kwargs)
    333 v1, p1 = velocity, pressure
    334 # PDE at current point
--> 335 rhs1 = pde(v1, **pde_aux_kwargs) - p1.gradient(at=v1.sampled_at, order=pressure_order)
    336 v2_old = velocity + (dt / 2) * rhs1
    337 v2, delta_p = make_incompressible(v2_old, solve=pressure_solve, order=pressure_order)

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phi/field/_field.py:547, in Field.gradient(self, boundary, at, dims, stack_dim, order, implicit, scheme, upwind, gradient_extrapolation)
    545 """Alias for `phi.field.spatial_gradient`"""
    546 from ._field_math import spatial_gradient
--> 547 return spatial_gradient(self, boundary=boundary, at=at, dims=dims, stack_dim=stack_dim, order=order, implicit=implicit, scheme=scheme, upwind=upwind, gradient_extrapolation=gradient_extrapolation)

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phi/field/_field_math.py:244, in spatial_gradient(field, boundary, at, dims, stack_dim, order, implicit, implicitness, scheme, upwind, gradient_extrapolation)
    240         assert stack_dim.name == 'vector'
    241         return stagger(field, lambda lower, upper: (upper - lower) / field.dx.vector.as_dual(), gradient_extrapolation)
    243 result_components = [
--> 244     perform_finite_difference_operation(field.values, dim, 1, field.dx.vector[dim], field.extrapolation,
    245                                         gradient_extrapolation, at, order, implicit, implicitness)
    246     for dim in field.shape.only(grad_dims).names]
    248 if at == 'center':
    249     result = field.with_values(math.stack(result_components, stack_dim))

File /opt/hostedtoolcache/Python/3.12.7/x64/lib/python3.12/site-packages/phi/field/_field_math.py:409, in perform_finite_difference_operation(field, dim, differentiation_order, dx, ext, output_ext, at, order, implicit, implicitness)
    407     standard_mask = CenteredGrid(0, resolution=field.shape.non_batch)  # ToDo ed is this okay with batch dimensions?
    408 else:
--> 409     standard_mask = CenteredGrid(0, resolution=field.shape.non_batch + spatial(
    410         **{dim: sum(output_ext.valid_outer_faces(dim)) - 1}))
    412 output_valid_ext = extrapolation.combine_sides(**{dim: tuple(
    413     extrapolation.ONE if valid_tuple else extrapolation.ZERO for valid_tuple in
    414     output_ext.valid_outer_faces(dim)) for dim in field.shape.spatial.names})
    415 output_valid_mask = standard_mask.with_extrapolation(output_valid_ext)

TypeError: unsupported operand type(s) for +: 'Shape' and 'Shape'

Further Reading

The Kolmogorov flow notebebook shows higher-order fluid flow with forcing.

For a comparison of various schemes in both accuracy and performance is given here.

Coupling between centered and staggered fields can be seen in the smoke plume notebook and Python script.