Newer
Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
"""
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 27.
When choosing "two", "four" or "four_2", make sure to adjust the number of
particles in "simulate.py" correspondingly.
"""
import matplotlib.pyplot as plt
import numpy as np
from jobs.src.bht_algorithm_2D import MainApp, Particle
class GalaxyCollisionSimulation:
"""
Simulation class for galaxy collision using Barnes-Hut and Euler method
or Velocity Verlet
"""
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_2') # 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'):
"""
Function to initialize particles with different initial conditions
"""
if initial == 'two':
# Generate two particles without initial velocities
particles = [Particle(60, 50, 1000), Particle(40, 50, 1000)]
particles[0].vx, particles[0].vy = 0, 0
particles[1].vx, particles[1].vy = 0, 0
#particles=particles[:2]
return particles
elif initial == 'random':
# Generate random particles within a specified range with random initial velocities
particles = []
for _ in range(self.num_particles):
x = np.random.uniform(0, 100)
y = np.random.uniform(0, 100)
vx = np.random.uniform(-5, 5)
vy = np.random.uniform(-5, 5)
#mass = np.random.uniform(10, 100)
mass = 1000
particle = Particle(x, y, mass)
particle.vx = vx
particle.vy = vy
particles.append(particle)
return particles
elif initial == 'random_2':
# 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)
vx = 0
vy = 0
#mass = np.random.uniform(10, 100)
mass = 1000
particle = Particle(x, y, mass)
particle.vx = vx
particle.vy = vy
particles.append(particle)
return particles
elif initial == 'four':
# Generate four particles with initial velocities
particles = [Particle(60, 50, 1000), Particle(40, 50, 1000), Particle(50, 60, 1000), Particle(50, 40, 1000)]
particles[0].vx, particles[0].vy = 0, 3
particles[1].vx, particles[1].vy = 0, -3
particles[2].vx, particles[2].vy = -3, 0
particles[3].vx, particles[3].vy = 3, 0
#particles=particles[:2]
return particles
elif initial == 'four_2':
# Generate four particles without initial velocities
particles = [Particle(60, 50, 1000), Particle(40, 50, 1000), Particle(50, 60, 1000), Particle(50, 40, 1000)]
particles[0].vx, particles[0].vy = 0, 0
particles[1].vx, particles[1].vy = 0, 0
particles[2].vx, particles[2].vy = 0, 0
particles[3].vx, particles[3].vy = 0, 0
#particles=particles[:2]
return particles
def simulate(self, num_steps, time_step):
"""
Function to simulate the galaxy collision for a certain number of step
"""
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):
n_particles, particles = self.barnes_hut.rootNode.particles_in_subtree()
print(f'===Quadtree at step {step}, n_particles:{n_particles}===')
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):
"""
Function to 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):
"""
Function to 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):
"""
Function to 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):
"""
Function to 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):
"""
Function to 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'Day {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=300)
#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()
if __name__ == "__main__":
sim = GalaxyCollisionSimulation(num_particles=3)
sim.simulate(num_steps=10000, time_step=0.001)
sim.display_snapshots(200, fix_axes=True)