Fluid simulation using Navier-Stokes equation in Python

fluidSim0.py

import pygame as pg
import colorsys
from time import time
from copy import deepcopy
from pprint import pprint

# COLORS

BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
FLUID = (0, 0, 255)
RED = (255, 0, 0)

# END COLORS

# SETTINGS

N = 100
visc = 1

# END SETTINGS

# GRID VARIABLES
size = (N+2)**2

dens = [0] * size
dens_prev = [0] * size

u = [0] * size
u_prev = [0] * size

v = [0] * size
v_prev = [0] * size

obstacles = [0] * size

# END GRID VARIABLES


# HELPER FUNCTIONS


def clip(lower, value, upper):
    return min(upper, max(lower, value))


def toggleObstacle(x, y):
    obstacles[IX(x, y)] = 1 if obstacles[IX(x, y)] == 0 else 0


def getNeighbors(arr, x, y):
    dirs = [(1, 0), (-1, 0), (0, 1), (0, -1)]
    return [
        arr[IX(x + dir[0], y + dir[1])]
        for dir in dirs
        if obstacles[IX(x + dir[0], y + dir[1])] != 1
    ]

# END HELPER FUNCTIONS

# SIMULATION FUNCTIONS


def diffuse(N, b, x, x0, diff, dt):
    a = dt*diff*N*N

    for _ in range(20):
        for i in range(1, N+1):
            for j in range(1, N+1):
                neighbors = getNeighbors(x, i, j)
                x[IX(i, j)] = (
                    x0[IX(i, j)]
                    + a*(
                        # x[IX(i-1, j)]
                        # + x[IX(i+1, j)]
                        # + x[IX(i, j-1)]
                        # + x[IX(i, j+1)]
                        sum(neighbors)
                    )
                ) / (1+len(neighbors)*a)
        set_bnd(N, b, x)


def advect(N, b, d, d0, u, v, dt):
    dt0 = dt*N
    for i in range(1, N+1):
        for j in range(1, N+1):
            x = clip(0.5, i - dt0*u[IX(i, j)], N+0.5)
            y = clip(0.5, j - dt0*u[IX(i, j)], N+0.5)
            i0 = int(x)
            j0 = int(y)
            i1 = i0 + 1
            j1 = j0 + 1
            s1 = x-i0
            s0 = 1-s1
            t1 = y-j0
            t0 = 1-t1
            d[IX(i, j)] = s0*(
                t0*d0[IX(i0, j0)]
                + t1*d0[IX(i0, j1)]
            ) + s1*(
                t0*d0[IX(i1, j0)]
                + t1*d0[IX(i1, j1)]
            )
    set_bnd(N, b, d)


def project(N, u, v, p, div):
    h = 1.0/N
    for i in range(1, N+1):
        for j in range(1, N+1):
            div[IX(i, j)] = -0.5*h*(
                u[IX(i+1, j)]
                - u[IX(i-1, j)]
                + v[IX(i, j+1)]
                - v[IX(i, j-1)]
            )
            p[IX(i, j)] = 0
    set_bnd(N, 0, div)
    set_bnd(N, 0, p)

    for _ in range(20):
        for i in range(1, N+1):
            for j in range(1, N+1):
                p[IX(i, j)] = (
                    div[IX(i, j)]
                    + p[IX(i-1, j)]
                    + p[IX(i+1, j)]
                    + p[IX(i, j-1)]
                    + p[IX(i, j+1)]
                )/4
        set_bnd(N, 0, p)

    for i in range(1, N+1):
        for j in range(1, N+1):
            u[IX(i, j)] -= 0.5*(p[IX(i+1, j)]-p[IX(i-1, j)])/h
            v[IX(i, j)] -= 0.5*(p[IX(i, j+1)]-p[IX(i, j-1)])/h


def set_bnd(N, b, x):
    for i in range(1, N+1):
        x[IX(0, i)] = -x[IX(1, i)] if b == 1 else x[IX(1, i)]
        x[IX(N+1, i)] = -x[IX(N, i)] if b == 1 else x[IX(N, i)]
        x[IX(i, 0)] = -x[IX(i, 1)] if b == 2 else x[IX(i, 1)]
        x[IX(i, N+1)] = -x[IX(i, N)] if b == 2 else x[IX(i, N)]


def densityStep(N, x, x0, u, v, diff, dt):
    diffuse(N, 0, x, x0, diff, dt)
    advect(N, 0, x0, x, u, v, dt)


def velStep(N, u, v, u0, v0, visc, dt):
    # add_source ( N, u, u0, dt ); add_source ( N, v, v0, dt );
    # SWAP ( u0, u )
    diffuse(N, 1, u0, u, visc, dt)
    # SWAP ( v0, v )
    diffuse(N, 2, v0, v, visc, dt)
    project(N, u0, v0, u, v)
    # SWAP ( u0, u ); SWAP ( v0, v );
    advect(N, 1, u, u0, u0, v0, dt)
    advect(N, 2, v, v0, u0, v0, dt)
    project(N, u, v, u0, v0)


def addSource(x, y):
    dens[IX(x, y)] = 255
    dens_prev[IX(x, y)] = 255
    u[IX(x, y)] = 255
    v[IX(x, y)] = 255


# END SIMULATION FUNCTIONS

# UI


def draw(dens, obs, surface):
    w = surface.get_width()
    h = surface.get_height()
    for i in range(1, N+1):
        for j in range(1, N+1):
            color = RED if obs[IX(i, j)] == 1 else tuple(
                dens[IX(i, j)] for _ in range(3)
            )
            pg.draw.rect(
                surface,
                color,
                (
                    (i-1)*(w/N),
                    (j-1)*(h/N),
                    (w/N),
                    (h/N)
                )
            )

# END UI


pg.font.init()
roboto = pg.font.SysFont('Roboto', 30)


def IX(x, y):
    return x + (N+2)*y


def isZero(x):
    if x == 0:
        return 1
    else:
        return x


class displayDraw():
    def __init__(self, display: pg.Surface):
        self.display = display

    def __enter__(self):
        self.display.fill(BLACK)

    def __exit__(self, *args):
        pg.display.update()


display = pg.display.set_mode((600, 600))

sources = [(2, 2)]

for i in range(5, 15):
    toggleObstacle(i, 5)
    toggleObstacle(5, i)

for i in range(5, 15):
    toggleObstacle(i, 14)

for i in range(2, 6):
    toggleObstacle(i, 5)

deltaT = 0
frame = 0
while True:
    frame += 1
    frameStartTime = time()
    with displayDraw(display):
        draw(dens, obstacles, display)
        statsLabel = roboto.render(
            f"FPS: {1/isZero(deltaT):0>6.1f}",
            True,
            (255, 0, 0)
        )
        display.blit(statsLabel, (0, 0))
    # print(dens)
    for s in sources:
        addSource(*s)
    velStep(N, u, v, u_prev, v_prev, visc, deltaT)
    densityStep(N, dens, dens_prev, u, v, 0.1, deltaT)

    dens_prev = deepcopy(dens)
    # print(dens)
    # if frame >= 10:
    #     break
    deltaT = time() - frameStartTime