Waves¶

Google Collab Book

This example simulates waves induced by a moving sphere on a 2D grid. It is a direct port of this example from warp.

In [1]:
%pip install --quiet phiflow
from phi.jax.flow import *
# from phi.flow import *  # If JAX is not installed. You can use phi.torch or phi.tf as well.
from tqdm.notebook import trange

We define two functions, wave_displace to model the interactions between the obstacle and the height field, and step to advance the simulation.

In [12]:
def wave_displace(sphere: Sphere, *fields: Field, mag=.5, t=-math.PI * 0.5):
    return [field.where(sphere, mag * math.sin(t), f) for f in fields]

@jit_compile
def step(h_c, h_p, time, dt=1/60./16, k_speed=1.0, k_damp=0.0):
    sphere = Sphere(center=h_c.bounds.center + math.rotate_vector(vec(x=0, y=-12.8/3), time), radius=1.)
    h_c, h_p = wave_displace(sphere, h_c, h_p)
    h_n = 2.0 * h_c - h_p + dt * dt * (k_speed * h_c.laplace() - k_damp * (h_c - h_p))  # wave_solve
    return h_n, h_c, time + dt

Now we can run the simulation on a $128\times128$ grid, performing 16 sub-steps per frame.

In [15]:
h_initial = CenteredGrid(x=128, y=128, bounds=Box(x=12.8, y=12.8), boundary=ZERO_GRADIENT)
h_trj, *_ = iterate(step, batch(time=300), h_initial, h_initial, 0, substeps=16, range=trange)
  0%|          | 0/300 [00:00<?, ?it/s]
In [17]:
plot(h_trj.time[::3], animate='time', size=(5, 4))
Out[17]:
Your browser does not support the video tag.