Extract hard sphere gas from particle playground

5 files changed, 545 insertions, 0 deletions

diff --git a/gas.py b/gas.py
new file mode 100644
index 0000000..78612df
--- /dev/null
+++ b/gas.py
@@ -0,0 +1,164 @@
+from OpenGL.GL import *
+from OpenGL.GLUT import *
+from OpenGL.GL import shaders
+
+from pyrr import matrix44
+
+import numpy as np
+
+from particles import GasFlow, HardSphereSetup, grid_of_random_velocity_particles
+
+def on_timer(t):
+    glutTimerFunc(t, on_timer, t)
+    glutPostRedisplay()
+
+def get_projection(width, height):
+    world_height = 1.0
+    world_width = world_height / height * width
+    pixels_per_unit = height / world_height
+
+    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])
+
+    return (np.matmul(translation, projection), pixels_per_unit)
+
+def on_reshape(width, height):
+    global projection, pixels_per_unit
+    glViewport(0,0,width,height)
+    projection, pixels_per_unit = get_projection(width, height)
+
+glutInit()
+glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
+glutInitWindowPosition(0, 0)
+window = glutCreateWindow("gas")
+glutTimerFunc(10, on_timer, 10)
+
+fragment_shader = shaders.compileShader("""
+#version 430
+
+in vec3 color;
+
+void main(){
+    if (length(gl_PointCoord - vec2(0.5)) > 0.5) {
+        discard;
+    }
+
+    gl_FragColor = vec4(color.xyz, 0.0);
+}""", GL_FRAGMENT_SHADER)
+
+particle_shader = shaders.compileShader("""
+#version 430
+
+layout (location=0) in vec2 particles;
+
+out vec3 color;
+
+uniform mat4 projection;
+
+void main() {
+    gl_Position = projection * vec4(particles, 0., 1.);
+    color = vec3(0.0);
+}""", GL_VERTEX_SHADER)
+
+particle_program = shaders.compileProgram(particle_shader, fragment_shader)
+
+background_fragment_shader = shaders.compileShader("""
+#version 430
+
+in vec3 color;
+
+void main(){
+    gl_FragColor = vec4(color.xyz, 0.0);
+}""", GL_FRAGMENT_SHADER)
+
+background_vertex_shader = shaders.compileShader("""
+#version 430
+
+in vec3 vertex;
+
+out vec3 color;
+
+uniform mat4 projection;
+uniform vec3 face_color;
+
+void main() {
+    gl_Position = projection * vec4(vertex, 1.);
+    color = face_color;
+}""", GL_VERTEX_SHADER)
+background_program = shaders.compileProgram(background_vertex_shader, background_fragment_shader)
+
+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')
+
+class View:
+    def __init__(self, gas):
+        self.gas = gas
+        self.energy = 0
+        self.tracer = [ ]
+
+    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 draw_trace(self):
+        glUniform3f(background_color_id, 1., 0., 0.)
+        glLineWidth(2)
+        glBegin(GL_LINE_STRIP)
+        for v in self.tracer:
+            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)
+
+        shaders.glUseProgram(background_program)
+        glUniformMatrix4fv(background_projection_id, 1, False, np.asfortranarray(projection))
+        glUniform3f(background_color_id, 1., 1., 1.)
+        glBegin(GL_TRIANGLE_STRIP)
+        glVertex(0.,0.,0.)
+        glVertex(1.,0.,0.)
+        glVertex(0.,1.,0.)
+        glVertex(1.,1.,0.)
+        glEnd()
+
+        if trace:
+            self.draw_trace()
+
+        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()
+
+        glutSwapBuffers()
+
+        if trace:
+            self.update_trace()
+
+        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
+
+grid_width = 10
+radius = 0.005
+char_u = 1
+trace = True
+
+config = HardSphereSetup(radius, char_u, *grid_of_random_velocity_particles(grid_width, radius, char_u))
+gas = GasFlow(config, opengl = True)
+
+view = View(gas)
+
+glutDisplayFunc(view.on_display)
+glutReshapeFunc(on_reshape)
+
+glutMainLoop()
diff --git a/gasplot.py b/gasplot.py
new file mode 100644
index 0000000..8adfe4d
--- /dev/null
+++ b/gasplot.py
@@ -0,0 +1,55 @@
+import numpy as np
+import scipy.stats as stats
+import scipy.constants as const
+from scipy.optimize import minimize
+
+import matplotlib
+import matplotlib.pyplot as plt
+
+from particles import GasFlow, HardSphereSetup, grid_of_random_velocity_particles
+
+grid_width = 30
+radius = 0.002
+char_u = 1120
+
+config = HardSphereSetup(radius, char_u, *grid_of_random_velocity_particles(grid_width, radius, char_u))
+gas = GasFlow(config)
+
+m_nitrogen = 0.028 / const.N_A
+
+def plot(step, velocities):
+    velocities = np.array([np.linalg.norm(v) for v in velocities])
+    maxwellian = stats.maxwell.fit(velocities)
+
+    print("T = %.0f K; u_mean = %.0f [m/s]; energy = %.05f" % ((maxwellian[1]**2 / const.k * m_nitrogen, stats.maxwell.mean(*maxwellian), np.sum([x*2 for x in velocities]))))
+
+    plt.figure()
+
+    plt.ylim(0, 0.003)
+    plt.ylabel('Probability')
+
+    plt.xlim(0, 1.2*char_u)
+    plt.xlabel('Velocity magnitude [m/s]')
+
+    plt.hist(velocities, bins=50, density=True, alpha=0.5, label='Simulated velocities')
+
+    xs = np.linspace(0, 1.2*char_u, 100)
+    plt.plot(xs, stats.maxwell.pdf(xs, *maxwellian), label='Maxwell-Boltzmann distribution')
+
+    plt.legend(loc='upper right')
+
+    plt.savefig("result/%04d.png" % step)
+    plt.close()
+
+def simulate(n_steps, section):
+    for i in range(0, int(n_steps / section)):
+        print("Plot step %d." % (i * section))
+
+        velocities = gas.get_velocities()
+
+        for j in range(0,section):
+            gas.evolve()
+
+        plot(i, velocities)
+
+simulate(10000, 500)
diff --git a/kernel.cl b/kernel.cl
new file mode 100644
index 0000000..8bbe28a
--- /dev/null
+++ b/kernel.cl
@@ -0,0 +1,153 @@
+#define RADIUS      $radius
+#define N_PARTICLES $n_particles
+#define DELTA_T     $delta_t
+
+typedef float  scalar_t;
+typedef float2 vec_t;
+
+#define SCALAR_MAX FLT_MAX
+
+__constant vec_t wall_normals[4] = {
+	(vec_t)(-1.,0.),
+	(vec_t)( 1.,0.),
+	(vec_t)(0., 1.),
+	(vec_t)(0.,-1.),
+};
+
+__constant scalar_t wall_loc[4] = {
+	-RADIUS, 1.-RADIUS, 1.-RADIUS, -RADIUS
+};
+
+scalar_t solve_wall_collision(vec_t n, scalar_t loc, vec_t p, vec_t v) {
+	if (dot(n,v) > 0.) {
+		vec_t wall_v = dot(n,v) * n;
+		return (loc - dot(p,n)) / dot(wall_v,n);
+	} else {
+		return SCALAR_MAX;
+	}
+}
+
+scalar_t pos_min_or_infty(scalar_t t0, scalar_t t1) {
+	if (t0 >= 0.) {
+		if (t1 >= 0.) {
+			return min(min(t0, t1), SCALAR_MAX);
+		} else {
+			return min(t0, SCALAR_MAX);
+		}
+	} else {
+		if (t1 >= 0.) {
+			return min(t1, SCALAR_MAX);
+		} else {
+			return SCALAR_MAX;
+		}
+	}
+}
+
+scalar_t solve_collision(vec_t p, vec_t v, vec_t p_, vec_t v_) {
+	scalar_t a = dot(v-v_, v-v_);
+	scalar_t b  = 2.*dot(p-p_, v-v_);
+	scalar_t c = dot(p-p_, p-p_) - 4.*RADIUS*RADIUS;
+	scalar_t d = b*b - 4.*a*c;
+
+	if (d >= 0.) {
+		scalar_t t0 = (-b + sqrt(d))/(2.*a);
+		scalar_t t1 = (-b - sqrt(d))/(2.*a);
+
+		return pos_min_or_infty(t0, t1);
+	} else {
+		return SCALAR_MAX;
+	}
+}
+
+__kernel void evolve(__global vec_t* pos_a,
+                     __global vec_t* vel_a,
+                     __global vec_t* pos_b,
+                     __global vec_t* vel_b,
+                     __global volatile unsigned int* last_collide)
+{
+	unsigned int i = get_global_id(0);
+
+	vec_t p = pos_a[i];
+	vec_t v = vel_a[i];
+
+	unsigned int jParticle = N_PARTICLES;
+	scalar_t min2intersect = SCALAR_MAX;
+
+	for (unsigned int iParticle=0; iParticle < N_PARTICLES; ++iParticle) {
+		if (iParticle != i && !(last_collide[i] == iParticle && last_collide[iParticle] == i)) {
+			vec_t p_ = pos_a[iParticle];
+			vec_t v_ = vel_a[iParticle];
+
+			scalar_t time2intersect = solve_collision(p, v, p_, v_);
+
+			if (time2intersect < min2intersect) {
+				min2intersect = time2intersect;
+				jParticle = iParticle;
+			}
+		}
+	}
+
+	unsigned int jWall = N_PARTICLES;
+	scalar_t min2wall = SCALAR_MAX;
+
+	for (unsigned int iWall=0; iWall < 4; ++iWall) {
+		scalar_t time2wall = solve_wall_collision(wall_normals[iWall], wall_loc[iWall], p, v);
+		if (time2wall < min2wall) {
+			min2wall = time2wall;
+			jWall = iWall;
+		}
+	}
+
+	if (min2intersect < DELTA_T) {
+		if (min2wall < min2intersect) {
+			p += min2wall * v;
+			v -= 2*dot(v,wall_normals[jWall])*wall_normals[jWall];
+			p += (DELTA_T - min2wall) * v;
+			last_collide[i] = N_PARTICLES;
+
+			pos_b[i] = p;
+			vel_b[i] = v;
+		} else {
+			vec_t p_ = pos_a[jParticle];
+			vec_t v_ = vel_a[jParticle];
+
+			p  += min2intersect * v;
+			p_ += min2intersect * v_;
+
+			vec_t omega = normalize(p - p_);
+
+			while (dot(omega,omega) < 1.0) {
+				omega *= 1.001;
+			}
+
+			v  -= dot(vel_a[i] - vel_a[jParticle], omega) * omega;
+			v_ -= dot(vel_a[jParticle] - vel_a[i], omega) * omega;
+
+			if (i < jParticle) {
+				p  += (DELTA_T - min2intersect) * v;
+				p_ += (DELTA_T - min2intersect) * v_;
+
+				pos_b[i] = p;
+				vel_b[i] = v;
+
+				pos_b[jParticle] = p_;
+				vel_b[jParticle] = v_;
+
+				last_collide[i] = jParticle;
+				last_collide[jParticle] = i;
+			}
+		}
+	} else {
+		if (min2wall < DELTA_T) {
+			p += min2wall * v;
+			v -= 2*dot(v,wall_normals[jWall])*wall_normals[jWall];
+			p += (DELTA_T - min2wall) * v;
+			last_collide[i] = N_PARTICLES;
+		} else {
+			p += DELTA_T * v;
+		}
+
+		pos_b[i] = p;
+		vel_b[i] = v;
+	}
+}
diff --git a/particles.py b/particles.py
new file mode 100644
index 0000000..0f05786
--- /dev/null
+++ b/particles.py
@@ -0,0 +1,129 @@
+import numpy as np
+
+import pyopencl as cl
+mf = cl.mem_flags
+from pyopencl.tools import get_gl_sharing_context_properties
+
+import OpenGL.GL as gl
+from OpenGL.arrays import vbo
+
+from string import Template
+
+class HardSphereSetup:
+    def __init__(self, radius, char_u, position, velocity):
+        self.radius = radius
+        self.char_u = char_u
+        self.position = position
+        self.velocity = velocity
+        self.n_particles = len(position[:,0])
+
+def build_kernel(delta_t, n_particles, radius):
+    with open('kernel.cl', 'r') as kernel_src:
+        return Template(kernel_src.read()).substitute(
+            delta_t     = delta_t,
+            n_particles = n_particles,
+            radius      = radius)
+
+def grid_of_random_velocity_particles(width, radius, u_scale):
+    np_position = np.ndarray((width**2, 2))
+    np_velocity = np.ndarray((width**2, 2))
+
+    grid = np.meshgrid(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_velocity[:,0] = u_scale*(-0.5 + np.random.random_sample((width**2,)))
+    np_velocity[:,1] = u_scale*(-0.5 + np.random.random_sample((width**2,)))
+
+    return (np_position, np_velocity)
+
+class GasFlow:
+    def setup_cl(self):
+        self.platform = cl.get_platforms()[0]
+        if self.opengl:
+            self.context = cl.Context(
+                properties=[(cl.context_properties.PLATFORM, self.platform)] + get_gl_sharing_context_properties())
+        else:
+            self.context = cl.Context(
+                properties=[(cl.context_properties.PLATFORM, self.platform)])
+        self.queue = cl.CommandQueue(self.context)
+
+        self.program = cl.Program(self.context, self.kernel_src).build(
+            '-cl-single-precision-constant -cl-opt-disable')
+
+        self.cl_particle_velocity_a = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_particle_velocity)
+        self.cl_particle_velocity_b = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_particle_velocity)
+
+        if self.opengl:
+            self.gl_particle_position_a = vbo.VBO(data=self.np_particle_position, usage=gl.GL_DYNAMIC_DRAW, target=gl.GL_ARRAY_BUFFER)
+            self.gl_particle_position_a.bind()
+            self.gl_particle_position_b = vbo.VBO(data=self.np_particle_position, usage=gl.GL_DYNAMIC_DRAW, target=gl.GL_ARRAY_BUFFER)
+            self.gl_particle_position_b.bind()
+
+            self.cl_particle_position_a = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_particle_position_a))
+            self.cl_particle_position_b = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_particle_position_b))
+        else:
+            self.cl_particle_position_a = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_particle_position)
+            self.cl_particle_position_b = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_particle_position)
+
+        self.cl_last_collide = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_last_collide)
+
+    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.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.kernel_src = build_kernel(self.t_scale*setup.radius/setup.char_u, self.n_particles, setup.radius)
+
+        self.setup_cl()
+
+        self.tick = True
+
+    def evolve(self):
+        if self.opengl:
+            cl.enqueue_acquire_gl_objects(self.queue, [self.cl_particle_position_a, self.cl_particle_position_b])
+
+        if self.tick:
+            self.tick = False
+            kernelargs = (self.cl_particle_position_a, self.cl_particle_velocity_a, self.cl_particle_position_b, self.cl_particle_velocity_b, self.cl_last_collide)
+        else:
+            self.tick = True
+            kernelargs = (self.cl_particle_position_b, self.cl_particle_velocity_b, self.cl_particle_position_a, self.cl_particle_velocity_a, self.cl_last_collide)
+
+        self.program.evolve(self.queue, (self.n_particles,), None, *(kernelargs))
+
+    def gl_draw_particles(self):
+        gl.glEnableClientState(gl.GL_VERTEX_ARRAY)
+
+        if self.tick:
+            self.gl_particle_position_b.bind()
+            gl.glVertexPointer(2, gl.GL_FLOAT, 0, self.gl_particle_position_b)
+        else:
+            self.gl_particle_position_a.bind()
+            gl.glVertexPointer(2, gl.GL_FLOAT, 0, self.gl_particle_position_a)
+
+        gl.glDrawArrays(gl.GL_POINTS, 0, self.n_particles)
+        gl.glDisableClientState(gl.GL_VERTEX_ARRAY)
+
+    def get_velocities(self):
+        if self.tick:
+            cl.enqueue_copy(self.queue, self.np_particle_velocity, self.cl_particle_velocity_b).wait()
+            return self.np_particle_velocity
+        else:
+            cl.enqueue_copy(self.queue, self.np_particle_velocity, self.cl_particle_velocity_a).wait()
+            return self.np_particle_velocity
+
+    def get_positions(self):
+        if self.tick:
+            cl.enqueue_copy(self.queue, self.np_particle_position, self.cl_particle_position_b).wait()
+            return self.np_particle_position
+        else:
+            cl.enqueue_copy(self.queue, self.np_particle_position, self.cl_particle_position_a).wait()
+            return self.np_particle_position
diff --git a/shell.nix b/shell.nix
new file mode 100644
index 0000000..9122012
--- /dev/null
+++ b/shell.nix
@@ -0,0 +1,44 @@
+{ pkgs ? import <nixpkgs> { }, ... }:
+
+pkgs.stdenvNoCC.mkDerivation rec {
+  name = "pycl-env";
+  env = pkgs.buildEnv { name = name; paths = buildInputs; };
+
+  buildInputs = let
+    custom-python = let
+      packageOverrides = self: super: {
+        pyopencl = super.pyopencl.overridePythonAttrs(old: rec {
+          buildInputs = with pkgs; [
+            opencl-headers ocl-icd python37Packages.pybind11
+            libGLU_combined
+          ];
+        # Enable OpenGL integration and fix build
+          preBuild = ''
+            python configure.py --cl-enable-gl
+            export HOME=/tmp/pyopencl
+          '';
+        });
+      };
+    in pkgs.python3.override { inherit packageOverrides; };
+
+    local-python = custom-python.withPackages (python-packages: with python-packages; [
+      numpy
+      sympy
+      scipy
+      pyopencl setuptools
+      pyopengl pyrr
+      matplotlib
+    ]);
+
+  in [
+    local-python
+    pkgs.opencl-info
+    pkgs.universal-ctags
+  ];
+
+  shellHook = ''
+    export NIX_SHELL_NAME="${name}"
+    export PYOPENCL_COMPILER_OUTPUT=1
+    export PYTHONPATH="$PWD:$PYTHONPATH"
+  '';
+}