3D obstacles not working with jit-compile

[KarlisFre]

Hi, 3D obstacles still not work with jit-compile.

from phi.tf.flow import *

object_geometry = Box(x=(40,60), y=(40, 60), z=(40, 60))
OBSTACLE = Obstacle(object_geometry, velocity=(0.,1.,0.))
velocity = StaggeredGrid((0, 0, 0), 0, x=32, y=32, z=32, bounds=Box(x=100, y=100, z=100))
smoke = CenteredGrid(0, ZERO_GRADIENT, x=32, y=32, z=32, bounds=Box(x=100, y=100, z=100))
INFLOW = 0.2 * resample(Sphere(x=50, y=50, z=10, radius=5), to=smoke, soft=True)
pressure = None


@jit_compile  # Only for PyTorch, TensorFlow and Jax
def step(v, s, p,obstacle, dt=1.):
    s = advect.mac_cormack(s, v, dt) + INFLOW
    buoyancy = resample(s * (0, 0, 0.1), to=v)
    v = advect.semi_lagrangian(v, v, dt) + buoyancy * dt
    v, p = fluid.make_incompressible(v, (obstacle), Solve('auto', 1e-5, x0=p))
    return v, s, p, obstacle


for _ in view(smoke, velocity, 'pressure', play=True, namespace=globals()).range(warmup=1):
    velocity, smoke, pressure, OBSTACLE = step(velocity, smoke, pressure, OBSTACLE)

It gives two errors:

"is not" with a literal. Did you mean "!="?
NotImplementedError: spatial_rank=3 not yet implemented

Also, at some point I would like to add object rotations. How this should be done in 3D? The obstacles have angularvelocity field but what is its meaning in 3D?
Thanks a lot!

          Hi, I pushed a fix to `2.4-develop`, https://github.com/tum-pbs/PhiFlow/commit/103e8be441a7fdb876973aefeb86c8cca09548ff .

pip uninstall phiflow
pip install git+https://github.com/tum-pbs/PhiFlow@2.4-develop

Originally posted by @holl- in #131 (comment)

[holl-]

I see, TensorFlow's jit compilation runs into the error even though it isn't actually executed.
The code should run with PyTorch's and Jax' jit already but I've pushed a fix for TensorFlow now.
You need to specify Obstacle(object_geometry, velocity=(0.,1.,0.), angular_velocity=(0, 0, 0)), however.
Could you try again with the latest 2.4-develop version?

[KarlisFre]

Yes, this works, Thanks!
But when I try rotated objects there are some more missing parts:

from phi.tf.flow import *

object_geometry = Box(x=(40,60), y=(40, 60), z=(40, 60))
object_geometry = object_geometry.rotated((0.0, 0.1, 0.0))
OBSTACLE = Obstacle(object_geometry, velocity=(0.,1.,0.), angular_velocity=(0.,0., 0.))

velocity = StaggeredGrid((0, 0, 0), 0, x=32, y=32, z=32, bounds=Box(x=100, y=100, z=100))
smoke = CenteredGrid(0, ZERO_GRADIENT, x=32, y=32, z=32, bounds=Box(x=100, y=100, z=100))
INFLOW = 0.2 * resample(Sphere(x=50, y=50, z=10, radius=5), to=smoke, soft=True)
pressure = None


@jit_compile  # Only for PyTorch, TensorFlow and Jax
def step(v, s, p,obstacle, dt=1.):
    s = advect.mac_cormack(s, v, dt) + INFLOW
    buoyancy = resample(s * (0, 0, 0.1), to=v)
    v = advect.semi_lagrangian(v, v, dt) + buoyancy * dt
    v, p = fluid.make_incompressible(v, (obstacle), Solve('auto', 1e-5, x0=p))
    return v, s, p, obstacle


for _ in view(smoke, velocity, 'pressure', play=True, namespace=globals()).range(warmup=1):
    velocity, smoke, pressure, OBSTACLE = step(velocity, smoke, pressure, OBSTACLE)

File "C:\Users\karlis\anaconda3\lib\site-packages\phi\physics\fluid.py", line 107, in make_incompressible *
accessible = CenteredGrid(~union([obs.geometry for obs in obstacles]), accessible_extrapolation, velocity.bounds, velocity.resolution)
File "C:\Users\karlis\anaconda3\lib\site-packages\phi\field_grid.py", line 201, in init **
values = reduce_sample(values, elements)
File "C:\Users\karlis\anaconda3\lib\site-packages\phi\field_field.py", line 402, in reduce_sample
return field._sample(geometry, **kwargs)
File "C:\Users\karlis\anaconda3\lib\site-packages\phi\field_point_cloud.py", line 156, in _sample
return math.where(self.elements.lies_inside(geometry.center), self._values[idx0], outside)
File "C:\Users\karlis\anaconda3\lib\site-packages\phi\geom_geom.py", line 392, in lies_inside
return ~self.geometry.lies_inside(location)
File "C:\Users\karlis\anaconda3\lib\site-packages\phi\geom_transform.py", line 64, in lies_inside
return self.geometry.lies_inside(self.global_to_child(location))
File "C:\Users\karlis\anaconda3\lib\site-packages\phi\geom_transform.py", line 52, in global_to_child
raise NotImplementedError('not yet implemented') # ToDo apply angle

[holl-]

I've now implemented 3D rotation, a2a5f66 .
Note that this is experimental and has not been thoroughly tested yet. Please check that your rotated obstacle works as intended.

[KarlisFre]

Many tanks for your quick fixes, it seems to work! I will start working with 3D rotations in a few weeks and then test it more carefully. One thing I noticed that composing several 3D rotations probably will not work correctly since the angles are summed. A proper solution would be to multiply the rotation matrices.

It should be evident in the Rotating Bar, if ported to 3D, demo when rotation is applied iteratively:
obstacle = obstacle.copied_with(geometry=obstacle.geometry.rotated(-obstacle.angular_velocity * DT)) # rotate bar

[holl-]

Right, we would need a function like combine_rotations...
If you know some reference code, let me know.

[KarlisFre]

A typical solution is to use matrices instead of Euler angles. The RotatedGeometry class would store the rotation matrix, the same currently is used for rotation:

        s1, s2, s3 = math.sin(angle).vector
        c1, c2, c3 = math.cos(angle).vector
        matrix = wrap([[c1*c2, c1*s2*s3 - s1*c3, c1*s2*c3 + s1*s3],
                       [s1*c2, s1*s2*s3 + c1*c3, s1*s2*c3 - c1*s3],
                       [-s2, c2*s3, c2*c3]],

And when combining rotations, we need to multiply the both matrices together.

Combining Euler angles is problematic and may lead to degeneracy: https://en.wikipedia.org/wiki/Gimbal_lock

[holl-]

In future versions, we will use rotation matrices directly. For now, you can add rotations by multiplying matrices and then getting back the Euler angles. To do this, I've added math.rotation_angles()

new_angles = math.rotation_angles(math.rotation_matrix(old_angles) @ math.rotation_matrix(delta_rotation))