-
Notifications
You must be signed in to change notification settings - Fork 26
/
Copy pathphysics_engine.py
executable file
·110 lines (101 loc) · 3.35 KB
/
physics_engine.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
105
106
107
108
109
110
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import argparse
import sys
import os
import numpy as np
import time
from math import sin, cos, radians, pi
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.animation as manimation
import cv2
# 1000 time steps
total_state=1000;
# 5 features on the state [mass,x,y,x_vel,y_vel]
fea_num=5;
# G
#G = 6.67428e-11;
G=10**5;
# time step
diff_t=0.001;
def init(total_state,n_body,fea_num,orbit):
data=np.zeros((total_state,n_body,fea_num),dtype=float);
if(orbit):
data[0][0][0]=100;
data[0][0][1:5]=0.0;
for i in range(1,n_body):
data[0][i][0]=np.random.rand()*8.98+0.02;
distance=np.random.rand()*90.0+10.0;
theta=np.random.rand()*360;
theta_rad = pi/2 - radians(theta);
data[0][i][1]=distance*cos(theta_rad);
data[0][i][2]=distance*sin(theta_rad);
data[0][i][3]=-1*data[0][i][2]/norm(data[0][i][1:3])*(G*data[0][0][0]/norm(data[0][i][1:3])**2)*distance/1000;
data[0][i][4]=data[0][i][1]/norm(data[0][i][1:3])*(G*data[0][0][0]/norm(data[0][i][1:3])**2)*distance/1000;
#data[0][i][3]=np.random.rand()*10.0-5.0;
#data[0][i][4]=np.random.rand()*10.0-5.0;
else:
for i in range(n_body):
data[0][i][0]=np.random.rand()*8.98+0.02;
distance=np.random.rand()*90.0+10.0;
theta=np.random.rand()*360;
theta_rad = pi/2 - radians(theta);
data[0][i][1]=distance*cos(theta_rad);
data[0][i][2]=distance*sin(theta_rad);
data[0][i][3]=np.random.rand()*6.0-3.0;
data[0][i][4]=np.random.rand()*6.0-3.0;
return data;
def norm(x):
return np.sqrt(np.sum(x**2));
def get_f(reciever,sender):
diff=sender[1:3]-reciever[1:3];
distance=norm(diff);
if(distance<1):
distance=1;
return G*reciever[0]*sender[0]/(distance**3)*diff;
def calc(cur_state,n_body):
next_state=np.zeros((n_body,fea_num),dtype=float);
f_mat=np.zeros((n_body,n_body,2),dtype=float);
f_sum=np.zeros((n_body,2),dtype=float);
acc=np.zeros((n_body,2),dtype=float);
for i in range(n_body):
for j in range(i+1,n_body):
if(j!=i):
f=get_f(cur_state[i][:3],cur_state[j][:3]);
f_mat[i,j]+=f;
f_mat[j,i]-=f;
f_sum[i]=np.sum(f_mat[i],axis=0);
acc[i]=f_sum[i]/cur_state[i][0];
next_state[i][0]=cur_state[i][0];
next_state[i][3:5]=cur_state[i][3:5]+acc[i]*diff_t;
next_state[i][1:3]=cur_state[i][1:3]+next_state[i][3:5]*diff_t;
return next_state;
def gen(n_body,orbit):
# initialization on just first state
data=init(total_state,n_body,fea_num,orbit);
for i in range(1,total_state):
data[i]=calc(data[i-1],n_body);
return data;
def make_video(xy,filename):
os.system("rm -rf pics/*");
FFMpegWriter = manimation.writers['ffmpeg']
metadata = dict(title='Movie Test', artist='Matplotlib',
comment='Movie support!')
writer = FFMpegWriter(fps=15, metadata=metadata)
fig = plt.figure()
plt.xlim(-200, 200)
plt.ylim(-200, 200)
fig_num=len(xy);
color=['ro','bo','go','ko','yo','mo','co'];
with writer.saving(fig, filename, len(xy)):
for i in range(len(xy)):
for j in range(len(xy[0])):
plt.plot(xy[i,j,1],xy[i,j,0],color[j%len(color)]);
writer.grab_frame();
if __name__=='__main__':
data=gen(3,True);
xy=data[:,:,1:3];
make_video(xy,"test.mp4");