New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WIP:Jitting newton #413

Merged
merged 6 commits into from Jul 20, 2018

Conversation

2 participants
@nikita-astronaut
Copy link
Contributor

nikita-astronaut commented Jul 19, 2018

First we need to figure out the root of evil...

nikita-astronaut added some commits Jul 19, 2018

@codecov

This comment has been minimized.

Copy link

codecov bot commented Jul 19, 2018

Codecov Report

Merging #413 into master will decrease coverage by 0.42%.
The diff coverage is 90.81%.

Impacted file tree graph

@@            Coverage Diff             @@
##           master     #413      +/-   ##
==========================================
- Coverage   85.31%   84.88%   -0.43%     
==========================================
  Files          39       41       +2     
  Lines        1947     1992      +45     
  Branches      149      152       +3     
==========================================
+ Hits         1661     1691      +30     
- Misses        249      261      +12     
- Partials       37       40       +3
Impacted Files Coverage Δ
src/poliastro/integrators.py 90.29% <ø> (-0.08%) ⬇️
src/poliastro/core/jit.py 100% <ø> (ø)
...oliastro/twobody/thrust/change_ecc_quasioptimal.py 100% <100%> (ø) ⬆️
src/poliastro/twobody/orbit.py 95.62% <100%> (ø) ⬆️
src/poliastro/hyper.py 100% <100%> (ø) ⬆️
src/poliastro/iod/vallado.py 86% <100%> (ø) ⬆️
src/poliastro/twobody/perturbations.py 100% <100%> (ø) ⬆️
src/poliastro/twobody/thrust/change_inc_ecc.py 100% <100%> (ø) ⬆️
src/poliastro/util.py 100% <100%> (ø) ⬆️
src/poliastro/twobody/thrust/change_argp.py 100% <100%> (ø) ⬆️
... and 11 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 533b429...8b5ab77. Read the comment docs.

@Juanlu001
Copy link
Member

Juanlu001 left a comment

This is in the right path, but I suggested more performance tricks. Bear in mind that creating a Quantity object is super slow, so the more we can avoid it the better.

