Add some structure to view implementation

1 files changed, 126 insertions, 106 deletions

diff --git a/gas.py b/gas.py
index 78612df..de047a1 100644
--- a/gas.py
+++ b/gas.py
@@ -8,157 +8,177 @@ import numpy as np
 
 from particles import GasFlow, HardSphereSetup, grid_of_random_velocity_particles
 
-def on_timer(t):
-    glutTimerFunc(t, on_timer, t)
-    glutPostRedisplay()
+glutInit()
+glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
+glutInitWindowPosition(0, 0)
+window = glutCreateWindow("gas")
 
-def get_projection(width, height):
-    world_height = 1.0
-    world_width = world_height / height * width
-    pixels_per_unit = height / world_height
+class Shader:
+    def __init__(self, vertex_src, fragment_src, uniform):
+        self.program = shaders.compileProgram(
+            shaders.compileShader(vertex_src, GL_VERTEX_SHADER),
+            shaders.compileShader(fragment_src, GL_FRAGMENT_SHADER))
+        self.uniform = { }
+        for name in uniform:
+            self.uniform[name] = shaders.glGetUniformLocation(self.program, name)
 
-    projection  = matrix44.create_orthogonal_projection(-world_width/2, world_width/2, -world_height/2, world_height/2, -1, 1)
-    translation = matrix44.create_from_translation([-1.0/2, -1.0/2, 0])
+    def use(self):
+        shaders.glUseProgram(self.program)
 
-    return (np.matmul(translation, projection), pixels_per_unit)
+particle_shader = Shader(
+    fragment_src = """
+        #version 430
 
-def on_reshape(width, height):
-    global projection, pixels_per_unit
-    glViewport(0,0,width,height)
-    projection, pixels_per_unit = get_projection(width, height)
+        in vec3 color;
 
-glutInit()
-glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
-glutInitWindowPosition(0, 0)
-window = glutCreateWindow("gas")
-glutTimerFunc(10, on_timer, 10)
+        void main(){
+            if (length(gl_PointCoord - vec2(0.5)) > 0.5) {
+                discard;
+            }
 
-fragment_shader = shaders.compileShader("""
-#version 430
+            gl_FragColor = vec4(color.xyz, 0.0);
+        }""",
+    vertex_src = """
+        #version 430
 
-in vec3 color;
+        layout (location=0) in vec2 particles;
 
-void main(){
-    if (length(gl_PointCoord - vec2(0.5)) > 0.5) {
-        discard;
-    }
+        out vec3 color;
 
-    gl_FragColor = vec4(color.xyz, 0.0);
-}""", GL_FRAGMENT_SHADER)
+        uniform mat4 projection;
 
-particle_shader = shaders.compileShader("""
-#version 430
+        void main() {
+            gl_Position = projection * vec4(particles, 0., 1.);
+            color = vec3(0.0);
+        }""",
+    uniform = ['projection']
+)
 
-layout (location=0) in vec2 particles;
+decoration_shader = Shader(
+    fragment_src = """
+        #version 430
 
-out vec3 color;
+        in vec3 color;
 
-uniform mat4 projection;
+        void main(){
+            gl_FragColor = vec4(color.xyz, 0.0);
+        }""",
+    vertex_src = """
+        #version 430
 
-void main() {
-    gl_Position = projection * vec4(particles, 0., 1.);
-    color = vec3(0.0);
-}""", GL_VERTEX_SHADER)
+        in vec3 vertex;
 
-particle_program = shaders.compileProgram(particle_shader, fragment_shader)
+        out vec3 color;
 
-background_fragment_shader = shaders.compileShader("""
-#version 430
+        uniform mat4 projection;
+        uniform vec3 face_color;
 
-in vec3 color;
+        void main() {
+            gl_Position = projection * vec4(vertex, 1.);
+            color = face_color;
+        }""",
+    uniform = ['projection', 'face_color']
+)
 
-void main(){
-    gl_FragColor = vec4(color.xyz, 0.0);
-}""", GL_FRAGMENT_SHADER)
+class View:
+    def __init__(self, gas, decorations):
+        self.gas = gas
+        self.decorations = decorations
 
-background_vertex_shader = shaders.compileShader("""
-#version 430
+    def reshape(self, width, height):
+        glViewport(0,0,width,height)
 
-in vec3 vertex;
+        world_height = 1.0
+        world_width = world_height / height * width
 
-out vec3 color;
+        projection  = matrix44.create_orthogonal_projection(-world_width/2, world_width/2, -world_height/2, world_height/2, -1, 1)
+        translation = matrix44.create_from_translation([-1.0/2, -1.0/2, 0])
 
-uniform mat4 projection;
-uniform vec3 face_color;
+        self.pixels_per_unit = height / world_height
+        self.projection = np.matmul(translation, projection)
 
-void main() {
-    gl_Position = projection * vec4(vertex, 1.);
-    color = face_color;
-}""", GL_VERTEX_SHADER)
-background_program = shaders.compileProgram(background_vertex_shader, background_fragment_shader)
+    def display(self):
+        glClearColor(0.4,0.4,0.4,1.)
+        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
 
