-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
104 lines (90 loc) · 4.09 KB
/
model.py
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
#
# SIMULATE ORBITS OF A PARTICLE SYSTEM USING NEWTONS LAW OF GRAVITY.
#
# O(h^2)
#
# sebastian angermund
import weakref
import numpy as np
import multiprocessing as mp
from multiprocessing import get_context
class Particle:
use_multiprocessing = False
timestep = 2 # timestep index
h = 3 # timestep length [seconds]
N = 300000 # N*timestep = simulation length [seconds]
G = 6.67e-11 # gravitational constant [SI units]
def __init__(self, index, mass, x_pos, y_pos, x_vel, y_vel):
self.index = index
self.mass = mass
self.k_array[self.index] = np.power(self.h, 2) * self.G * self.mass
self.particle_array[self.index*2, 0] = x_pos
self.particle_array[self.index*2, 1] = x_pos + (self.h * x_vel)
self.particle_array[self.index*2+1, 0] = y_pos
self.particle_array[self.index*2+1, 1] = y_pos + (self.h * y_vel)
if self.use_multiprocessing:
self._instances.append(self)
else:
self._instances.add(weakref.ref(self))
@classmethod
def _set_size(cls, num_particles=3):
cls.num_particles = num_particles
cls.particle_array = np.zeros((2*num_particles, cls.N), dtype=float)
cls.k_array = np.zeros(num_particles)
if num_particles > 2500:
cls.use_multiprocessing = True
cls._instances = []
else:
cls._instances = set()
@classmethod
def _r_update_func(cls, x_step, y_step, delta_x, delta_y, k_list):
r_sq = np.square(delta_x) + np.square(delta_y)
r_abs = np.sqrt(r_sq)
cos = np.divide(np.array(delta_x), r_abs)
sin = np.divide(np.array(delta_y), r_abs)
rx = 2 * x_step[0] - (x_step[1] + np.dot(np.divide(np.array(k_list), r_sq), cos))
ry = 2 * y_step[0] - (y_step[1] + np.dot(np.divide(np.array(k_list), r_sq), sin))
return rx, ry
@classmethod
def _update_positions(cls):
for ref_1 in cls._instances:
obj_1 = ref_1()
x_relevant = cls.particle_array[obj_1.index*2, cls.timestep-1]
y_relevant = cls.particle_array[obj_1.index*2+1, cls.timestep-1]
delta_array = np.delete(
cls.particle_array[:, cls.timestep-1],
(obj_1.index*2, obj_1.index*2+1),
axis=0,
)
k_list = np.delete(cls.k_array, (obj_1.index))
delta_x = x_relevant - delta_array[::2]
delta_y = y_relevant - delta_array[1::2]
x_step = [x_relevant, cls.particle_array[obj_1.index*2, cls.timestep-2]]
y_step = [y_relevant, cls.particle_array[obj_1.index*2+1, cls.timestep-2]]
cls.particle_array[obj_1.index*2, cls.timestep], cls.particle_array[obj_1.index*2+1, cls.timestep] \
= cls._r_update_func(x_step, y_step, delta_x, delta_y, k_list)
@classmethod
def _update_positions_parallel(cls, instance):
x_relevant = cls.particle_array[instance.index*2, cls.timestep-1]
y_relevant = cls.particle_array[instance.index*2+1, cls.timestep-1]
delta_array = np.delete(
cls.particle_array[:, cls.timestep-1],
(instance.index*2, instance.index*2+1),
axis=0,
)
k_list = np.delete(cls.k_array, (instance.index))
delta_x = x_relevant - delta_array[::2]
delta_y = y_relevant - delta_array[1::2]
x_step = [x_relevant, cls.particle_array[instance.index*2, cls.timestep-2]]
y_step = [y_relevant, cls.particle_array[instance.index*2+1, cls.timestep-2]]
return cls._r_update_func(x_step, y_step, delta_x, delta_y, k_list)
@classmethod
def run_update(cls, cores):
if cls.use_multiprocessing:
with mp.get_context("fork").Pool(cores) as p:
xy_list = p.map(cls._update_positions_parallel, cls._instances)
for index, instance in enumerate(cls._instances, 0):
cls.particle_array[instance.index*2, cls.timestep], cls.particle_array[instance.index*2+1, cls.timestep] = xy_list[index]
else:
cls._update_positions()
cls.timestep += 1