@@ -150,7 +35,7 @@ def D_to_nu(D, ecc):
"Robust resolution of Kepler’s equation in all eccentricity regimes."
Celestial Mechanics and Dynamical Astronomy 116, no. 1 (2013): 21-34.
"""
return 2.0 * np.arctan(D)
return D_to_nu_fast(D.to(u.rad).value, ecc.to(u.one).value) * u.rad

This comment has been minimized.

@Juanlu001

Juanlu001 Jul 19, 2018

Member

You can do ecc.value directly

@@ -346,7 +346,8 @@ def _sample(self, time_values, method=mean_motion):

def _generate_time_values(self, nu_vals):
# Subtract current anomaly to start from the desired point
M_vals = nu_to_M(nu_vals, self.ecc) - nu_to_M(self.nu, self.ecc)
M_vals = np.array([nu_to_M(nu_val, self.ecc).value -

This comment has been minimized.

@Juanlu001

Juanlu001 Jul 19, 2018

Member

This for loop in a list comprehension might be a bit slow, wouldn't it be better to use the @vectorize decorator in the corresponding angles functions instead? https://numba.pydata.org/numba-doc/0.39.0/user/vectorize.html#the-vectorize-decorator (https://numba.pydata.org/numba-doc/0.39.0/reference/jit-compilation.html#numba.vectorize)

M = M0 + tof * np.sqrt(k / np.abs(a ** 3)) * u.rad
nu = M_to_nu(M, ecc)
M = M0 + tof * np.sqrt(k / np.abs(a ** 3)) * u.rad
nu = M_to_nu(M, ecc * u.one)

This comment has been minimized.

@Juanlu001

Juanlu001 Jul 19, 2018

Member

Here you could use the fast version like M_to_nu(M.to(u.rad).value, ecc)

M = M0 + tof * np.sqrt(k / 2.0 / (q ** 3))
nu = M_to_nu(M, ecc)
M = M0 + tof * np.sqrt(k / 2.0 / (q ** 3)) * u.rad
nu = M_to_nu(M, ecc * u.one)

This comment has been minimized.

@Juanlu001

Juanlu001 Jul 19, 2018

Member

Same here!

@@ -140,23 +140,21 @@ def mean_motion(orbit, tof, **kwargs):
p, ecc, inc, raan, argp, nu0 = rv2coe(k, r0, v0)

# get the initial mean anomaly
M0 = nu_to_M(nu0, ecc)
M0 = nu_to_M(nu0 * u.rad, ecc * u.one)

This comment has been minimized.

@Juanlu001

Juanlu001 Jul 19, 2018

Member

Instead of creating units here, why not use nu_to_M from poliastro.core.angles now? 😉

@Juanlu001

This comment has been minimized.

Copy link
Member

Juanlu001 commented Jul 19, 2018

I gave it a rough try (I have a very bad connection at this moment so this took a while) but if you manage to @jit mean_motion we can cut the time by half. Here's my patch:

From 0299ce68f84c6b83743a369724764c7acb64794d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juan=20Luis=20Cano=20Rodr=C3=ADguez?= <juanlu001@gmail.com>
Date: Thu, 19 Jul 2018 19:33:47 +0000
Subject: [PATCH] WIP: jit mean_motion

---
 src/poliastro/core/util.py           | 59 ++++++++++++++++++++++++++++
 src/poliastro/twobody/classical.py   | 17 ++++----
 src/poliastro/twobody/propagation.py | 21 ++++++----
 src/poliastro/twobody/rv.py          | 28 +++++++++----
 4 files changed, 103 insertions(+), 22 deletions(-)
 create mode 100644 src/poliastro/core/util.py

diff --git src/poliastro/core/util.py src/poliastro/core/util.py
new file mode 100644
index 0000000..5f48538
--- /dev/null
+++ src/poliastro/core/util.py
@@ -0,0 +1,59 @@
+import numpy as np
+from numpy import sin, cos
+
+from poliastro.core.jit import jit
+
+@jit
+def rotate(vec, ax, angle):
+    """Rotates the coordinate system around axis x, y or z a CCW angle.
+    Parameters
+    ----------
+    vec : ndarray
+        Dimension 3 vector.
+    ax : int
+        Axis to be rotated.
+    angle : float
+        Angle of rotation (rad).
+    Notes
+    -----
+    This performs a so-called active or alibi transformation: rotates the
+    vector while the coordinate system remains unchanged. To do the opposite
+    operation (passive or alias transformation) call the function as
+    `rotate(vec, ax, -angle)` or use the convenience function `transform`,
+    see `[1]_`.
+    References
+    ----------
+    .. [1] http://en.wikipedia.org/wiki/Rotation_matrix#Ambiguities
+    """
+    assert vec.shape == (3,)
+
+    rot = np.eye(3)
+    if ax == 0:
+        sl = slice(1, 3, 1)
+    elif ax == 1:
+        sl = slice(0, 3, 2)
+    elif ax == 2:
+        sl = slice(0, 2, 1)
+    else:
+        raise ValueError("Invalid axis: must be one of 0, 1 or 2")
+
+    rot[sl, sl] = np.array((
+        (cos(angle), -sin(angle)),
+        (sin(angle), cos(angle))
+    ))
+    if ax == 1:
+        rot = rot.T
+
+    # np.dot() arguments must all have the same dtype
+    return np.dot(rot, vec.astype(rot.dtype))
+
+
+@jit
+def transform(vec, ax, angle):
+    """Rotates a coordinate system around axis a positive right-handed angle.
+    Notes
+    -----
+    This is a convenience function, equivalent to `rotate(vec, ax, -angle)`.
+    Refer to the documentation of that function for further information.
+    """
+    return rotate(vec, ax, -angle)
diff --git src/poliastro/twobody/classical.py src/poliastro/twobody/classical.py
index f6ecf0e..3966691 100644
--- src/poliastro/twobody/classical.py
+++ src/poliastro/twobody/classical.py
@@ -5,7 +5,8 @@ import numpy as np
 from numpy import sin, cos, sqrt
 from astropy import units as u
 
-from poliastro.util import transform
+from poliastro.core.jit import jit
+from poliastro.core.util import transform
 
 from poliastro.twobody import rv
 from poliastro.twobody import equinoctial
@@ -13,6 +14,7 @@ from poliastro.twobody import equinoctial
 from ._base import BaseState
 
 
+@jit
 def rv_pqw(k, p, ecc, nu):
     """Returns r and v vectors in perifocal frame.
 
@@ -22,6 +24,7 @@ def rv_pqw(k, p, ecc, nu):
     return r_pqw, v_pqw
 
 
