-
Notifications
You must be signed in to change notification settings - Fork 0
/
read_orbits.py
executable file
·123 lines (102 loc) · 3.79 KB
/
read_orbits.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
'''
Read NASA planetary orbits, download from https://ssd.jpl.nasa.gov/horizons.cgi#top
with format X, Y, Z, VX, VY, VZ format them, and return two data vectors, one
containing positions and one containing accelerations.
For the future: Create a class that contains orbits, names and masses. Will
improve appearance.
Written by Pablo Lemos (UCL)
Nov 2020
'''
from data.solar_system_names import *
from simulator.base_classes import *
import os
import pickle
def read_orbit(name, path):
''' Reads the data for a single orbit
Parameters:
-----------
name: str
the name of the body
path: string
the path to the orbit file
Returns:
--------
orbit_data : np.array
a numpy array containing the positions and accelerations for the body
'''
try:
orbit = np.loadtxt(os.path.join(path, name + '.txt'), usecols=[2, 3, 4, 5, 6, 7],
unpack=True, delimiter=',')
except IndexError:
orbit = np.genfromtxt(os.path.join(path, name + '.txt'), usecols=[2, 3, 4, 5, 6, 7],
unpack=True, delimiter=',',
skip_header=22, skip_footer=31
)
return orbit.T
def main(nplanets=0,
path=None,
read_data=True):
''' Reads the data files and returns a numpy array with the orbits
Parameters:
-----------
nplanets: int
the number of planets to be used. If 0, use all the planets.
Defaults to 0
path: string
the path to the orbit files. It should contain a folder: 'barycenter'
read_data: bool
whether to read data files, or just masses and names, Defauls to true
Returns:
--------
orbits_data : np.array()
a numpy array containing the positions and accelerations for each
body.
masses : np.array()
the masses of all bodies
names_ls : ls
the names of all bodies
'''
if not path:
dir_path = os.path.dirname(os.path.realpath(__file__))
path = os.path.join(dir_path, '../data/')
print('Reading data in Solar System barycenter reference frame')
path = os.path.join(path, 'barycenter')
# If using the default value for nplanets, we use all 8
if nplanets == 0:
nplanets = 8
sun = Body(name='sun', mass=sun_mass)
if read_data:
# Read the sun's orbit
orbit_sun = read_orbit('sun', path)
# Create a list that will contain all the orbits
sun.add_trajectory(orbit_sun)
# Create a list that will contain all the bodies
bodies = [sun]
for i in range(nplanets):
name = planet_names[i]
planet = Body(mass=planet_masses[i], name=name)
if read_data:
print('Reading data for', name)
orbit = read_orbit(name, path)
planet.add_trajectory(orbit)
bodies.append(planet)
#Ignore Moons for now because of time concerns
#if name in planets_with_moons:
# j = planets_with_moons.index(name)
# for (name_moon, mass_moon) in zip(moon_names[j], moon_masses[j]):
# moon = Body(mass=mass_moon, name=name_moon)
# if read_data:
# print('Reading data for', name_moon)
# orbit = read_orbit(name_moon, path)
# moon.add_trajectory(orbit)
# bodies.append(moon)
return StarSystem(bodies)
if __name__ == "__main__":
system = main(nplanets=8,path='./planets_data/' )
dir_path = os.path.dirname(os.path.realpath(__file__))
#path = os.path.join(dir_path, '../data/debug.pkl')
print("Writing StarSystem as pickle file")
path = os.path.join(dir_path, './planets_data/solar_system_data.pkl')
file = open(path, 'wb')
pickle.dump(system, file)