# poliastro/poliastro

src/poliastro/core/elements.py

### Summary

B
5 hrs
###### Test Coverage
"""This module contains a set of functions that can be used
to convert between different elements that define the orbit
of a body.
"""

import numpy as np
from numpy import cross
from numpy.core.umath import cos, sin, sqrt

from poliastro.core.angles import E_to_nu, F_to_nu
from poliastro.core.util import norm, rotation_matrix

from ._jit import jit

@jit
def rv_pqw(k, p, ecc, nu):
r"""Returns r and v vectors in perifocal frame.

.. math::

\vec{r} = \frac{h^2}{\mu}\frac{1}{1 + e\cos(\theta)}\begin{bmatrix}
\cos(\theta)\\
\sin(\theta)\\
0
\end{bmatrix} \\\\\\

\vec{v} = \frac{h^2}{\mu}\begin{bmatrix}
-\sin(\theta)\\
e+\cos(\theta)\\
0
\end{bmatrix}

Parameters
----------
k : float
Standard gravitational parameter (km^3 / s^2).
p : float
Semi-latus rectum or parameter (km).
ecc : float
Eccentricity.
nu: float

Returns
-------

r: ndarray
Position. Dimension 3 vector
v: ndarray
Velocity. Dimension 3 vector

Examples
--------
>>> from poliastro.constants import GM_earth
>>> k = GM_earth #Earth gravitational parameter

>>> ecc = 0.3 #Eccentricity
>>> h = 60000e6 #Angular momentum of the orbit [m^2]/[s]
>>> p = h**2 / k #Parameter of the orbit
>>> r, v = rv_pqw(k, p, ecc, nu)

>>> #Printing the results
r = [-5312706.25105345  9201877.15251336    0] [m]
v = [-5753.30180931 -1328.66813933  0] [m]/[s]

Note
----

These formulas can be checked at Curtis 3rd. Edition, page 110. Also the
example proposed is 2.11 of Curtis 3rd Edition book.
"""
r_pqw = (np.array([cos(nu), sin(nu), 0 * nu]) * p / (1 + ecc * cos(nu))).T
v_pqw = (np.array([-sin(nu), (ecc + cos(nu)), 0]) * sqrt(k / p)).T
return r_pqw, v_pqw

@jit
def coe_rotation_matrix(inc, raan, argp):
"""Create a rotation matrix for coe transformation
"""
r = rotation_matrix(raan, 2)
r = r @ rotation_matrix(inc, 0)
r = r @ rotation_matrix(argp, 2)
return r

@jit
def coe2rv(k, p, ecc, inc, raan, argp, nu):
r"""Converts from classical orbital to state vectors.

Classical orbital elements are converted into position and velocity
vectors by rv_pqw algorithm. A rotation matrix is applied to position
and velocity vectors to get them expressed in terms of an IJK basis.

.. math::
\begin{align}
\vec{r}_{IJK} &= [ROT3(-\Omega)][ROT1(-i)][ROT3(-\omega)]\vec{r}_{PQW}
= \left [ \frac{IJK}{PQW} \right ]\vec{r}_{PQW}\\
\vec{v}_{IJK} &= [ROT3(-\Omega)][ROT1(-i)][ROT3(-\omega)]\vec{v}_{PQW}
= \left [ \frac{IJK}{PQW} \right ]\vec{v}_{PQW}\\
\end{align}

Previous rotations (3-1-3) can be expressed in terms of a single rotation matrix:

.. math::
\left [ \frac{IJK}{PQW} \right ]

.. math::
\begin{bmatrix}
\cos(\Omega)\cos(\omega) - \sin(\Omega)\sin(\omega)\cos(i) & -\cos(\Omega)\sin(\omega) - \sin(\Omega)\cos(\omega)\cos(i) & \sin(\Omega)\sin(i)\\
\sin(\Omega)\cos(\omega) + \cos(\Omega)\sin(\omega)\cos(i) & -\sin(\Omega)\sin(\omega) + \cos(\Omega)\cos(\omega)\cos(i) & -\cos(\Omega)\sin(i)\\
\sin(\omega)\sin(i) & \cos(\omega)\sin(i) & \cos(i)
\end{bmatrix}

Parameters
----------
k : float
Standard gravitational parameter (km^3 / s^2).
p : float
Semi-latus rectum or parameter (km).
ecc : float
Eccentricity.
inc : float
omega : float
argp : float
nu : float

Returns
-------
r_ijk: np.array
Position vector in basis ijk
v_ijk: np.array
Velocity vector in basis ijk
"""
r_pqw, v_pqw = rv_pqw(k, p, ecc, nu)
rm = coe_rotation_matrix(inc, raan, argp)

r_ijk = rm @ r_pqw
v_ijk = rm @ v_pqw

return r_ijk, v_ijk

@jit
def coe2mee(p, ecc, inc, raan, argp, nu):
r"""Converts from classical orbital elements to modified equinoctial
orbital elements.

The definition of the modified equinoctial orbital elements is taken from
[Walker, 1985].

The modified equinoctial orbital elements are a set of orbital elements that are useful for
trajectory analysis and optimization. They are valid for circular, elliptic, and hyperbolic
orbits. These direct modified equinoctial equations exhibit no singularity for zero
eccentricity and orbital inclinations equal to 0 and 90 degrees. However, two of the
components are singular for an orbital inclination of 180 degrees.

.. math::
\begin{align}
p &= a(1-e^2) \\
f &= e\cos(\omega + \Omega) \\
g &= e\sin(\omega + \Omega) \\
h &= \tan(\frac{i}{2})\cos(\Omega) \\
k &= \tan(\frac{i}{2})\sin(\Omega) \\
L &= \Omega + \omega + \theta \\
\end{align}

Parameters
----------
k : float
Standard gravitational parameter (km^3 / s^2).
p : float
Semi-latus rectum or parameter (km).
ecc : float
Eccentricity.
inc : float
omega : float
argp : float
nu : float

Returns
-------
p: float
Semi-latus rectum or parameter
f: float
Equinoctial parameter f
g: float
Equinoctial parameter g
h: float
Equinoctial parameter h
k: float
Equinoctial parameter k
L: float
Longitude

Note
-----
The conversion equations are taken directly from the original paper.

"""
lonper = raan + argp
f = ecc * np.cos(lonper)
g = ecc * np.sin(lonper)
# TODO: Check polar case (see [Walker, 1985])
h = np.tan(inc / 2) * np.cos(raan)
k = np.tan(inc / 2) * np.sin(raan)
L = lonper + nu
return p, f, g, h, k, L

@jit
def rv2coe(k, r, v, tol=1e-8):
r"""Converts from vectors to classical orbital elements.

1. First the angular momentum is computed:
.. math::
\vec{h} = \vec{r} \times \vec{v}

2. With it the eccentricity can be solved:
.. math::
\begin{align}
\vec{e} &= \frac{1}{\mu}\left [ \left ( v^{2} - \frac{\mu}{r}\right ) \vec{r}  - (\vec{r} \cdot \vec{v})\vec{v} \right ] \\
e &= \sqrt{\vec{e}\cdot\vec{e}} \\
\end{align}

3. The node vector line is solved:
.. math::
\begin{align}
\vec{N} &= \vec{k} \times \vec{h} \\
N &= \sqrt{\vec{N}\cdot\vec{N}}
\end{align}

4. The rigth ascension node is computed:
.. math::
\Omega = \left\{ \begin{array}{lcc}
cos^{-1}{\left ( \frac{N_{x}}{N} \right )} &   if  & N_{y} \geq  0 \\
\\ 360^{o} -cos^{-1}{\left ( \frac{N_{x}}{N} \right )} &  if & N_{y} < 0 \\
\end{array}
\right.

5. The argument of perigee:
.. math::
\omega  = \left\{ \begin{array}{lcc}
cos^{-1}{\left ( \frac{\vec{N}\vec{e}}{Ne} \right )} &   if  & e_{z} \geq  0 \\
\\ 360^{o} -cos^{-1}{\left ( \frac{\vec{N}\vec{e}}{Ne} \right )} &  if & e_{z} < 0 \\
\end{array}
\right.

6. And finally the true anomaly:
.. math::
\nu  = \left\{ \begin{array}{lcc}
cos^{-1}{\left ( \frac{\vec{e}\vec{r}}{er} \right )} &   if  & v_{r} \geq  0 \\
\\ 360^{o} -cos^{-1}{\left ( \frac{\vec{e}\vec{r}}{er} \right )} &  if & v_{r} < 0 \\
\end{array}
\right.

Parameters
----------
k : float
Standard gravitational parameter (km^3 / s^2)
r : array
Position vector (km)
v : array
Velocity vector (km / s)
tol : float, optional
Tolerance for eccentricity and inclination checks, default to 1e-8

Returns
-------
p : float
Semi-latus rectum of parameter (km)
ecc: float
Eccentricity
inc: float
raan: float
Right ascension of the ascending nod (rad)
argp: float
nu: float

Examples
--------
>>> from poliastro.constants import GM_earth
>>> from astropy import units as u
>>> k = GM_earth.to(u.km ** 3 / u.s ** 2).value  # Earth gravitational parameter
>>> r = np.array([-6045., -3490., 2500.])
>>> v = np.array([-3.457, 6.618, 2.533])
>>> p, ecc, inc, raan, argp, nu = rv2coe(k, r, v)
>>> print("p:", p, "[km]")
p: 8530.47436396927 [km]
>>> print("ecc:", ecc)
ecc: 0.17121118195416898
inc: 153.2492285182475 [deg]
raan: 255.27928533439618 [deg]
argp: 20.068139973005366 [deg]
nu: 28.445804984192122 [deg]

Note
----
This example is a real exercise from Orbital Mechanics for Engineering
students by Howard D.Curtis. This exercise is 4.3 of 3rd. Edition, page 200.
"""

