Skip to content

Commit 5d2c4f7

Browse files
moeyensjmoeyensj
committed
Clean up tests_kepler, make use of to_pandas()
1 parent 06a43fa commit 5d2c4f7

File tree

1 file changed

+114
-81
lines changed

1 file changed

+114
-81
lines changed

‎thor/orbits/tests/test_kepler.py

+114-81
Original file line numberDiff line numberDiff line change
@@ -10,114 +10,147 @@
1010
from ..kepler import _convertKeplerianToCartesian
1111

1212
MU = c.G * c.M_SUN
13+
TOL = 1e-15
14+
MAX_ITER = 100
15+
16+
EPOCHS = Time(
17+
["{}-02-02T00:00:00.000".format(i) for i in range(1993, 2050)],
18+
format="isot",
19+
scale="tdb"
20+
)
21+
ISO_EPOCHS = Time(
22+
["{}-02-02T00:00:00.000".format(i) for i in range(2017, 2022)],
23+
format="isot",
24+
scale="tdb"
25+
)
26+
TARGETS = [
27+
"Amor",
28+
"Eros",
29+
"Eugenia",
30+
"Ceres",
31+
]
32+
ISO_TARGETS = [
33+
"1I/2017 U1", # Oumuamua
34+
"C/2019 Q4" # Borisov
35+
]
36+
1337

1438
def test_convertOrbitalElements():
15-
epochs = Time(["{}-02-02T00:00:00.000".format(i) for i in range(1993, 2050)], format="isot", scale="tdb")
16-
targets = [
17-
"Amor",
18-
"Eros",
19-
"Eugenia",
20-
"Ceres"
21-
]
22-
for name in targets:
23-
target = Horizons(id=name, epochs=epochs.mjd)
24-
vectors = target.vectors()
25-
vectors = np.array(vectors["x", "y", "z", "vx", "vy", "vz"])
26-
vectors = vectors.view("float64").reshape(vectors.shape + (-1,))
39+
for name in TARGETS:
40+
target = Horizons(id=name, epochs=EPOCHS.mjd)
41+
42+
vectors = target.vectors().to_pandas()
43+
vectors = vectors[["x", "y", "z", "vx", "vy", "vz"]].values
44+
45+
elements = target.elements().to_pandas()
46+
elements = elements[["a", "e", "incl", "Omega", "w", "M", "nu"]].values
2747

28-
elements = target.elements()
29-
elements = np.array(elements["a", "e", "incl", "Omega", "w", "M", "nu"])
30-
elements = elements.view("float64").reshape(elements.shape + (-1,))
48+
np.testing.assert_allclose(
49+
elements[:, :6],
50+
convertOrbitalElements(
51+
vectors,
52+
"cartesian",
53+
"keplerian",
54+
max_iter=MAX_ITER,
55+
tol=TOL,
56+
)
57+
)
3158

32-
np.testing.assert_allclose(elements[:, :6], convertOrbitalElements(vectors, "cartesian", "keplerian"))
33-
np.testing.assert_allclose(vectors, convertOrbitalElements(elements[:, :6], "keplerian", "cartesian"))
59+
np.testing.assert_allclose(
60+
vectors,
61+
convertOrbitalElements(
62+
elements[:, :6],
63+
"keplerian",
64+
"cartesian",
65+
max_iter=MAX_ITER,
66+
tol=TOL,
67+
)
68+
)
69+
return
3470

3571
def test_convertCartesianToKeplerian_elliptical():
36-
epochs = Time(["{}-02-02T00:00:00.000".format(i) for i in range(1993, 2050)], format="isot", scale="tdb")
37-
targets = [
38-
"Amor",
39-
"Eros",
40-
"Eugenia",
41-
"Ceres",
42-
]
43-
44-
for name in targets:
45-
target = Horizons(id=name, epochs=epochs.mjd)
46-
vectors = target.vectors()
47-
vectors = np.array(vectors["x", "y", "z", "vx", "vy", "vz"])
48-
vectors = vectors.view("float64").reshape(vectors.shape + (-1,))
72+
for name in TARGETS:
73+
target = Horizons(id=name, epochs=EPOCHS.mjd)
74+
vectors = target.vectors().to_pandas()
75+
vectors = vectors[["x", "y", "z", "vx", "vy", "vz"]].values
4976

50-
elements = target.elements()
51-
elements = np.array(elements["a", "q", "e", "incl", "Omega", "w", "M", "nu"])
52-
elements = elements.view("float64").reshape(elements.shape + (-1,))
77+
elements = target.elements().to_pandas()
78+
elements = elements[["a", "q", "e", "incl", "Omega", "w", "M", "nu"]].values
5379

5480
for v, e in zip(vectors, elements):
55-
np.testing.assert_allclose(_convertCartesianToKeplerian(v.reshape(1, -1), mu=MU), e.reshape(1, -1))
81+
np.testing.assert_allclose(
82+
_convertCartesianToKeplerian(
83+
v.reshape(1, -1),
84+
mu=MU,
85+
),
86+
e.reshape(1, -1)
87+
)
88+
return
5689

5790
def test_convertCartesianToKeplerian_parabolic():
5891
warnings.warn("Need to implement and test parabolic conversions!!!")
92+
return
5993

