forked from skulumani/asteroid_dumbbell
/
relative_driver.py
233 lines (187 loc) · 10 KB
/
relative_driver.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
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
223
224
225
226
227
228
229
230
231
232
233
from __future__ import absolute_import, division, print_function, unicode_literals
import dynamics.asteroid as asteroid
import dynamics.dumbbell as dumbbell
import kinematics.attitude as attitude
from visualization import plotting
from eom_comparison import transform
import pdb
import numpy as np
from scipy import integrate
import argparse
# valid initial condtions for castalia
# these are in the asteroid fixed frame
periodic_pos = np.array([1.495746722510590,0.000001002669660,0.006129720493607])
periodic_vel = np.array([0.000000302161724,-0.000899607989820,-0.000000013286327])
AbsTol = 1e-9
RelTol = 1e-9
def hamilton_eoms_driver(initial_state, time, ast, dum):
"""ODEINT Hamilton Relative EOMs driver
This function will simulate the relative hamiltonian equations of motion using
scipy.integrate.odeint and then convert them into the inertial and
asteroid frames.
Parameters
----------
initial_state - (18,) numpy ndarray representing the initial state of the dumbbell
pos - initial_state[0:3] in km position of the dumbbell with respect to the
asteroid and defined in the asteroid frame
lin_mom - initial_state[3:6] in Jkm/sec is the linear momentum of dumbbell wrt the asteroid
and defined in the asteroid frame
R - initial_state[6:15] rotation matrix which converts vectors from the dumbbell frame to
the asteroid frame
ang_mom - state[15:18] rad/sec angular momentum of the dumbbell wrt inertial frame and
defined in the asteroid frame
time - (n,) numpy ndarray representing the simulation time span
ast - asteroid object instance
dum - dumbbell object instance
Returns
-------
time - (n,) numpy ndarray time vector. Same as the input
inertial_state - (n,18) ndarray with the simulation results represented in the inertial frame
pos - inertial_state[0:3] in km position of the dumbbell with respect to the
asteroid and defined in the inertial frame
vel - inertial_state[3:6] in km/sec is the velocity of dumbbell wrt the asteroid
and defined in the inertial frame
R - inertial_state[6:15] rotation matrix which converts vectors from the dumbbell frame to
the inertial frame
w - inertial_state[15:18] rad/sec angular velocity of the dumbbell wrt inertial frame and
defined in the inertial frame
asteroid_state - (n,18) ndarray with the simulation results represented in the asteroid frame
pos - asteroid_state[0:3] in km position of the dumbbell with respect to the
asteroid and defined in the asteroid frame
vel - asteroid_state[3:6] in km/sec is the velocity of dumbbell wrt the asteroid
and defined in the asteroid frame
R - asteroid_state[6:15] rotation matrix which converts vectors from the dumbbell frame to
the asteroid frame
w - asteroid_state[15:18] rad/sec angular velocity of the dumbbell wrt inertial frame and
defined in the asteroid frame
ast_state - (n,18) ndarray with the simulation results represented in the asteroid frame
pos - asteroid_state[0:3] in km position of the dumbbell with respect to the
asteroid and defined in the inertial frame
vel - asteroid_state[3:6] in km/sec is the velocity of dumbbell wrt the asteroid
and defined in the inertial frame
R - asteroid_state[6:15] rotation matrix which converts vectors from the dumbbell frame to
the inertial frame
w - asteroid_state[15:18] rad/sec angular momentum of the dumbbell wrt inertial frame and
defined in the asteroid frame
This output is exactly what dum.eoms_hamilton_relative will output without any alduteration
"""
ast_state = integrate.odeint(dum.eoms_hamilton_relative, initial_state, time, args=(ast,), atol=AbsTol, rtol=RelTol)
# convert to inertial and asteroid frames
inertial_state = transform.eoms_hamilton_relative_to_inertial(time, ast_state, ast, dum)
asteroid_state = transform.eoms_hamilton_relative_to_asteroid(time, ast_state, ast, dum)
return (time, inertial_state, asteroid_state, ast_state)
def compute_energy(file_name):
with np.load(file_name, allow_pickle=True) as data:
ast_name = data['ast_name']
num_faces = data['num_faces']
time = data['time']
asteroid_state = data['asteroid_state']
dum = data['dum'][()]
ast = data['ast'][()]
KE, PE = dum.relative_energy(time,asteroid_state,ast)
return KE, PE
def hamilton_eoms_energy_behavior(initial_state, time, ast, dum):
"""See \Delta E for varying tolerances of the ODE function
"""
tol_array = np.logspace(-4, -12, 9)
time_dict = {}
ast_dict = {}
inertial_state_dict = {}
asteroid_state_dict = {}
KE_dict = {}
PE_dict = {}
for tol in tol_array:
print('Tolerance - %4.2e' % tol)
(_, inertial_state, asteroid_state, ast_state) = hamilton_eoms_driver(initial_state, time, ast, dum)
KE, PE = dum.relative_energy(time, asteroid_state, ast)
time_dict[str(tol)] = time
KE_dict[str(tol)] = KE
PE_dict[str(tol)] = PE
asteroid_state_dict[str(tol)] = ast_state
inertial_state_dict[str(tol)] = inertial_state
asteroid_state_dict[str(tol)] = asteroid_state
return (time_dict, KE_dict, PE_dict, asteroid_state_dict, inertial_state_dict, asteroid_state_dict)
def relative_sim_plotter(file_name, mode):
"""Plot the simulation results
"""
# load the data
with np.load(file_name, allow_pickle=True) as data:
if mode == 0:
time = data['time']
KE = data['KE']
PE = data['PE']
# compute the
e_fig = plotting.plt.figure()
plotting.plot_energy(time,KE, PE, e_fig)
# plotting.animate_relative_trajectory(time, state, ast, dum, file_name[:-4])
elif mode == 1:
time_dict = data['time_dict'][()]
KE_dict = data['KE_dict'][()]
PE_dict = data['PE_dict'][()]
e_fig = plotting.plt.figure()
for tol in time_dict:
E = KE_dict[tol] + PE_dict[tol]
Ediff = np.absolute(E - E[0])
plotting.plt.plot(time_dict[tol], Ediff, label=tol)
print("Tol - %s, DeltaE - %4e" % (tol, E[-1]))
plotting.plt.legend()
else:
print("Incorrect mode.")
return 1
plotting.plt.show()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Relative EOMs simulator for a dumbbell around an asteroid')
parser.add_argument('ast_name', help="String - asteroid name: castalia or itokawa", type=str)
parser.add_argument('num_faces', help='Integer - Number of faces in polyhedron model', type=int)
parser.add_argument('tf', help='Float - Terminal time for simulation', type=float)
parser.add_argument('num_steps', help='Float - Number of steps in integration', type=float)
parser.add_argument('file_name', help='String - Filename for npz archive', type=str)
parser.add_argument("-m", "--mode", type=int, choices=[0, 1],
help="Choose which relative energy mode to run. 0 - relative energy, 1 - \Delta E behavior")
args = parser.parse_args()
print("Starting the simulation...")
# instantiate the asteroid and dumbbell objects
ast = asteroid.Asteroid(args.ast_name, args.num_faces)
dum = dumbbell.Dumbbell(m1=500, m2=500, l=0.003)
# initialize simulation parameters
time = np.linspace(0, args.tf, args.num_steps)
initial_pos = periodic_pos
initial_R = attitude.rot2(0).reshape(9)
initial_w = np.array([0.01, 0.01, 0.01])
initial_lin_mom = (dum.m1 + dum.m2) * (periodic_vel + attitude.hat_map(ast.omega*np.array([0,0,1])).dot(initial_pos))
initial_ang_mom = initial_R.reshape((3,3)).dot(dum.J).dot(initial_w)
initial_ham_state = np.hstack((initial_pos, initial_lin_mom, initial_R, initial_ang_mom))
print("This will run a long simulation with the following parameters!")
print(" ast_name - %s" % args.ast_name)
print(" num_faces - %s" % args.num_faces)
print(" tf - %s" % args.tf)
print(" num_steps - %s" % args.num_steps)
print(" file_name - %s" % args.file_name)
print("")
if args.mode == 0:
(_, inertial_state, asteroid_state, ast_state) = hamilton_eoms_driver(initial_ham_state, time, ast, dum)
print("Finished the simulation...")
print("Saving to npz file")
np.savez(args.file_name,inertial_state=inertial_state, time=time, asteroid_state=asteroid_state,
ast_state=ast_state, ast_name=args.ast_name, num_steps=args.num_steps,tf=args.tf, num_faces=args.num_faces,
ast=ast, dum=dum)
print("Now working to calculate the KE/PE of the simulation!")
KE, PE = compute_energy(args.file_name)
print("Finished with energy computations!")
np.savez('hamilton_relative_energy_' + args.file_name, args.file_name,inertial_state=inertial_state, time=time, asteroid_state=asteroid_state,
ast_state=ast_state, ast_name=args.ast_name, num_steps=args.num_steps,tf=args.tf, num_faces=args.num_faces,
ast=ast, dum=dum, KE=KE, PE=PE)
print("All finished!")
elif args.mode == 1:
print("We're running many simulations wiht varying tolerances.")
(time_dict, KE_dict, PE_dict, ast_state_dict, inertial_state_dict, asteroid_state_dict) \
= hamilton_eoms_energy_behavior(initial_ham_state, time, ast, dum)
print("Finished with simulations. Now saving to data file")
np.savez('hamilton_relative_energy_behavior_' + args.file_name, KE_dict=KE_dict,
PE_dict=PE_dict, time_dict=time_dict, ast_name=args.ast_name, num_steps=args.num_steps,
tf=args.tf, num_faces=args.num_faces, ast_state_dict=ast_state_dict,
inertial_state_dict=inertial_state_dict, asteroid_state_dict=asteroid_state_dict,
ast=ast, dum=dum)
print("All finished!")
else:
print("Missing mode argument.")