h = cross(r, v)
n = cross([0, 0, 1], h)
e = ((v.dot(v) - k / (norm(r))) * r - r.dot(v) * v) / k
ecc = norm(e)
p = h.dot(h) / k
inc = np.arccos(h[2] / norm(h))

circular = ecc < tol
equatorial = abs(inc) < tol

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(cross(e, r)) / norm(h), r.dot(e))
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(cross(h, n)) / norm(h), r.dot(n))
elif equatorial and circular:
raan = 0
argp = 0
nu = np.arctan2(r[1], r[0]) % (2 * np.pi)  # True longitude
else:
a = p / (1 - (ecc ** 2))
ka = k * a
if a > 0:
e_se = r.dot(v) / sqrt(ka)
e_ce = norm(r) * (norm(v) ** 2) / k - 1
nu = E_to_nu(np.arctan2(e_se, e_ce), ecc)
else:
e_sh = r.dot(v) / sqrt(-ka)
e_ch = norm(r) * (norm(v) ** 2) / k - 1
nu = F_to_nu(np.log((e_ch + e_sh) / (e_ch - e_sh)) / 2, ecc)

raan = np.arctan2(n[1], n[0]) % (2 * np.pi)
px = r.dot(n)
py = r.dot(cross(h, n)) / norm(h)
argp = (np.arctan2(py, px) - nu) % (2 * np.pi)

nu = (nu + np.pi) % (2 * np.pi) - np.pi

return p, ecc, inc, raan, argp, nu

@jit
def mee2coe(p, f, g, h, k, L):
r"""Converts from modified equinoctial orbital elements to classical
orbital elements.

The definition of the modified equinoctial orbital elements is taken from
[Walker, 1985].

.. math::

\begin{align}
p &= a(1 - e^{2})\\
e &= \sqrt{f^{2} + g^{2}}\\
i &= 2\arctan{(\sqrt{h^{2} + k^{2}})}\\
raan &= atan2(k, h) \pmod{2\pi}\\
argp &= (atan2(g, f) - raan) \pmod{2\pi}\\
nu &= (L - atan2(g, f)) \pmod{2\pi}\\
\end{align}

Parameters
----------
p: float
Semi-latus rectum
f: float
Equinoctial parameter p
g: float
Equinoctial parameter g
h: float
Equinoctial parameter h
k: float
Equinoctial parameter k
L: float
Longitude

Returns
-------
p: float
Semi-latus rectum
ecc: float
Eccentricity of the orbit
inc: float
Inclination of the orbit
raan: float
RAAN of orbit
argp: float
Argument of the periapsis
nu: float
True anomaly

Note
-----
The conversion is always safe because arctan2 works also for 0, 0
arguments.

"""
ecc = np.sqrt(f ** 2 + g ** 2)
inc = 2 * np.arctan(np.sqrt(h ** 2 + k ** 2))
lonper = np.arctan2(g, f)
raan = np.arctan2(k, h) % (2 * np.pi)
argp = (lonper - raan) % (2 * np.pi)
nu = (L - lonper) % (2 * np.pi)
return p, ecc, inc, raan, argp, nu