Permalink
Browse files

Remove unit checks from *State classes

This leads to huge performance improvements, as `quantity_input` is very slow.
  • Loading branch information...
Juanlu001 committed Sep 13, 2016
1 parent bb4c1a5 commit ed931adbe2c2c5c3f0467d9cc800f62377c6e02a
@@ -5,4 +5,4 @@ norecursedirs =
build
python_files =
test_*.py
addopts = --benchmark-autosave
#addopts = --benchmark-autosave
@@ -99,11 +99,8 @@ class ClassicalState(BaseState):
"""State defined by its classical orbital elements.
"""
@u.quantity_input(a=u.m, ecc=u.one, inc=u.rad, raan=u.rad, argp=u.rad, nu=u.rad)
def __init__(self, attractor, a, ecc, inc, raan, argp, nu):
super(ClassicalState, self).__init__(attractor)
if ecc == 1.0:
raise ValueError("For parabolic orbits use Orbit.parabolic instead")
self._a = a
self._ecc = ecc
self._inc = inc
@@ -78,8 +78,7 @@ def from_classical(cls, attractor, a, ecc, inc, raan, argp, nu, epoch=J2000):
"""
if ecc == 1.0 * u.one:
raise ValueError("For parabolic orbits use "
"State.parabolic instead")
raise ValueError("For parabolic orbits use Orbit.parabolic instead")
elif not 0 * u.deg <= inc <= 180 * u.deg:
raise ValueError("Inclination must be between 0 and 180 degrees")

@@ -12,7 +12,7 @@
from poliastro.stumpff import c2, c3


def func_twobody(t0, u, k, ad):
def func_twobody(t0, u_, k, ad):
"""Differential equation for the initial value two body problem.
This function follows Cowell's formulation.
@@ -21,17 +21,17 @@ def func_twobody(t0, u, k, ad):
----------
t0 : float
Time.
u : ndarray
u_ : ndarray
Six component state vector [x, y, z, vx, vy, vz] (km, km/s).
k : float
Standard gravitational parameter.
ad : function(t0, u, k)
Non Keplerian acceleration (km/s2).
"""
ax, ay, az = ad(t0, u, k)
ax, ay, az = ad(t0, u_, k)

x, y, z, vx, vy, vz = u
x, y, z, vx, vy, vz = u_
r3 = (x**2 + y**2 + z**2)**1.5

du = np.array([
@@ -84,7 +84,7 @@ def cowell(k, r0, v0, tof, rtol=1e-10, *, ad=None, callback=None, nsteps=1000):

# Set the non Keplerian acceleration
if ad is None:
ad = lambda t0, u, k: (0, 0, 0)
ad = lambda t0, u_, k_: (0, 0, 0)

# Set the integrator
rr = ode(func_twobody).set_integrator('dop853', rtol=rtol, nsteps=nsteps)
@@ -67,7 +67,6 @@ class RVState(BaseState):
"""State defined by its position and velocity vectors.
"""
@u.quantity_input(r=u.m, v=u.m / u.s)
def __init__(self, attractor, r, v):
super(RVState, self).__init__(attractor)
self._r = r
@@ -20,6 +20,38 @@ def test_default_time_for_new_state():
assert ss.epoch == expected_epoch


def test_state_raises_unitserror_if_elements_units_are_wrong():
_d = 1.0 * u.AU # Unused distance
_ = 0.5 * u.one # Unused dimensionless value
_a = 1.0 * u.deg # Unused angle
wrong_angle = 1.0 * u.AU
with pytest.raises(u.UnitsError) as excinfo:
ss = Orbit.from_classical(Sun, _d, _, _a, _a, _a, wrong_angle)
assert ("UnitsError: Argument 'nu' to function 'from_classical' must be in units convertible to 'rad'."
in excinfo.exconly())


def test_state_raises_unitserror_if_rv_units_are_wrong():
_d = [1.0, 0.0, 0.0] * u.AU
wrong_v = [0.0, 1.0e-6, 0.0] * u.AU
with pytest.raises(u.UnitsError) as excinfo:
ss = Orbit.from_vectors(Sun, _d, wrong_v)
assert ("UnitsError: Argument 'v' to function 'from_vectors' must be in units convertible to 'm / s'."
in excinfo.exconly())


def test_parabolic_elements_fail_early():
attractor = Earth
ecc = 1.0 * u.one
_d = 1.0 * u.AU # Unused distance
_ = 0.5 * u.one # Unused dimensionless value
_a = 1.0 * u.deg # Unused angle
with pytest.raises(ValueError) as excinfo:
ss = Orbit.from_classical(attractor, _d, ecc, _a, _a, _a, _a)
assert ("ValueError: For parabolic orbits use Orbit.parabolic instead"
in excinfo.exconly())


def test_bad_inclination_raises_exception():
_d = 1.0 * u.AU # Unused distance
_ = 0.5 * u.one # Unused dimensionless value
@@ -48,45 +48,13 @@ def test_state_has_individual_elements():
assert ss.nu == nu


def test_state_raises_unitserror_if_elements_units_are_wrong():
_d = 1.0 * u.AU # Unused distance
_ = 0.5 * u.one # Unused dimensionless value
_a = 1.0 * u.deg # Unused angle
wrong_angle = 1.0 * u.AU
with pytest.raises(u.UnitsError) as excinfo:
ss = ClassicalState(Sun, _d, _, _a, _a, _a, wrong_angle)
assert ("UnitsError: Argument 'nu' to function '__init__' must be in units convertible to 'rad'."
in excinfo.exconly())


def test_state_has_rv_given_in_constructor():
r = [1.0, 0.0, 0.0] * u.AU
v = [0.0, 1.0e-6, 0.0] * u.AU / u.s
ss = RVState(Sun, r, v)
assert ss.rv() == (r, v)


def test_state_raises_unitserror_if_rv_units_are_wrong():
_d = [1.0, 0.0, 0.0] * u.AU
wrong_v = [0.0, 1.0e-6, 0.0] * u.AU
with pytest.raises(u.UnitsError) as excinfo:
ss = RVState(Sun, _d, wrong_v)
assert ("UnitsError: Argument 'v' to function '__init__' must be in units convertible to 'm / s'."
in excinfo.exconly())


def test_parabolic_elements_fail_early():
attractor = Earth
ecc = 1.0 * u.one
_d = 1.0 * u.AU # Unused distance
_ = 0.5 * u.one # Unused dimensionless value
_a = 1.0 * u.deg # Unused angle
with pytest.raises(ValueError) as excinfo:
ss = ClassicalState(attractor, _d, ecc, _a, _a, _a, _a)
assert ("ValueError: For parabolic orbits use Orbit.parabolic instead"
in excinfo.exconly())


def test_perigee_and_apogee():
expected_r_a = 500 * u.km
expected_r_p = 300 * u.km

0 comments on commit ed931ad

Please sign in to comment.