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)})

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[1], line 4
      1 from tqdm.notebook import trange
      2 from phi.jax.flow import *  # imports sub-modules + core classes
----> 4 velocity = StaggeredGrid(Noise(), 'periodic', x=64, y=96)
      5 plot({"velocity": velocity, "vorticity": field.curl(velocity)})

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_grid.py:158, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_)
    156             assert set(resolution_from_staggered_tensor(values, extrapolation)) == set(resolution), f"Failed to create StaggeredGrid: values {values.shape} do not match given resolution {resolution} for extrapolation {extrapolation}. See https://tum-pbs.github.io/PhiFlow/Staggered_Grids.html"
    157 elif isinstance(values, (Geometry, Field, FieldInitializer)):
--> 158     values = sample(values, elements, at='face', boundary=extrapolation, dot_face_normal=elements)
    159 elif callable(values):
    160     values = sample_function(values, elements, 'face', extrapolation)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_resample.py:157, in sample(field, geometry, at, boundary, dot_face_normal, **kwargs)
    155         raise NotImplementedError
    156 elif field.is_grid and field.is_centered:
--> 157     return sample_grid_at_faces(field, geometry, boundary, **kwargs)
    158 elif field.is_grid and field.is_staggered:
    159     faces = math.unstack(slice_off_constant_faces(geometry.faces, geometry.boundary_faces, boundary), dual)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_resample.py:278, in sample_grid_at_faces(field, elements, extrapolation, order, implicit)
    276 if isinstance(elements, UniformGrid):
    277     sampled = {dim: sample_grid_at_centers(field, g, order=order, implicit=implicit) for dim, g in elements.staggered_cells(extrapolation).items()}
--> 278     return math.stack(sampled, dual(elements.face_shape))
    279 raise NotImplementedError

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:174, in stack(values, dim, expand_values, simplify, layout_non_matching, **kwargs)
    172 for v in values:
    173     if hasattr(v, '__stack__'):
--> 174         result = v.__stack__(values, dim, **kwargs)
    175         if result is not NotImplemented:
    176             if DEBUG_CHECKS:

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:515, in Tensor.__stack__(values, dim, **_kwargs)
    513     return layout_.__stack__(values, dim, **_kwargs)
    514 from ._ops import stack_tensors
--> 515 return stack_tensors(values, dim)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_ops.py:902, in stack_tensors(values, dim)
    900     return values[0]
    901 values = [wrap(v) for v in values]
--> 902 values = [squeeze(v, dim) for v in values]
    903 values = cast_same(*values)
    904 # --- sparse to dense ---

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:524, in squeeze(x, dims)
    522 if not dims:
    523     return x
--> 524 assert dims.volume == 1, f"Cannot squeeze non-singleton dims {dims} from {x}"
    525 return x[{d: 0 for d in dims.names}]

AssertionError: Cannot squeeze non-singleton dims (~vectorᵈ=x,y) from (~vectorᵈ=x,y, xˢ=64, yˢ=96) 0.00e+00 ± 9.9e-01 (-4e+00...3e+00)

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, rank_deficiency=0))
    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)})

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[2], line 8
      5     v, p = fluid.make_incompressible(v, (), Solve(x0=p, rank_deficiency=0))
      6     return v, p
----> 8 velocity0, pressure0 = fluid.make_incompressible(velocity)
      9 velocity1, pressure1 = operator_split_step(velocity0, None, dt=1.)
     10 plot({'initial vorticity': field.curl(velocity0), 'after step': field.curl(velocity1)})

NameError: name 'velocity' is not defined

@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, rank_deficiency=0))
    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)})

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[3], line 8
      5     v, p = fluid.make_incompressible(v, (), Solve(x0=p, rank_deficiency=0))
      6     return v, p
