""" GPU version of https://github.com/salt-die/navier-stokes Interact with the window using to add pressure and momentum. * Left mouse button (including drag) pokes the surface * Right mouse button places walls """ import random from pathlib import Path import numpy as np from pyrr import matrix44 import moderngl_window from moderngl_window import geometry external_flow = .35 viscosity = .018 # Is it odd that negative viscosity still works? rho = 1.06 # Density damping = 0.999 class NavierStokes2D(moderngl_window.WindowConfig): title = "Navier Stokes 2D" resource_dir = (Path(__file__) / '../../resources').absolute() # window_size = 3440, 1440 # window_size = 512, 512 aspect_ratio = None # Aspect ratio should change with window size resizable = False def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # various vars self.m_proj = matrix44.create_orthogonal_projection( 0, self.wnd.buffer_width, 0, self.wnd.buffer_height, -1, 1, dtype='f4', ) size = self.wnd.buffer_size # Geometry self.quad_fs = geometry.quad_fs() sprite_size = 9 * 4 self.drop_geometry = geometry.quad_2d(size=(sprite_size, sprite_size)) self.wall_geometry = geometry.quad_2d(size=(sprite_size, sprite_size)) # Framebuffers self.momentum_texture_1 = self.ctx.texture(size, 1, dtype='f4') self.momentum_texture_1.repeat_x = True self.momentum_texture_1.repeat_y = True self.momentum_texture_2 = self.ctx.texture(size, 1, dtype='f4') self.momentum_texture_2.repeat_x = True self.momentum_texture_2.repeat_y = True self.pressure_texture_1 = self.ctx.texture(size, 1, dtype='f4') self.pressure_texture_2 = self.ctx.texture(size, 1, dtype='f4') self.walls_texture = self.ctx.texture(size, 1, dtype='f4') self.difference_texture = self.ctx.texture(size, 1, dtype='f4') self.momentum_fbo_1 = self.ctx.framebuffer(color_attachments=[self.momentum_texture_1]) self.momentum_fbo_2 = self.ctx.framebuffer(color_attachments=[self.momentum_texture_2]) self.pressure_fbo_1 = self.ctx.framebuffer(color_attachments=[self.pressure_texture_1]) self.pressure_fbo_2 = self.ctx.framebuffer(color_attachments=[self.pressure_texture_2]) self.walls_fbo = self.ctx.framebuffer(color_attachments=[self.walls_texture]) self.difference_fbo = self.ctx.framebuffer(color_attachments=[self.difference_texture]) # programs self.texture_prog = self.load_program('programs/navier-stokes/texture.glsl') self.drop_prog = self.load_program('programs/navier-stokes/drop.glsl') self.drop_prog['m_proj'].write(self.m_proj.tobytes()) self.combine_prog = self.load_program('programs/navier-stokes/combine.glsl') self.combine_prog['pressure_texture'].value = 0 self.combine_prog['wall_texture'].value = 1 self.momentum_prog = self.load_program('programs/navier-stokes/momentum.glsl') self.momentum_prog['rho'].value = rho self.momentum_prog['damping'].value = damping self.momentum_prog['external_flow'].value = external_flow self.momentum_prog['momentum_texture'].value = 0 self.momentum_prog['pressure_texture'].value = 1 self.momentum_prog['walls_texture'].value = 3 self.flow_prog = self.load_program('programs/navier-stokes/flow.glsl') self.flow_prog['external_flow'].value = external_flow self.poisson_prog = self.load_program('programs/navier-stokes/poisson.glsl') self.pressure_prog = self.load_program('programs/navier-stokes/pressure.glsl') self.pressure_prog['rho'].value = rho self.pressure_prog['damping'].value = damping self.pressure_prog['pressure_texture'].value = 0 self.pressure_prog['difference_texture'].value = 1 self.pressure_prog['walls_texture'].value = 3 self.reset() def reset(self): self.momentum_fbo_2.clear() self.pressure_fbo_2.clear() self.walls_fbo.clear() def drop(self, x, y): rng = 40 for i in range(10): pos = ( x + int((random.random() - 0.5) * rng), self.wnd.buffer_size[1] - y + int((random.random() - 0.5) * rng) ) # Render drop into pressure texture self.pressure_fbo_2.use() self.drop_prog['pos'].value = pos self.drop_prog['force'].value = random.random() * 2.0 self.drop_prog['write_value'].value = 1.0 # self.drop.use() self.drop_geometry.render(self.drop_prog) self.momentum_fbo_2.use() self.drop_prog['write_value'].value = 0.0 self.drop_geometry.render(self.drop_prog) def wall(self, x, y): pos = x, self.wnd.buffer_height - y self.walls_fbo.use() self.drop_prog['pos'].value = pos self.drop_prog['force'].value = 1.0 self.drop_prog['write_value'].value = 1.0 self.drop_geometry.render(self.drop_prog) def render(self, time, frame_time): # calculate momentum self.momentum_fbo_1.use() self.momentum_texture_2.use(location=0) self.pressure_texture_2.use(location=1) self.walls_texture.use(location=3) self.quad_fs.render(self.momentum_prog) # External Flow self.momentum_swap() self.momentum_fbo_1.use() self.momentum_texture_2.use() self.quad_fs.render(self.flow_prog) # change in momentum # .. poisson momentum self.difference_fbo.use() self.momentum_texture_1.use() self.quad_fs.render(self.poisson_prog) # .. pressure self.pressure_fbo_1.use() self.pressure_texture_2.use(location=0) self.difference_texture.use(location=1) self.walls_texture.use(location=3) self.quad_fs.render(self.pressure_prog) # Render final result self.wnd.fbo.use() self.pressure_texture_1.use(location=0) self.walls_texture.use(location=1) self.quad_fs.render(self.combine_prog) self.momentum_swap() self.pressure_swap() def momentum_swap(self): self.momentum_texture_1, self.momentum_texture_2 = self.momentum_texture_2, self.momentum_texture_1 self.momentum_fbo_1, self.momentum_fbo_2 = self.momentum_fbo_2, self.momentum_fbo_1 def pressure_swap(self): self.pressure_texture_1, self.pressure_texture_2 = self.pressure_texture_2, self.pressure_texture_1 self.pressure_fbo_1, self.pressure_fbo_2 = self.pressure_fbo_2, self.pressure_fbo_1 def mouse_press_event(self, x, y, button): if button == self.wnd.mouse.left: self.drop(x, y) elif button == self.wnd.mouse.right: self.wall(x, y) def mouse_drag_event(self, x, y, dx, dy): if self.wnd.mouse_states.left: self.drop(x, y) elif self.wnd.mouse_states.right: self.wall(x, y) def key_event(self, key, action, modifiers): keys = self.wnd.keys if action == keys.ACTION_PRESS: if key == keys.R: self.reset() if __name__ == '__main__': moderngl_window.run_window_config(NavierStokes2D)