Skip to content
Snippets Groups Projects
Commit 280ec952 authored by jakut77's avatar jakut77
Browse files

Add dynamic scaling in BHT alg, clean up and usability improvements for Simulation class

parent 4f7d9f22
Branches
Tags
1 merge request!12BHT algorithm and Galaxy Collision Sim improvements
Pipeline #59664 passed
Stage: test-jobs
Hide whitespace changes
Inline Side-by-side
.gitignore
+ 2
− 1
View file @ 280ec952
/__pycache__
*.pyc
/playground
\ No newline at end of file
/playground
*.DS_Store
\ No newline at end of file
jobs/src/bht_algorithm_3D.py
+ 23
− 12
View file @ 280ec952
......@@ -9,6 +9,12 @@ class MainApp:
def __init__(self):
"""Initialize the MainApp with a root node."""
self.root_x = 0
self.root_y = 0
self.root_z = 0
self.root_width = 100
self.root_height = 100
self.root_depth = 100
self.rootNode = TreeNode(x=0, y=0, z=0, width=100, height=100, depth=100)
def BuildTree(self, particles):
......@@ -24,17 +30,21 @@ class MainApp:
def ResetTree(self, particles):
"""Reset the Octtree 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=100, height=100, depth=100)
if not all(self.rootNode.contains(particle) for particle in particles):
# if a particle exits the boundary, increase the size of the boundary
self.root_x -= self.root_width / 2
self.root_y -= self.root_height / 2
self.root_z -= self.root_depth / 2
self.root_width *= 2
self.root_height *= 2
self.root_depth *= 2
self.rootNode = TreeNode(x=self.root_x,
y=self.root_y,
z=self.root_z,
width=self.root_width,
height=self.root_height,
depth=self.root_depth)
class Particle:
......@@ -311,6 +321,7 @@ class TreeNode:
"""
#G = 6.674 * (10 ** -11) # Gravitational constant
G = 1
#G = 4 * np.pi**2 # AU^3 / m / yr^2
dx = particle2.x - particle1.x
dy = particle2.y - particle1.y
......@@ -363,4 +374,4 @@ class TreeNode:
self.mass = total_mass
else:
self.center_of_mass = np.array([0.0, 0.0, 0.0])
self.mass = 0
\ No newline at end of file
self.mass = 0
tasks/src/ff_simulation_3D.py
+ 59
− 38
View file @ 280ec952
......@@ -2,11 +2,11 @@
This code contains a class to simulate the collision of a Galaoxy with a given
number of particles/planets. To choose the initial conditions regadring initial
positions, velocities and sometimes the number of particles, please check the
function "initialize_particles". To choose the initial coditions desired, please
sure to type the correct keyword for the variable "initial" in the command
self.particles = self.initialize_particles(initial='random_2') in line 28.
When choosing "two", "four" or "four_2", make sure to adjust the number of
particles in "simulate.py" correspondingly.
function "initialize_particles". To choose initial state of the system,
specify the "initial" variable and number of particles. Possible initializations:
"random", "random_v", "two", "four", "four_v". The initializations marked with v
generate particles with initial velocities. The "num_particles" parameter is only
relevant when choosing the "random" option.
"""
import matplotlib.pyplot as plt
......@@ -22,10 +22,10 @@ class GalaxyCollisionSimulation:
"""
def __init__(self, num_particles):
def __init__(self, initial, num_particles):
# Initialize the simulation with a given number of particles
self.num_particles = num_particles
self.particles = self.initialize_particles(initial='random_2') # Generate particles
self.particles = self.initialize_particles(initial=initial) # 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
......@@ -42,8 +42,8 @@ class GalaxyCollisionSimulation:
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.barnes_hut.rootNode.center_of_mass[2])
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'):
......@@ -59,7 +59,7 @@ class GalaxyCollisionSimulation:
#particles=particles[:2]
return particles
elif initial == 'random':
elif initial == 'random_v':
# Generate random particles within a specified range with random initial velocities
particles = []
for _ in range(self.num_particles):
......@@ -69,7 +69,6 @@ class GalaxyCollisionSimulation:
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
......@@ -79,27 +78,21 @@ class GalaxyCollisionSimulation:
return particles
elif initial == 'random_2':
elif initial == 'random':
# Generate random particles within a specified range without random initial velocities
particles = []
for _ in range(self.num_particles):
x = np.random.uniform(0, 100)
y = np.random.uniform(0, 100)
z = np.random.uniform(0, 100)
vx = 0
vy = 0
vz = 0
#mass = np.random.uniform(10, 100)
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':
elif initial == 'four_v':
# Generate four particles with initial velocities
particles = [
Particle(60, 50, 40, 1000),
......@@ -111,10 +104,9 @@ class GalaxyCollisionSimulation:
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
elif initial == 'four_2':
elif initial == 'four':
# Generate four particles without initial velocities
particles = [
Particle(60, 50, 40, 1000),
......@@ -122,14 +114,9 @@ class GalaxyCollisionSimulation:
Particle(50, 60, 70, 1000),
Particle(50, 40, 30, 1000)
]
particles[0].vx, particles[0].vy, particles[0].vz = 0, 0, 0
particles[1].vx, particles[1].vy, particles[1].vz = 0, 0, 0
particles[2].vx, particles[2].vy, particles[2].vz = 0, 0, 0
particles[3].vx, particles[3].vy, particles[3].vz = 0, 0, 0
#particles=particles[:2]
return particles
def simulate(self, num_steps, time_step):
def simulate(self, num_steps, time_step, print_tree):
"""
Function to simulate the galaxy collision for a certain number of step
"""
......@@ -137,16 +124,15 @@ class GalaxyCollisionSimulation:
# 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):
n_particles, particles = self.barnes_hut.rootNode.particles_in_subtree()
print(f'===Octtree at step {step}, n_particles:{n_particles}===')
if print_tree and step in np.arange(0, num_steps, 1000):
print(f'===Octtree 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],
self.barnes_hut.rootNode.center_of_mass[2])
self.system_center_of_mass[0].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]))
# 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))
......@@ -236,8 +222,8 @@ class GalaxyCollisionSimulation:
fx, fy, fz = self.particle_forces[particle_index][step]
#mass = self.particles[particle_index].mass
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=300)
c_x, c_y, c_z = self.system_center_of_mass[0][step]
ax.scatter(c_x, c_y, c_z, c='orange', s=40)
#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)}'
#)
......@@ -245,8 +231,43 @@ class GalaxyCollisionSimulation:
plt.show()
def plot_trajectory(self, update_interval=100):
"""
Method to visualize particles trajectories as a movie
"""
n_time_step = len(self.particle_positions[0])
n_bodies = len(self.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[0][i]
ax.scatter(c_x, c_y, c_z, c='orange', s=40)
ax.set_xlim(0, 100)
ax.set_ylim(0, 100)
ax.set_zlim(0, 100)
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.simulate(num_steps=10000, time_step=0.001)
sim.display_snapshots(200, fix_axes=True)
sim = GalaxyCollisionSimulation(initial='random_v', num_particles=10)
sim.simulate(num_steps=10000, time_step=0.001, print_tree=True)
#sim.display_snapshots(20, fix_axes=True)
sim.plot_trajectory()
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment