ΦFlow
A modern python computational fluid dynamics library for ML research
June 24, 2022 — April 17, 2023
A useful differentiable PDE solver.
ΦFlow: A differentiable PDE solving framework for machine learning (Holl et al. 2020):
- Variety of built-in PDE operations with a focus on fluid phenomena, allowing for concise formulation of simulations.
- Tight integration with PyTorch, Jax, and TensorFlow for straightforward neural network training with fully differentiable simulations that can run on the GPU.
- Flexible, easy-to-use web interface featuring live visualisations and interactive controls that can affect simulations or network training on the fly.
- Object-oriented, vectorised design for expressive code, ease of use, flexibility, and extensibility.
- Reusable simulation code, independent of backend and dimensionality, i.e. the exact same code can run a 2D fluid sim using NumPy and a 3D fluid sim on the GPU using TensorFlow or PyTorch.
- High-level linear equation solver with automated sparse matrix generation.
Phiflow seems to have fewer elaborate PDEs built-in than Mantaflow but has more intimate and flexible ML integration and more active development. It’s featured in various papers from the TUM group (Holl, Thuerey, and Koltun 2020; Um and Holl 2021; Um et al. 2021).
The documentation is dispersed and confusing. Crucial facts are scattered across video tutorials, idiosyncratic and rapidly outdated API docs, tutorials, demos, and manuals. This is not to disrespect Phillip’s amazing work on this package. The fact that there is documentation places this product well ahead of many academic research competitors; I’m merely setting expectations. It’s a lovely package, in its quirky, hipster, and artisanal way.
It reinvents a few wheels while trying to be helpful and sometimes clashes with the needs of ML, with lots of opinionated design choices. Pet peeve: providing a unified API over various toolkits, which makes 80% of PDE tasks easy and the remaining 20% confusing. I’m trying to figure out how easy it is to manually stitch together PDEs and NNs and propagate gradients; I spend much time working around the convenient generic wrappers to get what I actually need.
Most of the documentation on this page is about the Phiflow v2 API, tested on Phiflow 2.1.4 2.2.7 2.5.1.
1 Documentation central
The YouTube tutorials are useful. The “main” docs site is guide-like. API reference docs are stored separately: phi API documentation.
NB: Now (v2.5) a substantial part of the functionality has been broken into a generic numerical library, PhiML, so part of the documentation is filed there.
2 Arrays in Phiflow
A.k.a. Tensors, which are wrappers around the tensor objects in whatever math backend Phiflow is using.
Arrays have mandatory names which are used in broadcasting rules. This is best explained in the video tutorial, Working with Tensors. Broadcasting between objects with “spatial”, “batch”, “instance” and “channel” dimensions is largely automatic. I guess “temporal” is implied somehow? Maybe via Scene objects? I do not use those.
This system is neat and flexible, but unfortunately, I spend much time working around it because I need to drop into the specialised mathematical tools of specific NN toolkits to get things done. Ultimately, I still need to choose an interpretation for my tensors, so PhiFlow’s attempt is valiant but ultimately burdensome.
Documentation was previously at phi.math API documentation. Now that mathematical functionality has moved to PhiML, refer to phiml.math API documentation.
3 Fields in Phiflow
A sampling grid plus some sample data at grid locations gives us a SampledField. We could imagine fields defined by some general mathematical function. The documentation references an AnalyticField, but it appears not to be implemented, so consider everything as SampledField for now. CenteredGrid objects are straightforward.
Slightly fancy feature — Grids are not necessarily in the centre but may be offset to show values at the boundaries. This is called the Staggered Grid and is useful for velocity fields.
3.1 Sampling
Fields can be sampled in a few ways. I am trying to learn which are differentiable. Field.at is probably usually what I want?
But there are options. This also works:
inf_field = field.CenteredGrid(
field, x=16, y=16, bounds=field.bounds,
extrapolation=field.extrapolation)Sampling overview explains more.
Various sample methods take us from a Field to a Tensor. math.downsample2x and math.sample_subgrid address special grid relations. field.sample and field.reduce_sample seem to accept arbitrary geometries (?), which is useful although I am still confused about how to specify useful Geometries.
4 Actual physics
The actual physics part is what Phiflow makes easy.
- phi.physics API documentation:
- Fluids Tutorial: Introduction to core classes and fluid-related functions.
- Overview: Domains, built-in physics functions.
- Functions for Fluid Simulations: Advection, projection, diffusion
5 Optimisation
ΦFlow supports two types of optimisation with different APIs. One is the ΦFlow native optimisation, which optimises Fields and ΦFlow Tensor objects, copying the scipy.minimize interface.
The other is the NN-style SGD training, which optimises NN parameters. This looks like a normal SGD training loop, as per
As usual with the scipy.minimize style system, there is limited scope to see what is happening during the optimisation. There is an example showing how to do that better in the Physics-based Deep Learning textbook Burgers Optimization with a Differentiable Physics Gradient, although it uses an outdated record_gradients API. A shorter but more modern example is in the cookbook.
6 ML
The API is idiosyncratic and rapidly evolving. Best explained through examples:
7 Visualisation
Messy. They created a lovely UI for controlling simulations interactively, but it is messy and unsatisfactory because python GUIs suck and also jupyter GUIs suck, but trying to serve two sucky masters is tedious.
As per the advice of lead developer Phillip Holl, I ignore the entire Vis system which only works from command-line scripts. I plot inside jupyter notebooks for now. If I wanted something more sophisticated I might use the ΦFlow Web Interface which integrates with dash. I am not sure why they do not just use one of the fairly standard tools for ML experiment tracking and visualisation such as tensorboard or whatever. The developers of those tools have already experienced the many irritations of trying to do this stuff interactively and found workarounds.
8 Efficiency/GPUs
I find it congenial to use the excessive ‘Kill it with fire’ method to moderate ΦFlow’s ambition regarding GPUs. In some versions, it is too reluctant to use CUDA, in others too averse. e.g. in 2.2.7, even if I am not using CUDA, math.seed will still try to use the GPU which is rude behaviour on shared machines.
Keep it simple by avoiding ambiguity:
GPU mode:
os.environ.pop("CUDA_VISIBLE_DEVICES", None) # Dangerous on shared machines!
TORCH.set_default_device('GPU')
PHI_DEVICE = 'GPU' # pass this to Phiflow functions
DEVICE = torch.device('cuda') # pass this to pytorch functionsCPU mode:
9 Useful examples
10 Data storage
I do not use the native Phiflow system, since I store everything in hdf5. But it is nice that the API is documented: