Extract visualization setup into module

7 files changed, 134 insertions, 119 deletions

diff --git a/boltzgas/__init__.py b/boltzgas/__init__.py
index 3b623dd..504dc1a 100644
--- a/boltzgas/__init__.py
+++ b/boltzgas/__init__.py
@@ -1,4 +1,5 @@
 import boltzgas.initial_condition
 import boltzgas.visual
+import boltzgas.visualizer
 
 from boltzgas.simulation import HardSphereSetup, HardSphereSimulation
diff --git a/boltzgas/simulation.py b/boltzgas/simulation.py
index 6af9069..d818bab 100644
--- a/boltzgas/simulation.py
+++ b/boltzgas/simulation.py
@@ -25,7 +25,27 @@ def build_kernel(delta_t, n_particles, radius):
             radius      = radius)
 
 class HardSphereSimulation:
-    def setup_cl(self):
+    def __init__(self, setup, opengl = False, t_scale = 1.0):
+        self.np_particle_position = setup.position.astype(np.float32)
+        self.np_particle_velocity = setup.velocity.astype(np.float32)
+
+        self.n_particles = setup.n_particles
+        self.radius      = setup.radius
+        self.char_u      = setup.char_u
+
+        self.opengl = opengl
+        self.t_scale = t_scale
+
+        self.np_last_collide = np.ndarray((self.n_particles, 1), dtype=np.uint32)
+        self.np_last_collide[:,0] = self.n_particles
+
+        self.np_particle_velocity_norms = np.ndarray((self.n_particles, 1), dtype=np.float32)
+
+        self.kernel_src = build_kernel(self.t_scale*self.radius/self.char_u, self.n_particles, self.radius)
+
+        self.tick = True
+
+    def setup(self):
         self.platform = cl.get_platforms()[0]
         if self.opengl:
             self.context = cl.Context(
@@ -56,27 +76,6 @@ class HardSphereSimulation:
         self.cl_last_collide = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_last_collide)
         self.cl_particle_velocity_norms = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_particle_velocity_norms)
 
-    def __init__(self, setup, opengl = False, t_scale = 1.0):
-        self.np_particle_position = setup.position.astype(np.float32)
-        self.np_particle_velocity = setup.velocity.astype(np.float32)
-
-        self.n_particles = setup.n_particles
-        self.radius      = setup.radius
-        self.char_u      = setup.char_u
-
-        self.opengl = opengl
-        self.t_scale = t_scale
-
-        self.np_last_collide = np.ndarray((self.n_particles, 1), dtype=np.uint32)
-        self.np_last_collide[:,0] = self.n_particles
-
-        self.np_particle_velocity_norms = np.ndarray((self.n_particles, 1), dtype=np.float32)
-
-        self.kernel_src = build_kernel(self.t_scale*self.radius/self.char_u, self.n_particles, self.radius)
-
-        self.setup_cl()
-
-        self.tick = True
 
     def evolve(self):
         if self.opengl:
diff --git a/boltzgas/visual/histogram.py b/boltzgas/visual/histogram.py
index fc1e853..84bba34 100644
--- a/boltzgas/visual/histogram.py
+++ b/boltzgas/visual/histogram.py
@@ -62,6 +62,7 @@ class VelocityHistogram:
             self.origin[0] + self.extend[0], self.origin[1] + self.extend[1], 1., 0.
         ], dtype=np.float32)
 
+    def setup(self):
         self.vao = glGenVertexArrays(1)
         self.vbo = glGenBuffers(1)
 
diff --git a/boltzgas/visual/tracer.py b/boltzgas/visual/tracer.py
index a5b2203..09b058e 100644
--- a/boltzgas/visual/tracer.py
+++ b/boltzgas/visual/tracer.py
@@ -6,6 +6,9 @@ class Tracer:
         self.iParticle = iParticle
         self.trace = [ ]
 
+    def setup(self):
+        return
+
     def update(self):
         position = self.gas.get_positions()[self.iParticle]
         self.trace.append((position[0], position[1]))
@@ -17,4 +20,3 @@ class Tracer:
         for v in self.trace:
             glVertex(*v, 0.)
         glEnd()
