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

这里的Demo是以前用cpp写的一个2D测试改的,taichi用cpu后端要比写原生cpp慢一些,但是gpu后端就快多了,基于grid的计算性能还是可以的。stream function solver是最快的;vortex-in-cell solver慢一些,可能和实现有关,解了两次scalar poisson;Biot-Savart solver就慢多了,用vortex particles+fast multiple还是可以的,这里直接range遍历了。

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

        # add vortex
        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
        # add vortex
        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
def linear_solver_step(SOR_factor: ti.f32, mask: ti.template(),
                       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():
    radius = 100
    patch_size = 4
    num_patches = radius // patch_size
    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
def advect_vorticity(dt_substep: ti.f32):
    # 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):
        advect_vorticity(dt_substep)

        # 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()
10 Likes

vort
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)

        # add vortex ring
        rr = sq(d - s) + sq(y - ss)
        mag = ti.exp(- rr / (2.0 * sq(sigma))) * t / d
        wx += mag * z
        wz += -mag * x
        # add vortex ring
        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
def linear_solver_step(SOR_factor: ti.f32, mask: ti.template(),
                       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
def advect_vorticity(dt_substep: ti.f32):
    # 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

        grad_u = ti.Matrix.cols([u_dx, u_dy, u_dz])

        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):
        advect_vorticity(dt_substep)
        # omega.swap() # not working
        copy_buffer(omega.cur, omega.nxt)

    solve_velocity_streamfunc()


@ti.func
def phong_shading(normal):
    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())
            shading = phong_shading(normal)
            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

哇这个有点厉害!

你好!搭车问个小白问题:通常所说的计算机视觉的流体模拟计算(fluid simulation solver for computer graphics)和工程或者物理上求解NS方程的模拟计算有什么区别?
谢谢!

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

4 Likes

谢谢!所以其实只是在求解的细节侧重上有所不同,基本的方法都是一致的。。特别Taichi里支持了32位和64位的浮点数,其实只要去设计高精度的求解过程,框架都是可以共用的。何况工程实际其实也并不是“那么”的讲究精度,和实际实验能够一定程度的吻合就已经很有用了。

我是从事工程流体计算的背景,最近才发现其实计算机视觉领域的流体计算做得非常深入,特别是对计算速度的优化,对程序代码的设计理念,很多都非常值得工程领域借鉴学习。

1 Like

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

2 Likes