Adapt examples to new material interface

8 files changed, 89 insertions, 75 deletions

diff --git a/channel_2d_gl_interop.py b/channel_2d_gl_interop.py
index 3db64df..7c78af5 100644
--- a/channel_2d_gl_interop.py
+++ b/channel_2d_gl_interop.py
@@ -13,41 +13,31 @@ from OpenGL.GL import shaders
 
 screen_x = 1920
 screen_y = 1200
-pixels_per_cell   = 2
+pixels_per_cell   = 4
 updates_per_frame = 80
 
 inflow = 0.02
 relaxation_time = 0.51
 
+def get_channel_material_map(geometry):
+    return [
+        (lambda x, y: x > 0 and x < geometry.size_x-1 and y > 0 and y < geometry.size_y-1, 1), # bulk fluid
+        (lambda x, y: x == 1,                 3), # inflow
+        (lambda x, y: x == geometry.size_x-2, 4), # outflow
+        (lambda x, y: y == 1,                 2), # bottom
+        (lambda x, y: y == geometry.size_y-2, 2), # top
+        (lambda x, y: x == 0 or x == geometry.size_x-1 or y == 0 or y == geometry.size_y-1, 0) # ghost cells
+    ]
+
 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) ]
     rs = [ geometry.size_x//100 for i, y in enumerate(ys) ]
     return list(zip(xs, ys, rs))
 
-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
-
-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
+def get_obstacles_material_map(obstacles):
+    indicator = lambda ox, oy, r: (lambda x, y: (ox - x)**2 + (oy - y)**2 < r*r)
+    return [ (indicator(ox, oy, r), 2) for ox, oy, r in obstacles ]
 
 boundary = Template("""
     if ( m == 2 ) {
@@ -151,11 +141,10 @@ lattice = Lattice(
     opengl       = True
 )
 
-set_bulk(lattice)
-set_walls(lattice)
-
-for obstacle in get_obstacles(lattice.geometry):
-    set_obstacle(lattice, *obstacle)
+lattice.apply_material_map(
+    get_channel_material_map(lattice.geometry))
+lattice.apply_material_map(
+    get_obstacles_material_map(get_obstacles(lattice.geometry)))
 
 lattice.sync_material()
 
diff --git a/implosion.py b/implosion.py
index 245f2bb..b95dffd 100644
--- a/implosion.py
+++ b/implosion.py
@@ -25,11 +25,12 @@ def generate_moment_plots(lattice, moments):
         plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
         plt.savefig("result/implosion_%02d.png" % i, bbox_inches='tight', pad_inches=0)
 
-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
+def get_box_material_map(geometry):
+    return [
+        (lambda x, y: x > 0 and x < geometry.size_x-1 and y > 0 and y < geometry.size_y-1,  1), # bulk fluid
+        (lambda x, y: x == 1 or y == 1 or x == geometry.size_x-2 or y == geometry.size_y-2, 2), # walls
+        (lambda x, y: x == 0 or x == geometry.size_x-1 or y == 0 or y == geometry.size_y-1, 0)  # ghost cells
+    ]
 
 pop_eq = """
     if ( sqrt(pow(get_global_id(0) - ${geometry.size_x//2}.f, 2.f)
@@ -73,7 +74,9 @@ lattice = Lattice(
     pop_eq_src   = pop_eq,
     boundary_src = boundary)
 
-lattice.setup_geometry(box)
+lattice.apply_material_map(
+    get_box_material_map(lattice.geometry))
+lattice.sync_material()
 
 print("Starting simulation using %d cells...\n" % lattice.geometry.volume)
 
diff --git a/ldc_2d.py b/ldc_2d.py
index 5fae3fe..8869c35 100644
--- a/ldc_2d.py
+++ b/ldc_2d.py
@@ -29,13 +29,13 @@ def generate_moment_plots(lattice, moments):
         plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
         plt.savefig("result/ldc_2d_%02d.png" % i, bbox_inches='tight', pad_inches=0)
 
-def cavity(geometry, x, y):
-    if x == 1 or y == 1 or x == geometry.size_x-2:
-        return 2
-    elif y == geometry.size_y-2:
-        return 3
-    else:
-        return 1
+def get_cavity_material_map(geometry):
+    return [
+        (lambda x, y: x > 0 and x < geometry.size_x-1 and y > 0 and y < geometry.size_y-1,  1), # bulk fluid
+        (lambda x, y: x == 1 or y == 1 or x == geometry.size_x-2,                           2), # left, right, bottom walls
+        (lambda x, y: y == geometry.size_y-2,                                               3), # lid
+        (lambda x, y: x == 0 or x == geometry.size_x-1 or y == 0 or y == geometry.size_y-1, 0)  # ghost cells
+    ]
 
 boundary = Template("""
     if ( m == 2 ) {
@@ -72,7 +72,9 @@ lattice = Lattice(
 
     boundary_src = boundary)
 
-lattice.setup_geometry(cavity)
+lattice.apply_material_map(
+    get_cavity_material_map(lattice.geometry))
+lattice.sync_material()
 
 print("Starting simulation using %d cells...\n" % lattice.geometry.volume)
 
diff --git a/ldc_2d_benchmark.py b/ldc_2d_benchmark.py
index df697ea..88041b8 100644
--- a/ldc_2d_benchmark.py
+++ b/ldc_2d_benchmark.py
@@ -15,13 +15,13 @@ relaxation_time = 0.52
 def MLUPS(cells, steps, time):
     return cells * steps / time * 1e-6
 
-def cavity(geometry, x, y):
-    if x == 1 or y == 1 or x == geometry.size_x-2:
-        return 2
-    elif y == geometry.size_y-2:
-        return 3
-    else:
-        return 1
+def get_cavity_material_map(geometry):
+    return [
+        (lambda x, y: x > 0 and x < geometry.size_x-1 and y > 0 and y < geometry.size_y-1,  1), # bulk fluid
+        (lambda x, y: x == 1 or y == 1 or x == geometry.size_x-2,                           2), # left, right, bottom walls
+        (lambda x, y: y == geometry.size_y-2,                                               3), # lid
+        (lambda x, y: x == 0 or x == geometry.size_x-1 or y == 0 or y == geometry.size_y-1, 0)  # ghost cells
+    ]
 
 boundary = Template("""
     if ( m == 2 ) {
@@ -89,7 +89,9 @@ for size, layout, precision, opti, align in base_2_configs + align_configs + pad
         moments = lbm.moments(optimize = opti),
         collide = lbm.bgk(f_eq = lbm.equilibrium(), tau = relaxation_time, optimize = opti),
         boundary_src = boundary)
-    lattice.setup_geometry(cavity)
+    lattice.apply_material_map(
+        get_cavity_material_map(lattice.geometry))
+    lattice.sync_material()
 
     nUpdates = 500
     nStat = 100
diff --git a/ldc_2d_gl_interop.py b/ldc_2d_gl_interop.py
index e5ea866..ec94a75 100644
--- a/ldc_2d_gl_interop.py
+++ b/ldc_2d_gl_interop.py
@@ -19,13 +19,13 @@ updates_per_frame = 200
 lid_speed = 0.1
 relaxation_time = 0.515
 
-def cavity(geometry, x, y):
-    if x == 1 or y == 1 or x == geometry.size_x-2:
-        return 2
-    elif y == geometry.size_y-2:
-        return 3
-    else:
-        return 1
+def get_cavity_material_map(geometry):
+    return [
+        (lambda x, y: x > 0 and x < geometry.size_x-1 and y > 0 and y < geometry.size_y-1,  1), # bulk fluid
+        (lambda x, y: x == 1 or y == 1 or x == geometry.size_x-2,                           2), # left, right, bottom walls
+        (lambda x, y: y == geometry.size_y-2,                                               3), # lid
+        (lambda x, y: x == 0 or x == geometry.size_x-1 or y == 0 or y == geometry.size_y-1, 0)  # ghost cells
+    ]
 
 boundary = Template("""
     if ( m == 2 ) {
@@ -126,7 +126,9 @@ lattice = Lattice(
     opengl       = True
 )
 
-lattice.setup_geometry(cavity)
+lattice.apply_material_map(
+    get_cavity_material_map(lattice.geometry))
+lattice.sync_material()
 
 projection = get_projection()
 
diff --git a/ldc_3d.py b/ldc_3d.py
index a41f0ae..b0ed568 100644
--- a/ldc_3d.py
+++ b/ldc_3d.py
@@ -55,13 +55,17 @@ def generate_moment_plots(lattice, moments):
 
         plt.savefig("result/ldc_3d_%02d.png" % i, bbox_inches='tight', pad_inches=0)
 
-def cavity(geometry, x, y, z):
-    if x == 1 or y == 1 or z == 1 or x == geometry.size_x-2 or y == geometry.size_y-2:
-        return 2
-    elif z == geometry.size_z-2:
-        return 3
-    else:
-        return 1
+def get_cavity_material_map(geometry):
+    return [
+        (lambda x, y, z: x > 0 and x < geometry.size_x-1 and
+                         y > 0 and y < geometry.size_y-1 and
+                         z > 0 and z < geometry.size_z-1,                                                1), # bulk fluid
+        (lambda x, y, z: x == 1 or y == 1 or z == 1 or x == geometry.size_x-2 or y == geometry.size_y-2, 2), # walls
+        (lambda x, y, z: z == geometry.size_z-2,                                                         3), # lid
+        (lambda x, y, z: x == 0 or x == geometry.size_x-1 or
+                         y == 0 or y == geometry.size_y-1 or
+                         z == 0 or z == geometry.size_z-1,                                               0)  # ghost cells
+    ]
 
 boundary = """
     if ( m == 2 ) {
@@ -94,7 +98,9 @@ lattice = Lattice(
 
     boundary_src = boundary)
 
-lattice.setup_geometry(cavity)
+lattice.apply_material_map(
+    get_cavity_material_map(lattice.geometry))
+lattice.sync_material()
 
 print("Starting simulation using %d cells...\n" % lattice.geometry.volume)
 
diff --git a/ldc_3d_benchmark.py b/ldc_3d_benchmark.py
index 2075b32..4383add 100644
--- a/ldc_3d_benchmark.py
+++ b/ldc_3d_benchmark.py
@@ -16,13 +16,17 @@ relaxation_time = 0.52
 def MLUPS(cells, steps, time):
     return cells * steps / time * 1e-6
 
-def cavity(geometry, x, y, z):
-    if x == 1 or y == 1 or z == 1 or x == geometry.size_x-2 or y == geometry.size_y-2:
-        return 2
-    elif z == geometry.size_z-2:
-        return 3
-    else:
-        return 1
+def get_cavity_material_map(geometry):
+    return [
+        (lambda x, y, z: x > 0 and x < geometry.size_x-1 and
+                         y > 0 and y < geometry.size_y-1 and
+                         z > 0 and z < geometry.size_z-1,                                                1), # bulk fluid
+        (lambda x, y, z: x == 1 or y == 1 or z == 1 or x == geometry.size_x-2 or y == geometry.size_y-2, 2), # walls
+        (lambda x, y, z: z == geometry.size_z-2,                                                         3), # lid
+        (lambda x, y, z: x == 0 or x == geometry.size_x-1 or
+                         y == 0 or y == geometry.size_y-1 or
+                         z == 0 or z == geometry.size_z-1,                                               0)  # ghost cells
+    ]
 
 boundary = Template("""
     if ( m == 2 ) {
@@ -91,7 +95,9 @@ for size, layout, descriptor, precision, opti, align in base_2_configs + align_c
         moments = lbm.moments(optimize = opti),
         collide = lbm.bgk(f_eq = lbm.equilibrium(), tau = relaxation_time, optimize = opti),
         boundary_src = boundary)
-    lattice.setup_geometry(cavity)
+    lattice.apply_material_map(
+        get_cavity_material_map(lattice.geometry))
+    lattice.sync_material()
 
     nUpdates = 500
     nStat = 100
diff --git a/simulation.py b/simulation.py
index badd2bb..1f9a6c1 100644
--- a/simulation.py
+++ b/simulation.py
@@ -160,6 +160,10 @@ class Lattice:
 
         self.material = numpy.ndarray(shape=(self.memory.volume, 1), dtype=numpy.int32)
 
+    def apply_material_map(self, material_map):
+        for indicator, material in material_map:
+            self.material[[indicator(*idx) for idx in self.memory.cells()]] = material
+
     def sync_material(self):
         cl.enqueue_copy(self.queue, self.memory.cl_material, self.material).wait();