Vortex Method Demo

vortex method这类方法基于旋度形式的NS方程，适合模拟vortex ring之类的旋涡高度集中的现象。以及和stream function结合可以很好地保证divergence-free。SIGGRAPH 2015有一篇有趣的paper，A Stream Function Solver for Liquid Simulations，用的stream function模拟multiphase，不久前作者终于开源了，https://github.com/ryichando/shiokaze

``````import taichi as ti
import math

# ti.init(arch=ti.cpu)
ti.init(arch=ti.gpu)

N = 512
Nx = N
Ny = N
nTotal = Nx * Ny

dx = 1.0 / N
dt = 0.2 / N

ux = ti.var(dt=ti.f32, shape=(Nx, Ny))  # velocity x
uy = ti.var(dt=ti.f32, shape=(Nx, Ny))  # velocity y
omega = ti.var(dt=ti.f32, shape=(Nx, Ny))  # vorticity
omega_temp = ti.var(dt=ti.f32, shape=(Nx, Ny))  # vorticity
sigma = ti.var(dt=ti.f32, shape=(Nx, Ny))  # stream function

color_buffer = ti.Vector(3, dt=ti.f32, shape=(Nx, Ny))

solver_type = 0  # 0: stream func sovler, 1: vortex in cell solver, 2: biot-savart solver
solver_names = ["stream function solver",
"vortex-in-cell solver", "Biot-Savart solver"]

class ColorMap:
def __init__(self, h, wl, wr, c):
self.h = h
self.wl = wl
self.wr = wr
self.c = c

@ti.func
def clamp(self, x):
return max(0.0, min(1.0, x))

@ti.func
def map(self, x):
w = 0.0
if x < self.c:
w = self.wl
else:
w = self.wr
return self.clamp((w-abs(self.clamp(x)-self.c))/w*self.h)

jetR = ColorMap(1.5, .37, .37, .75)
jetG = ColorMap(1.5, .37, .37, .5)
jetB = ColorMap(1.5, .37, .37, .25)

@ti.func
def color_map(c):
return ti.Vector([jetR.map(c),
jetG.map(c),
jetB.map(c)])

@ti.kernel
def fill_velocity():
for i, j in ux:
_ux = ux[i, j]
_uy = uy[i, j]

color_buffer[i, j] = color_map(ti.sqrt(_ux * _ux + _uy * _uy))

@ti.kernel
def fill_vorticity():
for i, j in omega:
c = omega[i, j]
c = 0.25 + c * 0.01
color_buffer[i, j] = color_map(c)

@ti.kernel
def fill_streamfunc():
for i, j in sigma:
c = sigma[i, j] * 5.0 + 0.5
color_buffer[i, j] = color_map(c)

# -------------------------------------

def swap_buffers(a, b):
a, b = b, a

@ti.kernel
def copy_buffer(dst: ti.template(), src: ti.template()):
for i, j in src:
dst[i, j] = src[i, j]

# [-0.5, 0.5]
@ti.func
def xCoord(i):
return ((i + 0.5) / Nx) - 0.5

# [-0.5, 0.5]
@ti.func
def yCoord(j):
return ((j + 0.5) / Ny) - 0.5

@ti.func
def ix(i, j):
return i + j * Nx

@ti.func
def mod(a, b):
return int(a % b)

@ti.func
def sq(x):
return x ** 2

@ti.kernel
def initialize_vortices():
for i, j in ux:
s = 0.16
x = xCoord(i) * 2.0
y = yCoord(j) * 2.0
x0 = x - s
x1 = x + s
ux_ = 0.0
uy_ = 0.0

rr1 = sq(x0) + sq(y)
ux_ += - y * ti.exp(- rr1 / (2.0 * sq(0.12))) * 25.0
uy_ += x0 * ti.exp(- rr1 / (2.0 * sq(0.12))) * 25.0
rr2 = sq(x1) + sq(y)
ux_ += - y * ti.exp(- rr2 / (2.0 * sq(0.12))) * 25.0
uy_ += x1 * ti.exp(- rr2 / (2.0 * sq(0.12))) * 25.0

