pythonfluid-dynamicstaichi

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!

Solution

  • There are two main errors in your code (at least, that I've spotted so far).

    Firstly, you do NOT need to offset x,y from i,j by 0.5 in advect_density - you would introduce a bias that tends to move things in the wrong direction. (Similarly in set_velocity_vortex and add_density.)

    Secondly, in your original velocity it should be (with j replacing i for the logical test):

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

    I've made those changes in your code below. I've also switched to the vortex velocity field (because it's more interesting!) and switched off diffusion (because the diffusivity that you set is very high - it smears things out too rapidly to see what is going on).

    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


##################################
# Parametros de la simulacion
##################################
res = 512
h = 1.0 / res
dt = 0.03

# Coeficiente de difusion
k = 0.00003

# Numero de iteraciones Jacobi
JACOBI_ITERS = 32

# Parametros 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 j < third_height:  # Upper third
            vx = 1.0  # Velocity to the right (positive x-direction)
            vy = 0.0  # No vertical velocity
        elif j < 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 vortice simple
##################################
@ti.kernel
def set_velocity_vortex(v: ti.template()):
    for i, j in v.cur:
        x = i
        y = j
        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 raton
##################################
@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
        cy = j
        # Distancia al raton
        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
#  - Interpolacion 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]

    # Interpolacion
    return (
        d00 * (1 - sx) * (1 - sy)
        + d10 * sx * (1 - sy)
        + d01 * (1 - sx) * sy
        + d11 * sx * sy
    )


# ##################################
# # advect
# #  - Adveccion Semi-Lagrangiana
# #    dens_in -> dens_out usando velocidad (vel)
# ##################################
# @ti.kernel

@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 pequena, nos saltamos la interpolacion
        if abs(vx) + abs(vy) < epsilon:
            # Usamos indice entero => exactamente dens_in[i,j]
            dens_out[i, j] = dens_in[i, j]
        else:
            # Semi-Lagrangian normal
            x = i
            y = j

            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) anadir densidad
    add_density(dens.cur, input_data)

    # 2) difundir con Jacobi
#   diffuse_jacobi(dens.cur, dens.nxt, k)
#   dens.swap()  # Swap despues de la difusion para que dens.cur tenga el resultado

    # 3) advect con velocity
    advect_density(dens.cur, dens.nxt, vel.cur)  # Usamos vel.cur para la adveccion
    dens.swap()  # Swap despues de la adveccion 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

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

        # Simulacion
        if not paused:
            step(input_data)

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


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

    enter image description here