-particle_projection_id   = shaders.glGetUniformLocation(particle_program, 'projection')
-background_projection_id = shaders.glGetUniformLocation(background_program, 'projection')
-background_color_id = shaders.glGetUniformLocation(background_program, 'face_color')
+        decoration_shader.use()
+        glUniformMatrix4fv(decoration_shader.uniform['projection'], 1, False, np.asfortranarray(self.projection))
+        for decoration in self.decorations:
+            decoration.display(decoration_shader.uniform)
 
-class View:
-    def __init__(self, gas):
+        particle_shader.use()
+        glUniformMatrix4fv(particle_shader.uniform['projection'], 1, False, np.asfortranarray(self.projection))
+        glEnable(GL_POINT_SPRITE)
+        glPointSize(2*radius*self.pixels_per_unit)
+        self.gas.gl_draw_particles()
+
+        glutSwapBuffers()
+
+class Tracer:
+    def __init__(self, gas, iParticle):
         self.gas = gas
-        self.energy = 0
-        self.tracer = [ ]
+        self.iParticle = iParticle
+        self.trace = [ ]
 
-    def update_trace(self):
-        positions = self.gas.get_positions()
-        self.tracer.append((positions[5][0],positions[5][1]))
-        if len(self.tracer) > 100:
-            self.tracer.pop(0)
+    def update(self):
+        position = self.gas.get_positions()[self.iParticle]
+        self.trace.append((position[0], position[1]))
 
-    def draw_trace(self):
-        glUniform3f(background_color_id, 1., 0., 0.)
+    def display(self, uniform):
+        glUniform3f(uniform['face_color'], 1., 0., 0.)
         glLineWidth(2)
         glBegin(GL_LINE_STRIP)
-        for v in self.tracer:
+        for v in self.trace:
             glVertex(*v, 0.)
         glEnd()
 
-    def on_display(self):
-        for i in range(0,10):
-            self.gas.evolve()
-
-        glClearColor(0.4,0.4,0.4,1.)
-        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
+class ColoredBox:
+    def __init__(self, origin, extend, color):
+        self.origin = origin
+        self.extend = extend
+        self.color = color
 
-        shaders.glUseProgram(background_program)
-        glUniformMatrix4fv(background_projection_id, 1, False, np.asfortranarray(projection))
-        glUniform3f(background_color_id, 1., 1., 1.)
+    def display(self, uniform):
+        glUniform3f(uniform['face_color'], *self.color)
         glBegin(GL_TRIANGLE_STRIP)
-        glVertex(0.,0.,0.)
-        glVertex(1.,0.,0.)
-        glVertex(0.,1.,0.)
-        glVertex(1.,1.,0.)
+        glVertex(self.origin[0],                  self.origin[1]                 , 0.)
+        glVertex(self.origin[0] + self.extend[0], self.origin[1]                 , 0.)
+        glVertex(self.origin[0]                 , self.origin[1] + self.extend[1], 0.)
+        glVertex(self.origin[0] + self.extend[1], self.origin[1] + self.extend[1], 0.)
         glEnd()
 
-        if trace:
-            self.draw_trace()
+grid_width = 30
+radius = 0.002
+char_u = 1120
 
-        shaders.glUseProgram(particle_program)
-        glUniformMatrix4fv(particle_projection_id, 1, False, np.asfortranarray(projection))
-        glEnable(GL_POINT_SPRITE)
-        glPointSize(2*radius*pixels_per_unit)
-        gas.gl_draw_particles()
+position, velocity = grid_of_random_velocity_particles(grid_width, radius, char_u)
+velocity[:,:] = 0
+velocity[0,0] = 10*char_u
+velocity[0,1] =  1*char_u
 
-        glutSwapBuffers()
+config = HardSphereSetup(radius, char_u, position, velocity)
+gas = GasFlow(config, opengl = True, t_scale = 1.0)
 
-        if trace:
-            self.update_trace()
+tracer = Tracer(gas, 5)
+view = View(gas, [ColoredBox([0,0], [1,1], (1,1,1)), tracer])
 
-        velocities = gas.get_velocities()
-        energy = np.sum(np.array([np.linalg.norm(v)**2 for v in velocities]))
-        if abs(energy - self.energy) > 1e-4:
-            print("energy = %.05f" % energy)
-            self.energy = energy
+def on_display():
+    for i in range(0,10):
+        gas.evolve()
 
-grid_width = 10
-radius = 0.005
-char_u = 1
-trace = True
+    tracer.update()
 
-config = HardSphereSetup(radius, char_u, *grid_of_random_velocity_particles(grid_width, radius, char_u))
-gas = GasFlow(config, opengl = True)
+    view.display()
+
+def on_reshape(width, height):
+    view.reshape(width, height)
 
-view = View(gas)
+def on_timer(t):
+    glutTimerFunc(t, on_timer, t)
+    glutPostRedisplay()
 
-glutDisplayFunc(view.on_display)
+glutDisplayFunc(on_display)
 glutReshapeFunc(on_reshape)
+glutTimerFunc(10, on_timer, 10)
 
 glutMainLoop()