ux[i, j] = ux_
uy[i, j] = uy_

for i, j in omega:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)

omega[i, j] = \
(uy[i_p, j] - uy[i_n, j]) / (2.0 * dx) \
- (ux[i, j_p] - ux[i, j_n]) / (2.0 * dx)

# solve L*x=-rhs using periodic boundary
# using red-black ordering for successive over-relaxation/Gauss-Seidel linear solver
@ti.kernel
x: ti.template(), rhs: ti.template()):
for i, j in x:
if mod(i + j, 2) == mask:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)

x_update = (
- rhs[i, j] * sq(dx)
- x[i_p, j]
- x[i_n, j]
- x[i, j_p]
- x[i, j_n]) * -0.25

x[i, j] = SOR_factor * x_update + (1.0 - SOR_factor) * x[i, j]

@ti.kernel
def fix_streamfunc():
# since velocity is the gradient of sigma (the stream function),
# the constant drift in sigma is harmless to velocity,
# and the field is eased for better visualization
sigmaMean = 0.0
for i, j in sigma:
sigmaMean += sigma[i, j]
sigmaMean /= nTotal
for i, j in sigma:
sigma[i, j] -= sigmaMean

def poisson_solver(num_iterations, x, rhs):
SOR_factor = 1.99  # 1.0 for Gauss-Seidel

for iters in range(num_iterations):
linear_solver_step(SOR_factor, iters % 2, x, rhs)

def solve_streamfunc(num_iterations):
# solve L*sigma=-omega using periodic boundary
poisson_solver(num_iterations, sigma, omega)

fix_streamfunc()

@ti.kernel
def reconstruct_velocity_sigma():
# u = curl(sigma)
# in 2D, u = grad x (0,0,sigma) = (par_y_sigma, -par_x_sigma, 0)
for i, j in sigma:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)

ux[i, j] = (sigma[i, j_p] - sigma[i, j_n]) / (2.0 * dx)  # par_y_sigma
uy[i, j] = - (sigma[i_p, j] - sigma[i_n, j]) / \
(2.0 * dx)  # -par_x_sigma

# approximate Biot-Savart law, slow
# using the 2D version of the Biot-Savart equation
@ti.kernel
def reconstruct_velocity_biot_savart():
patch_size = 4
d_patch_size = patch_size * patch_size * sq(dx)

for i, j in omega:
ux_ = 0.0
uy_ = 0.0

for m_, n_ in ti.ndrange((-num_patches, num_patches + 1),
(-num_patches, num_patches + 1)):
m = m_ * patch_size
n = n_ * patch_size
if (m != 0 or n != 0):
ii = mod(i + m + Nx, Nx)
jj = mod(j + n + Ny, Ny)

vorticity = omega[ii, jj]

# x - p
dist_x = - m * dx
dist_y = - n * dx

# omega \cross (x - p)
vx = vorticity * - dist_y
vy = vorticity * dist_x

rr = (dist_x * dist_x + dist_y * dist_y)

ux_ += vx / rr
uy_ += vy / rr

ux[i, j] = ux_ / (2.0 * math.pi) * d_patch_size
uy[i, j] = uy_ / (2.0 * math.pi) * d_patch_size

# vortex in cell method
# solve L*U=-curl(omega)
@ti.kernel
def compute_omega_curl_x(temp: ti.template()):
for i, j in omega:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)

temp[i, j] = (omega[i, j_p] - omega[i, j_n]) / \
(2.0 * dx)  # par_y_omega

@ti.kernel
def compute_omega_curl_y(temp: ti.template()):
for i, j in omega:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)

sigma[i, j] = - (omega[i_p, j] - omega[i_n, j]) / \
(2.0 * dx)  # -par_x_omega

