Extract geometry information

4 files changed, 95 insertions, 74 deletions

diff --git a/implosion.py b/implosion.py
index a3bf8a1..514af46 100644
--- a/implosion.py
+++ b/implosion.py
@@ -4,8 +4,8 @@ import time
 import matplotlib
 import matplotlib.pyplot as plt
 matplotlib.use('AGG')
+from lbm import Lattice, Geometry
 
-from lbm import Lattice
 
 import symbolic.D2Q9 as D2Q9
 
@@ -16,31 +16,31 @@ def generate_moment_plots(lattice, moments):
     for i, m in enumerate(moments):
         print("Generating plot %d of %d." % (i+1, len(moments)))
 
-        density = numpy.ndarray(shape=(lattice.nY-2, lattice.nX-2))
-        for y in range(1,lattice.nY-1):
-            for x in range(1,lattice.nX-1):
-                density[y-1,x-1] = m[0,lattice.idx(x,y)]
+        velocity = numpy.ndarray(shape=tuple(reversed(lattice.geometry.inner_span())))
+        for x, y in lattice.geometry.inner_cells():
+            velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.idx(x,y)]**2 + m[2,lattice.idx(x,y)]**2)
 
         plt.figure(figsize=(10, 10))
-        plt.imshow(density, origin='lower', vmin=0.2, vmax=2.0, cmap=plt.get_cmap('seismic'))
-        plt.savefig("result/density_" + str(i) + ".png", bbox_inches='tight', pad_inches=0)
+        plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
+        plt.savefig("result/implosion_" + str(i) + ".png", bbox_inches='tight', pad_inches=0)
 
-def box(nX, nY, x, y):
-    if x == 1 or y == 1 or x == nX-2 or y == nY-2:
+def box(geometry, x, y):
+    if x == 1 or y == 1 or x == geometry.size_x-2 or y == geometry.size_y-2:
         return 2
     else:
         return 1
 
 pop_eq = """
-    if ( sqrt(pow(get_global_id(0) - ${nX//2}.f, 2.f) + pow(get_global_id(1) - ${nY//2}.f, 2.f)) < ${nX//10} ) {
+    if ( sqrt(pow(get_global_id(0) - ${geometry.size_x//2}.f, 2.f)
+            + pow(get_global_id(1) - ${geometry.size_y//2}.f, 2.f)) < ${geometry.size_x//10} ) {
 % for i, w_i in enumerate(descriptor.w):
-        preshifted_f_a[${i*nCells}] = 1./24.f;
-        preshifted_f_b[${i*nCells}] = 1./24.f;
+        preshifted_f_a[${i*geometry.volume}] = 1./24.f;
+        preshifted_f_b[${i*geometry.volume}] = 1./24.f;
 % endfor
     } else {
 % for i, w_i in enumerate(descriptor.w):
-        preshifted_f_a[${i*nCells}] = ${w_i}.f;
-        preshifted_f_b[${i*nCells}] = ${w_i}.f;
+        preshifted_f_a[${i*geometry.volume}] = ${w_i}.f;
+        preshifted_f_b[${i*geometry.volume}] = ${w_i}.f;
 % endfor
 }"""
 
@@ -60,14 +60,17 @@ print("Initializing simulation...\n")
 
 lattice = Lattice(
     descriptor = D2Q9,
-    nX = 1024, nY = 1024,
+    geometry   = Geometry(1024, 1024),
+
     moments = D2Q9.moments(optimize = False),
     collide = D2Q9.bgk(tau = 0.8),
-    geometry = box,
+
     pop_eq_src   = pop_eq,
     boundary_src = boundary)
 
-print("Starting simulation using %d cells...\n" % lattice.nCells)
+lattice.setup_geometry(box)
+
+print("Starting simulation using %d cells...\n" % lattice.geometry.volume)
 
 lastStat = time.time()
 
@@ -76,7 +79,7 @@ for i in range(1,nUpdates+1):
 
     if i % nStat == 0:
         lattice.sync()
-        print("i = %4d; %3.0f MLUPS" % (i, MLUPS(lattice.nCells, nStat, time.time() - lastStat)))
+        print("i = %4d; %3.0f MLUPS" % (i, MLUPS(lattice.geometry.volume, nStat, time.time() - lastStat)))
         moments.append(lattice.get_moments())
         lastStat = time.time()
 
