Add basic Generator class

5 files changed, 36 insertions, 32 deletions

diff --git a/boltzgen.py b/boltzgen.py
index 6d410c3..27a34fa 100644
--- a/boltzgen.py
+++ b/boltzgen.py
@@ -9,16 +9,12 @@ argparser.add_argument('language', help = 'Target language (currently either "op
 args = argparser.parse_args()
 
 lbm = LBM(D2Q9)
-geometry = Geometry(32,32)
+generator = Generator(
+    descriptor = D2Q9,
+    moments    = lbm.moments(),
+    collision  = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.6))
 
-src = source(
-    args.language,
-    D2Q9,
-    lbm.moments(),
-    lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.6),
-    "",
-    'float',
-    geometry
-)
+geometry = Geometry(32,32)
 
+src = generator.kernel(args.language, 'float', geometry)
 print(src)
diff --git a/boltzgen/__init__.py b/boltzgen/__init__.py
index 8896c65..0cf365b 100644
--- a/boltzgen/__init__.py
+++ b/boltzgen/__init__.py
@@ -6,4 +6,4 @@ import lbm.model.D3Q27 as D3Q27
 
 from geometry import Geometry
 
-from kernel.generator import source
+from kernel.generator import Generator
diff --git a/boltzgen/kernel/generator.py b/boltzgen/kernel/generator.py
index 59654e2..5853561 100644
--- a/boltzgen/kernel/generator.py
+++ b/boltzgen/kernel/generator.py
@@ -3,23 +3,31 @@ import sympy
 from mako.template import Template
 from pathlib import Path
 
-def source(target, descriptor, moments, collide, boundary_src, float_type, geometry):
-    return Template(filename = str(Path(__file__).parent/('template/basic.' + target + '.mako'))).render(
-        descriptor = descriptor,
-        geometry   = geometry,
+class Generator:
+    def __init__(self, descriptor, moments, collision, boundary = ''):
+        self.descriptor = descriptor
+        self.moments    = moments
+        self.collision  = collision
+        self.boundary   = boundary
 
-        moments_subexpr    = moments[0],
-        moments_assignment = moments[1],
-        collide_subexpr    = collide[0],
-        collide_assignment = collide[1],
+    def kernel(self, target, precision, geometry):
+        return Template(filename = str(Path(__file__).parent/('template/basic.' + target + '.mako'))).render(
+            descriptor = self.descriptor,
+            geometry   = geometry,
 
-        float_type = float_type,
+            moments_subexpr    = self.moments[0],
+            moments_assignment = self.moments[1],
 
-        boundary_src = Template(boundary_src).render(
-            descriptor = descriptor,
-            geometry   = geometry,
-            float_type = float_type
-        ),
+            collision_subexpr    = self.collision[0],
+            collision_assignment = self.collision[1],
+
+            float_type = precision,
+
+            boundary_src = Template(self.boundary).render(
+                descriptor = self.descriptor,
+                geometry   = geometry,
+                float_type = precision
+            ),
 
-        ccode = sympy.ccode
-    )
+            ccode = sympy.ccode
+        )
diff --git a/boltzgen/kernel/template/basic.cpp.mako b/boltzgen/kernel/template/basic.cpp.mako
index 1879c35..ef3486f 100644
--- a/boltzgen/kernel/template/basic.cpp.mako
+++ b/boltzgen/kernel/template/basic.cpp.mako
@@ -52,15 +52,15 @@ void collide_and_stream(      ${float_type}* f_next,
     ${float_type} ${ccode(expr)}
 % endfor
 
-% for i, expr in enumerate(collide_subexpr):
+% for i, expr in enumerate(collision_subexpr):
     const ${float_type} ${expr[0]} = ${ccode(expr[1])};
 % endfor
 
-% for i, expr in enumerate(collide_assignment):
+% for i, expr in enumerate(collision_assignment):
     const ${float_type} ${ccode(expr)}
 % endfor
 
-% for i, expr in enumerate(collide_assignment):
+% for i, expr in enumerate(collision_assignment):
     preshifted_f_next[${pop_offset(i)}] = m*f_next_${i} + (1.0-m)*${descriptor.w[i].evalf()};
 % endfor
 }
diff --git a/boltzgen/kernel/template/basic.opencl.mako b/boltzgen/kernel/template/basic.opencl.mako
index 5ddf64c..3ecb16c 100644
--- a/boltzgen/kernel/template/basic.opencl.mako
+++ b/boltzgen/kernel/template/basic.opencl.mako
@@ -70,11 +70,11 @@ __kernel void collide_and_stream(__global ${float_type}* f_next,
 
   ${boundary_src}
 
-% for i, expr in enumerate(collide_subexpr):
+% for i, expr in enumerate(collision_subexpr):
     const ${float_type} ${expr[0]} = ${ccode(expr[1])};
 % endfor
 
-% for i, expr in enumerate(collide_assignment):
+% for i, expr in enumerate(collision_assignment):
     const ${float_type} ${ccode(expr)}
 % endfor