def reconstruct_velocity_vic():
compute_omega_curl_x(sigma)
poisson_solver(20, ux, sigma)

compute_omega_curl_y(sigma)
poisson_solver(20, uy, sigma)

@ti.kernel
# advection of omega by 2nd order upwind scheme
for i, j in omega:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)

i_p2 = mod(i + 2, Nx)
i_n2 = mod(i - 2 + Nx, Nx)
j_p2 = mod(j + 2, Ny)
j_n2 = mod(j - 2 + Ny, Ny)

ux_pos = max(ux[i, j], 0.0)
ux_neg = min(ux[i, j], 0.0)
uy_pos = max(uy[i, j], 0.0)
uy_neg = min(uy[i, j], 0.0)

# second order
omega_dx_pos = \
(- omega[i_p2, j] + 4.0 * omega[i_p, j] -
3.0 * omega[i, j]) / (2.0 * dx)
omega_dx_neg = \
(omega[i_n2, j] - 4.0 * omega[i_n, j] +
3.0 * omega[i, j]) / (2.0 * dx)
omega_dy_pos = \
(- omega[i, j_p2] + 4.0 * omega[i, j_p] -
3.0 * omega[i, j]) / (2.0 * dx)
omega_dy_neg = \
(omega[i, j_n2] - 4.0 * omega[i, j_n] +
3.0 * omega[i, j]) / (2.0 * dx)

omega_temp[i, j] = omega[i, j] - dt_substep * (
ux_pos * omega_dx_neg + ux_neg * omega_dx_pos +
uy_pos * omega_dy_neg + uy_neg * omega_dy_pos)

def solve_velocity_streamfunc():
solve_streamfunc(20)
reconstruct_velocity_sigma()

def solve_velocity_vic():
reconstruct_velocity_vic()

def solve_biot_savart():
reconstruct_velocity_biot_savart()

def initialize():
initialize_vortices()
solve_streamfunc(1000)

def run_iteration():
num_substeps = 5  # to deal with large CFL
dt_substep = dt / num_substeps
for iters in range(num_substeps):

# may use the swap style in the stable-fluids example
copy_buffer(omega, omega_temp)

if (solver_type == 0):
solve_velocity_streamfunc()
elif (solver_type == 1):
solve_velocity_vic()
else:
solve_biot_savart()

def toggle_solver():
global solver_type
solver_type = (solver_type + 1) % 3
if (solver_type == 0):
sigma.fill(0.0)

class Viewer:
def __init__(self, dump):
self.display_mode = 0
self.is_active = True
self.dump = dump
self.frame = 0

if self.dump:
result_dir = "./results"
self.video_manager = ti.VideoManager(
output_dir=result_dir, framerate=24, automatic_build=False)

def toggle(self):
self.display_mode = (self.display_mode + 1) % 3

def active(self):
return self.is_active

def draw(self, gui):
if self.display_mode == 0:
fill_vorticity()
display_name = "vorticity"
elif self.display_mode == 1:
fill_velocity()
display_name = "velocity"
else:
fill_streamfunc()
display_name = "stream function"

img = color_buffer.to_numpy()
gui.set_image(img)
gui.text(content=f"display: {display_name}",
pos=(0.0, 0.99), color=0xFFFFFF)

if self.dump:
self.video_manager.write_frame(img)
print(f"\rframe {self.frame} written", end="")
self.frame = self.frame + 1

if self.frame == 300:
self.video_manager.make_video(gif=True, mp4=True)
self.is_active = False

def main():
initialize()

viewer = Viewer(False)

gui = ti.GUI("vortex method", Nx, Ny)
while viewer.active():
for e in gui.get_events(ti.GUI.PRESS):
if e.key == ti.GUI.ESCAPE or e.key == 'q':
exit(0)
elif e.key == 'r':
initialize()
elif e.key == ' ':
viewer.toggle()
elif e.key == 's':
toggle_solver()

