docs: finished docstrings

Author: PhySimdev

BarnesHut.py

@@ -6,7 +6,17 @@
 
 
 class Node():
+    """Stores data for Octree nodes."""
+
     def __init__(self, middle, dimension):
+        """Method that sets up a node.
+
+        Method that sets up and declares variables for an octree node.
+
+        Args:
+            middle: Position of center of the position.
+            dimension: Length of sides of a node.
+        """
         self.particle   = None
         self.middle     = middle
         self.dimension  = dimension
@@ -26,6 +36,15 @@ def get_subnode(self, quad):
         return (self.subnodes[quad[0]][quad[1]][quad[2]])
 
     def create_subnode(self, quad):
+        """Method that creates a subnode.
+
+        Method that determines the middle and dimension of the subnode
+        of a specific quadrant of the node. Initializes a subnode and adds
+        that subnode to the nodes.
+
+        Args:
+            quad: Quadrant of node.
+        """
         dimension  = self.dimension / 2
 
         x, y, z = 1, 1, 1
@@ -36,9 +55,9 @@ def create_subnode(self, quad):
         if quad[2] == 1:
             z = -1
 
-        middle = [self.middle[0] + ((dimension / 2) * x),  # quad[0] == 0: value  1, right
-                  self.middle[1] + ((dimension / 2) * y),  # quad[1] == 0: value  1, front
-                  self.middle[2] + ((dimension / 2) * z)]  # quad[2] == 0: value  1, top
+        middle = [self.middle[0] + ((dimension / 2) * x),  # value  1, right
+                  self.middle[1] + ((dimension / 2) * y),  # value  1, front
+                  self.middle[2] + ((dimension / 2) * z)]  # value  1, top
         node       = Node(middle, dimension)
         self.subnodes[quad[0]][quad[1]][quad[2]] = node
         self.nodes.append(node)
@@ -54,6 +73,11 @@ def get_quad(self, point):
         return [x, y, z]
 
     def get_corners(self):
+        """Method that gets corners of a node.
+
+        Method that get corners of a node for visualization. Iterates through
+        the top and bottom for front and back for right and left.
+        """
         if self.corners is None:
             self.corners = []
             for x in [1, -1]:          # right or left
@@ -67,13 +91,22 @@ def get_corners(self):
 
     def in_bounds(self, point):
         val = False
-        if point[0] <= self.middle[0] + (self.dimension / 2) and point[0] >= self.middle[0] - (self.dimension / 2) and\
-           point[1] <= self.middle[1] + (self.dimension / 2) and point[1] >= self.middle[1] - (self.dimension / 2) and\
-           point[2] <= self.middle[2] + (self.dimension / 2) and point[2] >= self.middle[2] - (self.dimension / 2):
+        if point[0] <= self.middle[0] + (self.dimension / 2) and\
+           point[0] >= self.middle[0] - (self.dimension / 2) and\
+           point[1] <= self.middle[1] + (self.dimension / 2) and\
+           point[1] >= self.middle[1] - (self.dimension / 2) and\
+           point[2] <= self.middle[2] + (self.dimension / 2) and\
+           point[2] >= self.middle[2] - (self.dimension / 2):
             val = True
         return val
 
     def compute_mass_distribution(self):
+        """Method that calculates the mass distribution.
+
+        Method that calculates the mass distribution of the node based on
+        the mass posistions of the subnode weighted by weights of 
+        the subnodes.
+        """
         if self.particle is not None:
             self.center_of_mass = np.array([*self.particle.position])
             self.mass = self.particle.mass
@@ -91,41 +124,61 @@ def compute_mass_distribution(self):
 
 
 class Octree():
-    def __init__(self, particles, root_node, theta):
+    """Handles setup and calculations of the Barnes-Hut octree."""
+
+    def __init__(self, particles, root_node, theta, node_type):
+        """Method that sets up an octree.
+
+        Method that sets up the variables for the octree. Calls functions
+        for creation of the octree.
+
+        Args:
+            particles: List of particles that are inserted.
+            root_node: Root node of the octree.
+            theta: Theta that determines the accuracy of the simulations.
+        """
         self.theta = theta
         self.root_node = root_node
         self.particles = particles
+        self.node_type = node_type
         for particle in self.particles:
             self.insert_to_node(self.root_node, particle)
 
     def insert_to_node(self, node, particle):
