Merge pull request #7 from hsorby/octree

Add a simple octree implementation.

Author: hsorby

src/cmlibs/maths/octree.py

@@ -0,0 +1,106 @@
+"""
+Octree for searching for objects by coordinates
+"""
+import math
+
+from cmlibs.maths.vectorops import sub, dot
+
+
+class Octree:
+    """
+    Octree for searching for objects by coordinates
+    """
+
+    def __init__(self, tolerance=None):
+        """
+        :param tolerance: If supplied, tolerance to use, or None to compute as 1.0E-12.
+        """
+        self._tolerance = 1.0E-12 if tolerance is None else tolerance
+        self._coordinates = None
+        # Octree is either leaf with _object, or has 2**self._dimension children
+        self._object = None
+        # exactly 2^self._dimension children, cycling in lowest x index fastest
+        self._children = None
+
+    def _location_match(self, x):
+        return all([math.fabs(x[i] - self._coordinates[i]) < self._tolerance for i in range(3)])
+
+    def _child_index_lookup(self, x):
+        switch = [0 if x[i] < self._coordinates[i] else 1 for i in range(3)]
+        return switch[0] + 2 * switch[1] + 4 * switch[2]
+
+    def _sq_distance(self, x):
+        diff = sub(x, self._coordinates)
+        return dot(diff, diff)
+
+    def _find_object_by_coordinates(self, x, nearest=False):
+        """
+        Find the closest existing object with |x - ox| < tolerance.
+        :param x: 3 coordinates in a list.
+        :return: nearest distance, nearest object or None, None if none found.
+        """
+        if self._coordinates is not None and self._location_match(x):
+            return 0.0, self._object
+
+        if not nearest and self._children is None:
+            return None, None
+
+        if nearest and self._coordinates is None:
+            return math.inf, None
+
+        if nearest and self._children is None:
+            return 0.0, self._object
+
+        index = self._child_index_lookup(x)
+        sq_distance_, object_ = self._children[index]._find_object_by_coordinates(x, nearest)
+
+        if nearest:
+            sq_distance = self._sq_distance(x)
+            if sq_distance_ < sq_distance:
+                return sq_distance_, object_
+            else:
+                return sq_distance, self._object
+
+        return None, object_
+
+    def tolerance(self):
+        return self._tolerance
+
+    def find_object_by_coordinates(self, x):
+        """
+        Find the closest existing object with |x - ox| < tolerance.
+        :param x: 3 coordinates in a list.
+        :return: nearest object or None if not found.
+        """
+        distance, nearest_object = self._find_object_by_coordinates(x)
+        return nearest_object
+
+    def find_nearest_object_by_coordinates(self, x):
+        distance, nearest_object = self._find_object_by_coordinates(x, True)
+        return nearest_object
+
+    def insert_object_at_coordinates(self, x, obj):
+        """
+        Add object at coordinates to octree.
+
+        :param x: 3 coordinates in a list.
+        :param obj: object to store with coordinates.
+        """
+        if self._coordinates is None:
+            self._coordinates = x
+            self._object = obj
+        else:
+            if self._location_match(x):
+                self._object = obj
+            else:
+                index = self._child_index_lookup(x)
+                if self._children is None:
+                    self._children = [Octree() for _ in range(8)]
+
+                self._children[index].insert_object_at_coordinates(x, obj)
+
+    def __repr__(self):
+        if self._children is None:
+            return f'\tleaf {self._coordinates}\n'
+
+        return f'{self._coordinates} - \n' + ''.join([f'{c}' for c in self._children])

tests/test_octree.py

