Newer
Older
"""
This module contains functions for a simulation with force field
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
of gravitational n-body system
"""
import matplotlib.pyplot as plt
import numpy as np
from jobs.src.bht_algorithm_3D 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(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, particle.z)]
for i, particle in enumerate(self.particles)
} # Store particle positions for each time step
self.particle_velocities = {
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, 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.barnes_hut.rootNode.center_of_mass[2])
}
def initialize_particles(self, initial='random'):
if initial == 'two':
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
elif initial == 'random':
# Generate random particles within a specified range
particles = []
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
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):
# 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],
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))
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, 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, particle.fz = 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.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]}, {pos[2]})")
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 = 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)
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}')
for particle_index, positions in self.particle_positions.items():
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, 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()
if __name__ == "__main__":
sim = GalaxyCollisionSimulation(num_particles=10)
sim.simulate(num_steps=10000, time_step=0.001)
sim.display_snapshots(200, fix_axes=True)