----> 8 velocity0, pressure0 = fluid.make_incompressible(velocity)
      9 velocity1, pressure1 = operator_split_step(velocity0, None, dt=1.)
     10 plot({'initial vorticity': field.curl(velocity0), 'after step': field.curl(velocity1)})

NameError: name 'velocity' is not defined

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)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[4], line 1
----> 1 velocity_trj, pressure_trj = iterate(operator_split_step, batch(time=100), velocity0, pressure0, dt=1., range=trange)

NameError: name 'velocity0' is not defined

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)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[5], line 1
----> 1 plot(field.curl(velocity_trj), animate='time', same_scale=False)

NameError: name 'velocity_trj' is not defined

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 = CenteredGrid(Noise(vector='x,y'), 'periodic', x=64, y=96)
velocity0, pressure0 = fluid.make_incompressible(velocity, order=4)
velocity_trj, pressure_trj = iterate(rk4_step, batch(time=100), velocity0, pressure0, dt=.5, substeps=2, range=trange)
plot(field.curl(velocity_trj), animate='time', same_scale=False)

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[7], line 6
      3     return fluid.incompressible_rk4(momentum_equation, v, p, dt, pressure_order=4)
      5 velocity = CenteredGrid(Noise(vector='x,y'), 'periodic', x=64, y=96)
----> 6 velocity0, pressure0 = fluid.make_incompressible(velocity, order=4)
      7 velocity_trj, pressure_trj = iterate(rk4_step, batch(time=100), velocity0, pressure0, dt=.5, substeps=2, range=trange)
      8 plot(field.curl(velocity_trj), animate='time', same_scale=False)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/physics/fluid.py:137, in make_incompressible(velocity, obstacles, solve, active, order, correct_skew, wide_stencil)
    135     active = accessible.with_extrapolation(extrapolation.NONE) if active is None else active * accessible  # no pressure inside obstacles
    136     velocity = apply_boundary_conditions(velocity, obstacles)
--> 137 div = divergence(velocity, order=order)
    138 if active is not None:
    139     div *= active  # inactive cells must solvable

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field_math.py:635, in divergence(field, order, implicit, upwind, implicitness)
    631         return result
    633 else:
    634     components = [
--> 635         spatial_gradient(f, dims=dim, at='center', order=order, implicit=implicit, implicitness=implicitness,
    636                          stack_dim="sum:b").sum[0] for f, dim in zip(field.vector, field.shape.only(spatial).names)]
    638 return sum(components)

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

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field_math.py:419, in perform_finite_difference_operation(field, dim, differentiation_order, dx, ext, output_ext, at, order, implicit, implicitness)
    416 one_sided_ext = extrapolation.map(lambda e: extrapolation.ONE if is_one_sided(e) else extrapolation.ZERO, ext)
    417 one_sided_mask = standard_mask.with_extrapolation(one_sided_ext)
--> 419 result = apply_stencils(field, ext, output_ext, dx, base_values, base_shifts, at, dim,
    420                         masks=(ext_valid_masks, output_valid_mask, one_sided_mask),
    421                         stencil_tensors=expl_one_sided_stencil_tensor,
    422                         differencing_order=differentiation_order)
    424 if implicit:
    425     implicit.x0 = result

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_functional.py:468, in LinearFunction.__call__(self, *args, **kwargs)
    466 matrix, bias, (out_tree, out_tensors) = self._get_or_trace(key, args, aux_kwargs)
    467 result = matrix @ tensors[0] + bias
--> 468 out_tensors = list(out_tensors)
    469 out_tensors[0] = result
    470 return assemble_tree(out_tree, out_tensors)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:783, in Tensor.__iter__(self)
    781     return iter([self.native()])
    782 else:
--> 783     native = self.native([self.shape])
    784     return iter(native)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_lin_trace.py:113, in LinTracer.native(self, order, force_expand)
    112 def native(self, order: Union[str, tuple, list, Shape] = None, force_expand=True):
--> 113     raise NotImplementedError

NotImplementedError: 

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.