diff --git a/lbm.py b/lbm.py
index 5118f17..e4eaf62 100644
--- a/lbm.py
+++ b/lbm.py
@@ -6,13 +6,28 @@ import sympy
 
 from mako.template import Template
 
+class Geometry:
+    def __init__(self, size_x, size_y):
+        self.size_x = size_x
+        self.size_y = size_y
+        self.volume = size_x * size_y
+
+    def inner_cells(self):
+        for y in range(1,self.size_y-1):
+            for x in range(1,self.size_x-1):
+                yield x, y
+
+    def span(self):
+        return (self.size_x, self.size_y)
+
+    def inner_span(self):
+        return (self.size_x-2, self.size_y-2)
+
+
 class Lattice:
-    def __init__(self, descriptor, nX, nY, moments, collide, geometry, pop_eq_src = '', boundary_src = ''):
+    def __init__(self, descriptor, geometry, moments, collide, pop_eq_src = '', boundary_src = ''):
         self.descriptor = descriptor
-
-        self.nX = nX
-        self.nY = nY
-        self.nCells = nX * nY
+        self.geometry   = geometry
 
         self.moments = moments
         self.collide = collide
@@ -24,13 +39,12 @@ class Lattice:
         self.context  = cl.Context(properties=[(cl.context_properties.PLATFORM, self.platform)])
         self.queue = cl.CommandQueue(self.context)
 
-        self.np_material = numpy.ndarray(shape=(self.nCells, 1), dtype=numpy.int32)
-        self.setup_geometry(geometry)
+        self.np_material = numpy.ndarray(shape=(self.geometry.volume, 1), dtype=numpy.int32)
 
         self.tick = True
 
-        self.pop_size     = descriptor.q     * self.nCells * numpy.float32(0).nbytes
-        self.moments_size = (descriptor.d+1) * self.nCells * numpy.float32(0).nbytes
+        self.pop_size     = descriptor.q     * self.geometry.volume * numpy.float32(0).nbytes
+        self.moments_size = (descriptor.d+1) * self.geometry.volume * numpy.float32(0).nbytes
 
         self.cl_pop_a = cl.Buffer(self.context, mf.READ_WRITE, size=self.pop_size)
         self.cl_pop_b = cl.Buffer(self.context, mf.READ_WRITE, size=self.pop_size)
@@ -40,37 +54,35 @@ class Lattice:
 
         self.build_kernel()
 
-        self.program.equilibrilize(self.queue, (self.nX,self.nY), (32,1), self.cl_pop_a, self.cl_pop_b).wait()
+        self.program.equilibrilize(
+            self.queue, (self.geometry.size_x,self.geometry.size_y), (32,1), self.cl_pop_a, self.cl_pop_b).wait()
 
     def idx(self, x, y):
-        return y * self.nX + x;
+        return y * self.geometry.size_x + x;
+
+    def setup_geometry(self, material_at):
+        for x, y in self.geometry.inner_cells():
+            self.np_material[self.idx(x,y)] = material_at(self.geometry, x, y)
 
