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.13/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, 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)})

/opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/physics/fluid.py:155: RuntimeWarning: Rank deficiency >= 1 detected in linear solve. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).
  pressure = math.solve_linear(masked_laplace, div, solve, velocity.boundary, hard_bcs, active, wide_stencil=wide_stencil, order=order, implicit=None, upwind=None, correct_skew=correct_skew)
/opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/backend/_linalg.py:348: SparseEfficiencyWarning: spsolve requires A be CSC or CSR matrix format
  x = spsolve(matrix, y[batch])  # returns nan when diverges

---------------------------------------------------------------------------
Diverged                                  Traceback (most recent call last)
Cell In[2], line 8
      4     v = diffuse.explicit(v, viscosity, dt)
      5     v, p = fluid.make_incompressible(v, (), Solve(x0=p, rank_deficiency=0))
      6     return v, p
      7 
----> 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)})

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/physics/fluid.py:155, in make_incompressible(velocity, obstacles, solve, active, order, correct_skew, wide_stencil)
    153 if wide_stencil is None:
    154     wide_stencil = not velocity.is_staggered
--> 155 pressure = math.solve_linear(masked_laplace, div, solve, velocity.boundary, hard_bcs, active, wide_stencil=wide_stencil, order=order, implicit=None, upwind=None, correct_skew=correct_skew)
    156 # --- Subtract grad p ---
    157 grad_pressure = field.spatial_gradient(pressure, input_velocity.extrapolation, at=velocity.sampled_at, order=order, scheme='green-gauss')

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:694, in solve_linear(f, y, solve, grad_for_f, f_kwargs, *f_args, **f_kwargs_)
    691         return result  # must return exactly `x` so gradient isn't computed w.r.t. other quantities
    693     _matrix_solve = attach_gradient_solve(_matrix_solve_forward, auxiliary_args=f'is_backprop,solve{",matrix" if matrix.backend == NUMPY else ""}', matrix_adjoint=grad_for_f)
--> 694     return _matrix_solve(assemble_tree(y_tree, y_tensors), solve, matrix)
    695 else:  # Matrix-free solve
    696     from ._ops import cached

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:972, in CustomGradientFunction.__call__(self, *args, **kwargs)
    967             if len(self.traces) >= 8:
    968                 warnings.warn(f"""{self.__name__} has been traced {len(self.traces)} times.
    969 To avoid memory leaks, call {f_name(self.f)}.traces.clear(), {f_name(self.f)}.recorded_mappings.clear().
    970 Traces can be avoided by jit-compiling the code that calls custom gradient functions.
    971 """, RuntimeWarning, stacklevel=2)
--> 972         native_result = self.traces[key](*natives)  # With PyTorch + jit, this does not call forward_native every time
    973         output_key = match_output_signature(key, self.recorded_mappings, self)
    974         output_tensors = assemble_tensors(native_result, output_key.specs)

    [... skipping hidden 8 frame]

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:928, in CustomGradientFunction._trace.<locals>.forward_native(*natives)
    926 kwargs = assemble_tree(in_key.tree, in_tensors, attr_type=variable_attributes)
    927 ML_LOGGER.debug(f"Running forward pass of custom op {forward_native.__name__} given args {tuple(kwargs.keys())} containing {len(natives)} native tensors")
--> 928 result = self.f(**kwargs, **in_key.auxiliary_kwargs)  # Tensor or tuple/list of Tensors
    929 nest, out_tensors = disassemble_tree(result, cache=True, attr_type=variable_attributes)
    930 result_natives, result_shapes, specs = disassemble_tensors(out_tensors, expand=True)

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:690, in solve_linear.<locals>._matrix_solve_forward(y, solve, matrix, is_backprop)
    688     rx, ry, ex, ey = reduce_y, reduce_x, expand_y, expand_x
    689     # pattern_dims are already switched due to matrix transposition
--> 690 result = _linear_solve_forward(y, solve, nat_matrix, pattern_dims_in, pattern_dims_out, preconditioner, backend, is_backprop, rx, ry, ex, ey)
    691 return result

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:818, in _linear_solve_forward(y, solve, native_lin_op, pattern_dims_in, pattern_dims_out, preconditioner, backend, is_backprop, reduce_x, reduce_y, expand_x, expand_y)
    816 for tape in _SOLVE_TAPES:
    817     tape._add(solve, trj, result)
--> 818 result.convergence_check(is_backprop and 'TensorFlow' in backend.name)  # raises ConvergenceException
    819 return final_x

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:203, in SolveInfo.convergence_check(self, only_warn)
    201             warnings.warn(self.msg, ConvergenceWarning)
    202         else:
