JIT compile moving obstacles

[KarlisFre]

Many thanks for your great fluid simulation package!

Is it possible to use @math.jit_compile when simulation involves moving obstacles? Currently, it becomes much slower and gives a warning: WARNING:tensorflow:6 out of the last 6 calls to <function JitFunction._jit_compile..jit_f_native at 0x7fa7a14e31f0> triggered tf.function retracing. With static obstacles jit_compile works great and gives significant speedup. See code below which is adapted from the demo example:

from phi.tf.flow import *
DOMAIN = dict(x=30, y=30)
DT = 0.1

def move_obstacle(obs: Obstacle):
    if (obs.geometry.center[0]) > 35:
        new_geometry = Box(x=(-6, 0), y=(10, 16))
    else:
        new_geometry = obs.geometry.shifted([1. * DT, 0])
    return obs.copied_with(geometry=new_geometry)


obstacle = Obstacle(Box(x=(5, 11), y=(10, 16)), velocity=[1., 0], angular_velocity=tensor(0,))
velocity = StaggeredGrid(0, extrapolation.ZERO, **DOMAIN)
obstacle_mask = CenteredGrid(HardGeometryMask(obstacle.geometry), extrapolation.BOUNDARY, **DOMAIN)
pressure = None

@math.jit_compile
def step(velocity, obstacle):
    velocity = advect.mac_cormack(velocity, velocity, DT)
    velocity, pressure = fluid.make_incompressible(velocity, (obstacle,))
    return velocity, pressure

for _ in view(velocity, obstacle_mask, play=True, namespace=globals(), port=6006).range():
    obstacle = move_obstacle(obstacle)
    velocity, pressure = step(velocity, obstacle)
    fluid.masked_laplace.tracers.clear()  # we will need to retrace because the matrix changes each step. This is not needed when JIT-compiling the physics.
    obstacle_mask = HardGeometryMask(obstacle.geometry) @ pressure

[holl-]

Yes, this is supposed to work. However, the Obstacle class does not currently support function transformations. Try passing the geometry to the function and constructing the Obstacle inside.

[KarlisFre]

Unfortunately this does not compile:

Traceback (most recent call last):
File "/home/flow_control/moving_obstacle.py", line 15, in
velocity, pressure = step(velocity, object)
File "/usr/local/lib/python3.8/dist-packages/phi/math/_functional.py", line 194, in call
self.traces[key] = self._jit_compile(key)
File "/usr/local/lib/python3.8/dist-packages/phi/math/_functional.py", line 171, in _jit_compile
PHI_LOGGER.debug(f"Φ-jit: '{f_name(self.f)}' called with new key. shapes={[s.volume for s in in_key.shapes]}, args={in_key.tree}")
File "/usr/local/lib/python3.8/dist-packages/phi/geom/_box.py", line 298, in repr
if self.shape.non_channel.volume == 1:
AttributeError: 'NoneType' object has no attribute 'non_channel'

from phi.tf.flow import *

DOMAIN = dict(x=30, y=30)
DT = 0.1
object = Box(x=(5, 11), y=(10, 16))

@math.jit_compile
def step(velocity, object_geometry):
    obstacle = Obstacle(object_geometry, velocity=[1., 0], angular_velocity=tensor(0, ))
    velocity = advect.mac_cormack(velocity, velocity, DT)
    velocity, pressure = fluid.make_incompressible(velocity, (obstacle,))
    return velocity, pressure

velocity = StaggeredGrid(0, extrapolation.ZERO, **DOMAIN)
velocity, pressure = step(velocity, object)
obstacle_mask = HardGeometryMask(object) @ pressure

def move_obstacle(geometry):
    if (geometry.center[0]) > 35:
        new_geometry = Box(x=(-6, 0), y=(10, 16))
    else:
        new_geometry = geometry.shifted([1. * DT, 0])
    return new_geometry

for _ in view(velocity, obstacle_mask, play=True, namespace=globals(), port=6006).range():
    object = move_obstacle(object)
    velocity, pressure = step(velocity, object)
    fluid.masked_laplace.tracers.clear()  # we will need to retrace because the matrix changes each step. This is not needed when JIT-compiling the physics.
    obstacle_mask = HardGeometryMask(object) @ pressure

[holl-]

I've pushed a fix to 2.3-develop. Your original version should now work after removing the line fluid.masked_laplace.tracers.clear().
Install 2.3-develop by running

$ pip install --upgrade git+https://github.com/tum-pbs/PhiFlow@2.3-develop