diff --git a/tasks/bht/2D/bht.py b/tasks/bht/2D/bht.py
new file mode 100644
index 0000000000000000000000000000000000000000..30f345a2744dec06dd289082f5cb4f0e21e3c692
--- /dev/null
+++ b/tasks/bht/2D/bht.py
@@ -0,0 +1,316 @@
+import numpy as np
+
+
+class MainApp:
+
+ def __init__(self):
+ """Initialize the MainApp with a root node."""
+ self.rootNode = TreeNode(x=0, y=0, width=100, height=100)
+
+ 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=0, y=0, width=100, height=100)
+
+
+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
+
+ dx = particle2.x - particle1.x
+ dy = particle2.y - particle1.y
+ cutoff_radius = 5
+ 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
diff --git a/tasks/bht/2D/galaxysimulation.py b/tasks/bht/2D/galaxysimulation.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8aae983d9ba230a7e8addf3c65fe797b492af95
--- /dev/null
+++ b/tasks/bht/2D/galaxysimulation.py
@@ -0,0 +1,147 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from bht import MainApp, Particle
+
+
+# Simulation class for galaxy collision using Barnes-Hut and Euler method
+class GalaxyCollisionSimulation:
+
+ def __init__(self, num_particles):
+ # Initialize the simulation with a given number of particles
+ self.num_particles = num_particles
+ self.particles = self.initialize_particles(initial='random') # Generate particles
+ self.barnes_hut = MainApp() # Initialize Barnes-Hut algorithm instance
+ self.barnes_hut.BuildTree(self.particles) # Build the Barnes-Hut tree with particles
+ self.barnes_hut.rootNode.ComputeMassDistribution() #Compute the center of mass of the tree nodes
+ self.particle_positions = {
+ i: [(particle.x, particle.y)]
+ for i, particle in enumerate(self.particles)
+ } # Store particle positions for each time step
+ self.particle_velocities = {
+ i: [(particle.vx, particle.vy)]
+ for i, particle in enumerate(self.particles)
+ } # Store particle velocities for each time step
+ self.particle_forces = {
+ i: [(particle.fx, particle.fy)]
+ for i, particle in enumerate(self.particles)
+ } # Store particle velocities for each time step
+ self.system_center_of_mass = {
+ 0: (self.barnes_hut.rootNode.center_of_mass[0], self.barnes_hut.rootNode.center_of_mass[1])
+ }
+
+ def initialize_particles(self, initial='random'):
+
+ if initial == 'two':
+ particles = [Particle(80, 50, 1000), Particle(20, 50, 1000)]
+ #Particle(0.5, 0, 10000),
+ #Particle(-0.5, 0, 10000)]
+ particles[0].vx, particles[0].vy = 0, 3
+ particles[1].vx, particles[1].vy = 0, -3
+ #particles=particles[:2]
+ return particles
+
+ elif initial == 'random':
+ # Generate random particles within a specified range
+ particles = []
+ for _ in range(self.num_particles):
+ x = np.random.uniform(0, 100)
+ y = np.random.uniform(0, 100)
+ #vx = 0
+ #vy = 0
+ #mass = np.random.uniform(10, 100)
+ mass = 10000
+ particles.append(Particle(x, y, mass))
+ return particles
+
+ def simulate(self, num_steps, time_step):
+ # Simulate the galaxy collision for a certain number of steps
+ for step in range(num_steps):
+ # For each step, build the tree
+ self.barnes_hut.BuildTree(self.particles)
+ self.barnes_hut.rootNode.ComputeMassDistribution()
+ if step in np.arange(0, num_steps, 1000):
+ print(f'===Quadtree at step {step}===')
+ self.barnes_hut.rootNode.print_tree()
+ self.evaluate_forces() # Evaluate forces acting on each particle
+ self.integrate(time_step, 'velocity_verlet') # Integrate to update particle positions
+
+ self.system_center_of_mass[step] = (self.barnes_hut.rootNode.center_of_mass[0],
+ self.barnes_hut.rootNode.center_of_mass[1])
+ # Store particle positions, velocities and forces for each time step
+ for i, particle in enumerate(self.particles):
+ self.particle_positions[i].append((particle.x, particle.y))
+ self.particle_velocities[i].append((particle.vx, particle.vy))
+ self.particle_forces[i].append((particle.fx, particle.fy))
+
+ def evaluate_forces(self):
+ # Evaluate forces acting on each particle using Barnes-Hut algorithm
+ for particle in self.particles:
+ force = self.barnes_hut.rootNode.CalculateForceFromTree(particle)
+ particle.fx, particle.fy = force
+
+ def calculate_force(self, particle):
+ # Calculate force acting on a single particle
+ force = self.barnes_hut.rootNode.CalculateForceFromTree(particle)
+ particle.fx, particle.fy = force
+
+ def integrate(self, time_step, method):
+ # Integrate to update particle positions based on calculated forces
+ if method == 'explicit_euler':
+ for particle in self.particles:
+ particle.vx += particle.fx * time_step / particle.mass # Update velocity components
+ particle.vy += particle.fy * time_step / particle.mass
+ particle.x += particle.vx * time_step # Update particle positions using velocity
+ particle.y += particle.vy * time_step
+
+ elif method == 'velocity_verlet':
+ for particle in self.particles:
+ particle.vx += particle.fx * time_step / (2 * particle.mass) # First half-step
+ particle.vy += particle.fy * time_step / (2 * particle.mass)
+ particle.x += particle.vx * time_step # Full step for position
+ particle.y += particle.vy * time_step
+
+ self.calculate_force(particle) # Recalculate force for the particle
+
+ particle.vx += particle.fx * time_step / (2 * particle.mass) # Second half-step
+ particle.vy += particle.fy * time_step / (2 * particle.mass)
+
+ def print_particle_positions(self):
+ # Print the positions of all particles for each time step
+ for i, positions in self.particle_positions.items():
+ print(f"Particle {i + 1} Positions:")
+ for step, pos in enumerate(positions):
+ print(f" Time Step {step}: Position ({pos[0]}, {pos[1]})")
+
+ def display_snapshots(self, num_shots, fix_axes=True):
+ # Display pictures for each time step
+ num_time_steps = len(next(iter(self.particle_positions.values())))
+ print('===Quadtree at the end===')
+ self.barnes_hut.rootNode.print_tree() # print the tree
+ axes_type = "fixed" if fix_axes else "unfixed"
+ #Fixed axis limits
+ print(f"For {axes_type} axis limits, num_particles={self.num_particles}")
+ for step in np.arange(0, num_time_steps, int(num_time_steps / num_shots)):
+ fig, ax = plt.subplots()
+ if fix_axes:
+ ax.set_xlim(0, 100)
+ ax.set_ylim(0, 100)
+ ax.set_xlabel('X-axis')
+ ax.set_ylabel('Y-axis')
+ ax.set_title(f'Time Step {step}')
+ ax.grid()
+
+ print(f'Step {step}')
+ for particle_index, positions in self.particle_positions.items():
+ x, y = positions[step]
+ vx, vy = self.particle_velocities[particle_index][step]
+ fx, fy = self.particle_forces[particle_index][step]
+ #mass = self.particles[particle_index].mass
+ ax.scatter(x, y, c='blue', s=20, alpha=0.5) # Plot particle position for the current time step
+ c_x, c_y = self.system_center_of_mass[step]
+ ax.scatter(c_x, c_y, c='orange', s=20)
+ print(
+ f'i={particle_index}: x={round(x,2)}, y={round(y,2)}, vx={round(vx,2)}, vy={round(vy,2)}, fx={round(fx,2)}, fy={round(fy,2)}'
+ )
+ #print(y)
+
+ plt.show()
diff --git a/tasks/bht/2D/simulate.py b/tasks/bht/2D/simulate.py
new file mode 100644
index 0000000000000000000000000000000000000000..c703e4231fcaa037e6bbe49f5f9195b7020e1509
--- /dev/null
+++ b/tasks/bht/2D/simulate.py
@@ -0,0 +1,5 @@
+from galaxysimulation import GalaxyCollisionSimulation
+
+sim = GalaxyCollisionSimulation(num_particles=6)
+sim.simulate(num_steps=7000, time_step=0.001)
+sim.display_snapshots(200, fix_axes=True)
diff --git a/tasks/bht/__pycache__/bht_oct.cpython-311.pyc b/tasks/bht/__pycache__/bht_oct.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..040eedbf568386874364ac28f4d803d8393a67f5
Binary files /dev/null and b/tasks/bht/__pycache__/bht_oct.cpython-311.pyc differ
diff --git a/tasks/bht/__pycache__/galaxysimulation.cpython-311.pyc b/tasks/bht/__pycache__/galaxysimulation.cpython-311.pyc
index 217cfe8eeb559192f49de5628658625e9ff730ff..1a9755185c2906f2b2ed6f525202579068b0e3dc 100644
Binary files a/tasks/bht/__pycache__/galaxysimulation.cpython-311.pyc and b/tasks/bht/__pycache__/galaxysimulation.cpython-311.pyc differ
diff --git a/tasks/bht/bht_oct.py b/tasks/bht/bht_oct.py
new file mode 100644
index 0000000000000000000000000000000000000000..c55b43fc09d0607b4635256a619434c10a16625d
--- /dev/null
+++ b/tasks/bht/bht_oct.py
@@ -0,0 +1,362 @@
+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
diff --git a/tasks/bht/galaxysimulation.py b/tasks/bht/galaxysimulation.py
index f8aae983d9ba230a7e8addf3c65fe797b492af95..03d532fed339454cf6209117f15564d350e99f92 100644
--- a/tasks/bht/galaxysimulation.py
+++ b/tasks/bht/galaxysimulation.py
@@ -1,6 +1,6 @@
import matplotlib.pyplot as plt
import numpy as np
-from bht import MainApp, Particle
+from bht_oct import MainApp, Particle
# Simulation class for galaxy collision using Barnes-Hut and Euler method
@@ -10,33 +10,32 @@ class GalaxyCollisionSimulation:
# Initialize the simulation with a given number of particles
self.num_particles = num_particles
self.particles = self.initialize_particles(initial='random') # Generate particles
- self.barnes_hut = MainApp() # Initialize Barnes-Hut algorithm instance
+ self.barnes_hut = MainApp(self.particles) # Initialize Barnes-Hut algorithm instance
self.barnes_hut.BuildTree(self.particles) # Build the Barnes-Hut tree with particles
self.barnes_hut.rootNode.ComputeMassDistribution() #Compute the center of mass of the tree nodes
self.particle_positions = {
- i: [(particle.x, particle.y)]
+ i: [(particle.x, particle.y, particle.z)]
for i, particle in enumerate(self.particles)
} # Store particle positions for each time step
self.particle_velocities = {
- i: [(particle.vx, particle.vy)]
+ i: [(particle.vx, particle.vy, particle.vz)]
for i, particle in enumerate(self.particles)
} # Store particle velocities for each time step
self.particle_forces = {
- i: [(particle.fx, particle.fy)]
+ i: [(particle.fx, particle.fy, particle.fz)]
for i, particle in enumerate(self.particles)
} # Store particle velocities for each time step
self.system_center_of_mass = {
- 0: (self.barnes_hut.rootNode.center_of_mass[0], self.barnes_hut.rootNode.center_of_mass[1])
+ 0: (self.barnes_hut.rootNode.center_of_mass[0], self.barnes_hut.rootNode.center_of_mass[1],
+ self.barnes_hut.rootNode.center_of_mass[2])
}
def initialize_particles(self, initial='random'):
if initial == 'two':
- particles = [Particle(80, 50, 1000), Particle(20, 50, 1000)]
- #Particle(0.5, 0, 10000),
- #Particle(-0.5, 0, 10000)]
- particles[0].vx, particles[0].vy = 0, 3
- particles[1].vx, particles[1].vy = 0, -3
+ particles = [Particle(60, 50, 40, 1000), Particle(40, 50, 60, 1000)]
+ particles[0].vx, particles[0].vy, particles[0].vz = 5, 3, 3
+ particles[1].vx, particles[1].vy, particles[1].vz = 5, -3, 3
#particles=particles[:2]
return particles
@@ -44,13 +43,34 @@ class GalaxyCollisionSimulation:
# Generate random particles within a specified range
particles = []
for _ in range(self.num_particles):
- x = np.random.uniform(0, 100)
- y = np.random.uniform(0, 100)
- #vx = 0
- #vy = 0
+ x = np.random.uniform(50, 150)
+ y = np.random.uniform(50, 150)
+ z = np.random.uniform(50, 150)
+ vx = np.random.uniform(-5, 5)
+ vy = np.random.uniform(-5, 5)
+ vz = np.random.uniform(-5, 5)
#mass = np.random.uniform(10, 100)
- mass = 10000
- particles.append(Particle(x, y, mass))
+ mass = 1000
+ particle = Particle(x, y, z, mass)
+ particle.vx = vx
+ particle.vy = vy
+ particle.vz = vz
+ particles.append(particle)
+
+ return particles
+
+ elif initial == 'four':
+ particles = [
+ Particle(60, 50, 40, 1000),
+ Particle(40, 50, 60, 1000),
+ Particle(50, 60, 70, 1000),
+ Particle(50, 40, 30, 1000)
+ ]
+ particles[0].vx, particles[0].vy, particles[0].vz = 0, 3, 3
+ particles[1].vx, particles[1].vy, particles[1].vz = 0, -3, 3
+ particles[2].vx, particles[2].vy, particles[2].vz = -3, 0, -3
+ particles[3].vx, particles[3].vy, particles[3].vz = 3, 0, -3
+ #particles=particles[:2]
return particles
def simulate(self, num_steps, time_step):
@@ -66,23 +86,24 @@ class GalaxyCollisionSimulation:
self.integrate(time_step, 'velocity_verlet') # Integrate to update particle positions
self.system_center_of_mass[step] = (self.barnes_hut.rootNode.center_of_mass[0],
- self.barnes_hut.rootNode.center_of_mass[1])
+ self.barnes_hut.rootNode.center_of_mass[1],
+ self.barnes_hut.rootNode.center_of_mass[2])
# Store particle positions, velocities and forces for each time step
for i, particle in enumerate(self.particles):
- self.particle_positions[i].append((particle.x, particle.y))
- self.particle_velocities[i].append((particle.vx, particle.vy))
- self.particle_forces[i].append((particle.fx, particle.fy))
+ self.particle_positions[i].append((particle.x, particle.y, particle.z))
+ self.particle_velocities[i].append((particle.vx, particle.vy, particle.vz))
+ self.particle_forces[i].append((particle.fx, particle.fy, particle.fz))
def evaluate_forces(self):
# Evaluate forces acting on each particle using Barnes-Hut algorithm
for particle in self.particles:
force = self.barnes_hut.rootNode.CalculateForceFromTree(particle)
- particle.fx, particle.fy = force
+ particle.fx, particle.fy, particle.fz = force
def calculate_force(self, particle):
# Calculate force acting on a single particle
force = self.barnes_hut.rootNode.CalculateForceFromTree(particle)
- particle.fx, particle.fy = force
+ particle.fx, particle.fy, particle.fz = force
def integrate(self, time_step, method):
# Integrate to update particle positions based on calculated forces
@@ -90,27 +111,32 @@ class GalaxyCollisionSimulation:
for particle in self.particles:
particle.vx += particle.fx * time_step / particle.mass # Update velocity components
particle.vy += particle.fy * time_step / particle.mass
+ particle.vz += particle.fz * time_step / particle.mass
particle.x += particle.vx * time_step # Update particle positions using velocity
particle.y += particle.vy * time_step
+ particle.z += particle.vz * time_step
elif method == 'velocity_verlet':
for particle in self.particles:
particle.vx += particle.fx * time_step / (2 * particle.mass) # First half-step
particle.vy += particle.fy * time_step / (2 * particle.mass)
+ particle.vz += particle.fz * time_step / (2 * particle.mass)
particle.x += particle.vx * time_step # Full step for position
particle.y += particle.vy * time_step
+ particle.z += particle.vz * time_step
self.calculate_force(particle) # Recalculate force for the particle
particle.vx += particle.fx * time_step / (2 * particle.mass) # Second half-step
particle.vy += particle.fy * time_step / (2 * particle.mass)
+ particle.vz += particle.fz * time_step / (2 * particle.mass)
def print_particle_positions(self):
# Print the positions of all particles for each time step
for i, positions in self.particle_positions.items():
print(f"Particle {i + 1} Positions:")
for step, pos in enumerate(positions):
- print(f" Time Step {step}: Position ({pos[0]}, {pos[1]})")
+ print(f" Time Step {step}: Position ({pos[0]}, {pos[1]}, {pos[2]})")
def display_snapshots(self, num_shots, fix_axes=True):
# Display pictures for each time step
@@ -121,27 +147,30 @@ class GalaxyCollisionSimulation:
#Fixed axis limits
print(f"For {axes_type} axis limits, num_particles={self.num_particles}")
for step in np.arange(0, num_time_steps, int(num_time_steps / num_shots)):
- fig, ax = plt.subplots()
+ fig = plt.figure()
+ ax = fig.add_subplot(111, projection='3d')
if fix_axes:
- ax.set_xlim(0, 100)
- ax.set_ylim(0, 100)
+ ax.set_xlim(0, 200)
+ ax.set_ylim(0, 200)
+ ax.set_zlim(0, 200)
ax.set_xlabel('X-axis')
ax.set_ylabel('Y-axis')
+ ax.set_zlabel('Z-axis')
ax.set_title(f'Time Step {step}')
ax.grid()
- print(f'Step {step}')
+ #print(f'Step {step}')
for particle_index, positions in self.particle_positions.items():
- x, y = positions[step]
- vx, vy = self.particle_velocities[particle_index][step]
- fx, fy = self.particle_forces[particle_index][step]
+ x, y, z = positions[step]
+ vx, vy, vz = self.particle_velocities[particle_index][step]
+ fx, fy, fz = self.particle_forces[particle_index][step]
#mass = self.particles[particle_index].mass
- ax.scatter(x, y, c='blue', s=20, alpha=0.5) # Plot particle position for the current time step
- c_x, c_y = self.system_center_of_mass[step]
- ax.scatter(c_x, c_y, c='orange', s=20)
- print(
- f'i={particle_index}: x={round(x,2)}, y={round(y,2)}, vx={round(vx,2)}, vy={round(vy,2)}, fx={round(fx,2)}, fy={round(fy,2)}'
- )
+ ax.scatter(x, y, z, c='blue', s=20, alpha=0.5) # Plot particle position for the current time step
+ c_x, c_y, c_z = self.system_center_of_mass[step]
+ ax.scatter(c_x, c_y, c_z, c='orange', s=20)
+ #print(
+ #f'i={particle_index}: x={round(x,2)}, y={round(y,2)}, z={round(z,2)}, vx={round(vx,2)}, vy={round(vy,2)}, vz={round(vz,2)}, fx={round(fx,2)}, fy={round(fy,2)}, fz={round(fz,2)}'
+ #)
#print(y)
plt.show()
diff --git a/tasks/bht/print_tree.py b/tasks/bht/print_tree.py
index 9390f837d3d1473bf229d06c8d28a11970442e87..5ed8745915b53ad8063ea64ea90c57ef9651519f 100644
--- a/tasks/bht/print_tree.py
+++ b/tasks/bht/print_tree.py
@@ -1,5 +1,5 @@
import numpy as np
-from bht import MainApp, Particle
+from tasks.bht.bht_oct import MainApp, Particle
# Function to generate random particles
diff --git a/tasks/bht/simulate.py b/tasks/bht/simulate.py
index c703e4231fcaa037e6bbe49f5f9195b7020e1509..1de867ad187aee396b3b7ffb9f9a6251ff8b5bfb 100644
--- a/tasks/bht/simulate.py
+++ b/tasks/bht/simulate.py
@@ -1,5 +1,5 @@
from galaxysimulation import GalaxyCollisionSimulation
-sim = GalaxyCollisionSimulation(num_particles=6)
-sim.simulate(num_steps=7000, time_step=0.001)
-sim.display_snapshots(200, fix_axes=True)
+sim = GalaxyCollisionSimulation(num_particles=4)
+sim.simulate(num_steps=10000, time_step=0.001)
+sim.display_snapshots(200, fix_axes=True)
\ No newline at end of file
diff --git a/tasks/bht/system.py b/tasks/bht/system.py
new file mode 100644
index 0000000000000000000000000000000000000000..264dcfdf22572b1208902fe034f99542928a374c
--- /dev/null
+++ b/tasks/bht/system.py
@@ -0,0 +1,57 @@
+"""
+This module contains base class for a n-body gravitational system
+with 2 methods of simulations: direct and approximated force field.
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class GravitationalSystem:
+ r0: np.ndarray[float] # shape = (N,d)
+ v0: np.ndarray[float] # shape = (N,d)
+ m: np.ndarray[float] # shape = N
+ t: np.ndarray[float] # shape = n_time_steps
+ force: Callable
+ solver: Callable
+
+ def __post_init__(self):
+ "Checking dimensions of inputs"
+ assert self.r0.shape == self.v0.shape
+ assert self.m.shape[0] == self.r0.shape[0]
+
+ solvers = ['verlet']
+ assert self.solver.__name__ in solvers
+
+ def direct_simulation(self):
+ """
+ Using integrator 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)):
+ 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