Make it easier to exchange initial equilibration logic

3 files changed, 39 insertions, 22 deletions

diff --git a/D2Q9.py b/D2Q9.py
index 173f1b7..fbce5d6 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):
+    def __init__(self, nX, nY, tau, geometry, pop_eq_src = ''):
         self.nX = nX
         self.nY = nY
         self.nCells = nX * nY
@@ -26,6 +26,8 @@ class Lattice:
         self.np_material = numpy.ndarray(shape=(self.nCells, 1), dtype=numpy.int32)
         self.setup_geometry(geometry)
 
+        self.pop_eq_src = pop_eq_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)
 
@@ -53,7 +55,14 @@ class Lattice:
             collide_assignment = D2Q9.collide_opt[1],
             c     = D2Q9.c,
             w     = D2Q9.w,
-            ccode = sympy.ccode
+            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
+            )
         )
         self.program = cl.Program(self.context, program_src).build()
 
diff --git a/implosion.py b/implosion.py
index 75e9640..d2754c2 100644
--- a/implosion.py
+++ b/implosion.py
@@ -29,14 +29,27 @@ def box(nX, nY, x, y):
     else:
         return 1
 
-nUpdates = 1000
+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} ) {
+% for i, w_i in enumerate(w):
+        preshifted_f_a[${i*nCells}] = 1./24.f;
+        preshifted_f_b[${i*nCells}] = 1./24.f;
+% endfor
+    } else {
+% for i, w_i in enumerate(w):
+        preshifted_f_a[${i*nCells}] = ${w_i}.f;
+        preshifted_f_b[${i*nCells}] = ${w_i}.f;
+% endfor
+}"""
+
+nUpdates = 2000
 nStat    = 100
 
 moments = []
 
 print("Initializing simulation...\n")
 
-lattice = Lattice(nX = 1024, nY = 1024, tau = 0.8, geometry = box)
+lattice = Lattice(nX = 1024, nY = 1024, tau = 0.8, geometry = box, pop_eq_src = pop_eq)
 
 print("Starting simulation using %d cells...\n" % lattice.nCells)
 
diff --git a/template/kernel.mako b/template/kernel.mako
index c43a0dc..c76e5d9 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -1,9 +1,5 @@
 __constant float tau = ${tau};
 
-bool is_in_circle(float x, float y, float a, float b, float r) {
-    return sqrt(pow(x-a,2)+pow(y-b,2)) < r;
-}
-
 __kernel void equilibrilize(__global __write_only float* f_a,
                             __global __write_only float* f_b)
 {
@@ -12,20 +8,14 @@ __kernel void equilibrilize(__global __write_only float* f_a,
     __global __write_only float* preshifted_f_a = f_a + gid;
     __global __write_only float* preshifted_f_b = f_b + gid;
 
-    if (  is_in_circle(get_global_id(0), get_global_id(1), ${nX//4},    ${nY//4},    ${nX//10})
-       || is_in_circle(get_global_id(0), get_global_id(1), ${nX//4},    ${nY-nY//4}, ${nX//10})
-       || is_in_circle(get_global_id(0), get_global_id(1), ${nX-nX//4}, ${nY//4},    ${nX//10})
-       || is_in_circle(get_global_id(0), get_global_id(1), ${nX-nX//4}, ${nY-nY//4}, ${nX//10}) ) {
-% for i, w_i in enumerate(w):
-        preshifted_f_a[${i*nCells}] = 1./24.f;
-        preshifted_f_b[${i*nCells}] = 1./24.f;
-% endfor
-    } else {
-% for i, w_i in enumerate(w):
-        preshifted_f_a[${i*nCells}] = ${w_i}.f;
-        preshifted_f_b[${i*nCells}] = ${w_i}.f;
-% endfor
-    }
+% if pop_eq_src == '':
+%     for i, w_i in enumerate(w):
+    preshifted_f_a[${i*nCells}] = ${w_i}.f;
+    preshifted_f_b[${i*nCells}] = ${w_i}.f;
+%     endfor
+% else:
+    ${pop_eq_src}
+% endif
 }
 
 <%
@@ -71,6 +61,11 @@ __kernel void collide_and_stream(__global __write_only float* f_a,
         u_1 = 0.0;
     }
 
+    if ( m == 3 ) {
+        u_0 = 0.1;
+        u_1 = 0.0;
+    }
+
 % for i, expr in enumerate(collide_helper):
     const float ${expr[0]} = ${ccode(expr[1])};
 % endfor