taichi实现三维LBM

以下是使用 Taichi 实现 Lattice Boltzmann 方法三维模拟的示例代码，其中采用了 MRT 碰撞模型。代码的注释中包含了详细的解释和说明。

import taichi as ti

# 模拟参数
# 格子数
NX, NY, NZ = 64, 64, 64
# 时间步长
DT = 1.0
# 模拟总时间
T_TOTAL = 100
# 碰撞模型参数
S = 1.8
OMEGA = 1.0 / (3.0 * S + 0.5)
M = ti.Matrix([
    [ 1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0],
    [-2.0, -1.0, -1.0, -1.0, -1.0,  2.0,  2.0,  2.0],
    [ 0.0,  1.0,  0.0,  0.0,  0.0, -1.0,  0.0,  0.0],
    [ 0.0,  0.0,  1.0,  0.0,  0.0,  0.0, -1.0,  0.0],
    [ 0.0,  0.0,  0.0,  1.0,  0.0,  0.0,  0.0, -1.0],
    [ 1.0, -1.0,  1.0, -1.0,  1.0, -1.0,  1.0, -1.0],
    [ 0.0,  0.0,  0.0,  0.0,  1.0, -1.0,  1.0, -1.0],
    [ 0.0,  0.0,  0.0,  0.0,  0.0,  1.0, -1.0,  1.0],
])
M_INV = M.inverse()

# Taichi 的配置
ti.init(arch=ti.gpu)

# 格子的状态量，包括速度和密度
vel = ti.Vector.field(3, float, (NX, NY, NZ))
rho = ti.field(float, (NX, NY, NZ))
f = ti.Vector.field(8, float, (NX, NY, NZ))

# 初始化函数
@ti.kernel
def init():
    # 初始化速度和密度
    for i, j, k in vel:
        vel[i, j, k] = ti.Vector([0.0, 0.0, 0.0])
        rho[i, j, k] = 1.0
    # 初始化分布函数
    for i, j, k in f:
        f[i, j, k] = ti.Vector([OMEGA * rho[i, j, k]] * 8)

@ti.kernel
def collide():
# 计算宏观量
for i, j, k in rho:
rho[i, j, k] = f[i, j, k].sum()
vel[i, j, k] = ti.Vector.zero(float, 3)
for l in ti.static(range(8)):
vel[i, j, k] += f[i, j, k][l] * M[l]
vel[i, j, k] /= rho[i, j, k]
# 碰撞操作
for i, j, k in f:
# 计算速度差
delta_v = vel[i, j, k] - ti.Vector([0.0, 0.0, 0.0])
# 计算 MRT 变换后的分布函数
f_eq = ti.Vector.zero(float, 8)
f_tmp = ti.Vector.zero(float, 8)
for l in ti.static(range(8)):
f_eq[l] = rho[i, j, k] * OMEGA * (1.0 + M[l] @ delta_v / S + (M[l] @ delta_v)2 / (2.0 * S2) - delta_v.norm_sqr() / (2.0 * S**2))
f_tmp[l] = f[i, j, k][l]
f[i, j, k] = f_tmp @ M_INV
f[i, j, k] -= S * (f[i, j, k] - f_eq)

边界条件函数
@ti.kernel
def apply_bc():
# 对于每个面，设置其上的分布函数为当前格子的反向分布函数
  for i, j, k in ti.static(ti.ndrange(NX, NY, NZ)):
    if i == 0:
      f[i, j, k][0] = f[i, j, k][1]
      f[i, j, k][3] = f[i, j, k][4]
      f[i, j, k][7] = f[i, j, k][6]
    if i == NX - 1:
      f[i, j, k][1] = f[i, j, k][0]
      f[i, j, k][4] = f[i, j, k][3]
      f[i, j, k][6] = f[i, j, k][7]
    if j == 0:
      f[i, j, k][2] = f[i, j, k][5]
      f[i, j, k][5] = f[i, j, k][2]
      f[i, j, k][6] = f[i, j, k][7]
    if j == NY - 1:
      f[i, j, k][5] = f[i, j, k][2]
      f[i, j, k][2] = f[i, j, k][5]
      f[i, j, k][7] = f[i, j, k][6]
    if k == 0:
      f[i, j, k][0] = f[i, j, k][1]
      f[i, j, k][1] = f[i, j, k][0]
      f[i, j, k][2] = f[i, j, k][5]
    if k == NZ - 1:
      f[i, j, k][3] = f[i, j, k][7]