for iters in range(10):
run_iteration()

viewer.draw(gui)

gui.text(content=f"solver: {solver_names[solver_type]}",
pos=(0, 0.96), color=0xFFFFFF)

gui.text(content="r: reset simulation", pos=(0, 0.86), color=0xFFFFFF)
gui.text(content="q, esc: quit", pos=(0, 0.83), color=0xFFFFFF)
gui.text(content="space: toggle display mode",
pos=(0, 0.8), color=0xFFFFFF)

gui.show()

if __name__ == '__main__':
main()
``````
11 Likes

3D leapfrog

``````import taichi as ti
import math

# ti.init(arch=ti.cpu)
ti.init(arch=ti.gpu)

N = 256
Nx = N
Ny = N
Nz = N
nTotal = Nx * Ny * Nz

dx = 1.0 / N
dt = 0.5 / N
num_substeps = 2  # to deal with large CFL
one_over_six = 1.0 / 6.0
one_over_two_dx = 1.0 / (2.0 * dx)

vel = ti.Vector(3, dt=ti.f32, shape=(Nx, Ny, Nz))  # velocity
omega_0 = ti.Vector(3, dt=ti.f32, shape=(Nx, Ny, Nz))  # vorticity
omega_1 = ti.Vector(3, dt=ti.f32, shape=(Nx, Ny, Nz))  # vorticity
sigma = ti.Vector(3, dt=ti.f32, shape=(Nx, Ny, Nz))  # stream function
volume = ti.var(dt=ti.f32, shape=(Nx, Ny, Nz))  # stream function

class TexPair:
def __init__(self, cur, nxt):
self.cur = cur
self.nxt = nxt

def swap(self):
self.cur, self.nxt = self.nxt, self.cur

@ti.kernel
def copy_buffer(dst: ti.template(), src: ti.template()):
for I in ti.grouped(src):
dst[I] = src[I]

omega = TexPair(omega_0, omega_1)

color_buffer = ti.Vector(3, dt=ti.f32, shape=(Nx, Ny))

@ti.data_oriented
class ColorMap:
def __init__(self, h, wl, wr, c):
self.h = h
self.wl = wl
self.wr = wr
self.c = c

@ti.func
def clamp(self, x):
return ti.max(0.0, ti.min(1.0, x))

@ti.func
def map(self, x):
w = 0.0
if x < self.c:
w = self.wl
else:
w = self.wr
return self.clamp((w-ti.abs(self.clamp(x)-self.c))/w*self.h)

jetR = ColorMap(1.5, .37, .37, .75)
jetG = ColorMap(1.5, .37, .37, .5)
jetB = ColorMap(1.5, .37, .37, .25)

bwrR = ColorMap(1.0, .25, 1, .5)
bwrG = ColorMap(1.0, .5, .5, .5)
bwrB = ColorMap(1.0, 1, .25, .5)

@ti.func
def color_map(c):
# this works by chance, must use ti.func in ti.kernel
# return ti.Vector([jetR.map(c),
#                   jetG.map(c),
#                   jetB.map(c)])
return ti.Vector([bwrR.map(c),
bwrG.map(c),
bwrB.map(c)])

# -------------------------------------

# [-1, 1]
@ti.func
def xCoord(i):
return ((i + 0.5) / Nx) * 2.0 - 1.0

# [-1, 1]
@ti.func
def yCoord(j):
return ((j + 0.5) / Ny) * 2.0 - 1.0

# [-1, 1]
@ti.func
def zCoord(k):
return ((k + 0.5) / Nz) * 2.0 - 1.0

@ti.func
def mod(a, b):
return int(a % b)

@ti.func
def sq(x):
return x ** 2

@ti.kernel
def initialize_vorticity_leapfrog():
for i, j, k in omega.cur:
x = xCoord(i)
y = yCoord(j)
z = zCoord(k)

wx = 0.0
wy = 0.0
wz = 0.0

