基于Python与Taichi实现Jos Stam稳定流体模拟时出现异常密度漂移问题

根据你描述的密度向左下漂移的问题，结合代码分析，主要存在以下几个关键问题，下面逐一说明并给出修改方案：

一、核心问题分析 #

1. 速度场坐标轴对应错误（最可能导致漂移的原因） #

你在set_velocity_pattern中错误地将第一个维度i当作了垂直方向（y轴），但实际上你的场shape=(res, res)中，i对应水平x轴，j才是垂直y轴。这导致你定义的"上三分之一区域"实际是窗口的左三分之一区域，速度方向与预期完全不符，进而表现出异常漂移。

2. 扩散步骤的双缓冲逻辑混乱 #

你的diffuse_jacobi函数在每次Jacobi迭代后直接将结果复制回原场，再加上多余的dens.swap()操作，导致双缓冲完全失去意义，甚至可能破坏模拟状态的一致性。

3. 平流步骤的边界与坐标匹配问题 #

虽然Semi-Lagrangian回溯公式本身正确，但窗口坐标系（y轴从下到上）与模拟场的索引对应关系容易被忽略，可能导致垂直方向的隐性漂移。

二、修改后的完整代码方案 #

下面是修复后的关键代码部分，保留你的原有结构并修正上述问题：

############################
# etapa2.py
# Etapa 2: Densidad con Difusión (Jacobi) + Advección (Semi-Lagrangiana)
# 修复坐标轴对应关系、双缓冲逻辑与速度场定义
############################

import numpy as np
import taichi as ti

##################################
# Inicializar Taichi
##################################
arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda
ti.init(arch=arch)


##################################
# FieldPair para double buffering
##################################
class FieldPair:
    def __init__(self, current_field, next_field):
        self.cur = current_field
        self.nxt = next_field

    def swap(self):
        self.cur, self.nxt = self.nxt, self.cur


##################################
# Parámetros de la simulación
##################################
res = 512
h = 1.0 / res
dt = 0.03

# Coeficiente de difusión
k = 0.00003

# Número de iteraciones Jacobi
JACOBI_ITERS = 32

# Parámetros de fuente de densidad
s_dens = 10.0
s_radius = res / 15.0

# Umbral para considerar velocidad "nula"
epsilon = 1e-6

##################################
# Campos Taichi
##################################
density_1 = ti.field(float, shape=(res, res))
density_2 = ti.field(float, shape=(res, res))
dens = FieldPair(density_1, density_2)

velocity_1 = ti.Vector.field(2, dtype=ti.f32, shape=(res, res))
velocity_2 = ti.Vector.field(2, dtype=ti.f32, shape=(res, res))
vel = FieldPair(velocity_1, velocity_2)


############################
# 修复：速度场基于y轴（j）定义区域，匹配窗口垂直方向
############################
@ti.kernel
def set_velocity_pattern(v: ti.template()):
    for i, j in v.cur:
        third_height = res / 3.0
        vx = 0.0
        vy = 0.0

        # j是垂直y轴，j越大对应窗口越靠上
        if j > 2 * third_height:  # 上三分之一区域
            vx = 1.0  # 向右
        elif j > third_height:  # 中间三分之一区域
            vx = -1.0  # 向左
        else:  # 下三分之一区域
            vx = 1.0  # 向右

        v.cur[i, j] = ti.Vector([vx, vy])
        v.nxt[i, j] = ti.Vector([vx, vy])


##################################
# set_velocity_vortex
#  - Inicializa un vórtice simple
##################################
@ti.kernel
def set_velocity_vortex(v: ti.template()):
    for i, j in v.cur:
        x = i + 0.5
        y = j + 0.5
        cx = res * 0.5
        cy = res * 0.5
        dx = x - cx
        dy = y - cy
        # Campo tangencial
        scale = 2e-1
        vx = -dy * scale
        vy = dx * scale
        v.cur[i, j] = ti.Vector([vx, vy])
        v.nxt[i, j] = ti.Vector([vx, vy])


##################################
# add_sources
#  - Inyecta densidad alrededor del ratón
##################################
@ti.kernel
def add_density(dens_field: ti.template(), input_data: ti.types.ndarray()):
    for i, j in dens_field:
        densidad = input_data[2] * s_dens
        mx, my = input_data[0], input_data[1]
        # Centro de la celda
        cx = i + 0.5
        cy = j + 0.5
        # Distancia al ratón
        d2 = (cx - mx) ** 2 + (cy - my) ** 2
        dens_field[i, j] += dt * densidad * ti.exp(-6.0 * d2 / s_radius**2)


