[Python] Stable fluid simulation with Python and Taichi has strange drifting

I'm trying to implement Jos Stam's stable fuild in paper: https://graphics.cs.cmu.edu/nsp/course/15-464/Spring11/papers/StamFluidforGames.pdf I want to do it in 2D but do wraparound when you hit the borders. I kind of managed to get he first stages working but as soon as I add a static velocity field the densities start drifting diagonally (down-left). If you try it you can add density with right mouse button. Can anybody see what the issue could be in the code? I'm thinking it could be related to the bilinear interpolation (but not sure, the velocity field is not causing it, the move seems to be a miscalculation):

############################
# etapa2.py
# Etapa 2: Densidad con Difusión (Jacobi) + Advección (Semi-Lagrangiana)
# Modificado para usar FieldPair para la velocidad y kernel de inicialización en cero
############################

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)


############################
# Inicializa un patron de tercios el campo de velocidad
############################
@ti.kernel
def set_velocity_pattern(v: ti.template()):
for i, j in v.cur:
# Calculate the third of the window
third_height = res / 3.0
vx = 0.0
vy = 0.0

if i < third_height: # Upper third
vx = 1.0 # Velocity to the right (positive x-direction)
vy = 0.0 # No vertical velocity
elif i < 2 * third_height: # Middle third
vx = -1.0 # Velocity to the left (negative x-direction)
vy = 0.0 # No vertical velocity
else: # Bottom third
vx = 1.0 # Velocity to the right (positive x-direction)
vy = 0.0 # No vertical velocity

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)


##################################
# diffuse_jacobi
# - Usa iteraciones Jacobi para difundir
##################################
@ti.kernel
def diffuse_jacobi(dens_in: ti.template(), dens_out: ti.template(), diff: float):
a = diff * dt / (h * h)
for _ in ti.static(range(JACOBI_ITERS)):
# 1) Jacobi
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)

# 2) Copiar dens_out -> dens_in
for i, j in dens_in:
dens_in[i, j] = dens_out[i, j]


##################################
# 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
)


# ##################################
# # advect
# # - Advección Semi-Lagrangiana
# # dens_in -> dens_out usando velocidad (vel)
# ##################################
# @ti.kernel
# def advect(dens_in: ti.template(), dens_out: ti.template(), vel: ti.template()):
# for i, j in dens_in:
# # Centro de la celda
# x = i + 0.5
# y = j + 0.5

# # Velocidad en (i,j)
# vx, vy = vel[i, j]

# # Posición anterior (backtrace)
# # Nota: asumiendo velocity con "cells per frame" o similar
# # factor dt ya está contemplado en velocity scale
# x_old = x - vx * dt
# y_old = y - vy * dt

# # Hacemos bilinear interpolation
# dens_out[i, j] = bilerp(dens_in, x_old, y_old)


@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]

# Si la magnitud es muy pequeña, nos saltamos la interpolación
if abs(vx) + abs(vy) < epsilon:
# Usamos índice entero => exactamente dens_in[i,j]
dens_out[i, j] = dens_in[i, j]
else:
# Semi-Lagrangian normal
x = i + 0.5
y = j + 0.5

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_jacobi(dens.cur, dens.nxt, k)
dens.swap() # Swap después de la difusión para que dens.cur tenga el resultado

# 3) advect con velocity
advect_density(dens.cur, dens.nxt, vel.cur) # Usamos vel.cur para la advección
dens.swap() # Swap después de la advección para que dens.cur tenga el resultado


##################################
# 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()


Any hints appreciated!

Continue reading...