new geometry class. beginning to implement physics by defining a material class; each triangle will have a material linked to both of its sides

1 files changed, 110 insertions, 0 deletions

diff --git a/geometry.py b/geometry.py
new file mode 100644
index 0000000..d365141
--- /dev/null
+++ b/geometry.py
@@ -0,0 +1,110 @@
+import numpy as np
+from bvh import *
+
+class Geometry(object):
+    """
+    Geometry object.
+    """
+
+    def __init__(self):
+        self.materials = []
+        self.mesh = np.empty(0)
+        self.inside = np.empty(0)
+        self.outside = np.empty(0)
+
+    def add_mesh(self, mesh, inside, outside):
+        """
+        Add a triangle mesh.
+
+        Args:
+            - mesh: array of shape (n,3,3)
+                Array of vertices for a closed mesh with n triangles.
+
+            - inside: Material
+                Material inside the mesh.
+
+            - outside: Material
+                Material outside the mesh.
+         """
+
+        if len(mesh.shape) != 3 or mesh.shape[1] != 3 or mesh.shape[2] != 3:
+            raise Exception('shape mismatch')
+
+        if inside not in self.materials:
+            self.materials.append(inside)
+        if outside not in self.materials:
+            self.materials.append(outside)
+
+        self.inside.resize(self.inside.size+mesh.shape[0])
+        self.inside[-mesh.shape[0]:] = np.tile(self.materials.index(inside), mesh.shape[0])
+
+        self.outside.resize(self.outside.size+mesh.shape[0])
+        self.outside[-mesh.shape[0]:] = np.tile(self.materials.index(outside), mesh.shape[0])
+
+        self.mesh.resize((self.mesh.shape[0]+mesh.shape[0],3,3))
+        self.mesh[-mesh.shape[0]:] = mesh
+
+    def build(self, bits=16):
+        order, zvalues = morton_order(self.mesh, bits)
+
+
+        zvalues = zvalues[order]
+        self.mesh = self.mesh[order]
+        self.inside = self.inside[order]
+        self.outside = self.outside[order]
+
+        leafs = []
+
+        for z in sorted(set(zvalues)):
+            mask = (zvalues == z)
+            leafs.append(Leaf(self.mesh[mask], 3*np.where(mask)[0][0], z))
+
+        layers = []
+        layers.append(leafs)
+
+        while True:
+            zvalues = np.array([node.zvalue for node in layers[-1]])
+
+            bit_shifted_zvalues = zvalues >> 1
+
+            nodes = []
+            for z in sorted(set(bit_shifted_zvalues)):
+                mask = (bit_shifted_zvalues == z)
+                nodes.append(Node(list(compress(layers[-1], mask)), z))
+
+            layers.append(nodes)
+
+            if len(nodes) == 1:
+                break
+
+        layers.reverse()
+
+        nodes = []
+        for layer in layers:
+            for node in layer:
+                nodes.append(node)
+
+
+        self.lower_bound = np.empty((len(nodes),3))
+        self.upper_bound = np.empty((len(nodes),3))
+        self.child_map = np.empty(len(nodes))
+        self.child_len = np.empty(len(nodes))
+        self.first_leaf = -1
+
+        for i, node in enumerate(nodes):
+            self.lower_bound[i] = node.lower_bound
+            self.upper_bound[i] = node.upper_bound
+            
+            self.child_len[i] = node.size
+
+            if isinstance(node, Node):
+                for j, child in enumerate(nodes):
+                    if child is node.children[0]:
+                        self.child_map[i] = j
+                        break
+
+            if isinstance(node, Leaf):
+                self.child_map[i] = node.mesh_idx
+
+                if self.first_leaf == -1:
+                    self.first_leaf = i