Differentiable Fluid Simulations with ΦFlow¶
This notebook steps you through setting up fluid simulations and using TensorFlow's differentiation to optimize them.
Execute the cell below to install the ΦFlow Python package from GitHub.
This notebook steps you through setting up fluid simulations and using TensorFlow's differentiation to optimize them.
Execute the cell below to install the ΦFlow Python package from GitHub.
# !pip install --quiet phiflow
from phi.flow import *
ΦFlow is vectorized but object-oriented, i.e. data are represented by Python objects that internally use tensors.
First, we create grids for the quantities we want to simulate. For this example, we require a velocity field and a smoke density field. We sample the smoke field at the cell centers and the velocity in staggered form.
smoke = CenteredGrid(0, extrapolation.BOUNDARY, x=32, y=40, bounds=Box(x=32, y=40)) # sampled at cell centers
velocity = StaggeredGrid(0, extrapolation.ZERO, x=32, y=40, bounds=Box(x=32, y=40)) # sampled in staggered form at face centers
Additionally, we want to add more smoke every time step.
We create the INFLOW field from a circle (2D Sphere) which defines where hot smoke is emitted.
Furthermore, we are interested in running the simulation for different inflow locations.
ΦFlow supports data-parallell execution via batch dimensions.
When a quantity has a batch dimension with size n, operations involving that quantity will be performed n times simultaneously and the result will also have that batch dimension. Here we add the batch dimension inflow_loc.
For an overview of the dimension types, see the documentation or watch the introductory tutorial video.
INFLOW_LOCATION = tensor([(4, 5), (8, 5), (12, 5), (16, 5)], batch('inflow_loc'), channel(vector='x,y'))
INFLOW = 0.6 * CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=3), extrapolation.BOUNDARY, x=32, y=40, bounds=Box(x=32, y=40))
The created grids are instances of the class Grid.
Like tensors, grids also have the shape attribute which lists all batch, spatial and channel dimensions.
Shapes in ΦFlow store not only the sizes of the dimensions but also their names and types.
print(f"Smoke: {smoke.shape}")
print(f"Velocity: {velocity.shape}")
print(f"Inflow: {INFLOW.shape}")
print(f"Inflow, spatial only: {INFLOW.shape.spatial}")
Smoke: (xˢ=32, yˢ=40) Velocity: (xˢ=32, yˢ=40, vectorᶜ=x,y) Inflow: (inflow_locᵇ=4, xˢ=32, yˢ=40) Inflow, spatial only: (xˢ=32, yˢ=40)
The grid values can be accessed using the values property.
print(smoke.values)
print(velocity.values)
print(INFLOW.values)
(xˢ=32, yˢ=40) const 0.0 (~vectorᵈ=x,y, xˢ=~(x=31, y=32) int64, yˢ=~(x=40, y=39) int64) const 0.0 (inflow_locᵇ=4, xˢ=32, yˢ=40) 0.015 ± 0.094 (0e+00...6e-01)
Grids have many more properties which are documented here. Also note that the staggered grid has a non-uniform shape because the number of faces is not equal to the number of cells.
Next, let's do some physics!
Since the initial velocity is zero, we just add the inflow and the corresponding buoyancy force.
For the buoyancy force we use the factor (0, 0.5) to specify strength and direction.
Finally, we project the velocity field to make it incompressible.
Note that the @ operator is a shorthand for resampling a field at different points. Since smoke is sampled at cell centers and velocity at face centers, this conversion is necessary.
smoke += INFLOW
buoyancy_force = smoke * (0, 0.5) @ velocity
velocity += buoyancy_force
velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0))
vis.plot(smoke)
--------------------------------------------------------------------------- IncompatibleShapes Traceback (most recent call last) Cell In[6], line 4 1 smoke += INFLOW 2 buoyancy_force = smoke * (0, 0.5) @ velocity 3 velocity += buoyancy_force ----> 4 velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0)) 5 6 vis.plot(smoke) File /opt/hostedtoolcache/Python/3.12.13/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.13/x64/lib/python3.12/site-packages/phi/field/_field_math.py:618, in divergence(field, order, implicit, upwind, implicitness) 616 if order == 2: 617 if field.is_staggered: --> 618 field = bake_extrapolation(field) 619 components = [] 620 for dim in field.shape.spatial.names: File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_field_math.py:38, in bake_extrapolation(grid) 36 value = grid.values[dim] 37 padded.append(math.pad(value, {dim: (0 if lower else 1, 0 if upper else 1)}, grid.extrapolation[{'vector': dim}], bounds=grid.bounds)) ---> 38 return StaggeredGrid(math.stack(padded, dual(vector=grid.shape.spatial)), bounds=grid.bounds, extrapolation=math.extrapolation.NONE) 39 elif grid.is_grid: 40 return pad(grid, 1).with_extrapolation(math.extrapolation.NONE) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_grid.py:172, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_) 170 assert resolution.spatial_rank == elements.bounds.spatial_rank, f"Resolution {resolution} does not match bounds {elements.bounds}" 171 if 'vector' in values.shape: --> 172 values = rename_dims(values, 'vector', dual(vector=values.vector.item_names)) 173 assert values.shape.spatial_rank == elements.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}" 174 assert values.shape.spatial_rank == elements.bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}" File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:566, in rename_dims(value, dims, names, **kwargs) 564 # --- First try __replace_dims__ --- 565 if hasattr(value, '__replace_dims__'): --> 566 result = value.__replace_dims__(old_dims.names, new_dims, **kwargs) 567 if result is not NotImplemented: 568 return result File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:526, in Tensor.__replace_dims__(self, dims, new_dims, **kwargs) 525 def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Tensor': --> 526 return self._with_shape_replaced(replace_dims(self.shape, dims, new_dims)) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:3078, in replace_dims(shape, dims, new) 3076 def replace_dims(shape: Shape, dims: DimFilter, new: DimFilter): 3077 old_dims, new_dims = _shape_replace(shape, dims, new) -> 3078 return shape.replace(old_dims, new_dims) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2054, in MixedShape.replace(self, dims, new) 2052 dim_list = [d for n, d in self.dims.items() if n not in dims] 2053 dim_list.insert(i0, new) -> 2054 return concat_shapes_(*dim_list) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2830, in concat_shapes_(*shapes) 2828 dims = {dim.name: dim for dim in all_dims} 2829 if len(dims) != len(all_dims): -> 2830 raise IncompatibleShapes(f"Cannot concatenate shapes {list(shapes)}. Duplicate dimension names are not allowed.") 2831 return MixedShape(dims=dims, **by_type) IncompatibleShapes: Cannot concatenate shapes [(~vectorᵈ=x,y), (inflow_locᵇ=4), (xˢ=~(x=33, y=32) int64), (yˢ=~(x=40, y=41) int64), (~vectorᵈ=x,y)]. Duplicate dimension names are not allowed.
Let's run a longer simulation! Now we add the transport or advection operations to the simulation. ΦFlow provides multiple algorithms for advection. Here we use semi-Lagrangian advection for the velocity and MacCormack advection for the smoke distribution.
trajectory = [smoke]
for i in range(20):
print(i, end=' ')
smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW
buoyancy_force = smoke * (0, 0.5) @ velocity
velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force
velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0))
trajectory.append(smoke)
trajectory = field.stack(trajectory, batch('time'))
vis.plot(trajectory, animate='time')
0
--------------------------------------------------------------------------- AssertionError Traceback (most recent call last) Cell In[7], line 4 1 trajectory = [smoke] 2 for i in range(20): 3 print(i, end=' ') ----> 4 smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW 5 buoyancy_force = smoke * (0, 0.5) @ velocity 6 velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force 7 velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0)) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/physics/advect.py:204, in mac_cormack(field, velocity, dt, correction_strength, integrator) 183 def mac_cormack(field: Field, 184 velocity: Field, 185 dt: float, 186 correction_strength=1.0, 187 integrator=euler) -> Field: 188 """ 189 MacCormack advection uses a forward and backward lookup to determine the first-order error of semi-Lagrangian advection. 190 It then uses that error estimate to correct the field values. (...) 202 Advected field of type `type(field)` 203 """ --> 204 v0 = sample(velocity, field.geometry, at=field.sampled_at, boundary=field.boundary) 205 points_bwd = integrator(field, velocity, -dt, v0=v0) 206 points_fwd = integrator(field, velocity, dt, v0=v0) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_resample.py:138, in sample(field, geometry, at, boundary, dot_face_normal, **kwargs) 136 return sample_grid_at_centers(field, geometry, **kwargs) 137 elif field.is_grid and field.is_staggered: --> 138 return sample_staggered_grid(field, geometry, **kwargs) 139 elif field.is_mesh and field.is_staggered: 140 return math.finite_mean(field.values, dual) # ToDo weigh by face areas? File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_resample.py:296, in sample_staggered_grid(self, geometry, **kwargs) 294 c_sampled = sample_grid_at_centers(Field(c_grid, c_values, ext, sampled_at='center'), geometry, **kwargs) 295 components.append(c_sampled) --> 296 return math.stack(components, geometry.shape['vector']) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:174, in stack(values, dim, expand_values, simplify, layout_non_matching, allow_varying_labels, **kwargs) 172 for v in values: 173 if hasattr(v, '__stack__'): --> 174 result = v.__stack__(values, dim, allow_varying_labels=allow_varying_labels, **kwargs) 175 if result is not NotImplemented: 176 if DEBUG_CHECKS: File /opt/hostedtoolcache/Python/3.12.13/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, **_kwargs) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_ops.py:902, in stack_tensors(values, dim, allow_varying_labels) 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.13/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 (inflow_locᵇ=4, xˢ=32, yˢ=40, vectorᶜ=x,y) 0.004 ± 0.031 (0e+00...3e-01)
The simulation we just computed was using pure NumPy so all operations were non-differentiable.
To enable differentiability, we need to use either PyTorch, TensorFlow or Jax.
This can be achieved by changing the import statement to phi.tf.flow, phi.torch.flow or phi.jax.flow, respectively.
Tensors created after this import will be allocated using PyTorch / TensorFlow / Jax and operations on these will be executed with the corresponding backend.
These operations can make use of a GPU through CUDA if your configuration supports it.
# from phi.jax.flow import *
from phi.torch.flow import *
# from phi.tf.flow import *
We set up the simulation as before.
INFLOW_LOCATION = tensor([(4, 5), (8, 5), (12, 5), (16, 5)], batch('inflow_loc'), channel(vector='x,y'))
INFLOW = 0.6 * CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=3), extrapolation.BOUNDARY, x=32, y=40, bounds=Box(x=32, y=40))
We can verify that tensors are now backed by TensorFlow / PyTorch / Jax.
type(INFLOW.values.native(INFLOW.shape))
torch.Tensor
Note that tensors created with NumPy will keep using NumPy/SciPy operations unless a TensorFlow tensor is also passed to the same operation.
Let's look at how to get gradients from our simulation.
Say we want to optimize the initial velocities so that all simulations arrive at a final state that is similar to the right simulation where the inflow is located at (16, 5).
To achieve this, we define the loss function as $L = | D(s - s_r) |^2$ where $s$ denotes the smoke density and the function $D$ diffuses the difference to smoothen the gradients.
def simulate(smoke: CenteredGrid, velocity: StaggeredGrid):
for _ in range(20):
smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW
buoyancy_force = smoke * (0, 0.5) @ velocity
velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force
velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0))
loss = math.sum(field.l2_loss(diffuse.explicit(smoke - field.stop_gradient(smoke.inflow_loc[-1]), 1, 1, 10)))
return loss, smoke, velocity
Now it is important that the initial velocity has the inflow_loc dimension before we record the gradients.
initial_smoke = CenteredGrid(0, extrapolation.BOUNDARY, x=32, y=40, bounds=Box(x=32, y=40))
initial_velocity = StaggeredGrid(math.zeros(batch(inflow_loc=4)), extrapolation.ZERO, x=32, y=40, bounds=Box(x=32, y=40))
Finally, we use gradient_function() to obtain the gradient with respect to the initial velocity. Since the velocity is the second argument to the simulate() function, we pass wrt=[1].
sim_grad = field.gradient(simulate, wrt='velocity', get_output=False)
The argument get_output=False specifies that we are not interested in the actual output of the function. By setting it to True, we would also get the loss value and the final simulation state.
To evaluate the gradient, we simply call the gradient function with the same arguments as we would call the simulation.
velocity_grad = sim_grad(initial_smoke, initial_velocity)
vis.plot(velocity_grad)
--------------------------------------------------------------------------- IncompatibleShapes Traceback (most recent call last) Cell In[14], line 1 ----> 1 velocity_grad = sim_grad(initial_smoke, initial_velocity) 2 3 vis.plot(velocity_grad) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:650, in GradientFunction.__call__(self, *args, **kwargs) 648 if key not in self.traces: 649 self.traces[key] = self._trace_grad(key, wrt_natives) --> 650 native_result = self.traces[key](*natives) 651 output_key = match_output_signature(key, self.recorded_mappings, self) 652 jac_shape = output_key.shapes[0].non_batch # ToDo prepend this to all wrt shapes File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1018, in TorchBackend.jacobian.<locals>.eval_grad(*args) 1016 with torch.enable_grad(): 1017 args, wrt_args = self._prepare_graph_inputs(args, wrt) -> 1018 loss, output = f(*args) 1019 if np.prod(self.staticshape(loss)) == 1: 1020 assert loss.requires_grad, f"Failed to compute gradient because function output does not depend on any input. Inputs: {args}" File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:617, in GradientFunction._trace_grad.<locals>.f_native(*natives) 615 kwargs = assemble_tree(in_key.tree, in_tensors, attr_type=variable_attributes) 616 with functional_derivative_evaluation(order=1): --> 617 result = self.f(**kwargs) # Tensor or tuple/list of Tensors 618 loss = result[0] if isinstance(result, (tuple, list)) else result 619 if isinstance(loss, Tensor): Cell In[11], line 6, in simulate(smoke, velocity) 2 for _ in range(20): 3 smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW 4 buoyancy_force = smoke * (0, 0.5) @ velocity 5 velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force ----> 6 velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0)) 7 loss = math.sum(field.l2_loss(diffuse.explicit(smoke - field.stop_gradient(smoke.inflow_loc[-1]), 1, 1, 10))) 8 return loss, smoke, velocity File /opt/hostedtoolcache/Python/3.12.13/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.13/x64/lib/python3.12/site-packages/phi/field/_field_math.py:618, in divergence(field, order, implicit, upwind, implicitness) 616 if order == 2: 617 if field.is_staggered: --> 618 field = bake_extrapolation(field) 619 components = [] 620 for dim in field.shape.spatial.names: File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_field_math.py:38, in bake_extrapolation(grid) 36 value = grid.values[dim] 37 padded.append(math.pad(value, {dim: (0 if lower else 1, 0 if upper else 1)}, grid.extrapolation[{'vector': dim}], bounds=grid.bounds)) ---> 38 return StaggeredGrid(math.stack(padded, dual(vector=grid.shape.spatial)), bounds=grid.bounds, extrapolation=math.extrapolation.NONE) 39 elif grid.is_grid: 40 return pad(grid, 1).with_extrapolation(math.extrapolation.NONE) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_grid.py:172, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_) 170 assert resolution.spatial_rank == elements.bounds.spatial_rank, f"Resolution {resolution} does not match bounds {elements.bounds}" 171 if 'vector' in values.shape: --> 172 values = rename_dims(values, 'vector', dual(vector=values.vector.item_names)) 173 assert values.shape.spatial_rank == elements.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}" 174 assert values.shape.spatial_rank == elements.bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}" File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:566, in rename_dims(value, dims, names, **kwargs) 564 # --- First try __replace_dims__ --- 565 if hasattr(value, '__replace_dims__'): --> 566 result = value.__replace_dims__(old_dims.names, new_dims, **kwargs) 567 if result is not NotImplemented: 568 return result File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:526, in Tensor.__replace_dims__(self, dims, new_dims, **kwargs) 525 def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Tensor': --> 526 return self._with_shape_replaced(replace_dims(self.shape, dims, new_dims)) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:3078, in replace_dims(shape, dims, new) 3076 def replace_dims(shape: Shape, dims: DimFilter, new: DimFilter): 3077 old_dims, new_dims = _shape_replace(shape, dims, new) -> 3078 return shape.replace(old_dims, new_dims) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2054, in MixedShape.replace(self, dims, new) 2052 dim_list = [d for n, d in self.dims.items() if n not in dims] 2053 dim_list.insert(i0, new) -> 2054 return concat_shapes_(*dim_list) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2830, in concat_shapes_(*shapes) 2828 dims = {dim.name: dim for dim in all_dims} 2829 if len(dims) != len(all_dims): -> 2830 raise IncompatibleShapes(f"Cannot concatenate shapes {list(shapes)}. Duplicate dimension names are not allowed.") 2831 return MixedShape(dims=dims, **by_type) IncompatibleShapes: Cannot concatenate shapes [(~vectorᵈ=x,y), (inflow_locᵇ=4), (xˢ=~(x=33, y=32) int64), (yˢ=~(x=40, y=41) int64), (~vectorᵈ=x,y)]. Duplicate dimension names are not allowed.
With the gradient, we can easily perform basic gradient descent optimization. For more advanced optimization techniques and neural network training, see the optimization documentation.
print(f"Initial loss: {simulate(initial_smoke, initial_velocity)[0]}")
initial_velocity -= 0.01 * velocity_grad
print(f"Next loss: {simulate(initial_smoke, initial_velocity)[0]}")
--------------------------------------------------------------------------- IncompatibleShapes Traceback (most recent call last) Cell In[15], line 1 ----> 1 print(f"Initial loss: {simulate(initial_smoke, initial_velocity)[0]}") 2 initial_velocity -= 0.01 * velocity_grad 3 print(f"Next loss: {simulate(initial_smoke, initial_velocity)[0]}") Cell In[11], line 6, in simulate(smoke, velocity) 2 for _ in range(20): 3 smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW 4 buoyancy_force = smoke * (0, 0.5) @ velocity 5 velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force ----> 6 velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0)) 7 loss = math.sum(field.l2_loss(diffuse.explicit(smoke - field.stop_gradient(smoke.inflow_loc[-1]), 1, 1, 10))) 8 return loss, smoke, velocity File /opt/hostedtoolcache/Python/3.12.13/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.13/x64/lib/python3.12/site-packages/phi/field/_field_math.py:618, in divergence(field, order, implicit, upwind, implicitness) 616 if order == 2: 617 if field.is_staggered: --> 618 field = bake_extrapolation(field) 619 components = [] 620 for dim in field.shape.spatial.names: File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_field_math.py:38, in bake_extrapolation(grid) 36 value = grid.values[dim] 37 padded.append(math.pad(value, {dim: (0 if lower else 1, 0 if upper else 1)}, grid.extrapolation[{'vector': dim}], bounds=grid.bounds)) ---> 38 return StaggeredGrid(math.stack(padded, dual(vector=grid.shape.spatial)), bounds=grid.bounds, extrapolation=math.extrapolation.NONE) 39 elif grid.is_grid: 40 return pad(grid, 1).with_extrapolation(math.extrapolation.NONE) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_grid.py:172, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_) 170 assert resolution.spatial_rank == elements.bounds.spatial_rank, f"Resolution {resolution} does not match bounds {elements.bounds}" 171 if 'vector' in values.shape: --> 172 values = rename_dims(values, 'vector', dual(vector=values.vector.item_names)) 173 assert values.shape.spatial_rank == elements.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}" 174 assert values.shape.spatial_rank == elements.bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}" File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:566, in rename_dims(value, dims, names, **kwargs) 564 # --- First try __replace_dims__ --- 565 if hasattr(value, '__replace_dims__'): --> 566 result = value.__replace_dims__(old_dims.names, new_dims, **kwargs) 567 if result is not NotImplemented: 568 return result File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:526, in Tensor.__replace_dims__(self, dims, new_dims, **kwargs) 525 def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Tensor': --> 526 return self._with_shape_replaced(replace_dims(self.shape, dims, new_dims)) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:3078, in replace_dims(shape, dims, new) 3076 def replace_dims(shape: Shape, dims: DimFilter, new: DimFilter): 3077 old_dims, new_dims = _shape_replace(shape, dims, new) -> 3078 return shape.replace(old_dims, new_dims) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2054, in MixedShape.replace(self, dims, new) 2052 dim_list = [d for n, d in self.dims.items() if n not in dims] 2053 dim_list.insert(i0, new) -> 2054 return concat_shapes_(*dim_list) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2830, in concat_shapes_(*shapes) 2828 dims = {dim.name: dim for dim in all_dims} 2829 if len(dims) != len(all_dims): -> 2830 raise IncompatibleShapes(f"Cannot concatenate shapes {list(shapes)}. Duplicate dimension names are not allowed.") 2831 return MixedShape(dims=dims, **by_type) IncompatibleShapes: Cannot concatenate shapes [(~vectorᵈ=x,y), (inflow_locᵇ=4), (xˢ=~(x=33, y=32) int64), (yˢ=~(x=40, y=41) int64), (~vectorᵈ=x,y)]. Duplicate dimension names are not allowed.
sim_grad = field.gradient(simulate, wrt='velocity', get_output=True)
for opt_step in range(4):
(loss, final_smoke, _v), velocity_grad = sim_grad(initial_smoke, initial_velocity)
print(f"Step {opt_step}, loss: {loss}")
initial_velocity -= 0.01 * velocity_grad
--------------------------------------------------------------------------- IncompatibleShapes Traceback (most recent call last) Cell In[16], line 4 1 sim_grad = field.gradient(simulate, wrt='velocity', get_output=True) 2 3 for opt_step in range(4): ----> 4 (loss, final_smoke, _v), velocity_grad = sim_grad(initial_smoke, initial_velocity) 5 print(f"Step {opt_step}, loss: {loss}") 6 initial_velocity -= 0.01 * velocity_grad File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:650, in GradientFunction.__call__(self, *args, **kwargs) 648 if key not in self.traces: 649 self.traces[key] = self._trace_grad(key, wrt_natives) --> 650 native_result = self.traces[key](*natives) 651 output_key = match_output_signature(key, self.recorded_mappings, self) 652 jac_shape = output_key.shapes[0].non_batch # ToDo prepend this to all wrt shapes File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1018, in TorchBackend.jacobian.<locals>.eval_grad(*args) 1016 with torch.enable_grad(): 1017 args, wrt_args = self._prepare_graph_inputs(args, wrt) -> 1018 loss, output = f(*args) 1019 if np.prod(self.staticshape(loss)) == 1: 1020 assert loss.requires_grad, f"Failed to compute gradient because function output does not depend on any input. Inputs: {args}" File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_functional.py:617, in GradientFunction._trace_grad.<locals>.f_native(*natives) 615 kwargs = assemble_tree(in_key.tree, in_tensors, attr_type=variable_attributes) 616 with functional_derivative_evaluation(order=1): --> 617 result = self.f(**kwargs) # Tensor or tuple/list of Tensors 618 loss = result[0] if isinstance(result, (tuple, list)) else result 619 if isinstance(loss, Tensor): Cell In[11], line 6, in simulate(smoke, velocity) 2 for _ in range(20): 3 smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW 4 buoyancy_force = smoke * (0, 0.5) @ velocity 5 velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force ----> 6 velocity, _ = fluid.make_incompressible(velocity, (), Solve(rank_deficiency=0)) 7 loss = math.sum(field.l2_loss(diffuse.explicit(smoke - field.stop_gradient(smoke.inflow_loc[-1]), 1, 1, 10))) 8 return loss, smoke, velocity File /opt/hostedtoolcache/Python/3.12.13/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.13/x64/lib/python3.12/site-packages/phi/field/_field_math.py:618, in divergence(field, order, implicit, upwind, implicitness) 616 if order == 2: 617 if field.is_staggered: --> 618 field = bake_extrapolation(field) 619 components = [] 620 for dim in field.shape.spatial.names: File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_field_math.py:38, in bake_extrapolation(grid) 36 value = grid.values[dim] 37 padded.append(math.pad(value, {dim: (0 if lower else 1, 0 if upper else 1)}, grid.extrapolation[{'vector': dim}], bounds=grid.bounds)) ---> 38 return StaggeredGrid(math.stack(padded, dual(vector=grid.shape.spatial)), bounds=grid.bounds, extrapolation=math.extrapolation.NONE) 39 elif grid.is_grid: 40 return pad(grid, 1).with_extrapolation(math.extrapolation.NONE) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phi/field/_grid.py:172, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_) 170 assert resolution.spatial_rank == elements.bounds.spatial_rank, f"Resolution {resolution} does not match bounds {elements.bounds}" 171 if 'vector' in values.shape: --> 172 values = rename_dims(values, 'vector', dual(vector=values.vector.item_names)) 173 assert values.shape.spatial_rank == elements.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}" 174 assert values.shape.spatial_rank == elements.bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}" File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:566, in rename_dims(value, dims, names, **kwargs) 564 # --- First try __replace_dims__ --- 565 if hasattr(value, '__replace_dims__'): --> 566 result = value.__replace_dims__(old_dims.names, new_dims, **kwargs) 567 if result is not NotImplemented: 568 return result File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:526, in Tensor.__replace_dims__(self, dims, new_dims, **kwargs) 525 def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Tensor': --> 526 return self._with_shape_replaced(replace_dims(self.shape, dims, new_dims)) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:3078, in replace_dims(shape, dims, new) 3076 def replace_dims(shape: Shape, dims: DimFilter, new: DimFilter): 3077 old_dims, new_dims = _shape_replace(shape, dims, new) -> 3078 return shape.replace(old_dims, new_dims) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2054, in MixedShape.replace(self, dims, new) 2052 dim_list = [d for n, d in self.dims.items() if n not in dims] 2053 dim_list.insert(i0, new) -> 2054 return concat_shapes_(*dim_list) File /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2830, in concat_shapes_(*shapes) 2828 dims = {dim.name: dim for dim in all_dims} 2829 if len(dims) != len(all_dims): -> 2830 raise IncompatibleShapes(f"Cannot concatenate shapes {list(shapes)}. Duplicate dimension names are not allowed.") 2831 return MixedShape(dims=dims, **by_type) IncompatibleShapes: Cannot concatenate shapes [(~vectorᵈ=x,y), (inflow_locᵇ=4), (xˢ=~(x=33, y=32) int64), (yˢ=~(x=40, y=41) int64), (~vectorᵈ=x,y)]. Duplicate dimension names are not allowed.
This notebook provided an introduction to running fluid simulations in NumPy and TensorFlow. It demonstrated how to obtain simulation gradients which can be used to optimize physical variables or train neural networks.
The full ΦFlow documentation is available at https://tum-pbs.github.io/PhiFlow/.
Visit the playground to run ΦFlow code in an empty notebook.