s = 0.5
ss = -0.9
sigma = 0.012
t = 500.0

d = ti.sqrt(x * x + z * z)

rr = sq(d - s) + sq(y - ss)
mag = ti.exp(- rr / (2.0 * sq(sigma))) * t / d
wx += mag * z
wz += -mag * x
rr = sq(d - s * 0.7) + sq(y - ss)
mag = ti.exp(- rr / (2.0 * sq(sigma))) * t / d
wx += mag * z
wz += -mag * x

omega.cur[i, j, k] = ti.Vector([wx, wy, wz])

def initialize_vortices():
initialize_vorticity_leapfrog()
reconstruct_velocity()

# solve L*x=-rhs using periodic boundary
# using red-black ordering for successive over-relaxation/Gauss-Seidel linear solver
@ti.kernel
x: ti.template(), rhs: ti.template()):
for i, j, k in x:
if mod(i + j + k, 2) == mask:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)
k_p = mod(k + 1, Nz)
k_n = mod(k - 1 + Nz, Nz)

x_update = (
- rhs[i, j, k] * sq(dx)
- x[i_p, j, k]
- x[i_n, j, k]
- x[i, j_p, k]
- x[i, j_n, k]
- x[i, j, k_p]
- x[i, j, k_n]) * -one_over_six

x[i, j, k] = SOR_factor * x_update + \
(1.0 - SOR_factor) * x[i, j, k]

# @ti.kernel
# def fix_streamfunc():
#     # since velocity is the gradient of sigma (the stream function),
#     # the constant drift in sigma is harmless to velocity,
#     # and the field is eased for better visualization
#     sigmaMean = 0.0
#     for I in ti.grouped(sigma):
#         sigmaMean += sigma[I]
#     sigmaMean /= nTotal
#     for I in sigma:
#         sigma[I] -= sigmaMean

def poisson_solver(num_iterations, x, rhs):
SOR_factor = 1.99  # 1.0 for Gauss-Seidel

for iters in range(num_iterations):
linear_solver_step(SOR_factor, iters % 2, x, rhs)

def solve_streamfunc(num_iterations):
# solve L*sigma=-omega using periodic boundary
poisson_solver(num_iterations, sigma, omega.cur)

# fix_streamfunc()

@ti.kernel
def reconstruct_velocity():
# u = curl(sigma)
for i, j, k in sigma:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)
k_p = mod(k + 1, Nz)
k_n = mod(k - 1 + Nz, Nz)

vel[i, j, k][0] = ((sigma[i, j_p, k][2] - sigma[i, j_n, k][2]) -
(sigma[i, j, k_p][1] - sigma[i, j, k_n][1])) / (2.0 * dx)
vel[i, j, k][1] = ((sigma[i, j, k_p][0] - sigma[i, j, k_n][0]) -
(sigma[i_p, j, k][2] - sigma[i_n, j, k][2])) / (2.0 * dx)
vel[i, j, k][2] = ((sigma[i_p, j, k][1] - sigma[i_n, j, k][1]) -
(sigma[i, j_p, k][0] - sigma[i, j_n, k][0])) / (2.0 * dx)

@ti.kernel
# advection of omega by 2nd order upwind scheme
for i, j, k in omega.cur:
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)
k_p = mod(k + 1, Nz)
k_n = mod(k - 1 + Nz, Nz)

i_p2 = mod(i + 2, Nx)
i_n2 = mod(i - 2 + Nx, Nx)
j_p2 = mod(j + 2, Ny)
j_n2 = mod(j - 2 + Ny, Ny)
k_p2 = mod(k + 2, Nz)
k_n2 = mod(k - 2 + Nz, Nz)

ux_pos = max(vel[i, j, k][0], 0.0)
ux_neg = min(vel[i, j, k][0], 0.0)
uy_pos = max(vel[i, j, k][1], 0.0)
uy_neg = min(vel[i, j, k][1], 0.0)
uz_pos = max(vel[i, j, k][2], 0.0)
uz_neg = min(vel[i, j, k][2], 0.0)

