CallumDoesStuff
CallumDoesStuff

Reputation: 53

Python Quadtree won't insert values

Ive created a simple quad tree class in python that initialises a tree, inserts a few bodies and attempts to plot them, alongside the boundaries it has made. This works for anything in the Quadrant 2 (South West), and its following subdivisions, however, Whenever I try to add bodies4 and bodies5 which should show up in Quadrant 3 (South East), nothing shows up. However, the node for quadrant 3 is created, but the bodies aren't added. Ive used recursive methods to repeatedly insert and subdivide, but that hasn't seemed to help:

import matplotlib.pyplot as plt

class Body:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.isNode = False

class Node:
    def __init__(self, boundx, boundy):
        self.children = [None, None, None, None]
        self.isNode = True
        self.boundx = boundx
        self.boundy = boundy
        self.width = boundx * 2
        self.height = boundy * 2

    def insert(self, body, position):
        self.children[position] = body

class Area:
    def __init__(self, width, height):
        self.width = width
        self.height = height

class QuadTree:
    def __init__(self, bounds):
        self.size = 4
        self.nodes = [None, None, None, None]
        self.bounds = bounds

    def insert(self, body):
        target_quadrant = self.get_quadrant(body.x, body.y, self.bounds.width, self.bounds.height)
        if self.nodes[target_quadrant] is None:
            self.nodes[target_quadrant] = body
        elif self.nodes[target_quadrant].isNode == False:
            bodies = [self.nodes[target_quadrant], body]
            self.nodes[target_quadrant] = self.subdivide(bodies, self.bounds, target_quadrant)
        elif self.nodes[target_quadrant].isNode == True:
            self.recursive_insert(body, self.nodes[target_quadrant])

    def recursive_insert(self, body, node):
        target_quadrant = self.get_quadrant(body.x, body.y, node.boundx, node.boundy)
        if node.children[target_quadrant] is None:
            node.children[target_quadrant] = body
        elif node.children[target_quadrant].isNode == False:
            bodies = [node.children[target_quadrant], body]
            node.children[target_quadrant] = self.subdivide(bodies, node, target_quadrant)
        elif node.children[target_quadrant].isNode == True:
            self.recursive_insert(body, node.children[target_quadrant])

    def get_offset(self, newBound, target_quadrant):
        if target_quadrant == 0:
            offset_x = newBound.width / 2
            offset_y = newBound.height / 2
        elif target_quadrant == 1:
            offset_x = newBound.width / 2
            offset_y = 0
        elif target_quadrant == 2:
            offset_x = 0
            offset_y = 0
        elif target_quadrant == 3:
            offset_x = 0
            offset_y = newBound.height / 2
        return offset_x, offset_y

    def subdivide(self, bodies, newBound, target_quadrant):
        offset_x, offset_y = self.get_offset(newBound, target_quadrant)
        new_Node = Node(newBound.width / 2, newBound.height / 2)
        target_quadrants = []
        for body in bodies:
            target_quadrant = self.get_quadrant(body.x, body.y, newBound.width / 2, newBound.height / 2)
            if target_quadrant == "Out of bounds":
                continue
            if target_quadrant in target_quadrants:
                sub_quad_bodies = [b for b in bodies if self.get_quadrant(b.x, b.y, newBound.width / 2, newBound.height / 2) == target_quadrant]
                sub_bound = Area(newBound.width / 2, newBound.height / 2)
                new_Node.children[target_quadrant] = self.subdivide(sub_quad_bodies, sub_bound, target_quadrant)
            else:
                new_Node.insert(body, target_quadrant)
                target_quadrants.append(target_quadrant)
        return new_Node

    def get_quadrant(self, x, y, width, height):
        if x < 0 or y < 0 or x >= width or y >= height:
            return "Out of bounds"
        elif x < width / 2:
            if y < height / 2:
                return 2
            else:
                return 0
        else:
            if y < height / 2:
                return 3
            else:
                return 1

    def plot_bodies(self):
        plt.figure(figsize=(8, 8))
        plt.xlim(0, self.bounds.width)
        plt.ylim(0, self.bounds.height)
        self.plot_quadrants(self.bounds, self.nodes)
        for node in self.nodes:
            if node is not None:
                self.plot_bodies_recursive(node)
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.title('QuadTree Bodies')
        plt.show()

    def plot_quadrants(self, bounds, nodes):
        if bounds.width == 1 or bounds.height == 1:
            return
        else:
            plt.plot([bounds.width / 2, bounds.width / 2], [0, bounds.height], color='gray', linestyle='--')
            plt.plot([0, bounds.width], [bounds.height / 2, bounds.height / 2], color='gray', linestyle='--')
            for node in nodes:
                if node is not None and isinstance(node, Node):
                    self.plot_quadrants(Area(bounds.width / 2, bounds.height / 2), node.children)

    def plot_bodies_recursive(self, node):
        if node is not None:
            if not node.isNode:
                plt.scatter(node.x, node.y, color='red')
            else:
                for child in node.children:
                    self.plot_bodies_recursive(child)

if __name__ == "__main__":
    bound = Area(128, 128)
    body = Body(1, 1)
    body2 = Body(5, 6)
    body3 = Body(2, 4)
    body4 = Body(80, 12)
    body5 = Body(80, 60)
    tree = QuadTree(bound)
    tree.insert(body)
    tree.insert(body2)
    tree.insert(body3)
    tree.insert(body4)
    tree.insert(body5)
    tree.plot_bodies()

Upvotes: 0

Views: 21

Answers (0)

Related Questions