taichi语言进行SPH数据集导出

import taichi as ti
import numpy as np
import os

ti.init(arch=ti.gpu)

# 模拟参数
num_particles = 500
particle_mass = 1.0
dt = 1e-4
h = 0.1  # 平滑长度
rho0 = 1000.0  # 参考密度
k = 1e3  # 弹性模量
mu = 1e-1  # 粘性系数

# 粒子属性
particle_position = ti.Vector.field(2, dtype=float, shape=num_particles)
particle_velocity = ti.Vector.field(2, dtype=float, shape=num_particles)
particle_density = ti.field(dtype=float, shape=num_particles)

# GUI显示设置
gui = ti.GUI("SPH Simulation", res=(800, 800))
output_folder = "sph_dataset"
os.makedirs(output_folder, exist_ok=True)

# 初始化粒子位置和速度
@ti.kernel
def initialize_particles():
    for i in range(num_particles):
        particle_position[i] = [ti.random(), ti.random()]  # 随机初始化位置
        particle_velocity[i] = [0, 0]  # 初始速度为零

initialize_particles()

# SPH核函数
@ti.kernel
def kernel_func(r, h):
    q = r / h
    if 0 <= q < 1:
        return 15 / (7 * np.pi * h**2) * (2 / 3 - q**2 + 0.5 * q**3)
    elif 1 <= q < 2:
        return 15 / (7 * np.pi * h**2) * 1 / 6 * (2 - q)**3
    else:
        return 0

# 计算密度和压力力
@ti.kernel
def compute_density_and_pressure():
    for i in range(num_particles):
        particle_density[i] = 0.0
        for j in range(num_particles):
            r = particle_position[i] - particle_position[j]
            particle_density[i] += particle_mass * kernel_func(r.norm(), h)

compute_density_and_pressure()

# 计算压力和粘性力
@ti.kernel
def compute_forces():
    for i in range(num_particles):
        pressure_force = ti.Vector([0.0, 0.0])
        viscosity_force = ti.Vector([0.0, 0.0])

        for j in range(num_particles):
            if i != j:
                r = particle_position[i] - particle_position[j]
                q = r.norm() / h

                # 计算压力力
                pressure_force += -particle_mass * (particle_density[i] + particle_density[j]) / (2 * rho0) \
                                  * k * kernel_func(q, h) * r / q

                # 计算粘性力
                viscosity_force += mu * particle_mass / (particle_density[i] + particle_density[j]) \
                                  * (particle_velocity[j] - particle_velocity[i]) * kernel_func(q, h)

        # 更新速度
        particle_velocity[i] += (pressure_force + viscosity_force) * dt / particle_mass

compute_forces()

# 更新位置并导出数据集
@ti.kernel
def update_position_and_export():
    for i in range(num_particles):
        particle_position[i] += particle_velocity[i] * dt

    # 将粒子位置数据导出为数据集
    for i in range(num_particles):
        pos_x, pos_y = particle_position[i]
        with ti.static(output_folder=output_folder):
            with open(f'{output_folder}/particle_data.txt', 'a') as f:
                f.write(f'{pos_x} {pos_y}\n')

update_position_and_export()

# GUI可视化
while gui.running:
    gui.circles(particle_position.to_numpy(), color=0xFFFFFF, radius=3)
    gui.show()
    update_position_and_export()