Start to use vectorizable indexing for material initialization

`utility.ndindex` adds support for specifying the traversal order to `numpy.ndindex`.

3 files changed, 58 insertions, 28 deletions

diff --git a/channel_2d_gl_interop.py b/channel_2d_gl_interop.py
index f8257e6..3db64df 100644
--- a/channel_2d_gl_interop.py
+++ b/channel_2d_gl_interop.py
@@ -22,27 +22,32 @@ relaxation_time = 0.51
 def get_obstacles(geometry):
     ys = numpy.linspace(geometry.size_y//50, geometry.size_y-geometry.size_y//50, num = 20)
     xs = [ 50 for i, y in enumerate(ys) ]
-    return list(zip(xs, ys))
+    rs = [ geometry.size_x//100 for i, y in enumerate(ys) ]
+    return list(zip(xs, ys, rs))
 
-def is_obstacle(geometry, x, y):
-    for (ox,oy) in obstacles:
-        if numpy.sqrt((x - ox)**2 + (y - oy)**2) < geometry.size_x//100:
-            return True
+def set_obstacle(lattice, x, y, r):
+    circle = [(x - idx[0])**2 + (y - idx[1])**2 < r*r for idx in lattice.memory.cells()]
+    lattice.material[circle] = 2
 
-    else:
-        return False
-
-def channel(geometry, x, y):
-    if x == 1:
-        return 3
-    elif x == geometry.size_x-2:
-        return 4
-    elif y == 1 or y == geometry.size_y-2:
-        return 2
-    elif is_obstacle(geometry, x, y):
-        return 2
-    else:
-        return 1
+def set_walls(lattice):
+    inflow = [idx[0] == 1 for idx in lattice.memory.cells()]
+    lattice.material[inflow] = 3
+
+    outflow = [idx[0] == lattice.geometry.size_x-2 for idx in lattice.memory.cells()]
+    lattice.material[outflow] = 4
+
+    top = [idx[1] == lattice.geometry.size_y-2 for idx in lattice.memory.cells()]
+    lattice.material[top] = 2
+
+    bottom = [idx[1] == 1 for idx in lattice.memory.cells()]
+    lattice.material[bottom] = 2
+
+    outer = [idx[0] == 0 or idx[0] == lattice.geometry.size_x-1 or idx[1] == 0 or idx[1] == lattice.geometry.size_y-1 for idx in lattice.memory.cells()]
+    lattice.material[outer] = 0
+
+def set_bulk(lattice):
+    bulk = [idx[0] > 0 and idx[0] < lattice.geometry.size_x-1 and idx[1] > 0 and idx[1] < lattice.geometry.size_y-1 for idx in lattice.memory.cells()]
+    lattice.material[bulk] = 1
 
 boundary = Template("""
     if ( m == 2 ) {
@@ -146,8 +151,13 @@ lattice = Lattice(
     opengl       = True
 )
 
-obstacles = get_obstacles(lattice.geometry)
-lattice.setup_geometry(channel)
+set_bulk(lattice)
+set_walls(lattice)
+
+for obstacle in get_obstacles(lattice.geometry):
+    set_obstacle(lattice, *obstacle)
+
+lattice.sync_material()
 
 projection = get_projection()
 
diff --git a/simulation.py b/simulation.py
index 700b8b6..badd2bb 100644
--- a/simulation.py
+++ b/simulation.py
@@ -2,6 +2,8 @@ import pyopencl as cl
 mf = cl.mem_flags
 
 import numpy
+from utility.ndindex import ndindex
+
 import sympy
 
 from mako.template import Template
@@ -98,6 +100,14 @@ class Memory:
     def gid(self, x, y, z = 0):
         return z * (self.size_x*self.size_y) + y * self.size_x + x;
 
+    def size(self):
+        if self.size_z == 1:
+            return (self.size_x, self.size_y)
+        else:
+            return (self.size_x, self.size_y, self.size_z)
+
+    def cells(self):
+        return ndindex(self.size(), order='F')
 
 class Lattice:
     def __init__(self,
@@ -148,14 +158,10 @@ class Lattice:
         self.program.equilibrilize(
             self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_pop_b).wait()
 
-    def setup_geometry(self, material_at):
-        material = numpy.ndarray(shape=(self.memory.volume, 1), dtype=numpy.int32)
-        material[:,:] = 0
-
-        for idx in self.geometry.inner_cells():
-            material[self.memory.gid(*idx)] = material_at(self.geometry, *idx)
+        self.material = numpy.ndarray(shape=(self.memory.volume, 1), dtype=numpy.int32)
 
-        cl.enqueue_copy(self.queue, self.memory.cl_material, material).wait();
+    def sync_material(self):
+        cl.enqueue_copy(self.queue, self.memory.cl_material, self.material).wait();
 
     def build_kernel(self):
         program_src = Template(filename = str(Path(__file__).parent/'template/kernel.mako')).render(
diff --git a/utility/ndindex.py b/utility/ndindex.py
new file mode 100644
index 0000000..0c5ed9b
--- /dev/null
+++ b/utility/ndindex.py
@@ -0,0 +1,14 @@
+import numpy
+from numpy.lib.stride_tricks import as_strided
+import numpy.core.numeric as _nx
+
+class ndindex(numpy.ndindex):
+    pass
+
+    def __init__(self, *shape, order):
+        if len(shape) == 1 and isinstance(shape[0], tuple):
+            shape = shape[0]
+        x = as_strided(_nx.zeros(1), shape=shape,
+                       strides=_nx.zeros_like(shape))
+        self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],
+                              order=order)