Created
November 10, 2019 18:32
-
-
Save sam-lb/c81d1616b19c1aac14d81653b627b79a to your computer and use it in GitHub Desktop.
Python program that visualizes fluid flow through vector fields
This file contains hidden or 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
""" | |
Fluid particles travelling through a vector field | |
By Sam Brunacini | |
Posts at http://sambrunacini.com/blog/ | |
""" | |
import pygame; | |
import math; | |
from math import hypot, atan2, sin, cos, pi; | |
pygame.init(); | |
WIDTH, HEIGHT = 600, 600; | |
TITLE = "Fluid simulation"; | |
MAX_FPS = 30; | |
TWO_PI = 2 * pi; | |
QUARTER_PI = pi / 4; | |
WHITE = (255, 255, 255); | |
RED = (255, 0, 0); | |
BLACK = (0, 0, 0); | |
screen = pygame.display.set_mode((WIDTH, HEIGHT)); | |
pygame.display.set_caption(TITLE); | |
pygame.key.set_repeat(100, 50); | |
clock = pygame.time.Clock(); | |
x_units, y_units = 8, 8; | |
hwidth, hheight = WIDTH / 2, HEIGHT / 2; | |
x_scale, y_scale = WIDTH / x_units, HEIGHT / y_units; | |
x_stop = x_units / 2; x_start = -x_stop; | |
y_stop = y_units / 2; y_start = -y_stop; | |
def from_polar(r, theta): | |
""" convert a vector from polar form to rectangular form """ | |
return r * cos(theta), r * sin(theta); | |
def create_text_surface(text, size, color): | |
""" creates a surface with text on it """ | |
return pygame.font.SysFont("arial", size).render(text, 1, color); | |
def text(surface, position, text, size, color): | |
""" draw text to a surface """ | |
text = create_text_surface(text, size, color); | |
surface.blit(text, position); | |
def safe_drange(start, stop, step=1, precision=2): | |
scaler = 10 ** precision; | |
start, stop = start * scaler, stop * scaler; | |
step = step * scaler; | |
for i in range(int(start), int(stop), int(step)): | |
yield i / scaler; | |
def translate_and_scale(x, y): | |
""" scale the coordinates and put the origin in the center of the screen """ | |
return (x * x_scale) + hwidth, hheight - (y * y_scale) | |
def draw_arrow(A, B, color, surf, width=2): | |
""" draw an arrow """ | |
# efficiency does not matter in these calculations because they are only done once at the beginning. | |
dy, dx = A[1] - B[1], A[0] - B[0]; | |
angle = atan2(dy, dx); | |
if angle < 0: | |
mang = (angle + TWO_PI) / TWO_PI; | |
else: | |
mang = angle / TWO_PI; | |
value = 255 * mang; | |
color = (0, value, value); | |
dist = hypot(dx, dy) / 5; # divide by five so the arrow's prongs aren't as long as the tail | |
x1, y1 = from_polar(dist, angle + QUARTER_PI); | |
x2, y2 = from_polar(dist, angle - QUARTER_PI); | |
pygame.draw.line(surf, color, A, B, width); | |
pygame.draw.line(surf, color, B, (B[0] + x1, B[1] + y1), width); | |
pygame.draw.line(surf, color, B, (B[0] + x2, B[1] + y2), width); | |
class VectorField(): | |
def __init__(self, function, vecs_per_unit=2, color=RED): | |
self.function = function; | |
self.vecs_per_unit = vecs_per_unit; | |
self.step = 1 / self.vecs_per_unit; | |
self.color = color; | |
self._generate_vectors(); | |
def _generate_vectors(self): | |
self.vectors = []; | |
for x in safe_drange(x_start, x_stop, self.step): | |
for y in safe_drange(y_start, y_stop, self.step): | |
try: | |
dx, dy = self.function(x, y); | |
self.vectors.append((x, y, x + dx / 8, y + dy / 8)); | |
except ZeroDivisionError: | |
continue; | |
def draw(self, surf): | |
for vector in self.vectors: | |
draw_arrow(translate_and_scale(vector[0], vector[1]), translate_and_scale(vector[2], vector[3]), self.color, surf); | |
class Particle(): | |
def __init__(self, function, x, y, color=(255, 0, 0), speed_scale=35): | |
self.function = function; | |
self.color = color; | |
self.speed_scale = speed_scale; | |
self.initial_speed = 1 / self.speed_scale; | |
self.x, self.y = x, y; | |
def move(self): | |
try: | |
dx, dy = self.function(self.x, self.y); | |
except ZeroDivisionError: | |
dx, dy = 0, 0; | |
self.x += dx / self.speed_scale; | |
self.y += dy / self.speed_scale; | |
def draw(self): | |
center = translate_and_scale(self.x, self.y); | |
pygame.draw.circle(screen, self.color, (int(center[0]), int(center[1])), 2); | |
def update(self): | |
self.move(); | |
self.draw(); | |
#function = lambda x, y: (cos(x*x + y*y), sin(x*x + y*y)); # Attraction curve | |
#function = lambda x, y: (sin(x + cos(y)), sin(y)); # Attraction points | |
#function = lambda x, y: (cos(y)/2, x/2); # Swirl point (off center in y direction) | |
#function = lambda x, y: (sin(y), -x); # Swirl point (center) | |
function = lambda x, y: (sin(2*y), -x/2); # Multiple swirl points along the y axis | |
#function = lambda x, y: (math.atan(y - x), x*cos(x)); # Attraction points | |
#function = lambda x, y: (x-y, x+y); # Repelling point | |
#function = lambda x, y: (sin(x-y), sin(y+x)); # Repelling points and attraction points | |
#function = lambda x, y: (0, y + sin(x)); # Repelling curve sin(x) | |
def generate_particles(function): | |
v = VectorField(function, vecs_per_unit=4); | |
particles = []; | |
step = 0.1 * (((x_units + y_units) // 2) / 8); | |
for x in safe_drange(x_start, x_stop, step): | |
for y in safe_drange(y_start, y_stop, step): | |
particles.append(Particle(function, x, y)); | |
return v, particles; | |
v, particles = generate_particles(function); | |
inter = pygame.Surface((WIDTH, HEIGHT)); | |
inter.fill(WHITE); | |
v.draw(inter); | |
running = True; | |
while running: | |
screen.blit(inter, (0, 0)); | |
for p in particles: | |
p.update(); | |
pygame.draw.rect(screen, WHITE, (10, 10, 200, 28)); | |
text(screen, (10, 10), "Swirl points", 25, BLACK); | |
pygame.display.flip(); | |
clock.tick(MAX_FPS); | |
for event in pygame.event.get(): | |
if event.type == pygame.QUIT: | |
running = False; | |
break; | |
elif event.type == pygame.KEYDOWN: | |
if event.key == pygame.K_ESCAPE: | |
running = False; | |
break; | |
pygame.quit(); |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment