整个代码如下:
import taichi as ti
ti.init(arch=ti.gpu)
n_grid = 512
dt = 0.03
use_mc = True
shape = ti.field(ti.u8, shape=(n_grid, n_grid))
_dye_buffer = ti.field(float, shape=(n_grid, n_grid))
_new_dye_buffer = ti.field(float, shape=(n_grid, n_grid))
_new_dye_buffer_aux = ti.field(float, shape=(n_grid, n_grid))
_velocities = ti.Vector.field(2, float, shape=(n_grid, n_grid))
_new_velocities = ti.Vector.field(2, float, shape=(n_grid, n_grid))
@ti.kernel
def init_dens_vel():
for i,j in ti.ndrange(n_grid, n_grid):
_velocities[i,j] = ti.Vector([10.0 ,0.0])
if i>0 and i<80 and j>200 and j<400:
_dye_buffer[i,j] = 1.0
@ti.func
def lerp(vl, vr, frac):
# frac: [0.0, 1.0]
return vl + frac * (vr - vl)
@ti.func
def sample_min(qf, p):
u, v = p
s, t = u - 0.5, v - 0.5
# floor
iu, iv = ti.floor(s), ti.floor(t)
a = sample(qf, iu, iv)
b = sample(qf, iu + 1, iv)
c = sample(qf, iu, iv + 1)
d = sample(qf, iu + 1, iv + 1)
return min(a, b, c, d)
@ti.func
def sample_max(qf, p):
u, v = p
s, t = u - 0.5, v - 0.5
# floor
iu, iv = ti.floor(s), ti.floor(t)
a = sample(qf, iu, iv)
b = sample(qf, iu + 1, iv)
c = sample(qf, iu, iv + 1)
d = sample(qf, iu + 1, iv + 1)
return max(a, b, c, d)
@ti.func
def sample(qf, u, v):
I = ti.Vector([int(u), int(v)])
I = max(0, min(n_grid - 1, I))
return qf[I]
@ti.func
def bilerp(vf, p):
u, v = p
s, t = u - 0.5, v - 0.5
# floor
iu, iv = ti.floor(s), ti.floor(t)
# fract
fu, fv = s - iu, t - iv
a = sample(vf, iu, iv)
b = sample(vf, iu + 1, iv)
c = sample(vf, iu, iv + 1)
d = sample(vf, iu + 1, iv + 1)
return lerp(lerp(a, b, fu), lerp(c, d, fu), fv)
@ti.func
def backtrace(vf, p, dt):
v1 = bilerp(vf, p)
p1 = p - 0.5 * dt * v1
v2 = bilerp(vf, p1)
p2 = p - 0.75 * dt * v2
v3 = bilerp(vf, p2)
p -= dt * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3)
return p
@ti.kernel
def advect(vf: ti.template(), qf: ti.template(), new_qf: ti.template()):
for i, j in ti.ndrange(n_grid, n_grid):
p = ti.Vector([i, j]) + 0.5
p = backtrace(vf, p, dt)
new_qf[i, j] = bilerp(qf, p)
@ti.kernel
def maccormack(vf: ti.template()):
for i, j in ti.ndrange(n_grid, n_grid):
p = ti.Vector([i, j]) + 0.5
p = backtrace( vf, p, -dt )
_new_dye_buffer_aux[i, j] = bilerp(_new_dye_buffer, p)
for i, j in ti.ndrange(n_grid, n_grid):
_new_dye_buffer[i, j] += 0.5 * (_dye_buffer[i,j]-_new_dye_buffer_aux[i,j])
def step():
advect(_velocities, _velocities, _new_velocities)
advect(_velocities, _dye_buffer, _new_dye_buffer)
maccormack(_velocities)
swap_v()
swap_d()
def swap_v():
global _velocities
global _new_velocities
_velocities,_new_velocities = _new_velocities,_velocities
def swap_d():
global _dye_buffer
global _new_dye_buffer
_dye_buffer,_new_dye_buffer = _new_dye_buffer,_dye_buffer
init_dens_vel()
gui = ti.GUI('shape control', (n_grid, n_grid))
while True:
step()
gui.set_image(_dye_buffer)
gui.show()