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)