+        """Recursive method that inserts particles into the octree.
+
+        Recursive method that inserts bodies into the octree.
+        Checks if particle is in the current node to prevent bounds issues.
+        Determines the appropriate child node and gets that subnode.
+        If that subnode is empty insert the particle and stop.
+        If the child node point is a point node (one particle) turn it into 
+        a regional node by inserting both particles into it.
+        If the child node is a regional node insert the particle.
+
+        Args:
+            node: Quadrant of node.
+            particle: Simulation body.
+        """
         # check if point is in cuboid of present node
         if not node.in_bounds(particle.position) and not np.array_equal(particle.position, self.root_node.middle):
             print("error particle not in bounds")
             print(f"middle: {node.middle}, dimension: {node.dimension}, particle position: {particle.position}, type: {type(particle)}")
             return
 
-        # determine the appropriate child node
         quad = node.get_quad(particle.position)
         if node.get_subnode(quad) is None:
             node.create_subnode(quad)
         subnode = node.get_subnode(quad)
 
-        # if subnode is empty, insert point, stop insertion
         if subnode.particle is None and len(subnode.nodes) == 0:  # case empty node
             subnode.insert_particle(particle)
 
-        # If the child node is a point node, replace it with a region node.
-        # Call insert for the point that just got replaced.
-        # Set current node as the newly formed regionnode.
         elif subnode.particle is not None:  # case point node
             old_particle = subnode.particle
             subnode.insert_particle(None)
             self.insert_to_node(subnode, old_particle)
             self.insert_to_node(subnode, particle)
 
-        # If the child node is a point node, replace it with a region node.
-        # Call insert for the point that just got replaced. Set current node as the newly formed region node.
         elif subnode.particle is None and len(subnode.nodes) >= 1:  # case region node
             self.insert_to_node(subnode, particle)
 
@@ -143,15 +196,28 @@ def update_forces_collisions(self):
                 del self.collision_dic[particle]
 
     def calc_forces(self, node, particle):
-        # Gravitational force and collision between two particles
+        """Method that calculates the force on an octree particle.
+
+        Method that calculates the force on an octree particle by iterating
+        through the octree.
+        If the node is a point node that doesnt hold the body itelf, directly
+        calculate the forces.
+        If its a regional node and the dimension/distance ratio is smaller
+        than theta and the center of mass position is not the same as the
+        particle position, calculate the force between the node and
+        the particle.
+
+        Args:
+            node: Quadrant of node.
+            particle: Simulation body.
+        """
         if node.particle is not None and node.particle != particle:
             force, e_pot, distance = self.gravitational_force(particle, node, np.array([]), np.array([]))
             particle.force -= force
             particle.e_pot -= e_pot
             if distance < particle.radius + node.particle.radius and particle.mass > node.particle.mass:
                 self.collision_dic[particle].append(node.particle)
 
-        # find regional node were particle is not the center of mass (subnodes particle is particle)
         elif node.particle is None and not np.array_equal(particle.position, node.center_of_mass):
             distance = np.array([*particle.position]) - np.array([*node.center_of_mass])
             r = math.sqrt(np.dot(distance, distance))
@@ -165,6 +231,22 @@ def calc_forces(self, node, particle):
                     self.calc_forces(subnode, particle)
 
     def gravitational_force(self, particle, node, distance_vec, distance_val):  # can be ragional or point node
+        """Method that calculates the force between two particles.
+    
+        Method that calculates the force acted on the particle by
+        another particle or a node. Only calculates the distance and vector
+        did not have to be calculated for theta.
+    
+        Args:
+            particle: Simulation body.
+            node: Node of the octree.
+            distance_vec: Vector of distance between the bodies.
+            distance_val: Magnitude of distance betweent the bodies
+        
+        Returns:
+            The force, potential energy and distance between two bodies or 
+            the body and the node.
+        """
         force = np.array([0., 0., 0.])
         if len(distance_vec) == 0 and len(distance_val) == 0:
             distance = np.array([*particle.position]) - np.array([*node.center_of_mass])
@@ -179,14 +261,15 @@ def gravitational_force(self, particle, node, distance_vec, distance_val):  # ca
         force = (distance / distance_mag) * force_mag
         return force, e_pot, distance_mag
 
-    def get_all_nodes(self, node, lst):  # all point nodes, could include regional nodes
+    def get_all_nodes(self, node, lst):
 
         if node.particle is None and len(node.nodes) >= 1 or node.particle is not None:
             if len(node.nodes) >= 1:
-                # lst.append(node.get_corners())  # if regional are shown aswell
+                if self.node_type == "regional" or self.node_type == "both":
+                    lst.append(node.get_corners()) 
                 for subnode in node.nodes:
                     self.get_all_nodes(subnode, lst)
-            if node.particle is not None:
+            if node.particle is not None and (self.node_type == "point" or self.node_type == "both"):
                 lst.append(node.get_corners())
 
 

Simulation.py

@@ -21,6 +21,7 @@ def __init__(self, theta=1, rc=0, absolute_pos=True, focus_index=0):
             rc: Restitution coefficient for collisions.
             absolute_pos: Bool value to determine type of movement.
             focus_index: Index of the list focus_options form 0 to 2.
+            node_type: String that determines what nodes are displayed.
         """
         self.restitution_coefficient = rc
         self.focus_options  = ["none", "body", "cm"]
@@ -36,7 +37,7 @@ def __init__(self, theta=1, rc=0, absolute_pos=True, focus_index=0):
         self.total_e        = 0
         self.cm_pos         = np.array([0, 0, 0])
         self.cm_velo        = None
-
+            
         if focus_index >= 0 and focus_index < len(self.focus_options):
             self.focus_index = focus_index
         else:
@@ -100,7 +101,7 @@ def clear_trail(self):
         for body in self.bodies:
             body.trail = []
 
-    def calculate(self, timestep, draw_box):
+    def calculate(self, timestep, draw_box, node_type):
         """Method that calculates a simulation physics step.
 
         Method that calls functions for physics calculations. 
@@ -115,7 +116,7 @@ def calculate(self, timestep, draw_box):
         """
         if self.first:
             self.first = False
-            self.update_interactions()
+            self.update_interactions(node_type)
             for body in self.bodies:
                 body.acceleration = body.force / body.mass
 
@@ -124,7 +125,7 @@ def calculate(self, timestep, draw_box):
             body.update_velocity(timestep)
             body.update_position(timestep)
 
-        self.update_interactions()
+        self.update_interactions(node_type)
 
         self.tree_nodes = []
         if draw_box:
@@ -185,7 +186,7 @@ def get_data(self):
 
         return body_data, self.tree_nodes, system_data, self.cm_pos - default_pos
 
-    def update_interactions(self):
+    def update_interactions(self, node_type):
         center = self.get_focus_pos()
 
         largest_val = 0
@@ -200,7 +201,7 @@ def update_interactions(self):
 
         dimension = math.sqrt(((furthest_body.position[largest_index] - center[largest_index]) * 2.5)**2)
         root = Node(center, dimension)
-        self.tree = Octree(self.bodies, root, self.theta)
+        self.tree = Octree(self.bodies, root, self.theta, node_type)
         root.compute_mass_distribution()
         self.tree.update_forces_collisions()
         self.compute_collisions(self.tree.collision_dic)

main.py

@@ -67,7 +67,7 @@ def __init__(self):
         self.planet_size_factor = 80
         self.min_body_size = .5
         self.path_planet_color = True
-        self.draw_box       = False
+        self.draw_box       = True
         self.draw_trail     = True
         self.min_trail_size = .5
         self.trail_length   = 100000
@@ -84,6 +84,10 @@ def __init__(self):
         self.max_frame      = 3600
 
         self.draw_rot       = True
+        self.node_list_index= 1
+        self.node_list      = ["regional", "point", "both"]
+        self.node_type      = self.node_list[self.node_list_index]
+        self.grid_thickness = 2 * 10**9
         self.rot_cube_pos   = [.8, -.8]  # zwischen -1,1, scale of screen
         self.rot_cube_scale = 80
         self.rot_cube_scolor = "grey"
@@ -484,7 +488,7 @@ def draw_cube(self, qube):
             pos2 = qube[line[1]][0]
             if pos1 and pos2 is not None:
                 f = (qube[line[0]][1] + qube[line[1]][1]) / 2
-                self.pointer.pensize(self.path_size * f / 2)
+                self.pointer.pensize(self.grid_thickness * f / 2)
                 self.pointer.goto(pos1)
                 self.pointer.down()
                 self.pointer.goto(pos2)
@@ -728,7 +732,7 @@ def update_program(self):
             self.frame_count += 1
             t1 = time.time()
             for system in self.simulations:
-                system.calculate(self.timestep, self.draw_box)
+                system.calculate(self.timestep, self.draw_box, self.node_type)
             t2 = time.time()
             self.physics_time = t2 - t1