6094
def test_convertCartesianToKeplerian_hyperbolic():
61-
epochs = Time(["{}-02-02T00:00:00.000".format(i) for i in range(2017, 2023)], format="isot", scale="tdb")
62-
iso_targets = [
63-
"1I/2017 U1", #Oumuamua
64-
"C/2019 Q4" #Borisov
65-
]
66-
67-
for name in iso_targets:
68-
target = Horizons(id=name, epochs=epochs.mjd)
69-
vectors = target.vectors()
70-
vectors = np.array(vectors["x", "y", "z", "vx", "vy", "vz"])
71-
vectors = vectors.view("float64").reshape(vectors.shape + (-1,))
95+
for name in ISO_TARGETS:
96+
target = Horizons(id=name, epochs=ISO_EPOCHS.mjd)
97+
vectors = target.vectors().to_pandas()
98+
vectors = vectors[["x", "y", "z", "vx", "vy", "vz"]].values
7299

73-
elements = target.elements()
74-
elements = np.array(elements["a", "q", "e", "incl", "Omega", "w", "M", "nu"])
75-
elements = elements.view("float64").reshape(elements.shape + (-1,))
100+
elements = target.elements().to_pandas()
101+
elements = elements[["a", "q", "e", "incl", "Omega", "w", "M", "nu"]].values
76102

77103
for v, e in zip(vectors, elements):
78-
np.testing.assert_allclose(_convertCartesianToKeplerian(v.reshape(1, -1), mu=MU), e.reshape(1, -1))
104+
np.testing.assert_allclose(
105+
_convertCartesianToKeplerian(
106+
v.reshape(1, -1),
107+
mu=MU,
108+
),
109+
e.reshape(1, -1)
110+
)
111+
return
79112

80113
def test_convertKeplerianToCartesian_elliptical():
81-
epochs = Time(["{}-02-02T00:00:00.000".format(i) for i in range(1993, 2050)], format="isot", scale="tdb")
82-
targets = [
83-
"Amor",
84-
"Eros",
85-
"Eugenia",
86-
"Ceres"
87-
]
88-
89-
for name in targets:
90-
target = Horizons(id=name, epochs=epochs.mjd)
91-
vectors = target.vectors()
92-
vectors = np.array(vectors["x", "y", "z", "vx", "vy", "vz"])
93-
vectors = vectors.view("float64").reshape(vectors.shape + (-1,))
114+
for name in TARGETS:
115+
target = Horizons(id=name, epochs=EPOCHS.mjd)
116+
vectors = target.vectors().to_pandas()
117+
vectors = vectors[["x", "y", "z", "vx", "vy", "vz"]].values
94118

95-
elements = target.elements()
96-
elements = np.array(elements["a", "e", "incl", "Omega", "w", "M", "nu"])
97-
elements = elements.view("float64").reshape(elements.shape + (-1,))
119+
elements = target.elements().to_pandas()
120+
elements = elements[["a", "e", "incl", "Omega", "w", "M", "nu"]].values
98121

99122
for v, e in zip(vectors, elements):
100-
np.testing.assert_allclose(_convertKeplerianToCartesian(e.reshape(1, -1), mu=MU, max_iter=100, tol=1e-15), v.reshape(1, -1))
123+
np.testing.assert_allclose(
124+
_convertKeplerianToCartesian(
125+
e.reshape(1, -1),
126+
mu=MU,
127+
max_iter=MAX_ITER,
128+
tol=TOL,
129+
),
130+
v.reshape(1, -1)
131+
)
132+
return
101133

102134
def test_convertKeplerianToCartesian_parabolic():
103135
warnings.warn("Need to implement and test parabolic conversions!!!")
104136

105137
def test_convertKeplerianToCartesian_hyperbolic():
106-
epochs = Time(["{}-02-02T00:00:00.000".format(i) for i in range(2017, 2023)], format="isot", scale="tdb")
107-
iso_targets = [
108-
"1I/2017 U1", #Oumuamua
109-
"C/2019 Q4" #Borisov
110-
]
111-
112-
for name in iso_targets:
113-
target = Horizons(id=name, epochs=epochs.mjd)
114-
vectors = target.vectors()
115-
vectors = np.array(vectors["x", "y", "z", "vx", "vy", "vz"])
116-
vectors = vectors.view("float64").reshape(vectors.shape + (-1,))
138+
for name in ISO_TARGETS:
139+
target = Horizons(id=name, epochs=ISO_EPOCHS.mjd)
140+
vectors = target.vectors().to_pandas()
141+
vectors = vectors[["x", "y", "z", "vx", "vy", "vz"]].values
117142

118-
elements = target.elements()
119-
elements = np.array(elements["a", "e", "incl", "Omega", "w", "M"])
120-
elements = elements.view("float64").reshape(elements.shape + (-1,))
143+
elements = target.elements().to_pandas()
144+
elements = elements[["a", "e", "incl", "Omega", "w", "M", "nu"]].values
121145

122146
for v, e in zip(vectors, elements):
123-
np.testing.assert_allclose(_convertKeplerianToCartesian(e.reshape(1, -1), mu=MU, max_iter=100, tol=1e-15), v.reshape(1, -1))
147+
np.testing.assert_allclose(
148+
_convertKeplerianToCartesian(
149+
e.reshape(1, -1),
150+
mu=MU,
151+
max_iter=MAX_ITER,
152+
tol=TOL,
153+
),
154+
v.reshape(1, -1)
155+
)
156+
return

0 commit comments

Comments
 (0)