-    def setup_geometry(self, geometry):
-        for y in range(1,self.nY-1):
-            for x in range(1,self.nX-1):
-                self.np_material[self.idx(x,y)] = geometry(self.nX,self.nY,x,y)
+        cl.enqueue_copy(self.queue, self.cl_material, self.np_material).wait();
 
     def build_kernel(self):
         program_src = Template(filename = './template/kernel.mako').render(
             descriptor = self.descriptor,
+            geometry   = self.geometry,
 
-            nX     = self.nX,
-            nY     = self.nY,
-            nCells = self.nCells,
-
-            moments_helper     = self.moments[0],
+            moments_subexpr    = self.moments[0],
             moments_assignment = self.moments[1],
-            collide_helper     = self.collide[0],
+            collide_subexpr    = self.collide[0],
             collide_assignment = self.collide[1],
 
             pop_eq_src = Template(self.pop_eq_src).render(
                 descriptor = self.descriptor,
-                nX     = self.nX,
-                nY     = self.nY,
-                nCells = self.nCells
+                geometry   = self.geometry
             ),
             boundary_src = Template(self.boundary_src).render(
-                descriptor = self.descriptor
+                descriptor = self.descriptor,
+                geometry   = self.geometry
             ),
 
             ccode = sympy.ccode
@@ -80,19 +92,23 @@ class Lattice:
     def evolve(self):
         if self.tick:
             self.tick = False
-            self.program.collide_and_stream(self.queue, (self.nX,self.nY), (32,1), self.cl_pop_a, self.cl_pop_b, self.cl_material)
+            self.program.collide_and_stream(
+                self.queue, self.geometry.span(), (32,1), self.cl_pop_a, self.cl_pop_b, self.cl_material)
         else:
             self.tick = True
-            self.program.collide_and_stream(self.queue, (self.nX,self.nY), (32,1), self.cl_pop_b, self.cl_pop_a, self.cl_material)
+            self.program.collide_and_stream(
+                self.queue, self.geometry.span(), (32,1), self.cl_pop_b, self.cl_pop_a, self.cl_material)
 
     def sync(self):
         self.queue.finish()
 
     def get_moments(self):
-        moments = numpy.ndarray(shape=(self.descriptor.d+1, self.nCells), dtype=numpy.float32)
+        moments = numpy.ndarray(shape=(self.descriptor.d+1, self.geometry.volume), dtype=numpy.float32)
         if self.tick:
-            self.program.collect_moments(self.queue, (self.nX,self.nY), (32,1), self.cl_pop_b, self.cl_moments)
+            self.program.collect_moments(
+                self.queue, self.geometry.span(), (32,1), self.cl_pop_b, self.cl_moments)
         else:
-            self.program.collect_moments(self.queue, (self.nX,self.nY), (32,1), self.cl_pop_a, self.cl_moments)
+            self.program.collect_moments(
+                self.queue, self.geometry.span(), (32,1), self.cl_pop_a, self.cl_moments)
         cl.enqueue_copy(self.queue, moments, self.cl_moments).wait();
         return moments
diff --git a/lid_driven_cavity.py b/lid_driven_cavity.py
index 2bfdb7a..873adeb 100644
--- a/lid_driven_cavity.py
+++ b/lid_driven_cavity.py
@@ -5,7 +5,7 @@ import matplotlib
 import matplotlib.pyplot as plt
 matplotlib.use('AGG')
 
-from lbm import Lattice
+from lbm import Lattice, Geometry
 
 import symbolic.D2Q9 as D2Q9
 
@@ -16,19 +16,18 @@ def generate_moment_plots(lattice, moments):
     for i, m in enumerate(moments):
         print("Generating plot %d of %d." % (i+1, len(moments)))
 
-        velocity = numpy.ndarray(shape=(lattice.nY-2, lattice.nX-2))
-        for y in range(1,lattice.nY-1):
-            for x in range(1,lattice.nX-1):
-                velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.idx(x,y)]**2 + m[2,lattice.idx(x,y)]**2)
+        velocity = numpy.ndarray(shape=tuple(reversed(lattice.geometry.inner_span())))
+        for x, y in lattice.geometry.inner_cells():
+            velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.idx(x,y)]**2 + m[2,lattice.idx(x,y)]**2)
 
         plt.figure(figsize=(10, 10))
         plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
-        plt.savefig("result/velocity_" + str(i) + ".png", bbox_inches='tight', pad_inches=0)
+        plt.savefig("result/ldc_" + str(i) + ".png", bbox_inches='tight', pad_inches=0)
 
-def cavity(nX, nY, x, y):
-    if x == 1 or y == 1 or x == nX-2:
+def cavity(geometry, x, y):
+    if x == 1 or y == 1 or x == geometry.size_x-2:
         return 2
-    elif y == nY-2:
+    elif y == geometry.size_y-2:
         return 3
     else:
         return 1
@@ -53,13 +52,16 @@ print("Initializing simulation...\n")
 
 lattice = Lattice(
     descriptor = D2Q9,
-    nX = 256, nY = 256,
-    geometry = cavity,
+    geometry   = Geometry(256, 256),
+
     moments = D2Q9.moments(optimize = False),
     collide = D2Q9.bgk(tau = 0.56),
+
     boundary_src = boundary)
 
-print("Starting simulation using %d cells...\n" % lattice.nCells)
+lattice.setup_geometry(cavity)
+
+print("Starting simulation using %d cells...\n" % lattice.geometry.volume)
 
 lastStat = time.time()
 
@@ -68,7 +70,7 @@ for i in range(1,nUpdates+1):
 
     if i % nStat == 0:
         lattice.sync()
