Add kernel customization point for velocity boundaries

4 files changed, 27 insertions, 12 deletions

diff --git a/D2Q9.py b/D2Q9.py
index fbce5d6..88bd88c 100644
--- a/D2Q9.py
+++ b/D2Q9.py
@@ -12,7 +12,7 @@ class Lattice:
     def idx(self, x, y):
         return y * self.nX + x;
 
-    def __init__(self, nX, nY, tau, geometry, pop_eq_src = ''):
+    def __init__(self, nX, nY, tau, geometry, pop_eq_src = '', boundary_src = ''):
         self.nX = nX
         self.nY = nY
         self.nCells = nX * nY
@@ -27,6 +27,7 @@ class Lattice:
         self.setup_geometry(geometry)
 
         self.pop_eq_src = pop_eq_src
+        self.boundary_src = boundary_src
 
         self.cl_pop_a = cl.Buffer(self.context, mf.READ_WRITE, size=9*self.nCells*numpy.float32(0).nbytes)
         self.cl_pop_b = cl.Buffer(self.context, mf.READ_WRITE, size=9*self.nCells*numpy.float32(0).nbytes)
@@ -56,12 +57,16 @@ class Lattice:
             c     = D2Q9.c,
             w     = D2Q9.w,
             ccode = sympy.ccode,
+
             pop_eq_src = Template(self.pop_eq_src).render(
                 nX     = self.nX,
                 nY     = self.nY,
                 nCells = self.nCells,
                 c     = D2Q9.c,
                 w     = D2Q9.w
+            ),
+            boundary_src = Template(self.boundary_src).render(
+                d = D2Q9.d
             )
         )
         self.program = cl.Program(self.context, program_src).build()
diff --git a/implosion.py b/implosion.py
index d2754c2..e5e87bc 100644
--- a/implosion.py
+++ b/implosion.py
@@ -42,6 +42,13 @@ pop_eq = """
 % endfor
 }"""
 
+boundary = """
+    if ( m == 2 ) {
+        u_0 = 0.0;
+        u_1 = 0.0;
+    }
+"""
+
 nUpdates = 2000
 nStat    = 100
 
@@ -49,7 +56,7 @@ moments = []
 
 print("Initializing simulation...\n")
 
-lattice = Lattice(nX = 1024, nY = 1024, tau = 0.8, geometry = box, pop_eq_src = pop_eq)
+lattice = Lattice(nX = 1024, nY = 1024, tau = 0.8, geometry = box, pop_eq_src = pop_eq, boundary_src = boundary)
 
 print("Starting simulation using %d cells...\n" % lattice.nCells)
 
diff --git a/lid_driven_cavity.py b/lid_driven_cavity.py
index b8a59db..db51e73 100644
--- a/lid_driven_cavity.py
+++ b/lid_driven_cavity.py
@@ -31,6 +31,17 @@ def box(nX, nY, x, y):
     else:
         return 1
 
+boundary = """
+    if ( m == 2 ) {
+        u_0 = 0.0;
+        u_1 = 0.0;
+    }
+    if ( m == 3 ) {
+        u_0 = 0.1;
+        u_1 = 0.0;
+    }
+"""
+
 nUpdates = 100000
 nStat    = 5000
 
@@ -38,7 +49,7 @@ moments = []
 
 print("Initializing simulation...\n")
 
-lattice = Lattice(nX = 256, nY = 256, tau = 0.56, geometry = box)
+lattice = Lattice(nX = 256, nY = 256, tau = 0.56, geometry = box, boundary_src = boundary)
 
 print("Starting simulation using %d cells...\n" % lattice.nCells)
 
diff --git a/template/kernel.mako b/template/kernel.mako
index c76e5d9..b81d432 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -56,15 +56,7 @@ __kernel void collide_and_stream(__global __write_only float* f_a,
     float ${ccode(expr)}
 % endfor
 
-    if ( m == 2 ) {
-        u_0 = 0.0;
-        u_1 = 0.0;
-    }
-
-    if ( m == 3 ) {
-        u_0 = 0.1;
-        u_1 = 0.0;
-    }
+  ${boundary_src}
 
 % for i, expr in enumerate(collide_helper):
     const float ${expr[0]} = ${ccode(expr[1])};