+@jit
 def coe2rv(k, p, ecc, inc, raan, argp, nu):
     """Converts from classical orbital elements to vectors.
 
@@ -45,12 +48,12 @@ def coe2rv(k, p, ecc, inc, raan, argp, nu):
     """
     r_pqw, v_pqw = rv_pqw(k, p, ecc, nu)
 
-    r_ijk = transform(r_pqw, -argp, 'z', u.rad)
-    r_ijk = transform(r_ijk, -inc, 'x', u.rad)
-    r_ijk = transform(r_ijk, -raan, 'z', u.rad)
-    v_ijk = transform(v_pqw, -argp, 'z', u.rad)
-    v_ijk = transform(v_ijk, -inc, 'x', u.rad)
-    v_ijk = transform(v_ijk, -raan, 'z', u.rad)
+    r_ijk = transform(r_pqw, 2, -argp)
+    r_ijk = transform(r_ijk, 0, -inc)
+    r_ijk = transform(r_ijk, 2, -raan)
+    v_ijk = transform(v_pqw, 2, -argp)
+    v_ijk = transform(v_ijk, 0, -inc)
+    v_ijk = transform(v_ijk, 2, -raan)
 
     return r_ijk, v_ijk
 
diff --git src/poliastro/twobody/propagation.py src/poliastro/twobody/propagation.py
index e45e246..b891fb9 100644
--- src/poliastro/twobody/propagation.py
+++ src/poliastro/twobody/propagation.py
@@ -10,7 +10,7 @@ from poliastro.integrators import DOP835
 from astropy import units as u
 from poliastro.twobody.rv import rv2coe
 from poliastro.twobody.classical import coe2rv
-from poliastro.twobody.angles import nu_to_M, M_to_nu
+from poliastro.core.angles import nu_to_M as nu_to_M_fast, M_to_nu as M_to_nu_fast
 
 from poliastro.core.jit import jit
 from poliastro.stumpff import c2, c3
@@ -136,27 +136,32 @@ def mean_motion(orbit, tof, **kwargs):
     r0 = orbit.r.to(u.km).value
     v0 = orbit.v.to(u.km / u.s).value
 
+    return _mean_motion(k, r0, v0, tof)
+
+
+@jit
+def _mean_motion(k, r0, v0, tof):
     # get the initial true anomaly and orbit parameters that are constant over time
     p, ecc, inc, raan, argp, nu0 = rv2coe(k, r0, v0)
 
     # get the initial mean anomaly
-    M0 = nu_to_M(nu0 * u.rad, ecc * u.one)
+    M0 = nu_to_M_fast(nu0, ecc)
     # strong elliptic or strong hyperbolic orbits
     if np.abs(ecc - 1.0) > 1e-2:
         a = p / (1.0 - ecc ** 2)
         # given the initial mean anomaly, calculate mean anomaly
         # at the end, mean motion (n) equals sqrt(mu / |a^3|)
-        M = M0 + tof * np.sqrt(k / np.abs(a ** 3)) * u.rad
-        nu = M_to_nu(M, ecc * u.one)
+        M = M0 + tof * np.sqrt(k / np.abs(a ** 3))
+        nu = M_to_nu_fast(M, ecc)
 
     # near-parabolic orbit
     else:
         q = p * np.abs(1.0 - ecc) / np.abs(1.0 - ecc ** 2)
         # mean motion n = sqrt(mu / 2 q^3) for parabolic orbit
-        M = M0 + tof * np.sqrt(k / 2.0 / (q ** 3)) * u.rad
-        nu = M_to_nu(M, ecc * u.one)
-    with u.set_enabled_equivalencies(u.dimensionless_angles()):
-        return coe2rv(k, p, ecc, inc, raan, argp, nu)
+        M = M0 + tof * np.sqrt(k / 2.0 / (q ** 3))
+        nu = M_to_nu_fast(M, ecc)
+
+    return coe2rv(k, p, ecc, inc, raan, argp, nu)
 
 
 def kepler(orbit, tof, *, numiter=350, **kwargs):
diff --git src/poliastro/twobody/rv.py src/poliastro/twobody/rv.py
index a70de6f..ae8581d 100644
--- src/poliastro/twobody/rv.py
+++ src/poliastro/twobody/rv.py
@@ -4,13 +4,27 @@
 import numpy as np
 from astropy import units as u
 