# first order
# omega_dx_pos = (  omega.cur[i_p, j, k] - omega.cur[i, j, k]) / (dx)
# omega_dx_neg = (- omega.cur[i_n, j, k] + omega.cur[i, j, k]) / (dx)
# omega_dy_pos = (  omega.cur[i, j_p, k] - omega.cur[i, j, k]) / (dx)
# omega_dy_neg = (- omega.cur[i, j_n, k] + omega.cur[i, j, k]) / (dx)
# omega_dz_pos = (  omega.cur[i, j, k_p] - omega.cur[i, j, k]) / (dx)
# omega_dz_neg = (- omega.cur[i, j, k_n] + omega.cur[i, j, k]) / (dx)

# second order
omega_dx_pos = \
(- omega.cur[i_p2, j, k] + 4.0 * omega.cur[i_p, j, k] -
3.0 * omega.cur[i, j, k]) / (2.0 * dx)
omega_dx_neg = \
(omega.cur[i_n2, j, k] - 4.0 * omega.cur[i_n, j, k] +
3.0 * omega.cur[i, j, k]) / (2.0 * dx)
omega_dy_pos = \
(- omega.cur[i, j_p2, k] + 4.0 * omega.cur[i, j_p, k] -
3.0 * omega.cur[i, j, k]) / (2.0 * dx)
omega_dy_neg = \
(omega.cur[i, j_n2, k] - 4.0 * omega.cur[i, j_n, k] +
3.0 * omega.cur[i, j, k]) / (2.0 * dx)
omega_dz_pos = \
(- omega.cur[i, j, k_p2] + 4.0 * omega.cur[i, j, k_p] -
3.0 * omega.cur[i, j, k]) / (2.0 * dx)
omega_dz_neg = \
(omega.cur[i, j, k_n2] - 4.0 * omega.cur[i, j, k_n] +
3.0 * omega.cur[i, j, k]) / (2.0 * dx)

vort = omega.cur[i, j, k] - dt_substep * (
ux_pos * omega_dx_neg + ux_neg * omega_dx_pos +
uy_pos * omega_dy_neg + uy_neg * omega_dy_pos +
uz_pos * omega_dz_neg + uz_neg * omega_dz_pos)

# vortex stretching term
u_dx = (vel[i_p, j, k] - vel[i_n, j, k]) * one_over_two_dx
u_dy = (vel[i, j_p, k] - vel[i, j_n, k]) * one_over_two_dx
u_dz = (vel[i, j, k_p] - vel[i, j, k_n]) * one_over_two_dx

vort += grad_u @ omega.cur[i, j, k] * dt_substep

omega.nxt[i, j, k] = vort

def solve_velocity_streamfunc():
solve_streamfunc(20)
reconstruct_velocity()

def initialize():
initialize_vortices()
sigma.fill(ti.Vector([0.0, 0.0, 0.0]))
solve_streamfunc(1000)

def run_iteration():
dt_substep = dt / num_substeps
for iters in range(num_substeps):
# omega.swap() # not working
copy_buffer(omega.cur, omega.nxt)

solve_velocity_streamfunc()

@ti.func
light_dir = ti.Vector([1.0, 1.0, 1.0]).normalized()
eye_dir = ti.Vector([0.0, 0.0, 1.0])
half_dir = (light_dir + eye_dir).normalized()
return 0.1 + ti.abs(light_dir.dot(normal)) + 3.0 * ti.pow(ti.abs(half_dir.dot(normal)), 60.0)