--> 203             raise Diverged(self)
    204 if not self.converged.trajectory[-1].all:
    205     if NotConverged not in self.solve.suppress:

Diverged: Direct solution does not satisfy tolerance: norm(residual)=0.0003806906461250037

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

---------------------------------------------------------------------------
Diverged                                  Traceback (most recent call last)
Cell In[3], line 8
      4     v = diffuse.explicit(v, 0.1, dt)
      5     v, p = fluid.make_incompressible(v, (), Solve(x0=p, rank_deficiency=0))
      6     return v, p
      7 
----> 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)})

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/physics/fluid.py:155, in make_incompressible(velocity, obstacles, solve, active, order, correct_skew, wide_stencil)
    153 if wide_stencil is None:
    154     wide_stencil = not velocity.is_staggered
--> 155 pressure = math.solve_linear(masked_laplace, div, solve, velocity.boundary, hard_bcs, active, wide_stencil=wide_stencil, order=order, implicit=None, upwind=None, correct_skew=correct_skew)
    156 # --- Subtract grad p ---
    157 grad_pressure = field.spatial_gradient(pressure, input_velocity.extrapolation, at=velocity.sampled_at, order=order, scheme='green-gauss')

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:694, in solve_linear(f, y, solve, grad_for_f, f_kwargs, *f_args, **f_kwargs_)
    691         return result  # must return exactly `x` so gradient isn't computed w.r.t. other quantities
    693     _matrix_solve = attach_gradient_solve(_matrix_solve_forward, auxiliary_args=f'is_backprop,solve{",matrix" if matrix.backend == NUMPY else ""}', matrix_adjoint=grad_for_f)
--> 694     return _matrix_solve(assemble_tree(y_tree, y_tensors), solve, matrix)
    695 else:  # Matrix-free solve
    696     from ._ops import cached

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:972, in CustomGradientFunction.__call__(self, *args, **kwargs)
    967             if len(self.traces) >= 8:
    968                 warnings.warn(f"""{self.__name__} has been traced {len(self.traces)} times.
    969 To avoid memory leaks, call {f_name(self.f)}.traces.clear(), {f_name(self.f)}.recorded_mappings.clear().
    970 Traces can be avoided by jit-compiling the code that calls custom gradient functions.
    971 """, RuntimeWarning, stacklevel=2)
--> 972         native_result = self.traces[key](*natives)  # With PyTorch + jit, this does not call forward_native every time
    973         output_key = match_output_signature(key, self.recorded_mappings, self)
    974         output_tensors = assemble_tensors(native_result, output_key.specs)

    [... skipping hidden 8 frame]

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:928, in CustomGradientFunction._trace.<locals>.forward_native(*natives)
    926 kwargs = assemble_tree(in_key.tree, in_tensors, attr_type=variable_attributes)
    927 ML_LOGGER.debug(f"Running forward pass of custom op {forward_native.__name__} given args {tuple(kwargs.keys())} containing {len(natives)} native tensors")
--> 928 result = self.f(**kwargs, **in_key.auxiliary_kwargs)  # Tensor or tuple/list of Tensors
    929 nest, out_tensors = disassemble_tree(result, cache=True, attr_type=variable_attributes)
    930 result_natives, result_shapes, specs = disassemble_tensors(out_tensors, expand=True)

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:690, in solve_linear.<locals>._matrix_solve_forward(y, solve, matrix, is_backprop)
    688     rx, ry, ex, ey = reduce_y, reduce_x, expand_y, expand_x
    689     # pattern_dims are already switched due to matrix transposition
--> 690 result = _linear_solve_forward(y, solve, nat_matrix, pattern_dims_in, pattern_dims_out, preconditioner, backend, is_backprop, rx, ry, ex, ey)
    691 return result

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:818, in _linear_solve_forward(y, solve, native_lin_op, pattern_dims_in, pattern_dims_out, preconditioner, backend, is_backprop, reduce_x, reduce_y, expand_x, expand_y)
    816 for tape in _SOLVE_TAPES:
    817     tape._add(solve, trj, result)
--> 818 result.convergence_check(is_backprop and 'TensorFlow' in backend.name)  # raises ConvergenceException
    819 return final_x

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:203, in SolveInfo.convergence_check(self, only_warn)
    201             warnings.warn(self.msg, ConvergenceWarning)
    202         else:
