def update_force(self, b: Body): """ Start at the main node of the tree. Then, recursively go each branch Until either we reach an leaf node or we reach a node that is sufficiently far away that the external nodes would not matter much. """ if self.is_leaf(): b.accelerate(self.body) elif self.body != b and self.quadrant.diameter / self.body.distance_to( b) < 2: b.accelerate(self.body) else: for quadrant in [self.nw, self.ne, self.sw, self.se]: if quadrant is not None: quadrant.update_force(b)
def insert(self, b: Body): """ We have to populate the tree with bodies. We start at the current tree and recursively travel through the branches """ # If there's not a b there already, put the b there. if not self.body: self.body = b # If there's already a body there, but it's not an external node # combine the two bodies and figure out which quadrant of the # tree it should be located in. Then recursively update the nodes below elif not self.is_leaf(): # In this case, this means "add the acting forces together b.accelerate(self.body) northwest = self.quadrant.NW() if b.is_in(northwest): if self.nw is None: self.nw = BHTree(northwest) self.nw.insert(b) return northeast = self.quadrant.NE() if b.is_in(northeast): if self.ne is None: self.ne = BHTree(northeast) self.ne.insert(b) return southeast = self.quadrant.SE() if b.is_in(southeast): if self.se is None: self.se = BHTree(northeast) self.se.insert(b) return southwest = self.quadrant.SW() if b.is_in(southwest): if self.sw is None: self.sw = BHTree(southwest) self.sw.insert(b) return # If the node is a leaf and contains another body, create BHTrees # where the bodies should go, update the node, and end # (do not do anything recursively) else: c = self.body northwest = self.quadrant.NW() if c.is_in(northwest): if self.nw is None: self.nw = BHTree(northwest) self.insert(c) northeast = self.quadrant.NE() if c.is_in(northeast): if self.ne is None: self.ne = BHTree(northeast) self.insert(c) southwest = self.quadrant.SW() if c.is_in(southwest): if self.sw is None: self.sw = BHTree(southwest) self.insert(c) southeast = self.quadrant.SE() if c.is_in(southeast): if self.se is None: self.se = BHTree(southeast) self.insert(c)