Add basic Generator class

1 files changed, 24 insertions, 16 deletions

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
+        )