solid-funicular/newbody.py at main · spaceybread/solid-funicular · GitHub

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
import numpy as np
import math
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import config

# sep check function
def sp(cor):
    X = np.array(cor[0])
    Y = np.array(cor[1])
    for n in range(len(X)):
        for i in range(len(X)):
            if n != i:
                sep = math.sqrt((X[n] - X[i])**2 + (Y[n] - Y[i])**2)
                if sep < 0.09:
                    return True
    return False

# helper function
def r(cor, mass, G):
    X = np.array(cor[0])
    Y = np.array(cor[1])
    vX = np.array(cor[2])
    vY = np.array(cor[3])
    delt = np.array(cor)

    for n in range(len(X)):
        dvx, dvy = 0, 0

        for i in range(len(X)):
            if n != i:
                sep = math.sqrt((X[n] - X[i])**2 + (Y[n] - Y[i])**2)
                dvx -= (G * mass[i] * (X[n] - X[i])) / sep**3
                dvy -= (G * mass[i] * (Y[n] - Y[i])) / sep**3

        delt[2][n] = dvx
        delt[3][n] = dvy

    delt[0] = vX
    delt[1] = vY

    return delt

#step function
def step(cor, mass, dt, G):
     if sp(cor) == True: # if the objects collide, the simulation stops
        return cor
     k1 = dt * r(cor, mass, G)
     k2 = dt * r(cor + k1/2, mass, G)
     k3 = dt * r(cor + k2/2, mass, G)
     k4 = dt * r(cor + k3, mass, G)

     cor = cor + (k1 + 2*k2 + 2*k3 + k4)/6

     return cor

#init
m = config.m
x = config.x
y = config.y
vx = config.vx
vy = config.vy

cords = np.array([x,y,vx,vy])

time = np.linspace(0, 2*6.3259, 2*1001, endpoint=True)
dt = time[1] - time[0]

X = np.zeros((3, len(time)))
Y = np.zeros((3, len(time)))
VX = np.zeros((3, len(time)))
VY = np.zeros((3, len(time)))


#computing values
for _ in range(len(time)):
    cords = step(cords, m, dt, 1)
    X[:, _], Y[:, _], VX[:, _], VY[:, _] = cords[0], cords[1], cords[2], cords[3]
    print(cords)

#plotting
fig, ax = plt.subplots(figsize=(10, 4))

ax.set_xlim(-10.1, 10.1)
ax.set_ylim(-10.1, 10.1)
ax.set_aspect('equal')
ax.grid(True, alpha=0.5)
ax.minorticks_on()
ax.grid(True, alpha=0.3, which='minor')
ax.set_facecolor('black')

line1, = ax.plot([], [], lw=2, alpha=0.2)
line2, = ax.plot([], [], lw=2, alpha=0.2)
line3, = ax.plot([], [], lw=2, alpha=0.2)
scat1 = ax.scatter([], [], marker='o', c='red', s=80)
scat2 = ax.scatter([], [], marker='o', c='blue', s=80)
scat3 = ax.scatter([], [], marker='o', c='green', s=80)

def update(frame):
    i = frame

    line1.set_data(X[0, :i], Y[0, :i])
    line2.set_data(X[1, :i], Y[1, :i])
    line3.set_data(X[2, :i], Y[2, :i])

    scat1.set_offsets(np.c_[X[0, i], Y[0, i]])
    scat2.set_offsets(np.c_[X[1, i], Y[1, i]])
    scat3.set_offsets(np.c_[X[2, i], Y[2, i]])

    return line1, line2, line3, scat1, scat2, scat3

frames = len(time)
ani = FuncAnimation(fig, update, frames=frames, blit=True, interval=5)

plt.show()