-        print("i = %4d; %3.0f MLUPS" % (i, MLUPS(lattice.nCells, nStat, time.time() - lastStat)))
+        print("i = %4d; %3.0f MLUPS" % (i, MLUPS(lattice.geometry.volume, nStat, time.time() - lastStat)))
         moments.append(lattice.get_moments())
         lastStat = time.time()
 
diff --git a/template/kernel.mako b/template/kernel.mako
index 24a6f64..f3de3da 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -1,15 +1,15 @@
 __kernel void equilibrilize(__global __write_only float* f_a,
                             __global __write_only float* f_b)
 {
-    const unsigned int gid = get_global_id(1)*${nX} + get_global_id(0);
+    const unsigned int gid = get_global_id(1)*${geometry.size_x} + get_global_id(0);
 
     __global __write_only float* preshifted_f_a = f_a + gid;
     __global __write_only float* preshifted_f_b = f_b + gid;
 
 % if pop_eq_src == '':
 %     for i, w_i in enumerate(descriptor.w):
-    preshifted_f_a[${i*nCells}] = ${w_i}.f;
-    preshifted_f_b[${i*nCells}] = ${w_i}.f;
+    preshifted_f_a[${i*geometry.volume}] = ${w_i}.f;
+    preshifted_f_b[${i*geometry.volume}] = ${w_i}.f;
 %     endfor
 % else:
     ${pop_eq_src}
@@ -24,14 +24,14 @@ def neighbor_offset(c_i):
     if c_i[1] == 0:
         return c_i[0]
     else:
-        return c_i[1]*nX + c_i[0]
+        return c_i[1]*geometry.size_x + c_i[0]
 %>
 
 __kernel void collide_and_stream(__global __write_only float* f_a,
                                  __global __read_only  float* f_b,
                                  __global __read_only  int* material)
 {
-    const unsigned int gid = get_global_id(1)*${nX} + get_global_id(0);
+    const unsigned int gid = get_global_id(1)*${geometry.size_x} + get_global_id(0);
 
     const int m = material[gid];
 
@@ -43,10 +43,10 @@ __kernel void collide_and_stream(__global __write_only float* f_a,
     __global __read_only  float* preshifted_f_b = f_b + gid;
 
 % for i, c_i in enumerate(descriptor.c):
-    const float f_curr_${i} = preshifted_f_b[${direction_index(c_i)*nCells + neighbor_offset(-c_i)}];
+    const float f_curr_${i} = preshifted_f_b[${direction_index(c_i)*geometry.volume + neighbor_offset(-c_i)}];
 % endfor
 
-% for i, expr in enumerate(moments_helper):
+% for i, expr in enumerate(moments_subexpr):
     const float ${expr[0]} = ${ccode(expr[1])};
 % endfor
 
@@ -56,7 +56,7 @@ __kernel void collide_and_stream(__global __write_only float* f_a,
 
   ${boundary_src}
 
-% for i, expr in enumerate(collide_helper):
+% for i, expr in enumerate(collide_subexpr):
     const float ${expr[0]} = ${ccode(expr[1])};
 % endfor
 
@@ -65,26 +65,26 @@ __kernel void collide_and_stream(__global __write_only float* f_a,
 % endfor
 
 % for i in range(0,descriptor.q):
-    preshifted_f_a[${i*nCells}] = f_next_${i};
+    preshifted_f_a[${i*geometry.volume}] = f_next_${i};
 % endfor
 }
 
 __kernel void collect_moments(__global __read_only  float* f,
                               __global __write_only float* moments)
 {
-    const unsigned int gid = get_global_id(1)*${nX} + get_global_id(0);
+    const unsigned int gid = get_global_id(1)*${geometry.size_x} + get_global_id(0);
 
     __global __read_only float* preshifted_f = f + gid;
 
 % for i in range(0,descriptor.q):
-    const float f_curr_${i} = preshifted_f[${i*nCells}];
+    const float f_curr_${i} = preshifted_f[${i*geometry.volume}];
 % endfor
 
-% for i, expr in enumerate(moments_helper):
+% for i, expr in enumerate(moments_subexpr):
     const float ${expr[0]} = ${ccode(expr[1])};
 % endfor
 
 % for i, expr in enumerate(moments_assignment):
-    moments[${i*nCells} + gid] = ${ccode(expr.rhs)};
+    moments[${i*geometry.volume} + gid] = ${ccode(expr.rhs)};
 % endfor
 }