I am trying to create a simulation of particles using the Barnes Hut Algorithm. I am representing the particles as a structured array where rows represent particles and the columns represent attributes of the particles. I have implemented the basics of the quadtree structure and the force calculations. I am having a hang up on the updating of the quadtree. I would rather not redraw the whole quadtree each iteration and I have a basic idea of what is entailed with updating the structure. I read a paper that discussed a way of implementing this that was efficient when compared to other algorithms. https://doi.org/10.1016/j.ins.2012.06.028
import numpy as np
from particledata import ParticleData
dt = 1/60
class QuadtreeNode:
"""
A class that represents a node in a quadtree.
Parameters:
parent (QuadtreeNode): The parent node of the current node.
rect (tuple): A tuple representing the boundaries of the node's rectangle in the form (x0, x1, y0, y1).
particle_data (ParticleData): The class containing the array which houses the data for the particles.
Attributes:
parent (QuadtreeNode): The parent node of the current node.
x0 (float): The minimum x-coordinate of the node's rectangle.
x1 (float): The maximum x-coordinate of the node's rectangle.
y0 (float): The minimum y-coordinate of the node's rectangle.
y1 (float): The maximum y-coordinate of the node's rectangle.
particle_indices (dict): A dictionary of indices of the particles mapped to the node they are contained in.
children (dict): A dictionary containing references to the four child nodes of the current node (nw, ne, sw, se).
com (numpy.ndarray): The center of mass of the node.
total_mass (float): The total mass of the particles contained in the node.
"""
particle_indices = {}
def __init__(self, parent, rect: tuple, particle_data:ParticleData) -> None:
self.parent = parent
self.x0,self.x1,self.y0,self.y1 = rect
self.children = {"nw": None, "ne": None, "sw": None, "se": None}
self.particle_data = particle_data
self.particle = None
self.com = np.zeros(2)
self.total_mass = 0.0
self.theta = 0.5
def build_quadtree(self) -> None:
"""
Builds the quadtree structure by initializing the root node and recursively subdiving it.
Returns:
None
"""
self.particle_indices = {i: self for i in range(self.particle_data.num_particles)}
self._recursive_subdivision()
def update(self) -> None:
"""
Update the quadtree by removing particles that have moved out of the node and reinserting them.
Returns:
None
"""
self.particle_data = self.calculate_forces()
self.particle_data.integrate(dt)
for child_node in self.children.values():
if child_node is not None:
child_node.update()
def _get_root_node(self):
if self.parent is None:
return self
else:
return self.parent._get_root_node()
def _recursive_subdivision(self) -> None:
"""
Recursively subdivides the current node into four child nodes and assigns the particles to the appropriate child nodes. Is meant to be called on the root node of a quadtree structure.
Returns:
None
"""
if len(self.particle_indices) > 1:
x_mid = (self.x1 + self.x0) / 2
y_mid = (self.y1 + self.y0) / 2
self.children['nw'] = QuadtreeNode(self, (self.x0, x_mid, y_mid, self.y1))
self.children['ne'] = QuadtreeNode(self, (x_mid, self.x1, y_mid, self.y1))
self.children['sw'] = QuadtreeNode(self, (self.x0, x_mid, self.y0, y_mid))
self.children['se'] = QuadtreeNode(self, (x_mid, self.x1, self.y0, y_mid))
for particle_index in self.particle_indices:
particle = self.particle_data.get_particle(particle_index)
x, y = particle['position']
child = self._contains(x,y)
child.particle_indices[particle_index] = child
self.com += particle['position'] * particle['mass']
self.total_mass += particle['mass']
if self.total_mass > 0:
self.com /= self.total_mass
for child in self.children.values():
child._recursive_subdivision()
else:
particle_index=next(iter(self.particle_indices))
self.particle = self.particle_data.get_particle(particle_index)
def _contains(self, x: int, y: int) -> 'QuadtreeNode':
"""
Returns the child node that contains the given coordinates (x, y).
Parameters:
x : The x-coordinate of the point.
y : The y-coordinate of the point.
Returns:
QuadtreeNode: The child node that contains the given coordinates (x, y).
"""
root_node = self._get_root_node()
root_x0, root_x1, root_y0, root_y1 = root_node.x0, root_node.x1, root_node.y0, root_node.y1
if x < root_x0 or x > root_x1 or y < root_y0 or y > root_y1:
return None
x_mid = (self.x0 + self.x1) / 2
y_mid = (self.y0 + self.y1) / 2
if x <= x_mid:
if y >= y_mid:
return self.children['nw']
else:
return self.children['sw']
else:
if y >= y_mid:
return self.children['ne']
else:
return self.children['se']
def _compute_theta(self, particle_index) -> float:
"""
Calculate theta which is the ratio node_size / distance from particle to node center of mass.
Returns:
float: The value theta which determines the accuracy of the computation.
"""
node_width = self.x1 - self.x0
node_height = self.y1 - self.y0
node_size = max(node_width, node_height)
position_of_particle = self.particle_data.get_particle(particle_index)['position']
dx = self.com[0] - position_of_particle[0]
dy = self.com[1] - position_of_particle[1]
distance = np.sqrt(dx**2 + dy**2)
theta = node_size / distance
return theta
def remove_particle_data_from_parents(self, particle_index: int) -> None:
particle = self.particle_data.get_particle(particle_index)
if self.parent is None:
return
if self is None:
return
else:
self.parent.com -= ()
def _recursively_reinsert_particles(self, reinsert_particles: list[int]) -> None:
"""
Reinserts the removed particles into the correct nodes based on their new positions.
Parameters:
reinsert_particles (List[int]): A list of particle indices that need to be reinserted.
Returns:
None
"""
for particle_index in reinsert_particles:
particle = self.particle_data.get_particle(particle_index)
x, y = particle['position']
node = self
while True:
if all(child is None for child in node.children.values()):
break
child = node._contains(x, y)
if child is None:
break
node = child
node.particle_indices.append(particle_index)
def _update_properties(self) -> None:
"""
Update the properties of the current node based on the new particle positions.
Returns:
None
"""
self.com = np.zeros(2, dtype=float)
self.total_mass = 0
for particle_index in self.particle_indices:
particle = self.particle_data.get_particle(particle_index)
self.com += particle['position']
self.total_mass += 1
if self.total_mass > 0:
self.com /= self.total_mass
for child_node in self.children.values():
if child_node is not None:
child_node._update_properties(self.particle_data)
def calculate_forces(self):
"""
Calculate the forces acting on particles using the Barnes-Hut algorithm and the quadtree structure.
Returns:
ParticleData: The updated particle data with calculated forces.
"""
particle_data = self.particle_data
for particle_index in range(particle_data.num_particles):
particle = particle_data.get_particle(particle_index)
x, y = particle['position']
force = np.zeros(2, dtype=np.float32)
self._calculate_force(particle_index, x, y, force)
particle['force'] = force
return particle_data
def _calculate_force(self, particle_index: int, x: float, y:float, force:np.ndarray):
"""
Calculate the net force acting on a particle recursively by traversing the quadtree.
Parameters:
particle_index (int): The index of the particle for which to calculate the force.
x (float): The x-coordinate of the particle.
y (float): The y-coordinate of the particle.
force (numpy.ndarray): The array to store the calculated net force.
Returns:
None
"""
if self is None:
return
if self.particle_indices and len(self.particle_indices) == 1 and self.particle_indices[0] == particle_index:
return
if self._compute_theta(particle_index) < 0.5:
particle = self.particle_data.get_particle(particle_index)
dx = self.com[0] - x
dy = self.com[1] - y
r_squared = dx ** 2 + dy ** 2
force_magnitude = particle['mass'] / r_squared
force[0] += force_magnitude * dx
force[1] += force_magnitude * dy
else:
for child_node in self.children.values():
if child_node is not None:
child_node._calculate_force(particle_index, x, y, force)
def resolve_movement(self):
for particle_index, node in self.particle_indices.items():
particle = self.particle_data.get_particle(particle_index)
x,y = particle['position']
This is my implementation of the quadtree data structure. Below is the class that creates and handles the particledata.
import numpy as np
class ParticleData:
"""
Class representing particle data.
Parameters
----------
num_particles : int
The number of particles.
dtype : float32
The datatype for the values in the position, velocity, and force arrays.
Attributes
----------
num_particles : int
The number of particles.
particles : ndarray
Array containing particle data.
position : ndarray
Array representing particle positions.
velocity : ndarray
Array representing particle velocities.
force : ndarray
Array representing particle forces.
"""
def __init__(self, num_particles: int, dtype=np.float32) -> None:
self.num_particles = num_particles
dtype = [
('position', dtype, 2),
('velocity', dtype, 2),
('force', dtype, 2),
('mass', dtype, 1)
]
self.particles = np.zeros(num_particles, dtype)
self.position = self.particles['position']
self.velocity = self.particles['velocity']
self.force = self.particles['force']
self.mass = self.particle['mass']
def initialize_particles(self, min_position:float, max_position:float, min_velocity:float, max_velocity:float) -> None:
"""
Initialize particle positions and velocities with random values.
Parameters
----------
min_position : float
The minimum value for particle positions.
max_position : float
The maximum value for particle positions.
min_velocity : float
The minimum value for particle velocities.
max_velocity : float
The maximum value for particle velocities.
The bounds for these parameters is defined by the border height and width for animation, although it might make more sense to create a different distribution of particles then randomly across the whole screen.
"""
self.position[:, 0] = np.random.uniform(min_position, max_position + 1, size = self.num_particles)
self.position[:, 1] = np.random.uniform(min_position, max_position + 1, size = self.num_particles)
self.velocity[:, 0] = np.random.uniform(min_velocity, max_velocity, size=self.num_particles)
self.velocity[:, 1] = np.random.uniform(min_velocity, max_velocity, size=self.num_particles)
def get_particle(self, index:int) -> np.ndarray:
return self.particles[index]
def remove(self, remove_particles:list[int]) -> None:
self.particles = np.delete(self.particles, remove_particles)
self.num_particles = len(self.particles)
self.position=self.particles['position']
self.velocity=self.particles['velocity']
self.force=self.particles['force']
def integrate(self,dt) -> None:
"""
Perform integration to update the particle positions and velocities.
Parameters:
dt (float): The time step for the integration.
Returns:
None
"""
self.velocity += self.force
self.position += self.velocity * dt
This is my current implementation of the quadtree. It is very rough, and is my first attempt a larger more complex problem. I have conceived two possible ways to go about updating positions and corresponding data. The first assumes that all particles are moving in the quadtree. It uses the dictionary of leaves to check whether the new particle position is still contained in the current leaf, if it is not it traverses upwards until it is, then downards until it reaches a leaf node. Then there would have to be some form of merging and splitting due to overcrowded and empty nodes. The second is to insert each particle at the root node and traverse to the children where it is contained until a leaf node is reached or further subdivision is required. My question is how can I compare the efficiency of different approaches? I believe 1 approach might be more efficient than the other in certain circumstances, how might I determine those? I would also need to update the center of mass and total mass for the node, or it might be better to calculate those on the fly as the force is being calculated. Again, I have no idea how to maximize efficiency with this problem. Any advice on the code I have provided is welcomed, I am fairly new to programming and I have not set into stone habits to live by, so I am very welcome to advice on that end as well. Are there better resources that I have yet to find for doing these sorts of things? Due to Python being my first programming language there is a certain level of abstraction from the basics of computer science that I have, please keep this in mind. Thanks for taking the time to read my question.