fslbm.py

#!/usr/bin/env python

import taichi as ti
import numpy as np
from matplotlib import cm


# Basic Taichi settings
ti.init(ti.gpu)
nx = 2000
ny = 400
shape = (nx, ny)
steps = 1000000

######################
# Non-dimensional factors
# C_l      := 1000 (m)
# C_t      := 10^5 (s)
# C_{\rho} := 1.0 (kg/m^d)
######################
# C_u      := C_l/C_t (m/s)
# C_v      := C_rho*C_l/C_t
######################
C_l = 10**-3
C_t = 10**-5
C_rho = 1.0
C_u = C_l / C_t
C_v = C_rho*C_l/C_t
def densityToSI(x): return x*C_rho
def lengthToSI(x): return x*C_l
def velocityToSI(x): return x*C_u
def viscosityToSI(x): return x*C_v
def densityFromSI(x): return x/C_rho
def lengthFromSI(x): return x/C_l
def velocityFromSI(x): return x/C_u
def viscosityFromSI(x): return x/C_v


# Macroscopic values
density = ti.field(dtype=ti.f32, shape=shape)
velocity = ti.Vector.field(n=2, dtype=ti.f32, shape=shape)
velocity_old = ti.Vector.field(n=2, dtype=ti.f32, shape=shape)
velocity_range = (0.0, 0.015)
wall_velocity = ti.Vector([0.01, 0])
radius = 40
center = ti.Vector([nx//6, ny//2])
viscosity = 0.00001
blending_factor = 1.0


def printInfo():
    print(f"Scene size: {lengthToSI(nx)} x {lengthToSI(ny)} (m)")
    print(f"Re: {2*radius*wall_velocity[0]/viscosity}")
    print(f"Wall velocity: {velocityToSI(wall_velocity[0])} (m/s)")
    print(f"Radius: {lengthToSI(radius)} (m)")
    print(
        f"Velocity Range: [{velocityToSI(velocity_range[0])}, {velocityToSI(velocity_range[1])}] (m/s)")
    print(f"Kinematic Viscosity: {viscosityToSI(viscosity)} (kg/(m s))")
    print(f"Step per second: {1//C_t}")


# D2Q9 {velocity set}
w = ti.field(dtype=ti.f32, shape=9)
c = ti.Vector.field(n=2, dtype=ti.i32, shape=9)
w.from_numpy(np.array([4.0 / 9.0, 1.0 / 9.0, 1.0 / 9.0, 1.0 / 9.0, 1.0 / 9.0, 1.0 / 36.0,
                       1.0 / 36.0, 1.0 / 36.0, 1.0 / 36.0], dtype=np.float32))
c.from_numpy(np.array([[0, 0], [1, 0], [0, 1], [-1, 0], [0, -1], [1, 1],
                       [-1, 1], [-1, -1], [1, -1]], dtype=np.int32))


@ti.func
def f_eq(i, j, k):
    # Maxwell-Boltzmann Distribution with Isothermal assumption
    rho, u = density[i, j], velocity_old[i, j]
    uc, uu = u.dot(c[k]), u.dot(u)
    return w[k] * rho * \
        (1 + 3*uc + 4.5*uc**2 - 1.5*uu)


@ti.func
def f(i, j, k):
    # Since $\tau^{\star} = 1$, $f = f^{\mathrm{eq}}$
    i_ = i - c[k][0]
    j_ = j - c[k][1]
    assert 0 <= i_ < nx and 0 <= j_ < ny
    return f_eq(i_, j_, k)


@ti.kernel
def init():
    for i, j in ti.ndrange(nx, ny):
        density[i, j] = 1
        velocity_old[i, j] = ti.Vector([0.0, 0.0])


@ti.kernel
def L1():
    for i, j in ti.ndrange((1, nx-1), (1, ny-1)):
        rho, u = 0.0, ti.Vector([0.0, 0.0])
        for k in ti.static(range(9)):
            f_ = f(i, j, k)
            rho += f_
            u += c[k] * f_
        density[i, j] = rho
        velocity[i, j] += u / rho


@ti.kernel
def L2():
    for i, j in ti.ndrange((1, nx-1), (1, ny-1)):
        u = ti.Vector([0.0, 0.0])
        for k in ti.static([1, 2, 3, 4]):
            i_ = i - c[k][0]
            j_ = j - c[k][1]
            u += velocity_old[i_, j_]
        u = (viscosity - 1/6) * (u - 4 * velocity_old[i, j])
        velocity[i, j] += u


@ti.kernel
def boundary():
    for i, j in ti.ndrange(nx, ny):
        if i == 0:  # left wall
            velocity[i, j] = wall_velocity
        if j == 0 or j == ny-1:  # top wall
            velocity[i, j] = ti.Vector([0.0, 0.0])
        # if abs(i - center[0]) <= length and abs(j - center[1]) <= length:
        if (ti.Vector([i, j]) - center).norm() < radius:
            velocity[i, j] = ti.Vector([0.0, 0.0])


@ti.kernel
def blend():
    for i, j in ti.ndrange(nx, ny):
        velocity_old[i, j] = (1 - blending_factor) * \
            velocity_old[i, j] + blending_factor * velocity[i, j]


def render():
    image = np.linalg.norm(velocity_old.to_numpy(), axis=2)
    image = np.clip(image, velocity_range[0],
                    velocity_range[1]) / velocity_range[1]
    return image


gui = ti.GUI("FSLBM", shape)
printInfo()
init()
for i in range(steps):
    L1()  # L1 operator
    L2()  # L2 operator
    boundary()
    blend()
    if i % 100 == 0:
        gui.set_image(cm.plasma(render()))
        gui.show()
    velocity.fill(ti.Vector([0.0, 0.0]))