wobbl_tools/math/spvo.py

#!/usr/bin/python3

"""
SPVO = Sparse Voxel Octree

A Sparse Voxel Octree is a way of storing multidimensional data. (In this case it is 3D Data.)
It works similarly to a Sparse Voxel Matrix, but it is a bit compressed.
In a Sparse Voxel Octree, you have a root-node. The root node gets split up into 8 child nodes,
which can be all split up into 8 more child nodes which can be split up too and so on.
Thanks to this "compression", there are no zero values on places where there is just air unlike as in a normal matrix.
"""


class Node:
    def __init__(self, data=None):
        self.children = [None] * 8
        self.leaf = True

        self.data = data


class SparseVoxelOctree:
    def __init__(self):
        self.root_node = Node()
        self.origin = (0, 0, 0)

    def _get_child_index(self, x, y, z, size):
        """
        Get the index of a child.
        Index: The position of the node in the node parent's children list.
        """

        index = 0

        if x >= size // 2:  # x pos
            # there are only 2 possible x coordinate values in a 2x2 cube, so we use the 1st bit in index for x
            # xyz: the coordinates of the parent node
            index |= 1  # set 1st bit to 1
            x -= size // 2  # get parent x

        # and the same for y and z

        if y >= size // 2:  # y pos
            index |= 2
            y -= size // 2

        if z >= size // 2:  # z pos
            index |= 4
            z -= size // 2

        return index, x, y, z

    def _extend_tree(self, x, y, z):
        new_root = Node()
        new_root.leaf = False

        index, nx, ny, nz = self._get_child_index(x, y, z, 2)
        new_root.children[index] = self.root_node
        self.root_node = new_root

        self.origin = (
            self.origin[0] - (x >= 0) * 2 + 1,
            self.origin[1] - (y >= 0) * 2 + 1,
            self.origin[2] - (z >= 0) * 2 + 1
        )

    def _insert(self, node, x, y, z, size):
        while not size == 1:
            index, x, y, z = self._get_child_index(x, y, z, size)

            if node.children[index] is None:
                node.children[index] = Node()

            size //= 2

        node.leaf = True

    def insert(self, x, y, z):
        # move origin if necessary
        while not (0 <= x < 2 ** 30 and 0 <= y < 2 ** 30 and 0 <= z < 2 ** 30):
            self._extend_tree(x, y, z)

            x, y, z = (
                x - self.origin[0],
                y - self.origin[1],
                z - self.origin[2]
            )

        self._insert(self.root_node, x, y, z, 2 ** 30)