##################################
# 修复：Jacobi迭代改为分步骤双缓冲实现
##################################
@ti.kernel
def diffuse_jacobi_step(dens_in: ti.template(), dens_out: ti.template(), diff: float):
    a = diff * dt / (h * h)
    for i, j in dens_in:
        i_left = (i - 1 + res) % res
        i_right = (i + 1) % res
        j_down = (j - 1 + res) % res
        j_up = (j + 1) % res

        dens_out[i, j] = (
            dens_in[i, j]
            + a
            * (
                dens_in[i_left, j]
                + dens_in[i_right, j]
                + dens_in[i, j_down]
                + dens_in[i, j_up]
            )
        ) / (1.0 + 4.0 * a)

def diffuse(dens_pair: FieldPair, diff: float):
    # 迭代过程中始终用cur计算nxt，swap后cur保持最新状态
    for _ in range(JACOBI_ITERS):
        diffuse_jacobi_step(dens_pair.cur, dens_pair.nxt, diff)
        dens_pair.swap()


##################################
# bilerp
#  - Interpolación bilineal de dens_in
#  - Asume x,y en float con "wrap" en [0,res)
##################################
@ti.func
def bilerp(dens_in, x, y):
    # Indices base
    x0 = int(ti.floor(x)) % res
    y0 = int(ti.floor(y)) % res

    # Indices vecinos
    x1 = (x0 + 1) % res
    y1 = (y0 + 1) % res

    # Distancias fraccionarias
    sx = x - ti.floor(x)
    sy = y - ti.floor(y)

    # Valores en esquinas
    d00 = dens_in[x0, y0]
    d01 = dens_in[x0, y1]
    d10 = dens_in[x1, y0]
    d11 = dens_in[x1, y1]

    # Interpolación
    return (
        d00 * (1 - sx) * (1 - sy)
        + d10 * sx * (1 - sy)
        + d01 * (1 - sx) * sy
        + d11 * sx * sy
    )


@ti.kernel
def advect_density(dens_in: ti.template(), dens_out: ti.template(), vel: ti.template()):
    for i, j in dens_in:
        # Velocidad
        vx, vy = vel[i, j]

        if abs(vx) + abs(vy) < epsilon:
            dens_out[i, j] = dens_in[i, j]
        else:
            # Centro de la celda
            x = i + 0.5
            y = j + 0.5

            # 正确的Semi-Lagrangian回溯公式
            x_old = x - vx * dt
            y_old = y - vy * dt

            dens_out[i, j] = bilerp(dens_in, x_old, y_old)


##################################
# init
##################################
def init():
    dens.cur.fill(0)
    dens.nxt.fill(0)
    vel.cur.fill(0)
    vel.nxt.fill(0)
    # Inicializar la velocidad con un campo estatico
    # set_velocity_vortex(vel)
    set_velocity_pattern(vel)


##################################
# step
#  - add_sources -> diffuse -> advect
##################################
def step(input_data):
    # 1) añadir densidad
    add_density(dens.cur, input_data)

    # 2) difundir con Jacobi（内部已处理双缓冲）
    diffuse(dens, k)

    # 3) advect con velocity
    advect_density(dens.cur, dens.nxt, vel.cur)
    dens.swap()  # Swap后cur为平流后的结果


##################################
# main
##################################
def main():
    # Ventana Taichi
    window = ti.ui.Window("Etapa 2: Diffusion + Advection", (res, res), vsync=True)
    canvas = window.get_canvas()

    paused = False

    # Inicializar todo
    init()

    while window.running:
        # input_data = (mx, my, active)
        input_data = np.zeros(3, dtype=np.float32)

        # Controles
        if window.get_event(ti.ui.PRESS):
            e = window.event
            if e.key == ti.ui.ESCAPE:
                break
            elif e.key == "r":
                paused = False
                init()
            elif e.key == "p":
                paused = not paused

        # Ratón
        if window.is_pressed(ti.ui.RMB):
            mouse_xy = window.get_cursor_pos()
            input_data[0] = mouse_xy[0] * res
            input_data[1] = mouse_xy[1] * res
            input_data[2] = 1.0

        # Simulación
        if not paused:
            step(input_data)

        # Render
        canvas.set_image(dens.cur)
        window.show()


# Si se llama directamente el script
if __name__ == "__main__":
    main()

三、关键修复点说明 #

1. 速度场坐标轴修正：
   将set_velocity_pattern中的区域判断从i改为j，确保"上/中/下三分之一"对应窗口的垂直方向，速度方向与视觉预期一致。

2. Jacobi迭代双缓冲修复：
   拆分原diffuse_jacobi为diffuse_jacobi_step（单步迭代）和diffuse（多步循环+swap），严格遵循双缓冲逻辑，避免状态覆盖。

3. 移除多余Swap操作：
   删除扩散步骤后的无意义Swap，仅在平流步骤后保留必要的Swap，确保模拟状态的正确性。

4. 坐标系一致性验证：
   保持add_density与平流函数的坐标对应关系，确保鼠标点击位置与模拟场的坐标完全匹配。

运行修改后的代码后，密度的流动方向应与速度场一致，不会再出现无规律的左下漂移。

备注：内容来源于stack exchange，提问作者German