updated to Octree

tasks/bht/2D/bht.py

+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

tasks/bht/2D/galaxysimulation.py

+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()

tasks/bht/2D/simulate.py

+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)

tasks/bht/bht_oct.py

+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

tasks/bht/galaxysimulation.py

 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()

tasks/bht/print_tree.py

 import numpy as np
-from bht import MainApp, Particle
+from tasks.bht.bht_oct import MainApp, Particle
 
 
 # Function to generate random particles

tasks/bht/simulate.py

 from galaxysimulation import GalaxyCollisionSimulation
 
-sim = GalaxyCollisionSimulation(num_particles=6)
-sim.simulate(num_steps=7000, time_step=0.001)
+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

tasks/bht/system.py

+"""
+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