simple-gas-sim/main.py

114 lines
3.0 KiB
Python

import pygame
import pygame.gfxdraw
import gas_sim as gs
import numpy as np
import random
##### CONFIG #####
window_size = (1280, 720)
target_fps = 60
particle_radius_ws = 0.025
particle_compressibility = 0.1
particle_friction = 0.1
cursor_radius_ws = 0.15
substeps = 10
##################
# ws - world space
# ss - screen space
window_aspect = window_size[0] / window_size[1]
cursor_radius_ss = cursor_radius_ws * window_size[1]
particle_radius_ss = particle_radius_ws * window_size[1]
# pygame init
pygame.init()
window = pygame.display.set_mode(window_size)
clock = pygame.time.Clock()
# gas sim init
gas_sim = gs.GasSim(particle_radius_ws, particle_compressibility, particle_friction, substeps, window_aspect)
not_loaded_particles = True
dt = target_fps / 1000
# main loop
while True:
# events
for event in pygame.event.get():
if event.type == pygame.QUIT:
pygame.quit()
quit()
elif event.type == pygame.KEYDOWN:
# on s key press
if event.key == pygame.K_s:
# save particles to file
particles = tuple(pos for pos, _ in gas_sim.particles)
np.save("particles.npy", particles)
print("Saved particles to file")
# on l key press
elif event.key == pygame.K_l:
if not_loaded_particles:
# load particles from file
particles = np.load("particles.npy")
for particle in particles:
gas_sim.add_particle(particle)
print("Loaded particles from file")
not_loaded_particles = False
else:
print("Particles already loaded !!!")
# on left mouse down
if pygame.mouse.get_pressed()[0]:
mouse_pos_ss = pygame.mouse.get_pos()
mouse_pos_ws = (mouse_pos_ss[0] / window_size[1], mouse_pos_ss[1] / window_size[1])
# if not colliding with any particles
if not gas_sim.circle_particle_collide_test(mouse_pos_ws, particle_radius_ws):
gas_sim.add_particle(mouse_pos_ws)
# on right mouse down
elif pygame.mouse.get_pressed()[2]:
# apply wind force
mouse_pos_ss = pygame.mouse.get_pos()
mouse_pos_ws = (mouse_pos_ss[0] / window_size[1], mouse_pos_ss[1] / window_size[1])
gas_sim.circle_force(mouse_pos_ws, cursor_radius_ws, (0.2, random.uniform(-0.5, 0.5)), dt)
# update particles
gas_sim.update(dt)
window.fill((0, 0, 0))
# draw cursor
cursor_pos = pygame.mouse.get_pos()
pygame.gfxdraw.aacircle(window, cursor_pos[0], cursor_pos[1], int(cursor_radius_ss), (255, 255, 255))
# draw particles
for particle_pos, _ in gas_sim.particles:
particle_pos_ss = (int(particle_pos[0] * window_size[1]), int(particle_pos[1] * window_size[1]))
pygame.gfxdraw.aacircle(window, particle_pos_ss[0], particle_pos_ss[1], int(particle_radius_ss), (255, 255, 255))
# update
pygame.display.update()
dt = clock.tick(60) / 1000