Reenable 2D support

7 files changed, 60 insertions, 38 deletions

diff --git a/boltzgas/initial_condition.py b/boltzgas/initial_condition.py
index 2526566..f5fdcff 100644
--- a/boltzgas/initial_condition.py
+++ b/boltzgas/initial_condition.py
@@ -1,20 +1,27 @@
 import numpy as np
 
-def grid_of_random_velocity_particles(width, radius, u_scale):
-    np_position = np.ndarray((width**3, 4))
-    np_velocity = np.ndarray((width**3, 4))
+def grid_of_random_velocity_particles(d, width, radius, u_scale):
+    np_position = np.ndarray((width**d, d))
+    np_velocity = np.ndarray((width**d, d))
+
+    if d == 2:
+        grid = np.meshgrid(np.linspace(2*radius, 1-2*radius, width),
+                           np.linspace(2*radius, 1-2*radius, width))
+    elif d == 3:
+        grid = np.meshgrid(np.linspace(2*radius, 1-2*radius, width),
+                           np.linspace(2*radius, 1-2*radius, width),
+                           np.linspace(2*radius, 1-2*radius, width))
 
-    grid = np.meshgrid(np.linspace(2*radius, 1-2*radius, width),
-                       np.linspace(2*radius, 1-2*radius, width),
-                       np.linspace(2*radius, 1-2*radius, width))
     np_position[:,0] = grid[0].flatten()
     np_position[:,1] = grid[1].flatten()
-    np_position[:,2] = grid[2].flatten()
-    np_position[:,3] = 0
 
-    np_velocity[:,0] = u_scale*(-0.5 + np.random.random_sample((width**3,)))
-    np_velocity[:,1] = u_scale*(-0.5 + np.random.random_sample((width**3,)))
-    np_velocity[:,2] = u_scale*(-0.5 + np.random.random_sample((width**3,)))
-    np_velocity[:,3] = 0
+    if d == 3:
+        np_position[:,2] = grid[2].flatten()
+
+    np_velocity[:,0] = u_scale*(-0.5 + np.random.random_sample((width**d,)))
+    np_velocity[:,1] = u_scale*(-0.5 + np.random.random_sample((width**d,)))
+
+    if d == 3:
+        np_velocity[:,2] = u_scale*(-0.5 + np.random.random_sample((width**d,)))
 
     return (np_position, np_velocity)
diff --git a/boltzgas/kernel.cl b/boltzgas/kernel.cl
index aabd1af..b69df51 100644
--- a/boltzgas/kernel.cl
+++ b/boltzgas/kernel.cl
@@ -2,7 +2,9 @@
 #define N_PARTICLES $n_particles
 #define DELTA_T     $delta_t
 
+#if $dimension == 3
 #define ENABLE_3D
+#endif
 
 typedef float  scalar_t;
 
diff --git a/boltzgas/simulation.py b/boltzgas/simulation.py
index 905a04f..e523360 100644
--- a/boltzgas/simulation.py
+++ b/boltzgas/simulation.py
@@ -17,18 +17,16 @@ class HardSphereSetup:
         self.velocity = velocity
         self.n_particles = len(position[:,0])
 
-def build_kernel(delta_t, n_particles, radius):
+def build_kernel(dimension, delta_t, n_particles, radius):
     with open('boltzgas/kernel.cl', 'r') as kernel_src:
         return Template(kernel_src.read()).substitute(
             delta_t     = delta_t,
             n_particles = n_particles,
-            radius      = radius)
+            radius      = radius,
+            dimension   = dimension)
 
 class HardSphereSimulation:
     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
@@ -36,12 +34,25 @@ class HardSphereSimulation:
         self.opengl = opengl
         self.t_scale = t_scale
 
+        self.dimension = len(setup.position[0,:])
+
+        if self.dimension == 2:
+            self.np_particle_position = setup.position.astype(np.float32)
+            self.np_particle_velocity = setup.velocity.astype(np.float32)
+        elif self.dimension == 3:
+            self.np_particle_position = np.ndarray((self.n_particles, 4), dtype=np.float32)
+            self.np_particle_position[:,0:3] = setup.position
+            self.np_particle_velocity = np.ndarray((self.n_particles, 4), dtype=np.float32)
+            self.np_particle_velocity[:,0:3] = setup.velocity
+        else:
+            raise 'Dimension invalid'
+
         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.kernel_src = build_kernel(self.dimension, self.t_scale*self.radius/self.char_u, self.n_particles, self.radius)
 
         self.tick = True
 
@@ -95,10 +106,10 @@ class HardSphereSimulation:
 
         if self.tick:
             self.gl_particle_position_b.bind()
-            gl.glVertexPointer(4, gl.GL_FLOAT, 0, self.gl_particle_position_b)
+            gl.glVertexPointer(len(self.np_particle_position[0,:]), gl.GL_FLOAT, 0, self.gl_particle_position_b)
         else:
             self.gl_particle_position_a.bind()
-            gl.glVertexPointer(4, gl.GL_FLOAT, 0, self.gl_particle_position_a)
+            gl.glVertexPointer(len(self.np_particle_position[0,:]), gl.GL_FLOAT, 0, self.gl_particle_position_a)
 
         gl.glDrawArrays(gl.GL_POINTS, 0, self.n_particles)
         gl.glDisableClientState(gl.GL_VERTEX_ARRAY)
