diff --git a/jobs/src/bht_algorithm.py b/jobs/src/bht_algorithm.py
deleted file mode 100644
index c6630cae4f9ba184cc3789c4b41d35e86dee8ef3..0000000000000000000000000000000000000000
--- a/jobs/src/bht_algorithm.py
+++ /dev/null
@@ -1,3 +0,0 @@
-"""
-This module implements the Barnes Hut Tree Algorithm
-"""
\ No newline at end of file
diff --git a/jobs/src/bht_algorithm_2D.py b/jobs/src/bht_algorithm_2D.py
new file mode 100644
index 0000000000000000000000000000000000000000..b896fa8407a1030501c48f23aa15817a2e7339d8
--- /dev/null
+++ b/jobs/src/bht_algorithm_2D.py
@@ -0,0 +1,321 @@
+"""
+This module implements the Barnes Hut Tree (BHT) Algorithm for 2D data.
+"""
+
+import numpy as np
+
+
+class MainApp:
+
+ def __init__(self):
+ """Initialize the MainApp with a root node."""
+ self.rootNode = TreeNode(x=-1, y=-1, width=2, height=2)
+
+ def BuildTree(self, particles):
+ """Build the Quadtree by inserting particles.
+
+ Args:
+ particles (list): A list of Particle objects to be inserted into the Quadtree.
+ """
+ self.ResetTree() # Empty the tree
+ for particle in particles:
+ self.rootNode.insert(particle)
+
+ def ResetTree(self):
+ """Reset the Quadtree by reinitializing the root node."""
+ self.rootNode = TreeNode(x=-1, y=-1, width=2, height=2)
+
+
+class Particle:
+
+ def __init__(self, x, y, mass):
+ """Initialize a Particle with x and y coordinates and mass."""
+ self.x = x
+ self.y = y
+ self.mass = mass
+ self.vx = 0.0 # Velocity component in x direction
+ self.vy = 0.0 # Velocity component in y direction
+ self.fx = 0.0 # Force component in x direction
+ self.fy = 0.0 # Force component in y direction
+
+
+class TreeNode:
+
+ def __init__(self, x, y, width, height):
+ """Initialize a TreeNode representing a quadrant in the Quadtree.
+
+ Args:
+ x (float): x-coordinate of the node.
+ y (float): y-coordinate of the node.
+ width (float): Width of the node.
+ height (float): Height of the node.
+ """
+ self.x = x
+ self.y = y
+ self.width = width
+ self.height = height
+ self.particle = None # Particle contained in this node
+ self.center_of_mass = np.array([(x + width) / 2, (y + width) / 2])
+ self.total_mass = 0
+ self.children = np.empty(4, dtype=object) # Children nodes (SW, SE, NW, NE)
+
+ def contains(self, particle):
+ """Check if the particle is within the bounds of this node.
+
+ Args:
+ particle (Particle): The particle to be checked.
+
+ Returns:
+ bool: True if the particle is within the node's bounds, False otherwise.
+ """
+ return (self.x <= particle.x < self.x + self.width and self.y <= particle.y < self.y + self.height)
+
+ def insert(self, particle):
+ """Insert a particle into the Quadtree.
+
+ Args:
+ particle (Particle): The particle to be inserted.
+
+ Returns:
+ bool: True if the particle is inserted, False otherwise.
+ """
+ if not self.contains(particle):
+ return False # Particle doesn't belong in this node
+
+ if self.particle is None and all(child is None for child in self.children):
+ # If the node is empty and has no children, insert particle here.
+ self.particle = particle
+ #print(f'particle inserted: x={round(self.particle.x,2)}, y={round(self.particle.y,2)}')
+ return True # Particle inserted in an empty node
+
+ if all(child is None for child in self.children):
+ # If no children exist, create and insert both particles
+ self.subdivide()
+ self.insert(self.particle) # Reinsert existing particle
+ self.insert(particle) # Insert new particle
+ self.particle = None # Clear particle from this node
+ else:
+ # If the node has children, insert particle in the child node.
+ quad_index = self.get_quadrant(particle)
+ if self.children[quad_index] is None:
+ # Create a child node if it doesn't exist
+ self.children[quad_index] = TreeNode(self.x + (quad_index % 2) * (self.width / 2),
+ self.y + (quad_index // 2) * (self.height / 2), self.width / 2,
+ self.height / 2)
+ self.children[quad_index].insert(particle)
+
+ def subdivide(self):
+ """Subdivide the node into four quadrants."""
+ sub_width = self.width / 2
+ sub_height = self.height / 2
+ self.children[0] = TreeNode(self.x, self.y, sub_width, sub_height) # SW
+ self.children[1] = TreeNode(self.x + sub_width, self.y, sub_width, sub_height) # SE
+ self.children[2] = TreeNode(self.x, self.y + sub_height, sub_width, sub_height) # NW
+ self.children[3] = TreeNode(self.x + sub_width, self.y + sub_height, sub_width, sub_height) # NE
+
+ def get_quadrant(self, particle):
+ """Determine the quadrant index for a particle based on its position.
+
+ Args:
+ particle (Particle): The particle to determine the quadrant index for.
+
+ Returns:
+ int: Quadrant index (0 for SW, 1 for SE, 2 for NW, 3 for NE).
+ """
+ mid_x = self.x + self.width / 2
+ mid_y = self.y + self.height / 2
+ quad_index = (particle.x >= mid_x) + 2 * (particle.y >= mid_y)
+ return quad_index
+
+ def print_tree(self, depth=0):
+ """Print the structure of the Quadtree.
+
+ Args:
+ depth (int): Current depth in the tree (for indentation in print).
+ """
+ if self.particle:
+ print(
+ f"{' ' * depth}Particle at ({round(self.particle.x,2)}, {round(self.particle.y,2)}) in Node ({self.x}, {self.y}), size={self.width}"
+ )
+ else:
+ print(f"{' ' * depth}Node ({self.x}, {self.y}) - Width: {self.width}, Height: {self.height}")
+ for child in self.children:
+ if child:
+ child.print_tree(depth + 2)
+
+ def ComputeMassDistribution(self):
+ """Compute the mass distribution for the tree nodes.
+
+ This function calculates the total mass and the center of mass
+ for each node in the Quadtree. It's a recursive function that
+ computes the mass distribution starting from the current node.
+
+ Note:
+ This method modifies the 'mass' and 'center_of_mass' attributes
+ for each node in the Quadtree.
+
+ Returns:
+ None
+ """
+ if self.particle is not None:
+ # Node contains only one particle
+ self.center_of_mass = np.array([self.particle.x, self.particle.y])
+ self.total_mass = self.particle.mass
+ else:
+ # Multiple particles in node
+ total_mass = 0
+ center_of_mass_accumulator = np.array([0.0, 0.0])
+
+ for child in self.children:
+ if child is not None:
+ # Recursively compute mass distribution for child nodes
+ child.ComputeMassDistribution()
+ total_mass += child.total_mass
+ center_of_mass_accumulator += child.total_mass * child.center_of_mass
+
+ if total_mass > 0:
+ self.center_of_mass = center_of_mass_accumulator / total_mass
+ self.total_mass = total_mass
+ else:
+ # If total mass is 0 or no child nodes have mass, leave values as default
+ pass
+ #self.center_of_mass = np.array([(x+width)/2, (y+width)/2])
+ #self.total_mass = 0
+
+ def CalculateForceFromTree(self, target_particle, theta=1.0):
+ """Calculate the force on a target particle using the Barnes-Hut algorithm.
+
+ Args:
+ target_particle (Particle): The particle for which the force is calculated.
+ theta (float): The Barnes-Hut criterion for force approximation.
+
+ Returns:
+ np.ndarray: The total force acting on the target particle.
+ """
+
+ total_force = np.array([0.0, 0.0])
+
+ if self.particle is not None:
+ # Node contains only one particle
+ if self.particle != target_particle:
+ # Calculate gravitational force between target_particle and node's particle
+ force = self.GravitationalForce(target_particle, self.particle)
+ total_force += force
+ else:
+ if self.total_mass == 0:
+ return total_force
+
+ r = np.linalg.norm(np.array([target_particle.x, target_particle.y]) - self.center_of_mass)
+ d = max(self.width, self.height)
+
+ if d / r < theta:
+ # Calculate gravitational force between target_particle and "node particle" representing cluster
+ node_particle = Particle(self.center_of_mass[0], self.center_of_mass[1], self.total_mass)
+ force = self.GravitationalForce(target_particle, node_particle)
+ total_force += force
+ else:
+ for child in self.children:
+ if child is not None:
+ # Recursively calculate force from child nodes
+ if target_particle is not None: # Check if the target_particle is not None
+ force = child.CalculateForceFromTree(target_particle)
+ total_force += force
+
+ return total_force
+
+ def CalculateForce(self, target_particle, particle, theta=1.0):
+ """Calculate the gravitational force between two particles.
+
+ Args:
+ target_particle (Particle): The particle for which the force is calculated.
+ particle (Particle): The particle exerting the force.
+
+ Returns:
+ np.ndarray: The force vector acting on the target particle due to 'particle'.
+ """
+ force = np.array([0.0, 0.0])
+ print('function CalculateForce is called')
+ if self.particle is not None:
+ # Node contains only one particle
+ if self.particle != target_particle:
+ # Calculate gravitational force between target_particle and node's particle
+ force = self.GravitationalForce(target_particle, self.particle)
+ else:
+ if target_particle is not None and particle is not None: # Check if both particles are not None
+ r = np.linalg.norm(
+ np.array([target_particle.x, target_particle.y]) - np.array([particle.x, particle.y]))
+ d = max(self.width, self.height)
+
+ if d / r < theta:
+ # Calculate gravitational force between target_particle and particle
+ force = self.GravitationalForce(target_particle, particle)
+ else:
+ for child in self.children:
+ if child is not None:
+ # Recursively calculate force from child nodes
+ force += child.CalculateForce(target_particle, particle)
+ return force
+
+ def GravitationalForce(self, particle1, particle2):
+ """Calculate the gravitational force between two particles.
+
+ Args:
+ particle1 (Particle): First particle.
+ particle2 (Particle): Second particle.
+
+ Returns:
+ np.ndarray: The gravitational force vector between particle1 and particle2.
+ """
+ #G = 6.674 * (10 ** -11) # Gravitational constant
+ #G = 1
+ G = 4 * np.pi**2 # AU^3 / m / yr^2
+
+ dx = particle2.x - particle1.x
+ dy = particle2.y - particle1.y
+ cutoff_radius = 0
+ r = max(np.sqrt(dx**2 + dy**2), cutoff_radius)
+
+ force_magnitude = G * particle1.mass * particle2.mass / (r**2)
+ force_x = force_magnitude * (dx / r)
+ force_y = force_magnitude * (dy / r)
+
+ return np.array([force_x, force_y])
+
+ # Helper method to retrieve all particles in the subtree
+ def particles_in_subtree(self):
+ """Retrieve all particles in the subtree rooted at this node.
+
+ Returns:
+ list: A list of particles in the subtree rooted at this node.
+ """
+ particles = []
+ if self.particle is not None:
+ particles.append(self.particle)
+ else:
+ for child in self.children:
+ if child is not None:
+ particles.extend(child.particles_in_subtree())
+ return particles
+
+ def compute_center_of_mass(self):
+ """Compute the center of mass for the node."""
+ print('Function compute_center_of_mass is called')
+ if self.particle is not None:
+ self.center_of_mass = np.array([self.particle.x, self.particle.y])
+ self.mass = self.particle.mass
+ else:
+ total_mass = 0
+ center_of_mass_accumulator = np.array([0.0, 0.0])
+
+ for child in self.children:
+ if child is not None:
+ child.compute_center_of_mass()
+ total_mass += child.mass
+ center_of_mass_accumulator += child.mass * child.center_of_mass
+
+ if total_mass > 0:
+ self.center_of_mass = center_of_mass_accumulator / total_mass
+ self.mass = total_mass
+ else:
+ self.center_of_mass = np.array([0.0, 0.0])
+ self.mass = 0
\ No newline at end of file
diff --git a/jobs/src/bht_algorithm_3D.py b/jobs/src/bht_algorithm_3D.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ade62d4c16433743a372b7384e55b928140f61c
--- /dev/null
+++ b/jobs/src/bht_algorithm_3D.py
@@ -0,0 +1,366 @@
+"""
+This module implements the Barnes Hut Tree (BHT) Algorithm for 3D data.
+"""
+
+import numpy as np
+
+
+class MainApp:
+
+ def __init__(self, particles):
+ """Initialize the MainApp with a root node."""
+ self.rootNode = TreeNode(x=0, y=0, z=0, width=200, height=200, depth=200)
+
+ def BuildTree(self, particles):
+ """Build the Octree by inserting particles.
+
+ Args:
+ particles (list): A list of Particle objects to be inserted into the Octree.
+ """
+ self.ResetTree(particles) # Empty the tree
+ for particle in particles:
+ self.rootNode.insert(particle)
+
+ def ResetTree(self, particles):
+ """Reset the Quadtree by reinitializing the root node."""
+ # Define the size of the rootNode based on the positions of the particles
+ #min_x = min([particle.x for particle in particles])
+ #min_y = min([particle.y for particle in particles])
+ #min_z = min([particle.z for particle in particles])
+ #max_x = max([particle.x for particle in particles])
+ #max_y = max([particle.y for particle in particles])
+ #max_z = max([particle.z for particle in particles])
+ #root_height = max_y - min_y
+ #root_width = max_x - min_x
+ #root_depth = max_z - min_z
+ #self.rootNode = TreeNode(x=min_x, y=min_y, z=min_z, width=root_width, height=root_height, depth=root_depth)
+ self.rootNode = TreeNode(x=0, y=0, z=0, width=200, height=200, depth=200)
+
+
+class Particle:
+
+ def __init__(self, x, y, z, mass):
+ """Initialize a Particle with x, y, z coordinates and mass."""
+ self.x = x
+ self.y = y
+ self.z = z
+ self.mass = mass
+
+ self.vx = 0.0 # Velocity component in x direction
+ self.vy = 0.0 # Velocity component in y direction
+ self.vz = 0.0 # Velocity component in z direction
+
+ self.fx = 0.0 # Force component in x direction
+ self.fy = 0.0 # Force component in y direction
+ self.fz = 0.0 # Force component in z direction
+
+
+class TreeNode:
+
+ def __init__(self, x, y, z, width, height, depth):
+ """Initialize a TreeNode representing a octant in the Octree.
+
+ Args:
+ x (float): x-coordinate of the node.
+ y (float): y-coordinate of the node.
+ z (float): z-coordinate of the node.
+ width (float): Width of the node.
+ height (float): Height of the node.
+ depth (float): Depth of the node.
+ """
+ self.x = x
+ self.y = y
+ self.z = z
+ self.width = width
+ self.height = height
+ self.depth = depth
+ self.particle = None # Particle contained in this node
+ self.center_of_mass = np.array([(x + width) / 2, (y + height) / 2, (z + depth) / 2])
+ self.total_mass = 0 # Total mass contained in the node
+ self.children = np.empty(8, dtype=object) # Children nodes (F-SW, F-SE, F-NW, F-NE, B-SW, B-SE, B-NW, B-NE,)
+
+ def contains(self, particle):
+ """Check if the particle is within the bounds of this node.
+
+ Args:
+ particle (Particle): The particle to be checked.
+
+ Returns:
+ bool: True if the particle is within the node's bounds, False otherwise.
+ """
+ return (self.x <= particle.x <= self.x + self.width and self.y <= particle.y <= self.y + self.height
+ and self.z <= particle.z <= self.z + self.depth)
+
+ def insert(self, particle):
+ """Insert a particle into the Octree.
+
+ Args:
+ particle (Particle): The particle to be inserted.
+
+ Returns:
+ bool: True if the particle is inserted, False otherwise.
+ """
+ if not self.contains(particle):
+ return False # Particle doesn't belong in this node
+
+ if self.particle is None and all(child is None for child in self.children):
+ # If the node is empty and has no children, insert particle here.
+ self.particle = particle
+ #print(f'particle inserted: x={round(self.particle.x,2)}, y={round(self.particle.y,2)}, , z={round(self.particle.z,2)}')
+ return True # Particle inserted in an empty node
+
+ if all(child is None for child in self.children):
+ # If no children exist, create and insert both particles
+ self.subdivide()
+ self.insert(self.particle) # Reinsert existing particle into subnode
+ self.insert(particle) # Insert new particle
+ self.particle = None # Clear particle from this node
+ else:
+ # If the node has children, insert particle in the child node.
+ oct_index = self.get_octant(particle)
+ if self.children[oct_index] is None:
+ print('Missing node:', self.children)
+ # Create a child node if it doesn't exist
+ self.children[oct_index] = TreeNode(self.x + (oct_index % 2) * (self.width / 2),
+ self.y + (oct_index // 2) * (self.height / 2),
+ self.z + (oct_index // 2) * (self.depth / 2), self.width / 2,
+ self.height / 2, self.depth / 2)
+ self.children[oct_index].insert(particle)
+
+ def subdivide(self):
+ """Create the children of the node."""
+ sub_width = self.width / 2
+ sub_height = self.height / 2
+ sub_depth = self.depth / 2
+ self.children[0] = TreeNode(self.x, self.y, self.z, sub_width, sub_height, sub_depth) # B-SW
+ self.children[1] = TreeNode(self.x + sub_width, self.y, self.z, sub_width, sub_height, sub_depth) # B-SE
+ self.children[2] = TreeNode(self.x, self.y + sub_height, self.z, sub_width, sub_height, sub_depth) # B-NW
+ self.children[3] = TreeNode(self.x + sub_width, self.y + sub_height, self.z, sub_width, sub_height,
+ sub_depth) # B-NE
+ self.children[4] = TreeNode(self.x, self.y, self.z + sub_depth, sub_width, sub_height, sub_depth) # T-SW
+ self.children[5] = TreeNode(self.x + sub_width, self.y, self.z + sub_depth, sub_width, sub_height,
+ sub_depth) # T-SE
+ self.children[6] = TreeNode(self.x, self.y + sub_height, self.z + sub_depth, sub_width, sub_height,
+ sub_depth) # T-NW
+ self.children[7] = TreeNode(self.x + sub_width, self.y + sub_height, self.z + sub_depth, sub_width, sub_height,
+ sub_depth) # T-NE
+
+ def get_octant(self, particle):
+ """Determine the octant index for a particle based on its position.
+
+ Args:
+ particle (Particle): The particle to determine the octant index for.
+
+ Returns:
+ int: Octant index
+ B - Bottom, T - Top
+ (0 for B-SW, 1 for B-SE, 2 for B-NW, 3 for B-NE,
+ 4 for T-SW, 5 for T-SE, 6 for T-NW, 7 for T-NE).
+ """
+ # Determine separating planes
+ mid_x = self.x + self.width / 2
+ mid_y = self.y + self.height / 2
+ mid_z = self.z + self.depth / 2
+ quad_index = (particle.x >= mid_x) + 2 * (particle.y >= mid_y) + 4 * (particle.z >= mid_z)
+ return quad_index
+
+ def print_tree(self, depth_=0):
+ """Print the structure of the Octree.
+
+ Args:
+ depth_ (int): Current depth in the tree (for indentation in print).
+ """
+ if self.particle:
+ print(
+ f"{' ' * depth_}Particle at ({round(self.particle.x,2)}, {round(self.particle.y,2)}, {round(self.particle.z,2)}) in Node ({round(self.x,2)}, {round(self.y,2)}, {round(self.z,2)}), width={round(self.width,2)}, height={round(self.height,2)}, depth={round(self.depth,2)}"
+ )
+ else:
+ print(
+ f"{' ' * depth_}Node ({round(self.x,2)}, {round(self.y,2)}, {round(self.z,2)}) - Width: {round(self.width,2)}, Height: {round(self.height,2)}, Depth: {round(self.depth,2)}"
+ )
+ for child in self.children:
+ if child:
+ child.print_tree(depth_ + 2)
+
+ def ComputeMassDistribution(self):
+ """Compute the mass distribution for the tree nodes.
+
+ This function calculates the total mass and the center of mass
+ for each node in the Octree. It's a recursive function that
+ computes the mass distribution starting from the current node.
+
+ Note:
+ This method modifies the 'total_mass' and 'center_of_mass' attributes
+ for each node in the Octree.
+
+ Returns:
+ None
+ """
+ if self.particle is not None:
+ # Node contains only one particle
+ self.center_of_mass = np.array([self.particle.x, self.particle.y, self.particle.z])
+ self.total_mass = self.particle.mass
+ else:
+ # Multiple particles in node
+ total_mass = 0
+ center_of_mass_accumulator = np.array([0.0, 0.0, 0.0])
+
+ for child in self.children:
+ if child is not None:
+ # Recursively compute mass distribution for child nodes
+ child.ComputeMassDistribution()
+ total_mass += child.total_mass
+ center_of_mass_accumulator += child.total_mass * child.center_of_mass
+
+ if total_mass > 0:
+ self.center_of_mass = center_of_mass_accumulator / total_mass
+ self.total_mass = total_mass
+ else:
+ # If total mass is 0 or no child nodes have mass, leave values as default
+ pass
+ #self.center_of_mass = np.array([(x+width)/2, (y+height)/2, (z+depth)/2])
+ #self.total_mass = 0
+
+ def CalculateForceFromTree(self, target_particle, theta=1.0):
+ """Calculate the force on a target particle using the Barnes-Hut algorithm.
+
+ Args:
+ target_particle (Particle): The particle for which the force is calculated.
+ theta (float): The Barnes-Hut criterion for force approximation.
+
+ Returns:
+ np.ndarray: The total force acting on the target particle.
+ """
+
+ total_force = np.array([0.0, 0.0, 0.0])
+
+ if self.particle is not None:
+ # Node contains only one particle
+ if self.particle != target_particle:
+ # Calculate gravitational force between target_particle and node's particle
+ force = self.GravitationalForce(target_particle, self.particle)
+ total_force += force
+ else:
+ if self.total_mass == 0:
+ return total_force
+
+ r = np.linalg.norm(
+ np.array([target_particle.x, target_particle.y, target_particle.z]) - self.center_of_mass)
+ d = max(self.width, self.height, self.depth)
+
+ if d / r < theta:
+ # Calculate gravitational force between target_particle and "node particle" representing cluster
+
+ force = self.GravitationalForce(
+ target_particle,
+ Particle(self.center_of_mass[0], self.center_of_mass[1], self.center_of_mass[2], self.total_mass))
+ total_force += force
+ else:
+ for child in self.children:
+ if child is not None:
+ # Recursively calculate force from child nodes
+ if target_particle is not None: # Check if the target_particle is not None
+ force = child.CalculateForceFromTree(target_particle)
+ total_force += force
+
+ return total_force
+
+ def CalculateForce(self, target_particle, particle, theta=1.0):
+ """Calculate the gravitational force between two particles.
+
+ Args:
+ target_particle (Particle): The particle for which the force is calculated.
+ particle (Particle): The particle exerting the force.
+
+ Returns:
+ np.ndarray: The force vector acting on the target particle due to 'particle'.
+ """
+ force = np.array([0.0, 0.0, 0.0])
+ print('function CalculateForce is called')
+ if self.particle is not None:
+ # Node contains only one particle
+ if self.particle != target_particle:
+ # Calculate gravitational force between target_particle and node's particle
+ force = self.GravitationalForce(target_particle, self.particle)
+ else:
+ if target_particle is not None and particle is not None: # Check if both particles are not None
+ r = np.linalg.norm(
+ np.array([target_particle.x, target_particle.y, target_particle.z]) -
+ np.array([particle.x, particle.y, particle.z]))
+ d = max(self.width, self.height, self.depth)
+
+ if d / r < theta:
+ # Calculate gravitational force between target_particle and particle
+ force = self.GravitationalForce(target_particle, particle)
+ else:
+ for child in self.children:
+ if child is not None:
+ # Recursively calculate force from child nodes
+ force += child.CalculateForce(target_particle, particle)
+ return force
+
+ def GravitationalForce(self, particle1, particle2):
+ """Calculate the gravitational force between two particles.
+
+ Args:
+ particle1 (Particle): First particle.
+ particle2 (Particle): Second particle.
+
+ Returns:
+ np.ndarray: The gravitational force vector between particle1 and particle2.
+ """
+ #G = 6.674 * (10 ** -11) # Gravitational constant
+ G = 1
+
+ dx = particle2.x - particle1.x
+ dy = particle2.y - particle1.y
+ dz = particle2.z - particle1.z
+ cutoff_radius = 5
+ r = max(np.sqrt(dx**2 + dy**2 + dz**2), cutoff_radius)
+
+ force_magnitude = G * particle1.mass * particle2.mass / (r**2)
+ force_x = force_magnitude * (dx / r)
+ force_y = force_magnitude * (dy / r)
+ force_z = force_magnitude * (dz / r)
+
+ return np.array([force_x, force_y, force_z])
+
+ # Helper method to retrieve all particles in the subtree
+ def particles_in_subtree(self):
+ """Retrieve all particles in the subtree rooted at this node.
+
+ Returns:
+ list: A list of particles in the subtree rooted at this node.
+ """
+ particles = []
+
+ if self.particle is not None:
+ particles.append(self.particle)
+ else:
+ for child in self.children:
+ if child is not None:
+ particles.extend(child.particles_in_subtree())
+ return len(particles), particles
+
+ def compute_center_of_mass(self):
+ """Compute the center of mass for the node."""
+ print('Function compute_center_of_mass is called')
+ if self.particle is not None:
+ self.center_of_mass = np.array([self.particle.x, self.particle.y, self.particle.z])
+ self.mass = self.particle.mass
+ else:
+ total_mass = 0
+ center_of_mass_accumulator = np.array([0.0, 0.0, 0.0])
+
+ for child in self.children:
+ if child is not None:
+ child.compute_center_of_mass()
+ total_mass += child.mass
+ center_of_mass_accumulator += child.mass * child.center_of_mass
+
+ if total_mass > 0:
+ self.center_of_mass = center_of_mass_accumulator / total_mass
+ self.mass = total_mass
+ else:
+ self.center_of_mass = np.array([0.0, 0.0, 0.0])
+ self.mass = 0
\ No newline at end of file
diff --git a/jobs/src/system.py b/jobs/src/system.py
index 264dcfdf22572b1208902fe034f99542928a374c..3e008f00defd27d1c05c6864574d296f4335cc23 100644
--- a/jobs/src/system.py
+++ b/jobs/src/system.py
@@ -1,6 +1,5 @@
"""
-This module contains base class for a n-body gravitational system
-with 2 methods of simulations: direct and approximated force field.
+This module contains base class for a n-body gravitational system.
"""
from collections.abc import Callable
@@ -23,12 +22,12 @@ class GravitationalSystem:
assert self.r0.shape == self.v0.shape
assert self.m.shape[0] == self.r0.shape[0]
- solvers = ['verlet']
+ solvers = ['verlet', 'bht']
assert self.solver.__name__ in solvers
- def direct_simulation(self):
+ def simulation(self):
"""
- Using integrator to compute trajector in phase space
+ Using integrator to compute trajectory in phase space
"""
p = np.zeros((len(self.t), *self.v0.shape))
q = np.zeros((len(self.t), *self.r0.shape))
@@ -39,19 +38,4 @@ class GravitationalSystem:
dt = self.t[i] - self.t[i - 1]
p[i], q[i] = self.solver(self.force, q[i - 1], p[i - 1], self.m, dt)
- return self.t, p, q
-
- def force_field(self):
- """
- Using force field method to compute trajector in phase space
- """
- p = np.zeros((len(self.t), *self.v0.shape))
- q = np.zeros((len(self.t), *self.r0.shape))
- q[0] = self.r0
- p[0] = self.m * self.v0
- for i in range(1, len(self.t)):
- # something with dt = self.t[i] - self.t[i-1]
- # p = m*v
- pass
-
- return self.t, p, q
+ return self.t, p, q
\ No newline at end of file
diff --git a/jobs/tests/test_bht_algorithm.py b/jobs/tests/test_bht_algorithm.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/jobs/tests/test_bht_algorithm_2D.py b/jobs/tests/test_bht_algorithm_2D.py
new file mode 100644
index 0000000000000000000000000000000000000000..0318849e84fb25bdd49676bf5134d385994cd9d6
--- /dev/null
+++ b/jobs/tests/test_bht_algorithm_2D.py
@@ -0,0 +1,150 @@
+"""
+This unittest tests implementation of the BHT algorithm
+for the restricted three-body (sun-earth-moon) problem.
+The sun is fixed at the origin (center of mass). For
+simplicity, the moon is assumed to be in the eliptic plane,
+so that vectors can be treated as two-dimensional.
+"""
+
+import logging
+import unittest
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from jobs.src.integrators import verlet
+from jobs.src.system import GravitationalSystem
+from jobs.src.bht_algorithm_2D import MainApp, Particle
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# system settings
+R_SE = 1 # distance between Earth and the Sun, 1 astronomical unit (AU)
+R_L = 2.57e-3 * R_SE # distance between Earth and the Moon
+
+M_S = 1 # solar mass
+M_E = 3.00e-6 * M_S
+M_L = 3.69e-8 * M_S
+
+T_E = 1 # earth yr
+T_L = 27.3 / 365.3 * T_E
+
+G = 4 * np.pi**2 # AU^3 / m / yr^2
+
+# simulation settings
+T = 0.3 # yr
+n_order = 6
+dt = 4**(-n_order) # yr
+
+
+# force computation
+def force(q: np.ndarray) -> np.ndarray:
+ r1 = q[0:2] # Earth coordinates
+ r2 = q[2:4] # Moon coordinates
+
+ sun = Particle(0, 0, M_S)
+ earth = Particle(r1[0], r1[1], M_E)
+ moon = Particle(r2[0], r2[1], M_L)
+
+ particles = [sun, earth, moon]
+
+ barnes_hut = MainApp() # Initialize Barnes-Hut algorithm instance
+ barnes_hut.BuildTree(particles) # Build the Barnes-Hut tree with particles
+ barnes_hut.rootNode.ComputeMassDistribution() #Compute the center of mass of the tree nodes
+
+ f_earth = barnes_hut.rootNode.CalculateForceFromTree(earth)
+ f_moon = barnes_hut.rootNode.CalculateForceFromTree(moon)
+
+ return np.concatenate((f_earth, f_moon), axis=0)
+
+
+class IntegratorTest(unittest.TestCase):
+
+ def test_bht(self):
+ """
+ Test functionalities of velocity-Verlet algorithm
+ """
+ # vector of r0 and v0
+ x0 = np.array([
+ R_SE,
+ 0,
+ R_SE + R_L,
+ 0,
+ 0,
+ R_SE * 2 * np.pi / T_E,
+ 0,
+ 1 / M_E * M_E * R_SE * 2 * np.pi / T_E + 1 * R_L * 2 * np.pi / T_L,
+ ])
+
+ system = GravitationalSystem(r0=x0[:4],
+ v0=x0[4:],
+ m=np.array([M_E, M_E, M_L, M_L]),
+ t=np.linspace(0, T, int(T // dt)),
+ force=force,
+ solver=verlet)
+
+ t, p, q = system.simulation()
+
+ ## checking total energy conservation
+ H = np.linalg.norm(p[:,:2], axis=1)**2 / (2 * M_E) + np.linalg.norm(p[:,2:], axis=1)**2 / (2 * M_L) + \
+ -G * M_S * M_E / np.linalg.norm(q[:,:2], axis=1) - G * M_S * M_L / np.linalg.norm(q[:,2:], axis=1) + \
+ -G * M_E * M_L / np.linalg.norm(q[:,2:] - q[:,:2], axis=1)
+
+ self.assertTrue(np.greater(1e-7 + np.zeros(H.shape[0]), H - H[0]).all())
+
+ ## checking total linear momentum conservation
+ P = p[:, :2] + p[:, 2:]
+ self.assertTrue(np.greater(1e-10 + np.zeros(P[0].shape), P - P[0]).all())
+
+ ## checking total angular momentum conservation
+ L = np.cross(q[:, :2], p[:, :2]) + np.cross(q[:, 2:], p[:, 2:])
+ self.assertTrue(np.greater(1e-8 + np.zeros(L.shape[0]), L - L[0]).all())
+
+ ## checking error
+ dts = [dt, 2 * dt, 4 * dt]
+ errors = []
+ ts = []
+ for i in dts:
+ system = GravitationalSystem(r0=x0[:4],
+ v0=x0[4:],
+ m=np.array([M_E, M_E, M_L, M_L]),
+ t=np.linspace(0, T, int(T // i)),
+ force=force,
+ solver=verlet)
+
+ t, p_t, q_t = system.simulation()
+
+
+ H = np.linalg.norm(p_t[:,:2], axis=1)**2 / (2 * M_E) + np.linalg.norm(p_t[:,2:], axis=1)**2 / (2 * M_L) + \
+ -G * M_S * M_E / np.linalg.norm(q_t[:,:2], axis=1) - G * M_S * M_L / np.linalg.norm(q_t[:,2:], axis=1) + \
+ -G * M_E * M_L / np.linalg.norm(q_t[:,2:] - q_t[:,:2], axis=1)
+
+ errors.append((H - H[0]) / i**2)
+ ts.append(t)
+
+ plt.figure()
+ plt.plot(ts[0], errors[0], label="dt")
+ plt.plot(ts[1], errors[1], linestyle='--', label="2*dt")
+ plt.plot(ts[2], errors[2], linestyle=':', label="4*dt")
+ plt.xlabel("$t$")
+ plt.ylabel("$\delta E(t)/(\Delta t)^2$")
+ plt.legend()
+ plt.show()
+
+ ## checking time reversal: p -> -p
+ x0 = np.concatenate((q[-1, :], -1 * p[-1, :] / np.array([M_E, M_E, M_L, M_L])), axis=0)
+ system = GravitationalSystem(r0=x0[:4],
+ v0=x0[4:],
+ m=np.array([M_E, M_E, M_L, M_L]),
+ t=np.linspace(0, T, int(T // dt)),
+ force=force,
+ solver=verlet)
+
+ t, p_reverse, q_reverse = system.simulation()
+ self.assertTrue(np.greater(1e-10 + np.zeros(4), q_reverse[-1] - q[0]).all())
+ self.assertTrue(np.greater(1e-10 + np.zeros(4), p_reverse[-1] - p[0]).all())
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/jobs/tests/test_integrator.py b/jobs/tests/test_integrator.py
index f772a1376a1a0b3cfed84b736ba14eee4fbc5715..fbd7008b4b6fc50f031150cd410577a64e2cdaa8 100644
--- a/jobs/tests/test_integrator.py
+++ b/jobs/tests/test_integrator.py
@@ -73,7 +73,7 @@ class IntegratorTest(unittest.TestCase):
force=force,
solver=verlet)
- t, p, q = system.direct_simulation()
+ t, p, q = system.simulation()
## checking total energy conservation
H = np.linalg.norm(p[:,:2], axis=1)**2 / (2 * M_E) + np.linalg.norm(p[:,2:], axis=1)**2 / (2 * M_L) + \
@@ -102,7 +102,7 @@ class IntegratorTest(unittest.TestCase):
force=force,
solver=verlet)
- t, p_t, q_t = system.direct_simulation()
+ t, p_t, q_t = system.simulation()
H = np.linalg.norm(p_t[:,:2], axis=1)**2 / (2 * M_E) + np.linalg.norm(p_t[:,2:], axis=1)**2 / (2 * M_L) + \
@@ -130,7 +130,7 @@ class IntegratorTest(unittest.TestCase):
force=force,
solver=verlet)
- t, p_reverse, q_reverse = system.direct_simulation()
+ t, p_reverse, q_reverse = system.simulation()
self.assertTrue(np.greater(1e-10 + np.zeros(4), q_reverse[-1] - q[0]).all())
self.assertTrue(np.greater(1e-10 + np.zeros(4), p_reverse[-1] - p[0]).all())