--> 203             raise Diverged(self)
    204 if not self.converged.trajectory[-1].all:
    205     if NotConverged not in self.solve.suppress:

Diverged: Direct solution does not satisfy tolerance: norm(residual)=0.0003806906461250037

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)

---------------------------------------------------------------------------
Diverged                                  Traceback (most recent call last)
Cell In[7], line 6
      2 def rk4_step(v, p, dt):
      3     return fluid.incompressible_rk4(momentum_equation, v, p, dt, pressure_order=4)
      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.13/x64/lib/python3.12/site-packages/phi/physics/fluid.py:155, in make_incompressible(velocity, obstacles, solve, active, order, correct_skew, wide_stencil)
    153 if wide_stencil is None:
    154     wide_stencil = not velocity.is_staggered
--> 155 pressure = math.solve_linear(masked_laplace, div, solve, velocity.boundary, hard_bcs, active, wide_stencil=wide_stencil, order=order, implicit=None, upwind=None, correct_skew=correct_skew)
    156 # --- Subtract grad p ---
    157 grad_pressure = field.spatial_gradient(pressure, input_velocity.extrapolation, at=velocity.sampled_at, order=order, scheme='green-gauss')

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:694, in solve_linear(f, y, solve, grad_for_f, f_kwargs, *f_args, **f_kwargs_)
    691         return result  # must return exactly `x` so gradient isn't computed w.r.t. other quantities
    693     _matrix_solve = attach_gradient_solve(_matrix_solve_forward, auxiliary_args=f'is_backprop,solve{",matrix" if matrix.backend == NUMPY else ""}', matrix_adjoint=grad_for_f)
--> 694     return _matrix_solve(assemble_tree(y_tree, y_tensors), solve, matrix)
    695 else:  # Matrix-free solve
    696     from ._ops import cached

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:972, in CustomGradientFunction.__call__(self, *args, **kwargs)
    967             if len(self.traces) >= 8:
    968                 warnings.warn(f"""{self.__name__} has been traced {len(self.traces)} times.
    969 To avoid memory leaks, call {f_name(self.f)}.traces.clear(), {f_name(self.f)}.recorded_mappings.clear().
    970 Traces can be avoided by jit-compiling the code that calls custom gradient functions.
    971 """, RuntimeWarning, stacklevel=2)
--> 972         native_result = self.traces[key](*natives)  # With PyTorch + jit, this does not call forward_native every time
    973         output_key = match_output_signature(key, self.recorded_mappings, self)
    974         output_tensors = assemble_tensors(native_result, output_key.specs)

    [... skipping hidden 8 frame]

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:928, in CustomGradientFunction._trace.<locals>.forward_native(*natives)
    926 kwargs = assemble_tree(in_key.tree, in_tensors, attr_type=variable_attributes)
    927 ML_LOGGER.debug(f"Running forward pass of custom op {forward_native.__name__} given args {tuple(kwargs.keys())} containing {len(natives)} native tensors")
--> 928 result = self.f(**kwargs, **in_key.auxiliary_kwargs)  # Tensor or tuple/list of Tensors
    929 nest, out_tensors = disassemble_tree(result, cache=True, attr_type=variable_attributes)
    930 result_natives, result_shapes, specs = disassemble_tensors(out_tensors, expand=True)

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:690, in solve_linear.<locals>._matrix_solve_forward(y, solve, matrix, is_backprop)
    688     rx, ry, ex, ey = reduce_y, reduce_x, expand_y, expand_x
    689     # pattern_dims are already switched due to matrix transposition
--> 690 result = _linear_solve_forward(y, solve, nat_matrix, pattern_dims_in, pattern_dims_out, preconditioner, backend, is_backprop, rx, ry, ex, ey)
    691 return result

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:818, in _linear_solve_forward(y, solve, native_lin_op, pattern_dims_in, pattern_dims_out, preconditioner, backend, is_backprop, reduce_x, reduce_y, expand_x, expand_y)
    816 for tape in _SOLVE_TAPES:
    817     tape._add(solve, trj, result)
--> 818 result.convergence_check(is_backprop and 'TensorFlow' in backend.name)  # raises ConvergenceException
    819 return final_x

File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:203, in SolveInfo.convergence_check(self, only_warn)
    201             warnings.warn(self.msg, ConvergenceWarning)
    202         else:
--> 203             raise Diverged(self)
    204 if not self.converged.trajectory[-1].all:
    205     if NotConverged not in self.solve.suppress:

Diverged: Direct solution does not satisfy tolerance: norm(residual)=0.009055467322468758

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.