Add basic D2Q9 LBM

Ported the basic compustream structure

2 files changed, 167 insertions, 0 deletions

diff --git a/implosion.py b/implosion.py
new file mode 100644
index 0000000..0a10827
--- /dev/null
+++ b/implosion.py
@@ -0,0 +1,166 @@
+import pyopencl as cl
+mf = cl.mem_flags
+from pyopencl.tools import get_gl_sharing_context_properties
+
+from string import Template
+
+import numpy
+import matplotlib.pyplot as plt
+
+kernel = """
+float constant w[9] = {
+    1./36., 1./9., 1./36.,
+    1./9. , 4./9., 1./9. ,
+    1./36 , 1./9., 1./36.
+};
+
+uint2 cellAtGid(unsigned int gid)
+{
+    const int y = gid / $nX;
+    return (uint2)(gid - $nX*y, y);
+}
+
+unsigned int gidOfCell(int x, int y)
+{
+    return y * $nX + x;
+}
+
+unsigned int indexOfDirection(int i, int j) {
+    return 3*(i+1) + (j+1);
+}
+
+float pop(__global float* cell, int i, int j) {
+    return cell[indexOfDirection(i,j)];
+}
+
+float comp(int i, int j, float2 v) {
+    return i*v.x + j*v.y;
+}
+
+float sq(float x) {
+    return x*x;
+}
+
+float norm(float2 v) {
+    return sqrt(dot(v,v));
+}
+
+float density(__global float* cell) {
+    float d = 0.;
+    for ( int i = 0; i < 9; ++i ) {
+        d += cell[i];
+    }
+    return d;
+}
+
+float2 velocity(__global float* cell, float d)
+{
+    return (float2)(
+        (pop(cell,1,0) - pop(cell,-1,0) + pop(cell,1,1) - pop(cell,-1,-1) + pop(cell,1,-1) - pop(cell,-1,1)) / d,
+        (pop(cell,0,1) - pop(cell,0,-1) + pop(cell,1,1) - pop(cell,-1,-1) - pop(cell,1,-1) + pop(cell,-1,1)) / d
+    );
+}
+
+float equilibrium(float d, float2 v, int i, int j) {
+    return w[indexOfDirection(i,j)] * d * (1 + 3*comp(i,j,v) + 4.5*sq(comp(i,j,v)) - 1.5*sq(norm(v)));
+}
+
+__kernel void collide_and_stream(__global float* pop_a,
+                                 __global float* pop_b)
+{
+    const unsigned int gid = get_global_id(0);
+    const uint2 cell = cellAtGid(gid);
+
+    float  d = density(&pop_b[gid*9]);
+    float2 v = velocity(&pop_b[gid*9],d);
+
+    if ( cell.x >= 2 && cell.x <= $nX-3 && cell.y >= 2 && cell.y <= $nY-3 ) {
+        for ( int i = -1; i <= 1; ++i ) {
+            for ( int j = -1; j <= 1; ++j ) {
+                pop_a[gidOfCell(cell.x+i, cell.y+j)*9 + indexOfDirection(i,j)] =
+                    pop_b[gid*9 + indexOfDirection(i,j)] + $tau * (equilibrium(d,v,i,j) - pop_b[gid*9 + indexOfDirection(i,j)]);
+            }
+        }
+    }
+    else if ( ((cell.y == 1 || cell.y == $nY-2) && (cell.x > 0 && cell.x < $nX-1)) ||
+              ((cell.x == 1 || cell.x == $nX-2) && (cell.y > 0 && cell.y < $nY-1)) )
+    {
+        for ( int i = -1; i <= 1; ++i ) {
+            for ( int j = -1; j <= 1; ++j ) {
+                pop_a[gidOfCell(cell.x-i, cell.y-j)*9 + indexOfDirection(-i,-j)] =
+                    pop_b[gid*9 + indexOfDirection(i,j)] + $tau * (equilibrium(d,v,i,j) - pop_b[gid*9 + indexOfDirection(i,j)]);
+            }
+        }
+    }
+}"""
+
+class D2Q9_BGK_Lattice:
+    def idx(self, x, y):
+        return y * self.nX + x;
+
+    def __init__(self, nX, nY):
+        self.nX = nX
+        self.nY = nY
+        self.nCells = nX * nY
+        self.tick = True
+
+        self.platform = cl.get_platforms()[0]
+        self.context  = cl.Context(properties=[(cl.context_properties.PLATFORM, self.platform)])
+        self.queue = cl.CommandQueue(self.context)
+
+        self.np_pop_a = numpy.ndarray(shape=(self.nCells, 9), dtype=numpy.float32)
+        self.np_pop_b = numpy.ndarray(shape=(self.nCells, 9), dtype=numpy.float32)
+
+        self.np_pop_a[:,:] = [ 1./36., 1./9., 1./36., 1./9. , 4./9., 1./9. , 1./36 , 1./9., 1./36. ]
+        self.np_pop_b[:,:] = [ 1./36., 1./9., 1./36., 1./9. , 4./9., 1./9. , 1./36 , 1./9., 1./36. ]
+
+        for x in range(self.nX//3,self.nX-self.nX//3):
+            for y in range(self.nY//3,self.nY-self.nY//3):
+                self.np_pop_a[self.idx(x,y),:] = 1./24.
+                self.np_pop_b[self.idx(x,y),:] = 1./24.
+
+        self.cl_pop_a = cl.Buffer(self.context, mf.READ_WRITE | mf.COPY_HOST_PTR, hostbuf=self.np_pop_a)
+        self.cl_pop_b = cl.Buffer(self.context, mf.READ_WRITE | mf.COPY_HOST_PTR, hostbuf=self.np_pop_b)
+
+        self.update_kernel()
+
+    def update_kernel(self):
+        self.program = cl.Program(self.context, Template(kernel).substitute({
+            'nX' : self.nX,
+            'nY' : self.nY,
+            'tau': 0.56
+        })).build()
+
+    def evolve(self):
+        if self.tick:
+            self.tick = False
+            self.program.collide_and_stream(self.queue, (self.nCells,), None, self.cl_pop_a, self.cl_pop_b)
+            self.queue.finish()
+        else:
+            self.tick = True
+            self.program.collide_and_stream(self.queue, (self.nCells,), None, self.cl_pop_b, self.cl_pop_a)
+            self.queue.finish()
+
+    def show(self, i):
+        if self.tick:
+            cl.enqueue_copy(LBM.queue, LBM.np_pop_a, LBM.cl_pop_b).wait();
+        else:
+            cl.enqueue_copy(LBM.queue, LBM.np_pop_a, LBM.cl_pop_a).wait();
+
+        pop = numpy.ndarray(shape=(self.nX, self.nY))
+
+        for y in range(0,self.nY-1):
+            for x in range(0,self.nX-1):
+                pop[x,y] = numpy.sum(self.np_pop_a[self.idx(x,y),:])
+
+        plt.imshow(pop, vmin=0.2, vmax=2, cmap=plt.get_cmap("seismic"))
+        plt.savefig("result/density_" + str(i) + ".png")
+
+
+LBM = D2Q9_BGK_Lattice(1024, 1024)
+
+for i in range(0,10000):
+    if i % 100 == 0:
+        LBM.show(i)
+
+    LBM.evolve()
diff --git a/shell.nix b/shell.nix
index ae13444..bbb4b85 100644
--- a/shell.nix
+++ b/shell.nix
@@ -27,6 +27,7 @@ pkgs.stdenvNoCC.mkDerivation rec {
       pyopencl
       pyopengl
       pygobject3
+      matplotlib
     ]))
 
     pkgs.opencl-info