diff --git a/boltzgas/visual/camera.py b/boltzgas/visual/camera.py
index 44831f4..2bfbc1d 100644
--- a/boltzgas/visual/camera.py
+++ b/boltzgas/visual/camera.py
@@ -42,27 +42,27 @@ class Projection:
         return self.matrix
 
 class Rotation:
-    def __init__(self, shift, x = np.pi, z = np.pi):
+    def __init__(self, shift, x = np.pi, y = np.pi):
         self.matrix = matrix44.create_from_translation(shift),
         self.rotation_x = quaternion.Quaternion()
-        self.update(x,z)
+        self.update(x,y)
 
-    def shift(self, x, z):
+    def shift(self, x, y):
         self.matrix = np.matmul(
             self.matrix,
-            matrix44.create_from_translation([x,0,z])
+            matrix44.create_from_translation([x,y,0])
         )
         self.inverse_matrix = np.linalg.inv(self.matrix)
 
-    def update(self, x, z):
+    def update(self, x, y):
         rotation_x = quaternion.Quaternion(quaternion.create_from_eulers([x,0,0]))
-        rotation_z = self.rotation_x.conjugate.cross(
-                quaternion.Quaternion(quaternion.create_from_eulers([0,0,z])))
+        rotation_y = self.rotation_x.conjugate.cross(
+                quaternion.Quaternion(quaternion.create_from_eulers([0,y,0])))
         self.rotation_x = self.rotation_x.cross(rotation_x)
 
         self.matrix = np.matmul(
             self.matrix,
-            matrix44.create_from_quaternion(rotation_z.cross(self.rotation_x))
+            matrix44.create_from_quaternion(rotation_y.cross(self.rotation_x))
         )
         self.inverse_matrix = np.linalg.inv(self.matrix)
 
diff --git a/boltzgas/visual/view.py b/boltzgas/visual/view.py
index b4c1acd..bdc18e8 100644
--- a/boltzgas/visual/view.py
+++ b/boltzgas/visual/view.py
@@ -125,8 +125,8 @@ class View:
         self.particle_shader   = Shader(*particle_shader)
         self.decoration_shader = Shader(*decoration_shader)
 
-        self.camera_projection = Projection(distance = 6)
-        self.camera_rotation = Rotation([-1/2, -1/2, -1/2])
+        self.camera_projection = Projection(distance = 5)
+        self.camera_rotation = Rotation([-1/2, -1/2, -1/2], np.pi/2, 0)
         self.camera_pos = np.matmul([0,-self.camera_projection.distance,0,1], self.camera_rotation.get_inverse())
 
         self.mouse_monitors = [
diff --git a/random_velocities.py b/random_velocities.py
index cf18e6b..2f62267 100644
--- a/random_velocities.py
+++ b/random_velocities.py
@@ -4,18 +4,18 @@ from boltzgas import HardSphereSetup, HardSphereSimulation
 from boltzgas.initial_condition import grid_of_random_velocity_particles
 from boltzgas.visual import VelocityHistogram, WireBox, Tracer
 
+dimension = 3
 grid_width = 10
 radius = 0.005
 char_u = 1120
 
-position, velocity = grid_of_random_velocity_particles(grid_width, radius, char_u)
+position, velocity = grid_of_random_velocity_particles(dimension, grid_width, radius, char_u)
 
 config = HardSphereSetup(radius, char_u, position, velocity)
 gas = HardSphereSimulation(config, opengl = True, t_scale=0.5)
 
-#tracer = Tracer(gas, int((grid_width**2)/2+grid_width/2))
 histogram = VelocityHistogram(gas, [1.1,0], [1,1])
 
-instruments = [ WireBox(0,1,0,1,0,1), histogram ]
+instruments = [ WireBox(0,1,0,1,0,dimension == 3 if 1 else 0), histogram ]
 
 boltzgas.visualizer.simulate(config, gas, instruments)
diff --git a/zero_velocities.py b/zero_velocities.py
index d0de774..26a21f9 100644
--- a/zero_velocities.py
+++ b/zero_velocities.py
@@ -4,21 +4,23 @@ from boltzgas import HardSphereSetup, HardSphereSimulation
 from boltzgas.initial_condition import grid_of_random_velocity_particles
 from boltzgas.visual import VelocityHistogram, WireBox
 
-grid_width = 10
-radius = 0.005
+dimension = 2
+grid_width = 30
+radius = 0.002
 char_u = 1120
 
-position, velocity = grid_of_random_velocity_particles(grid_width, radius, char_u)
+position, velocity = grid_of_random_velocity_particles(dimension, grid_width, radius, char_u)
 velocity[:,:] = 0
 velocity[0,0] = 5*char_u
 velocity[0,1] = 4*char_u
-velocity[0,2] = 3*char_u
+if dimension == 3:
+    velocity[0,2] = 3*char_u
 
 config = HardSphereSetup(radius, char_u, position, velocity)
 gas = HardSphereSimulation(config, opengl = True, t_scale = 0.5)
 
 histogram = VelocityHistogram(gas, [1.1,0], [1,1])
 
-instruments = [ WireBox(0,1,0,1,0,1), histogram ]
+instruments = [ WireBox(0,1,0,1,0,dimension == 3 if 1 else 0), histogram ]
 
 boltzgas.visualizer.simulate(config, gas, instruments)