Created
November 3, 2021 20:28
-
-
Save Bertik23/d7ffa6e326483034b43d95800d15d5d8 to your computer and use it in GitHub Desktop.
Fluid simulation using Navier-Stokes equation in Python
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment