Replace some explicit dimension branching

4 files changed, 26 insertions, 37 deletions

diff --git a/implosion.py b/implosion.py
index 2ba831f..63c1950 100644
--- a/implosion.py
+++ b/implosion.py
@@ -17,9 +17,9 @@ 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=tuple(reversed(lattice.geometry.inner_span())))
+        velocity = numpy.ndarray(shape=tuple(reversed(lattice.geometry.inner_size())))
         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)
+            velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.gid(x,y)]**2 + m[2,lattice.gid(x,y)]**2)
 
         plt.figure(figsize=(10, 10))
         plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
diff --git a/ldc_2d.py b/ldc_2d.py
index 979f7be..40edf89 100644
--- a/ldc_2d.py
+++ b/ldc_2d.py
@@ -17,9 +17,9 @@ 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=tuple(reversed(lattice.geometry.inner_span())))
+        velocity = numpy.ndarray(shape=tuple(reversed(lattice.geometry.inner_size())))
         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)
+            velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.gid(x,y)]**2 + m[2,lattice.gid(x,y)]**2)
 
         plt.figure(figsize=(10, 10))
         plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
diff --git a/ldc_3d.py b/ldc_3d.py
index c008b06..a8660a3 100644
--- a/ldc_3d.py
+++ b/ldc_3d.py
@@ -17,13 +17,13 @@ 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=tuple(reversed(lattice.geometry.inner_span())))
+        velocity = numpy.ndarray(shape=tuple(reversed(lattice.geometry.inner_size())))
 
         # plot x-z-plane
         y = lattice.geometry.size_y//2
-        for z in range(1,lattice.geometry.size_z-1):
-            for x in range(1,lattice.geometry.size_x-1):
-                velocity[z-1,y,x-1] = numpy.sqrt(m[1,lattice.idx(x,y,z)]**2 + m[2,lattice.idx(x,y,z)]**2 + m[3,lattice.idx(x,y,z)]**2)
+        for z, x in numpy.ndindex(lattice.geometry.size_z-2, lattice.geometry.size_x-2):
+            gid = lattice.gid(x+1,y,z+1)
+            velocity[z,y,x] = numpy.sqrt(m[1,gid]**2 + m[2,gid]**2 + m[3,gid]**2)
 
         plt.figure(figsize=(20, 10))
 
@@ -32,12 +32,12 @@ def generate_moment_plots(lattice, moments):
 
         # plot y-z-plane
         x = lattice.geometry.size_x//2
-        for z in range(1,lattice.geometry.size_z-1):
-            for y in range(1,lattice.geometry.size_y-1):
-                velocity[z-1,y-1,x] = numpy.sqrt(m[1,lattice.idx(x,y,z)]**2 + m[2,lattice.idx(x,y,z)]**2 + m[3,lattice.idx(x,y,z)]**2)
+        for z, y in numpy.ndindex(lattice.geometry.size_z-2, lattice.geometry.size_y-2):
+            gid = lattice.gid(x,y+1,z+1)
+            velocity[z,y,x] = numpy.sqrt(m[1,gid]**2 + m[2,gid]**2 + m[3,gid]**2)
 
         plt.subplot(1, 2, 2)
-        plt.imshow(velocity[:,:,x], origin='lower', vmin=0.0, vmax=0.15, cmap=plt.get_cmap('seismic'))
+        plt.imshow(velocity[:,:,x], origin='lower', vmin=0.0, vmax=0.175, cmap=plt.get_cmap('seismic'))
 
         plt.savefig("result/ldc_3d_%02d.png" % i, bbox_inches='tight', pad_inches=0)
 
@@ -62,8 +62,8 @@ boundary = """
     }
 """
 
-nUpdates = 20000
-nStat    = 500
+nUpdates = 40000
+nStat    = 250
 
 moments = []
 
diff --git a/simulation.py b/simulation.py
index ba3945c..672003a 100644
--- a/simulation.py
+++ b/simulation.py
@@ -15,23 +15,16 @@ class Geometry:
         self.volume = size_x * size_y * size_z
 
     def inner_cells(self):
-        if self.size_z == 1:
-            for y in range(1,self.size_y-1):
-                for x in range(1,self.size_x-1):
-                    yield x, y
-        else:
-            for z in range(1,self.size_z-1):
-                for y in range(1,self.size_y-1):
-                    for x in range(1,self.size_x-1):
-                        yield x, y, z
+        for idx in numpy.ndindex(self.inner_size()):
+            yield tuple(map(lambda i: i + 1, idx))
 
-    def span(self):
+    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 inner_span(self):
+    def inner_size(self):
         if self.size_z == 1:
             return (self.size_x-2, self.size_y-2)
         else:
@@ -75,18 +68,14 @@ class Lattice:
         }.get((descriptor.d, descriptor.q), None)
 
         self.program.equilibrilize(
-            self.queue, self.geometry.span(), self.layout, self.cl_pop_a, self.cl_pop_b).wait()
+            self.queue, self.geometry.size(), self.layout, self.cl_pop_a, self.cl_pop_b).wait()
 
-    def idx(self, x, y, z = 0):
+    def gid(self, x, y, z = 0):
         return z * (self.geometry.size_x*self.geometry.size_y) + y * self.geometry.size_x + x;
 
     def setup_geometry(self, material_at):
-        if self.descriptor.d == 2:
-            for x, y in self.geometry.inner_cells():
-                self.np_material[self.idx(x,y)] = material_at(self.geometry, x, y)
-        elif self.descriptor.d == 3:
-            for x, y, z in self.geometry.inner_cells():
-                self.np_material[self.idx(x,y,z)] = material_at(self.geometry, x, y, z)
+        for idx in self.geometry.inner_cells():
+            self.np_material[self.gid(*idx)] = material_at(self.geometry, *idx)
 
         cl.enqueue_copy(self.queue, self.cl_material, self.np_material).wait();
 
@@ -117,11 +106,11 @@ class Lattice:
         if self.tick:
             self.tick = False
             self.program.collide_and_stream(
-                self.queue, self.geometry.span(), self.layout, self.cl_pop_a, self.cl_pop_b, self.cl_material)
+                self.queue, self.geometry.size(), self.layout, self.cl_pop_a, self.cl_pop_b, self.cl_material)
         else:
             self.tick = True
             self.program.collide_and_stream(
-                self.queue, self.geometry.span(), self.layout, self.cl_pop_b, self.cl_pop_a, self.cl_material)
+                self.queue, self.geometry.size(), self.layout, self.cl_pop_b, self.cl_pop_a, self.cl_material)
 
     def sync(self):
         self.queue.finish()
@@ -130,9 +119,9 @@ class Lattice:
         moments = numpy.ndarray(shape=(self.descriptor.d+1, self.geometry.volume), dtype=numpy.float32)
         if self.tick:
             self.program.collect_moments(
-                self.queue, self.geometry.span(), self.layout, self.cl_pop_b, self.cl_moments)
+                self.queue, self.geometry.size(), self.layout, self.cl_pop_b, self.cl_moments)
         else:
             self.program.collect_moments(
-                self.queue, self.geometry.span(), self.layout, self.cl_pop_a, self.cl_moments)
+                self.queue, self.geometry.size(), self.layout, self.cl_pop_a, self.cl_moments)
         cl.enqueue_copy(self.queue, moments, self.cl_moments).wait();
         return moments