-
diff --git a/boltzgas/visualizer.py b/boltzgas/visualizer.py
new file mode 100644
index 0000000..44cec7a
--- /dev/null
+++ b/boltzgas/visualizer.py
@@ -0,0 +1,94 @@
+from OpenGL.GL import *
+from OpenGL.GLUT import *
+
+from boltzgas.visual import View
+
+class SimulationController:
+    def __init__(self, gas, instruments, updates_per_frame):
+        self.running = False
+        self.gas = gas
+        self.instruments = instruments
+        self.updates_per_frame = updates_per_frame
+
+    def isRunning(self):
+        return self.running
+
+    def run(self):
+        self.running = True
+
+    def pause(self):
+        self.running = False
+
+    def evolve(self):
+        if self.running:
+            for i in range(0,self.updates_per_frame):
+                self.gas.evolve()
+
+            for instrument in self.instruments:
+                instrument.update()
+
+    def shutdown(self):
+        self.pause()
+
+        for instrument in self.instruments:
+            try:
+                instrument.shutdown()
+            except AttributeError:
+                return # Doesn't matter, shutdown is optional
+
+
+def make_display_handler(controller, view):
+    def on_display():
+        controller.evolve()
+        view.display()
+
+    return on_display
+
+def make_reshape_handler(view):
+    def on_reshape(width, height):
+        view.reshape(width, height)
+
+    return on_reshape
+
+def make_timer():
+    def on_timer(t):
+        glutTimerFunc(t, on_timer, t)
+        glutPostRedisplay()
+
+    return on_timer
+
+def make_keyboard_handler(controller):
+    def on_keyboard(key, x, y):
+        if controller.isRunning():
+            controller.pause()
+        else:
+            controller.run()
+
+    return on_keyboard
+
+def make_close_handler(controller):
+    def on_close():
+        controller.shutdown()
+
+    return on_close
+
+def simulate(config, gas, instruments, decorations, windows, updates_per_frame = 5):
+    glutInit()
+    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
+    glutInitWindowPosition(0, 0)
+    glutCreateWindow("BoltzGas")
+
+    gas.setup()
+    for instrument in instruments:
+        instrument.setup()
+
+    view = View(gas, decorations, windows)
+    controller = SimulationController(gas, instruments, updates_per_frame)
+
+    glutDisplayFunc(make_display_handler(controller, view))
+    glutReshapeFunc(make_reshape_handler(view))
+    glutTimerFunc(20, make_timer(), 20)
+    glutKeyboardFunc(make_keyboard_handler(controller))
+    glutCloseFunc(make_close_handler(controller))
+
+    glutMainLoop()
diff --git a/gas.py b/gas.py
index 6d4f927..135796a 100644
--- a/gas.py
+++ b/gas.py
@@ -1,108 +1,25 @@
-from OpenGL.GL import *
-from OpenGL.GLUT import *
-
-import numpy as np
+import boltzgas.visualizer
 
 from boltzgas import HardSphereSetup, HardSphereSimulation
 from boltzgas.initial_condition import grid_of_random_velocity_particles
-from boltzgas.visual import View, VelocityHistogram, Tracer, ColoredBox
+from boltzgas.visual import VelocityHistogram, ColoredBox
 
 grid_width = 40
 radius = 0.002
 char_u = 1120
 
-class SimulationController:
-    def __init__(self, gas, instruments):
-        self.running = False
-        self.gas = gas
-        self.instruments = instruments
-
-    def isRunning(self):
-        return self.running
-
-    def run(self):
-        self.running = True
-
-    def pause(self):
-        self.running = False
-
-    def evolve(self):
-        if self.running:
-            for i in range(0,5):
-                self.gas.evolve()
-
-            for instrument in self.instruments:
-                instrument.update()
-
-    def shutdown(self):
-        self.pause()
-
-        for instrument in self.instruments:
-            try:
-                instrument.shutdown()
-            except AttributeError:
-                return # Doesn't matter, shutdown is optional
-
-
-def make_display_handler(controller, view):
-    def on_display():
-        controller.evolve()
-        view.display()
-
-    return on_display
-
-def make_reshape_handler(view):
-    def on_reshape(width, height):
-        view.reshape(width, height)
-
-    return on_reshape
-
-def make_timer():
-    def on_timer(t):
-        glutTimerFunc(t, on_timer, t)
-        glutPostRedisplay()
-
-    return on_timer
-
-def make_keyboard_handler(controller):
-    def on_keyboard(key, x, y):
-        if controller.isRunning():
-            controller.pause()
-        else:
-            controller.run()
-
-    return on_keyboard
-
-def make_close_handler(controller):
-    def on_close():
-        controller.shutdown()
-
-    return on_close
-
-if __name__ == "__main__":
-    glutInit()
-    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
-    glutInitWindowPosition(0, 0)
-    glutCreateWindow("BoltzGas")
-
-    position, velocity = grid_of_random_velocity_particles(grid_width, radius, char_u)
-    velocity[:,:] = 0
-    velocity[0,0] = 10*char_u
-    velocity[0,1] = 4*char_u
-
-    config = HardSphereSetup(radius, char_u, position, velocity)
-    gas = HardSphereSimulation(config, opengl = True, t_scale = 0.5)
+position, velocity = grid_of_random_velocity_particles(grid_width, radius, char_u)
+velocity[:,:] = 0
+velocity[0,0] = 10*char_u
+velocity[0,1] = 4*char_u
 
-    tracer = Tracer(gas, 4)
-    histogram = VelocityHistogram(gas, [1.1,0], [1,1])
-    view = View(gas, [ ColoredBox([0,0], [1,1], (0.2,0.2,0.2)) ], [ histogram ])
+config = HardSphereSetup(radius, char_u, position, velocity)
+gas = HardSphereSimulation(config, opengl = True, t_scale = 0.5)
 
-    controller = SimulationController(gas, [ histogram ])
+histogram = VelocityHistogram(gas, [1.1,0], [1,1])
 
-    glutDisplayFunc(make_display_handler(controller, view))
-    glutReshapeFunc(make_reshape_handler(view))
-    glutTimerFunc(20, make_timer(), 20)
-    glutKeyboardFunc(make_keyboard_handler(controller))
-    glutCloseFunc(make_close_handler(controller))
+decorations = [ ColoredBox([0,0], [1,1], (0.2,0.2,0.2)) ]
+instruments = [ histogram ]
+windows = [ histogram ]
 
-    glutMainLoop()
+boltzgas.visualizer.simulate(config, gas, instruments, decorations, windows)
diff --git a/plot.py b/plot.py
index 422db21..07fb8b1 100644
--- a/plot.py
+++ b/plot.py
@@ -19,6 +19,7 @@ velocity[0,0] = 10.75*char_u
 velocity[0,1] = -.25*char_u
 config = HardSphereSetup(radius, char_u, position, velocity)
 gas = HardSphereSimulation(config)
+gas.setup()
 
 m_nitrogen = 0.028 / const.N_A