-from poliastro.util import norm
+from poliastro.core.jit import jit
+from poliastro.util import norm_fast as norm
 
 from poliastro.twobody import classical
 
 from ._base import BaseState
 
 
+@jit
+def cross(a, b):
+    """Computes cross product between two vectors"""
+    # np.cross is not supported in numba nopython mode, see
+    # https://github.com/numba/numba/issues/2978
+    return np.array((
+        a[1] * b[2] - a[2] * b[1],
+        a[2] * b[0] - a[0] * b[2],
+        a[0] * b[1] - a[1] * b[0]
+    ))
+
+
+@jit
 def rv2coe(k, r, v, tol=1e-8):
     """Converts from vectors to classical orbital elements.
 
@@ -26,8 +40,8 @@ def rv2coe(k, r, v, tol=1e-8):
         Tolerance for eccentricity and inclination checks, default to 1e-8.
 
     """
-    h = np.cross(r, v)
-    n = np.cross([0, 0, 1], h) / norm(h)
+    h = cross(r, v)
+    n = cross([0, 0, 1], h) / norm(h)
     e = ((v.dot(v) - k / (norm(r))) * r - r.dot(v) * v) / k
     ecc = norm(e)
     p = h.dot(h) / k
@@ -39,13 +53,13 @@ def rv2coe(k, r, v, tol=1e-8):
     if equatorial and not circular:
         raan = 0
         argp = np.arctan2(e[1], e[0]) % (2 * np.pi)  # Longitude of periapsis
-        nu = (np.arctan2(h.dot(np.cross(e, r)) / norm(h), r.dot(e)) %
+        nu = (np.arctan2(h.dot(cross(e, r)) / norm(h), r.dot(e)) %
               (2 * np.pi))
     elif not equatorial and circular:
         raan = np.arctan2(n[1], n[0]) % (2 * np.pi)
         argp = 0
         # Argument of latitude
-        nu = (np.arctan2(r.dot(np.cross(h, n)) / norm(h), r.dot(n)) %
+        nu = (np.arctan2(r.dot(cross(h, n)) / norm(h), r.dot(n)) %
               (2 * np.pi))
     elif equatorial and circular:
         raan = 0
@@ -53,9 +67,9 @@ def rv2coe(k, r, v, tol=1e-8):
         nu = np.arctan2(r[1], r[0]) % (2 * np.pi)  # True longitude
     else:
         raan = np.arctan2(n[1], n[0]) % (2 * np.pi)
-        argp = (np.arctan2(e.dot(np.cross(h, n)) / norm(h), e.dot(n)) %
+        argp = (np.arctan2(e.dot(cross(h, n)) / norm(h), e.dot(n)) %
                 (2 * np.pi))
-        nu = (np.arctan2(r.dot(np.cross(h, e)) / norm(h), r.dot(e))
+        nu = (np.arctan2(r.dot(cross(h, e)) / norm(h), r.dot(e))
               % (2 * np.pi))
 
     return p, ecc, inc, raan, argp, nu
-- 
2.17.1

nikita-astronaut added some commits Jul 20, 2018

@Juanlu001

This comment has been minimized.

Copy link
Member

Juanlu001 commented Jul 20, 2018

I'm merging this branch as it is, then I will make some changes on my side. Thanks a lot!

@Juanlu001 Juanlu001 merged commit 5d5c1b9 into poliastro:master Jul 20, 2018

9 of 10 checks passed

codeclimate 3 issues to fix
Details
ci/circleci: coverage Your tests passed on CircleCI!
Details
ci/circleci: docs Your tests passed on CircleCI!
Details
ci/circleci: quality Your tests passed on CircleCI!
Details
ci/circleci: test_py35 Your tests passed on CircleCI!
Details
ci/circleci: test_py36 Your tests passed on CircleCI!
Details
codecov/patch 90.81% of diff hit (target 85.31%)
Details
codecov/project Absolute coverage decreased by -0.42% but relative coverage increased by +5.5% compared to 533b429
Details
continuous-integration/appveyor/pr AppVeyor build succeeded
Details
continuous-integration/travis-ci/pr The Travis CI build passed
Details

@nikita-astronaut nikita-astronaut deleted the nikita-astronaut:jitting_newton_2 branch Jul 21, 2018

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment