327 lines
10 KiB
Python
327 lines
10 KiB
Python
"""
|
|
Stable 2D fluid simulation (Jos Stam) using Pygame + PyOpenGL + NumPy.
|
|
|
|
- Single-file demo.
|
|
- Left mouse: add density (smoke).
|
|
- Right mouse: add velocity (drag to create flow).
|
|
- Keys: C = clear, D = toggle display velocity, Esc or close window to quit.
|
|
|
|
Dependencies:
|
|
pip install pygame PyOpenGL numpy
|
|
|
|
Run:
|
|
python stable_fluid_pygame_opengl.py
|
|
|
|
Notes: This is a CPU solver on a grid (not GPU accelerated). Grid size defaults to 128x128 which is interactive on modern machines.
|
|
"""
|
|
|
|
import sys
|
|
import math
|
|
import numpy as np
|
|
import pygame
|
|
from pygame.locals import *
|
|
from OpenGL.GL import *
|
|
|
|
# ----------------------- Fluid solver (Stable Fluids, Stam) -----------------------
|
|
|
|
class StableFluid:
|
|
def __init__(self, N, diffusion, viscosity, dt):
|
|
self.N = N # grid size (N x N)
|
|
self.size = N + 2 # include boundary
|
|
self.dt = dt
|
|
self.diff = diffusion
|
|
self.visc = viscosity
|
|
|
|
# fields (with boundary padding)
|
|
self.s = np.zeros((self.size, self.size), dtype=np.float32) # temp
|
|
self.density = np.zeros((self.size, self.size), dtype=np.float32) # smoke
|
|
|
|
self.Vx = np.zeros((self.size, self.size), dtype=np.float32)
|
|
self.Vy = np.zeros((self.size, self.size), dtype=np.float32)
|
|
self.Vx0 = np.zeros((self.size, self.size), dtype=np.float32)
|
|
self.Vy0 = np.zeros((self.size, self.size), dtype=np.float32)
|
|
|
|
def add_density(self, x, y, amount):
|
|
self.density[y, x] += amount
|
|
|
|
def add_velocity(self, x, y, amount_x, amount_y):
|
|
self.Vx[y, x] += amount_x
|
|
self.Vy[y, x] += amount_y
|
|
|
|
def step(self):
|
|
N = self.N
|
|
diff = self.diff
|
|
visc = self.visc
|
|
dt = self.dt
|
|
|
|
self.diffuse(1, self.Vx0, self.Vx, visc, dt)
|
|
self.diffuse(2, self.Vy0, self.Vy, visc, dt)
|
|
|
|
self.project(self.Vx0, self.Vy0, self.Vx, self.Vy)
|
|
|
|
self.advect(1, self.Vx, self.Vx0, self.Vx0, self.Vy0, dt)
|
|
self.advect(2, self.Vy, self.Vy0, self.Vx0, self.Vy0, dt)
|
|
|
|
self.project(self.Vx, self.Vy, self.Vx0, self.Vy0)
|
|
|
|
self.diffuse(0, self.s, self.density, diff, dt)
|
|
self.advect(0, self.density, self.s, self.Vx, self.Vy, dt)
|
|
|
|
# helper routines
|
|
def set_bnd(self, b, x):
|
|
N = self.N
|
|
# edges
|
|
x[0, 1:N+1] = -x[1, 1:N+1] if b == 2 else x[1, 1:N+1]
|
|
x[N+1, 1:N+1] = -x[N, 1:N+1] if b == 2 else x[N, 1:N+1]
|
|
|
|
x[1:N+1, 0] = -x[1:N+1, 1] if b == 1 else x[1:N+1, 1]
|
|
x[1:N+1, N+1] = -x[1:N+1, N] if b == 1 else x[1:N+1, N]
|
|
|
|
# corners
|
|
x[0, 0] = 0.5 * (x[1, 0] + x[0, 1])
|
|
x[0, N+1] = 0.5 * (x[1, N+1] + x[0, N])
|
|
x[N+1, 0] = 0.5 * (x[N, 0] + x[N+1, 1])
|
|
x[N+1, N+1] = 0.5 * (x[N, N+1] + x[N+1, N])
|
|
|
|
def lin_solve(self, b, x, x0, a, c):
|
|
N = self.N
|
|
for _ in range(20): # Gauss-Seidel iterations
|
|
x[1:N+1, 1:N+1] = (
|
|
x0[1:N+1, 1:N+1] + a * (
|
|
x[0:N, 1:N+1] + x[2:N+2, 1:N+1] + x[1:N+1, 0:N] + x[1:N+1, 2:N+2]
|
|
)
|
|
) / c
|
|
self.set_bnd(b, x)
|
|
|
|
def diffuse(self, b, x, x0, diff, dt):
|
|
N = self.N
|
|
a = dt * diff * N * N
|
|
self.lin_solve(b, x, x0, a, 1 + 4 * a)
|
|
|
|
def advect(self, b, d, d0, velocX, velocY, dt):
|
|
N = self.N
|
|
dt0 = dt * N
|
|
for i in range(1, N+1):
|
|
for j in range(1, N+1):
|
|
x = i - dt0 * velocX[j, i]
|
|
y = j - dt0 * velocY[j, i]
|
|
|
|
if x < 0.5:
|
|
x = 0.5
|
|
if x > N + 0.5:
|
|
x = N + 0.5
|
|
i0 = int(math.floor(x))
|
|
i1 = i0 + 1
|
|
|
|
if y < 0.5:
|
|
y = 0.5
|
|
if y > N + 0.5:
|
|
y = N + 0.5
|
|
j0 = int(math.floor(y))
|
|
j1 = j0 + 1
|
|
|
|
s1 = x - i0
|
|
s0 = 1 - s1
|
|
t1 = y - j0
|
|
t0 = 1 - t1
|
|
|
|
d[j, i] = (
|
|
s0 * (t0 * d0[j0, i0] + t1 * d0[j1, i0]) +
|
|
s1 * (t0 * d0[j0, i1] + t1 * d0[j1, i1])
|
|
)
|
|
self.set_bnd(b, d)
|
|
|
|
def project(self, velocX, velocY, p, div):
|
|
N = self.N
|
|
div[1:N+1, 1:N+1] = -0.5 * (
|
|
velocX[1:N+1, 2:N+2] - velocX[1:N+1, 0:N] + velocY[2:N+2, 1:N+1] - velocY[0:N, 1:N+1]
|
|
) / N
|
|
p[1:N+1, 1:N+1] = 0
|
|
self.set_bnd(0, div)
|
|
self.set_bnd(0, p)
|
|
|
|
self.lin_solve(0, p, div, 1, 4)
|
|
|
|
velocX[1:N+1, 1:N+1] -= 0.5 * (p[1:N+1, 2:N+2] - p[1:N+1, 0:N]) * N
|
|
velocY[1:N+1, 1:N+1] -= 0.5 * (p[2:N+2, 1:N+1] - p[0:N, 1:N+1]) * N
|
|
|
|
self.set_bnd(1, velocX)
|
|
self.set_bnd(2, velocY)
|
|
|
|
# ----------------------- OpenGL texture renderer -----------------------
|
|
|
|
class GLRenderer:
|
|
def __init__(self, width, height, gridN):
|
|
self.width = width
|
|
self.height = height
|
|
self.gridN = gridN
|
|
self.tex_id = glGenTextures(1)
|
|
self.init_texture()
|
|
|
|
def init_texture(self):
|
|
glBindTexture(GL_TEXTURE_2D, self.tex_id)
|
|
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
|
|
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
|
|
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
|
|
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
|
|
|
|
# initial empty texture (RGBA)
|
|
data = np.zeros((self.gridN, self.gridN, 4), dtype=np.uint8)
|
|
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, self.gridN, self.gridN, 0, GL_RGBA, GL_UNSIGNED_BYTE, data)
|
|
|
|
def update_texture_from_density(self, density):
|
|
# density is (N+2)x(N+2), we upload only 1..N
|
|
N = self.gridN
|
|
den = density[1:N+1, 1:N+1]
|
|
# map density to 0-255
|
|
img = np.empty((N, N, 4), dtype=np.uint8)
|
|
# grayscale smoke map; store in rgb and alpha
|
|
clipped = np.clip(den * 255.0, 0, 255).astype(np.uint8)
|
|
img[:, :, 0] = clipped
|
|
img[:, :, 1] = clipped
|
|
img[:, :, 2] = clipped
|
|
img[:, :, 3] = clipped
|
|
|
|
glBindTexture(GL_TEXTURE_2D, self.tex_id)
|
|
glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, N, N, GL_RGBA, GL_UNSIGNED_BYTE, img)
|
|
|
|
def draw_fullscreen(self):
|
|
glEnable(GL_TEXTURE_2D)
|
|
glBindTexture(GL_TEXTURE_2D, self.tex_id)
|
|
glBegin(GL_QUADS)
|
|
glTexCoord2f(0.0, 0.0); glVertex2f(-1.0, -1.0)
|
|
glTexCoord2f(1.0, 0.0); glVertex2f(1.0, -1.0)
|
|
glTexCoord2f(1.0, 1.0); glVertex2f(1.0, 1.0)
|
|
glTexCoord2f(0.0, 1.0); glVertex2f(-1.0, 1.0)
|
|
glEnd()
|
|
glDisable(GL_TEXTURE_2D)
|
|
|
|
def draw_velocity(self, fluid):
|
|
# draw simple vector field lines
|
|
N = fluid.N
|
|
glPushAttrib(GL_CURRENT_BIT | GL_POINT_BIT | GL_LINE_BIT)
|
|
glPointSize(2.0)
|
|
glBegin(GL_LINES)
|
|
step = max(1, N // 32)
|
|
for i in range(1, N+1, step):
|
|
for j in range(1, N+1, step):
|
|
x = (i - 0.5) / N * 2 - 1
|
|
y = (j - 0.5) / N * 2 - 1
|
|
vx = fluid.Vx[j, i]
|
|
vy = fluid.Vy[j, i]
|
|
glVertex2f(x, y)
|
|
glVertex2f(x + vx * 0.02, y + vy * 0.02)
|
|
glEnd()
|
|
glPopAttrib()
|
|
|
|
# ----------------------- Main app -----------------------
|
|
|
|
def main():
|
|
# settings
|
|
WIN_W, WIN_H = 800, 800
|
|
GRID_N = 128
|
|
DIFFUSION = 0.0001
|
|
VISCOSITY = 0.0001
|
|
DT = 0.1
|
|
|
|
pygame.init()
|
|
pygame.display.set_mode((WIN_W, WIN_H), DOUBLEBUF | OPENGL)
|
|
pygame.display.set_caption('Stable Fluid (Stam) - Pygame + PyOpenGL')
|
|
|
|
# setup orthographic projection
|
|
glViewport(0, 0, WIN_W, WIN_H)
|
|
glMatrixMode(GL_PROJECTION)
|
|
glLoadIdentity()
|
|
glOrtho(-1, 1, -1, 1, -1, 1)
|
|
glMatrixMode(GL_MODELVIEW)
|
|
glLoadIdentity()
|
|
|
|
fluid = StableFluid(GRID_N, DIFFUSION, VISCOSITY, DT)
|
|
renderer = GLRenderer(WIN_W, WIN_H, GRID_N)
|
|
|
|
clock = pygame.time.Clock()
|
|
running = True
|
|
show_velocity = False
|
|
|
|
is_down_left = False
|
|
is_down_right = False
|
|
last_mouse = None
|
|
|
|
while running:
|
|
for event in pygame.event.get():
|
|
if event.type == QUIT:
|
|
running = False
|
|
elif event.type == KEYDOWN:
|
|
if event.key == K_ESCAPE:
|
|
running = False
|
|
elif event.key == K_c:
|
|
fluid.density.fill(0)
|
|
fluid.Vx.fill(0)
|
|
fluid.Vy.fill(0)
|
|
elif event.key == K_d:
|
|
show_velocity = not show_velocity
|
|
elif event.type == MOUSEBUTTONDOWN:
|
|
if event.button == 1:
|
|
is_down_left = True
|
|
elif event.button == 3:
|
|
is_down_right = True
|
|
last_mouse = pygame.mouse.get_pos()
|
|
elif event.type == MOUSEBUTTONUP:
|
|
if event.button == 1:
|
|
is_down_left = False
|
|
elif event.button == 3:
|
|
is_down_right = False
|
|
last_mouse = None
|
|
|
|
mx, my = pygame.mouse.get_pos()
|
|
# convert to grid coords (i,x) with padding
|
|
def screen_to_grid(px, py):
|
|
gx = int((px / WIN_W) * GRID_N) + 1
|
|
gy = int(((WIN_H - py) / WIN_H) * GRID_N) + 1
|
|
gx = max(1, min(GRID_N, gx))
|
|
gy = max(1, min(GRID_N, gy))
|
|
return gx, gy
|
|
|
|
if last_mouse is not None:
|
|
lx, ly = last_mouse
|
|
if is_down_left or is_down_right:
|
|
g1 = screen_to_grid(lx, ly)
|
|
g2 = screen_to_grid(mx, my)
|
|
gx1, gy1 = g1
|
|
gx2, gy2 = g2
|
|
# draw along line between points
|
|
steps = max(abs(gx2 - gx1), abs(gy2 - gy1), 1)
|
|
for s in range(steps + 1):
|
|
t = s / steps
|
|
gx = int(round(gx1 + (gx2 - gx1) * t))
|
|
gy = int(round(gy1 + (gy2 - gy1) * t))
|
|
if is_down_left:
|
|
fluid.add_density(gx, gy, 100.0 * DT)
|
|
if is_down_right:
|
|
vx = (gx2 - gx1) * 5.0
|
|
vy = (gy2 - gy1) * 5.0
|
|
fluid.add_velocity(gx, gy, vx, vy)
|
|
last_mouse = (mx, my)
|
|
|
|
# step simulation
|
|
fluid.step()
|
|
|
|
# render
|
|
glClear(GL_COLOR_BUFFER_BIT)
|
|
renderer.update_texture_from_density(fluid.density)
|
|
renderer.draw_fullscreen()
|
|
|
|
if show_velocity:
|
|
glColor3f(1.0, 0.0, 0.0)
|
|
renderer.draw_velocity(fluid)
|
|
glColor3f(1.0, 1.0, 1.0)
|
|
|
|
pygame.display.flip()
|
|
#
|
|
clock.tick(10)
|
|
|
|
pygame.quit()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|