import taichi as ti
import numpy as np
ti.init(arch=ti.gpu)
参数设置
dim = 3
n = 128 # 空间网格数
n_particles = 51200 # 粒子数
dt = 1e-3 # 时间步长
alpha = 0.5 # 速度场更新系数
dx = 1.0 / n # 空间步长（网格尺寸）
inv_dx = 1.0 / dx # 空间步长的倒数
temp_decay = 0.95 # 烟雾温度衰减因子
cavas_color = (0, 0, 0) # 画布颜色
定义 Taichi 字段
smoke_density = ti.Vector.field(dim, dtype=float, shape=(n, n, n)) # 烟雾密度
temperature = ti.field(dtype=float, shape=(n, n, n)) # 温度场
velocity = ti.Vector.field(dim, dtype=float, shape=(n, n, n)) # 速度场
new_velocity = ti.Vector.field(dim, dtype=float, shape=(n, n, n)) # 新速度场
pressure = ti.field(dtype=float, shape=(n, n, n)) # 压力场
particle_positions = ti.Vector.field(dim, dtype=float, shape=n_particles) # 粒子位置，有n*1个格，每格含一个三维向量
particle_velocities = ti.Vector.field(dim, dtype=float, shape=n_particles) # 粒子速度
@ti.kernel
def init_smoke():
for i, j, k in smoke:
smoke_density[i, j, k] = (0.0, 0.0, 0.0)
temperature[i, j, k] = 23.0
velocity[i, j, k] = ti.Vector([0.0, 0.0, 0.0]) # 整个烟雾场的速度
new_velocity[i, j, k] = ti.Vector([0.0, 0.0, 0.0])
pressure[i, j, k] = 0.0
for p in range(n_particles):
particle_positions[p] = ti.Vector([ti.random(), ti.random(), ti.random()]) # 粒子位置在0-1中随机
particle_velocities[p] = ti.Vector([0, 0, 0])
@ti.kernel
def step_smoke():
# 计算速度场
for i, j, k in velocity:
vel = velocity[i, j, k]
new_velocity[i, j, k] = vel - dt * alpha * pressure[i, j, k] / dx
# 更新速度场
for i, j, k in velocity:
velocity[i, j, k] = new_velocity[i, j, k]
# 更新烟雾
for i, j, k in smoke:
smoke[i, j, k] = ti.max(smoke[i, j, k] * temp_decay, 0.0)
# 更新粒子位置
for n in range(n_particles):
p_pos = particle_positions[n]
p_vel = particle_velocities[n]
new_p_pos = p_pos + dt * p_vel
new_p_vel = p_vel + dt * velocity[int(new_p_pos[0] * inv_dx), int(new_p_pos[1] * inv_dx), int(new_p_pos[2] * inv_dx)]
particle_positions[n] = new_p_pos
particle_velocities[n] = new_p_vel
int_grid_x = ti.floor(particle_positions[n].x / dx) # 获取每个粒子所处的网格
int_grid_y = ti.floor(particle_positions[n].y / dx)
int_grid_z = ti.floor(particle_positions[n].z / dx)
smoke_density[int(int_grid_x), int(int_grid_y), int(int_grid_z)] = [0.3, 0.3, 0.3] # 把粒子所在的网格的颜色设置为定值
初始化
init_smoke()
window = ti.ui.Window(“Test for 3D_smoke”, (768, 768))
cavas = window.get_canvas()
cavas.set_background_color(cavas_color) # 画布颜色
scene = ti.ui.Scene()
camera = ti.ui.Camera()
camera.position(2, 2, 1)
while window.running:
step_smoke()
camera.track_user_inputs(window, movement_speed=0.003, hold_key=ti.ui.RMB)
scene.set_camera(camera)
scene.ambient_light((0.8, 0.8, 0.8))
scene.point_light(pos=(0.5, 1.5, 1.5), color=(1, 1, 1))
scene.particles(particle_positions, color= (0.5, 0.5, 0.5), radius=0.000001)
cavas.scene(scene)
window.show()