@@ -0,0 +1,87 @@
+import unittest
+
+from cmlibs.maths.octree import Octree
+
+
+class OctTreeTestCase(unittest.TestCase):
+
+    def test_oct_tree_1(self):
+        pts = [[0.5002720342505129, 0.049999999999999996, -0.08501700857389402],
+               [0.2645697738549972, -0.35824829046386303, 0.2483163247594392]]
+
+        o = Octree()
+        o.insert_object_at_coordinates(pts[0], 1)
+        o.insert_object_at_coordinates(pts[1], 2)
+
+        self.assertEqual(1, o.find_object_by_coordinates(pts[0]))
+        self.assertEqual(2, o.find_object_by_coordinates(pts[1]))
+
+    def test_oct_tree_2(self):
+        pts = [[-0.2934372873144786, -0.40824829046386296, -0.08501700857389399],
+               [-0.30539112847838207, -0.47895555285030167, -0.15471008125537822],
+               [-0.24765910625478924, -0.4789555528487444, -0.23635549006530096],
+               [-0.17796723347655352, -0.40824829046386296, -0.24831632475943927],
+               [-0.16601339231265005, -0.3375410280774243, -0.178623252077955],
+               [-0.22374541453624286, -0.33754102807898145, -0.09697784326803233],
+               [0.17796723347655338, -0.408248290463863, 0.2483163247594392],
+               [0.24765910625759371, -0.47895555285030167, 0.2363554900632418],
+               [0.3053911284793884, -0.4789555528487444, 0.15471008125204752],
+               [0.29343728731447843, -0.408248290463863, 0.08501700857389394],
+               [0.22374541453343808, -0.3375410280774243, 0.09697784327009142],
+               [0.16601339231164347, -0.33754102807898156, 0.17862325208128554]]
+
+        o = Octree()
+        for i, pt in enumerate(pts):
+            o.insert_object_at_coordinates(pt, i + 1)
+
+        for i, pt in enumerate(pts):
+            self.assertEqual(i + 1, o.find_object_by_coordinates(pt))
+
+        nd1 = [[-0.04936372890878149, -0.08550048652106605, -0.03490542745605377],
+               [0.027648784957331453, -0.04277813171414004, -0.0914935840944021],
+               [0.07704471667898151, 0.04277813171695018, -0.05656538580855996],
+               [0.04936372890878152, 0.08550048652106605, 0.03490542745605376],
+               [-0.027648784957331436, 0.04277813171413982, 0.09149358409440221],
+               [-0.07704471667898151, -0.042778131716949955, 0.05656538580856002],
+               [0.04936372890878148, -0.08550048652106605, 0.03490542745605377],
+               [0.07704471667850743, -0.04277813171413998, -0.05656538581133083],
+               [0.027648784954560728, 0.04277813171695016, -0.09149358409392554],
+               [-0.04936372890878149, 0.0855004865210661, -0.03490542745605379],
+               [-0.07704471667850739, 0.04277813171413971, 0.05656538581133116],
+               [-0.027648784954560936, -0.04277813171694988, 0.09149358409392555]]
+
+        self.assertIsNone(o.find_object_by_coordinates(nd1[0]))
+        self.assertEqual(11, o.find_nearest_object_by_coordinates([0.22374541453343808, -0.337541028078, 0.09697784327009142]))
+        self.assertEqual(6, o.find_nearest_object_by_coordinates([-0.22374541453343808, -0.337541028078, -0.09697784327009142]))
+
+    def test_oct_tree_insert_existing(self):
+        pts = [[0.5002720342505129, 0.049999999999999996, -0.08501700857389402],
+               [0.2645697738549972, -0.35824829046386303, 0.2483163247594392]]
+
+        o = Octree()
+        o.insert_object_at_coordinates(pts[0], 1)
+        o.insert_object_at_coordinates(pts[1], 2)
+
+        self.assertEqual(1, o.find_object_by_coordinates(pts[0]))
+        self.assertEqual(2, o.find_object_by_coordinates(pts[1]))
+
+        o.insert_object_at_coordinates(pts[1], 3)
+
+        self.assertEqual(3, o.find_object_by_coordinates(pts[1]))
+
+    def test_oct_tree_nearest(self):
+        pts = [[0.5002720342505129, 0.049999999999999996, -0.08501700857389402],
+               [0.2645697738549972, -0.35824829046386303, 0.2483163247594392]]
+        find_pts = [[0.5002720342506, 0.049999999999999996, -0.08501700857389402],
+                    [0.2645697738549972, -0.35824829046386303, 0.24831632475944]]
+
+        o = Octree()
+        o.insert_object_at_coordinates(pts[0], 1)
+        o.insert_object_at_coordinates(pts[1], 2)
+
+        self.assertEqual(1, o.find_nearest_object_by_coordinates(find_pts[0]))
+        self.assertEqual(2, o.find_nearest_object_by_coordinates(find_pts[1]))
+
+
+if __name__ == "__main__":
+    unittest.main()