@ti.kernel
def volume_render(vol: ti.template(), fb: ti.template()):
for i, j in fb:
c = ti.Vector([0.0, 0.0, 0.0])
f = 1.0
for k in range(Nz):
i_p = mod(i + 1, Nx)
i_n = mod(i - 1 + Nx, Nx)
j_p = mod(j + 1, Ny)
j_n = mod(j - 1 + Ny, Ny)
k_p = mod(k + 1, Nz)
k_n = mod(k - 1 + Nz, Nz)
normal = -ti.Vector([
vol[i_p, j, k] - vol[i_n, j, k],
vol[i, j_p, k] - vol[i, j_n, k],
vol[i, j, k_p] - vol[i, j, k_n]])
normal *= 1.0 / (1e-8 + normal.norm())
v = ti.max(0.0, ti.min(1.0, vol[i, j, k] * 20.0))
c += color_map(v) * f * shading * v
f *= 1.0 - v * 0.05
fb[i, j] = c * 0.08

@ti.kernel
def fill_velocity():
for i, j in color_buffer:
color_buffer[i, j] = color_map(vel[i, j, Nz // 2].norm())

@ti.kernel
def fill_vorticity():
for i, j, k in omega.cur:
volume[i, j, k] = omega.cur[i, j, k].norm() * 0.002

@ti.kernel
def fill_streamfunc():
for i, j in color_buffer:
c = sigma[i, j, Nz // 2].norm() * 5.0 + 0.5
color_buffer[i, j] = color_map(c)

class Viewer:
def __init__(self, dump):
self.display_mode = 0
self.is_active = True
self.dump = dump
self.frame = 0

if self.dump:
result_dir = "./results"
self.video_manager = ti.VideoManager(
output_dir=result_dir, framerate=24, automatic_build=False)

def toggle(self):
self.display_mode = (self.display_mode + 1) % 3

def active(self):
return self.is_active

def draw(self, gui):
if self.display_mode == 0:
fill_vorticity()
color_buffer.fill(ti.Vector([0.0, 0.0, 0.0]))
volume_render(volume, color_buffer)
display_name = "vorticity"
elif self.display_mode == 1:
fill_velocity()
display_name = "velocity"
else:
fill_streamfunc()
display_name = "stream function"

img = color_buffer.to_numpy()
gui.set_image(img)
gui.text(content=f"display: {display_name}",
pos=(0.0, 0.99), color=0xFFFFFF)

if self.dump:
self.video_manager.write_frame(img)
print(f"\rframe {self.frame} written", end="")
self.frame = self.frame + 1

if self.frame == 200:
self.video_manager.make_video(gif=True, mp4=True)
self.is_active = False

def main():
initialize()

viewer = Viewer(False)

gui = ti.GUI("vortex method", Nx, Ny)
while viewer.active():
for e in gui.get_events(ti.GUI.PRESS):
if e.key == ti.GUI.ESCAPE or e.key == 'q':
exit(0)
elif e.key == 'r':
initialize()
elif e.key == ' ':
viewer.toggle()

for iters in range(10):
run_iteration()

viewer.draw(gui)

gui.text(content="r: reset simulation", pos=(0, 0.86), color=0xFFFFFF)
gui.text(content="q, esc: quit", pos=(0, 0.83), color=0xFFFFFF)
gui.text(content="space: toggle display mode",
pos=(0, 0.8), color=0xFFFFFF)

gui.show()

if __name__ == '__main__':
main()
``````
8 Likes

cg的一是比较注重视觉效果，模拟湍流比较多，二是精度要求没那么高。工程的cfd不追求视觉效果，而是希望求得准确的压力分布等用于受力分析，比较注重定常问题的求解，对于非定常的湍流问题也希望求得准确的平均值。但是因为都是以求解NS方程为主，两者的计算方法有很多共同点，例如基本都会使用有限差分之类的数值方法。

4 Likes

2 Likes

taichi的主要优势是运行环境简单，以及支持多后端和可微编程，可以快速验证计算量大的模型。用cpp可以进一步优化，这个3d计算用最快的opengl后端1.37s每帧的情况下，cpp的cuda实现用8x8x8的block未优化时间差不多，用shared memory优化在0.9s左右。

2 Likes