Add geometry argument

3 files changed, 22 insertions, 1 deletions

diff --git a/boltzgen.py b/boltzgen.py
index 6f8f20c..95cccd9 100644
--- a/boltzgen.py
+++ b/boltzgen.py
@@ -11,6 +11,7 @@ argparser.add_argument('language', help = 'Target language (currently either "cl
 argparser.add_argument('--lattice',   dest = 'lattice',   required = True, help = 'Lattice type (D2Q9, D3Q7, D3Q19, D3Q27)')
 argparser.add_argument('--layout',    dest = 'layout',    required = True, help = 'Memory layout ("AOS" or "SOA")')
 argparser.add_argument('--precision', dest = 'precision', required = True, help = 'Floating precision ("single" or "double")')
+argparser.add_argument('--geometry',  dest = 'geometry',  required = True, help = 'Size of the block geometry ("x:y(:z)")')
 
 args = argparser.parse_args()
 
@@ -22,7 +23,7 @@ generator = Generator(
     moments    = lbm.moments(),
     collision  = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.6))
 
-geometry = Geometry(1024,1024)
+geometry = Geometry.parse(args.geometry)
 
 src = generator.kernel(args.language, args.precision, args.layout, geometry)
 print(src)
diff --git a/boltzgen/geometry.py b/boltzgen/geometry.py
index 4b61982..193cbe6 100644
--- a/boltzgen/geometry.py
+++ b/boltzgen/geometry.py
@@ -1,3 +1,4 @@
+import re
 from utility.ndindex import ndindex
 
 class Geometry:
@@ -7,6 +8,22 @@ class Geometry:
         self.size_z = size_z
         self.volume = size_x * size_y * size_z
 
+    @classmethod
+    def parse(self, s):
+        args = re.search(r'([0-9]+):([0-9]+):?([0-9]+)?', s)
+        if args.group(3) == None:
+            return self(int(args.group(1)), int(args.group(2)))
+        else:
+            return self(int(args.group(1)), int(args.group(2)), int(args.group(3)))
+
+    def dimension(self):
+        if self.size_x > 1 and self.size_y > 1 and self.size_z == 1:
+            return 2
+        elif self.size_x > 1 and self.size_y > 1 and self.size_z > 1:
+            return 3
+        else:
+            raise Exception('Geometry malformed')
+
     def inner_cells(self):
         for idx in numpy.ndindex(self.inner_size()):
             yield tuple(map(lambda i: i + 1, idx))
diff --git a/boltzgen/kernel/generator.py b/boltzgen/kernel/generator.py
index d91ddd5..434d163 100644
--- a/boltzgen/kernel/generator.py
+++ b/boltzgen/kernel/generator.py
@@ -24,6 +24,9 @@ class Generator:
         else:
             layout_impl = layout_impl(self.descriptor, geometry)
 
+        if geometry.dimension() != self.descriptor.d:
+            raise Exception('Geometry dimension must match descriptor dimension')
+
         return Template(filename = str(template_path)).render(
             descriptor = self.descriptor,
             geometry   = geometry,