1 files changed, 24 insertions, 0 deletions
diff --git a/boltzgen/geometry.py b/boltzgen/geometry.py
new file mode 100644
index 0000000..4b61982
--- /dev/null
+++ b/boltzgen/geometry.py
@@ -0,0 +1,24 @@
+from utility.ndindex import ndindex
+
+class Geometry:
+    def __init__(self, size_x, size_y, size_z = 1):
+        self.size_x = size_x
+        self.size_y = size_y
+        self.size_z = size_z
+        self.volume = size_x * size_y * size_z
+
+    def inner_cells(self):
+        for idx in numpy.ndindex(self.inner_size()):
+            yield tuple(map(lambda i: i + 1, idx))
+
+    def size(self):
+        if self.size_z == 1:
+            return (self.size_x, self.size_y)
+        else:
+            return (self.size_x, self.size_y, self.size_z)
+
+    def inner_size(self):
+        if self.size_z == 1:
+            return (self.size_x-2, self.size_y-2)
+        else:
+            return (self.size_x-2, self.size_y-2, self.size_z-2)