Skip to content
Snippets Groups Projects
Commit 98e1e9a3 authored by jakut77's avatar jakut77
Browse files

some experiments + time measurements

parent 8595d3e5
Branches
Tags
No related merge requests found
/__pycache__
*.pyc
/playground
*.DS_Store
\ No newline at end of file
"""
This module implements the Barnes Hut Tree (BHT) Algorithm for 3D data.
"""
import numpy as np
class MainApp:
def __init__(self):
"""Initialize the MainApp with a root node."""
self.rootNode = TreeNode(x=0, y=0, z=0, width=20, height=20, depth=20)
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)
for particle in particles:
self.rootNode.insert(particle)
def ResetTree(self, particles):
"""Reset the Octree by reinitializing the root node."""
if all(self.rootNode.contains(particle) for particle in particles):
# If all particles are contained in the current tree, keep same dimensions
root_x = self.rootNode.x
root_y = self.rootNode.y
root_z = self.rootNode.z
root_width = self.rootNode.width
root_height = self.rootNode.height
root_depth = self.rootNode.depth
else:
# If a particle leaves the root node boundaries, resize the tree
print('New tree dimensions')
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])
x_m = (min_x + max_x) / 2
y_m = (min_y + max_y) / 2
z_m = (min_z + max_z) / 2
delta_x = max_x - min_x
delta_y = max_y - min_y
delta_z = max_z - min_z
# Recalculate dimensions
root_x = x_m - delta_x
root_y = y_m - delta_y
root_z = z_m - delta_z
root_dim = max(delta_x, delta_y, delta_z)
root_width = 2 * root_dim
root_height = 2 * root_dim
root_depth = 2 * root_dim
self.rootNode = TreeNode(x=root_x, y=root_y, z=root_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
......@@ -6,7 +6,7 @@ of gravitational n-body system
import matplotlib.pyplot as plt
import numpy as np
from jobs.src.bht_algorithm_3D import MainApp, Particle
from bht_algorithm_3D import MainApp, Particle
# Simulation class for galaxy collision using Barnes-Hut and Euler method
......@@ -16,7 +16,7 @@ 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(self.particles) # Initialize Barnes-Hut algorithm instance
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 = {
......@@ -31,15 +31,17 @@ class GalaxyCollisionSimulation:
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],
self.system_center_of_mass = [
(self.barnes_hut.rootNode.center_of_mass[0], self.barnes_hut.rootNode.center_of_mass[1],
self.barnes_hut.rootNode.center_of_mass[2])
}
]
self.treeDims = [(self.barnes_hut.rootNode.x, self.barnes_hut.rootNode.width, self.barnes_hut.rootNode.y,
self.barnes_hut.rootNode.height, self.barnes_hut.rootNode.z, self.barnes_hut.rootNode.depth)]
def initialize_particles(self, initial='random'):
if initial == 'two':
particles = [Particle(60, 50, 40, 1000), Particle(40, 50, 60, 1000)]
particles = [Particle(60, 49, 40, 1000), Particle(40, 51, 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]
......@@ -48,19 +50,19 @@ class GalaxyCollisionSimulation:
elif initial == 'random':
# Generate random particles within a specified range
particles = []
initial_v = True
for _ in range(self.num_particles):
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 = 1000
particle = Particle(x, y, z, mass)
particle.vx = vx
particle.vy = vy
particle.vz = vz
if initial_v:
particle.vx = np.random.uniform(-5, 5)
particle.vy = np.random.uniform(-5, 5)
particle.vz = np.random.uniform(-5, 5)
particles.append(particle)
return particles
......@@ -91,14 +93,17 @@ class GalaxyCollisionSimulation:
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],
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, particle.z))
self.particle_velocities[i].append((particle.vx, particle.vy, particle.vz))
self.particle_forces[i].append((particle.fx, particle.fy, particle.fz))
self.system_center_of_mass.append(
(self.barnes_hut.rootNode.center_of_mass[0], self.barnes_hut.rootNode.center_of_mass[1],
self.barnes_hut.rootNode.center_of_mass[2]))
self.treeDims.append(
(self.barnes_hut.rootNode.x, self.barnes_hut.rootNode.width, self.barnes_hut.rootNode.y,
self.barnes_hut.rootNode.height, self.barnes_hut.rootNode.z, self.barnes_hut.rootNode.depth))
def evaluate_forces(self):
# Evaluate forces acting on each particle using Barnes-Hut algorithm
......@@ -145,7 +150,9 @@ class GalaxyCollisionSimulation:
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
"""
Method to display particle positions as snapshots
"""
num_time_steps = len(next(iter(self.particle_positions.values())))
print('===Quadtree at the end===')
self.barnes_hut.rootNode.print_tree() # print the tree
......@@ -156,9 +163,12 @@ class GalaxyCollisionSimulation:
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
if fix_axes:
ax.set_xlim(0, 200)
ax.set_ylim(0, 200)
ax.set_zlim(0, 200)
x_min, width = self.treeDims[step][0], self.treeDims[step][1]
y_min, height = self.treeDims[step][2], self.treeDims[step][3]
z_min, depth = self.treeDims[step][4], self.treeDims[step][5]
ax.set_xlim(x_min, x_min + width)
ax.set_ylim(y_min, y_min + height)
ax.set_zlim(z_min, z_min + depth)
ax.set_xlabel('X-axis')
ax.set_ylabel('Y-axis')
ax.set_zlabel('Z-axis')
......@@ -181,8 +191,44 @@ class GalaxyCollisionSimulation:
plt.show()
def plot_trajectory(self, update_interval=100):
"""
Method to visualize particles trajectories as a slideshow
"""
n_time_step = len(self.particle_positions[0])
n_bodies = self.num_particles
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
for i in range(0, n_time_step, update_interval):
ax.clear()
for j in range(n_bodies):
body_traj = self.particle_positions[j][i]
ax.scatter(body_traj[0], body_traj[1], body_traj[2], c='blue', alpha=0.5)
c_x, c_y, c_z = self.system_center_of_mass[i]
ax.scatter(c_x, c_y, c_z, c='orange', s=40)
ax.set_xlim(0, 200)
ax.set_ylim(0, 200)
ax.set_zlim(0, 200)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.set_title(f"Time step: {i}")
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
ax.grid()
plt.pause(0.01)
if __name__ == "__main__":
sim = GalaxyCollisionSimulation(num_particles=10)
sim = GalaxyCollisionSimulation(num_particles=5)
sim.simulate(num_steps=10000, time_step=0.001)
sim.display_snapshots(200, fix_axes=True)
sim.plot_trajectory()
#sim.display_snapshots(20)
import numpy as np
import time
import matplotlib.pyplot as plt
if __name__ == "__main__":
N_particles = [2, 3, 10, 15, 25, 50, 100]
no_time_steps = [5, 10, 100, 500, 1000, 5000, 10000]
results = np.zeros((7, 7))
for i in range(len(N_particles)):
num_particles = N_particles[i]
for j in range(len(no_time_steps)):
num_steps = no_time_steps[j]
start_time = time.time()
#sim = GalaxyCollisionSimulation(num_particles=num_particles)
#sim.simulate(num_steps=num_steps, time_step=0.001)
execution_time = time.time() - start_time
print(f'N_PARTICLES={num_particles}, NUMSTEPS={num_steps}')
#results[i, j] = execution_time
results = [[
3.20506096e-03, 4.54711914e-03, 2.99420357e-02, 1.30662918e-01, 5.54497004e-01, 1.61694264e+00, 2.64192891e+00
], [1.98988914e-02, 5.33080101e-03, 4.81491089e-02, 2.01280832e-01, 3.57017279e-01, 1.68726897e+00, 4.41522312e+00],
[
4.58290577e-02, 4.86457348e-02, 4.17565823e-01, 1.45790005e+00, 3.72849512e+00, 1.69347689e+01,
3.80113401e+01
],
[
1.59001112e-01, 7.18071461e-02, 7.43870020e-01, 3.17859292e+00, 6.89351201e+00, 4.38741171e+01,
7.23900406e+01
],
[
2.58275032e-01, 1.59539938e-01, 1.69004893e+00, 8.19918489e+00, 1.74204440e+01, 9.41446784e+01,
1.98575519e+02
],
[
5.22159815e-01, 4.45141077e-01, 4.30359578e+00, 2.12155340e+01, 4.82974751e+01, 2.71531863e+02,
5.68442169e+02
],
[
1.30096102e+00, 1.02588534e+00, 9.99782109e+00, 4.98463769e+01, 1.01500050e+02, 7.60922490e+02,
1.57951990e+03
]]
fig, ax = plt.subplots()
im = plt.imshow(results)
ax.set_xticks(np.arange(len(no_time_steps)), labels=no_time_steps)
ax.set_yticks(np.arange(len(N_particles)), labels=N_particles)
ax.set_title("Time scaling for various particle numbers and time steps")
ax.set_xlabel('Num steps')
ax.set_ylabel('N particles')
plt.colorbar()
plt.show()
fig, ax = plt.subplots()
im = plt.imshow(np.log(results))
ax.set_xticks(np.arange(len(no_time_steps)), labels=no_time_steps)
ax.set_yticks(np.arange(len(N_particles)), labels=N_particles)
ax.set_title("Time scaling, logarithmic scale")
ax.set_xlabel('Num steps')
ax.set_ylabel('N particles')
plt.colorbar()
plt.show()
model_function = np.array([n * np.emath.logn(8, n) for n in N_particles])
n_squared = np.array([n**2 for n in N_particles])
for k in range(3, len(no_time_steps)):
data = np.transpose(results)[k]
coeff = data[-1] / model_function[-1]
rescaled_model = coeff * model_function
plt.plot(N_particles, rescaled_model, '--', label='a*NlogN')
plt.scatter(N_particles, data, alpha=0.3, label=f'{no_time_steps[k]} time steps')
print(model_function[-1], data[-1])
plt.plot(N_particles, n_squared, '--', label='N**2')
plt.ylim(0, 1600)
plt.xlabel('N particles')
plt.ylabel('Time complexity')
plt.title('Comparative plot for various simulation lengths')
plt.legend()
plt.show()
print('RESULTS')
#print(results)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment