# Source code for stonesoup.types.orbitalstate

# -*- coding: utf-8 -*-
import datetime
from enum import Enum
from typing import Mapping, Any

import numpy as np

from ..base import Property
from ..functions.orbital import keplerian_to_rv, tru_anom_from_mean_anom
from .array import StateVector
from .state import State, GaussianState
from .angle import Inclination, EclipticLongitude

[docs]class CoordinateSystem(Enum):
"""Enumerates the allowable coordinate systems. See OrbitalState help for full explanation of
what each of the elements does.
"""
CARTESIAN = "Cartesian"
KEPLERIAN = "Keplerian"
TLE = "TLE"
EQUINOCTIAL = "Equinoctial"

# To allow case insensitivity and the use of "TwoLineElement" as a string to mean TLE
@classmethod
def _missing_(cls, value):

for element in cls:
if element.value.lower() == value.lower():
return element
if element.value == "TLE" and value.lower() == "twolineelement":
return element

raise ValueError("%r is not a valid %s" % (value, cls.__name__))

[docs]class OrbitalState(State):
r"""The orbital state base type. This is the building block of Stone Soup's orbital inference
routines and follows the principle that you shouldn't have to care which parameterisation you
use. The class stores relevant information internally and undertakes whatever conversions are
necessary.

The :attr:state_vector is held as :math:[\mathbf{r}, \dot{\mathbf{r}}], the "Orbital State
Vector" (as traditionally understood in orbital mechanics), where :math:\mathbf{r} is the
(3D) Cartesian position in the primary-centered inertial frame, while :math:\dot{\mathbf{r}}
is the corresponding velocity vector. All other parameters are accessed via functions. Formulae
for conversions are generally found in, or derived from [1]_.

The gravitational parameter :math:\mu = GM can be defined. If left undefined it defaults to
that of the Earth, :math:3.986004418 \, (\pm \, 0.000000008) \times 10^{14} \mathrm{m}^3
\mathrm{s}^{−2}

The object is constructed from the input vector :math:X_{t_{0}} at epoch
:attr:State.timestamp, :math:t_0. The coordinates of :math:X_{t_{0}} are Cartesian
Earth-Centered Inertial (ECI) [m] by default, but may be selected via the "coordinates" keyword
by passing a :class:~.CoordinateSystem object, or an appropriate string. Allowable coordinate
systems are,

coordinates = "Cartesian", the input state vector is

.. math::

X_{t_0} = [r_x, r_y, r_z, \dot{r}_x, \dot{r}_y,
\dot{r}_z]^{T}

where :math:r_x, r_y, r_z are the Cartesian position coordinates in the Primary-Centered
Inertial frame and :math:\dot{r}_x, \dot{r}_y, \dot{r}_z are the corresponding velocity
coordinates.

coordinates = "Keplerian" (Keplarian elements), construct using input state vector:

.. math::

X_{t_0} = [e, a, i, \Omega, \omega, \theta]^{T} \\

where:
:math:e is the orbital eccentricity (unitless),
:math:a the semi-major axis ([length]),
:math:i the inclination (radian),
:math:\Omega is the longitude of the ascending node (radian),
:math:\omega the argument of periapsis (radian), and
:math:\theta the true anomaly (radian)

coordinates = "TLE" (Two-Line Elements [2]_), initiates using input vector

.. math::

X_{t_0} = [i, \Omega, e, \omega, M_0, n]^{T}

where :math:i the inclination (radian),
:math:\Omega is the longitude of the ascending node (radian),
:math:e is the orbital eccentricity (unitless),
:math:\omega the argument of perigee (radian),
:math:M_0 the mean anomaly (radian) and
:math:n the mean motion (radian / [time]).

This can also be constructed by passing state_vector=None and using the metadata. In this
instance the metadata must conform to the TLE standard format [2]_ and be included in the
metadata dictionary as 'line_1' and 'line_2'.

coordinates = "Equinoctial" (equinoctial elements [3]_),

.. math::

X_{t_0} = [a, h, k, p, q, \lambda]^{T} \\

where :math:a the semi-major axis ([length]),
:math:h is the horizontal component of the eccentricity (unitless),
:math:k is the vertical component of the eccentricity (unitless),
:math:q is the horizontal component of the inclination (radian),
:math:k is the vertical component of the inclination (radian),
:math:\lambda is the mean longitude (radian)

References
----------
.. [1] Curtis, H.D. 2010, Orbital Mechanics for Engineering Students (3rd Ed), Elsevier
Aerospace Engineering Series

.. [2] NASA, Definition of Two-line Element Set Coordinate System, [spaceflight.nasa.gov](
https://spaceflight.nasa.gov/realdata/sightings/SSapplications/Post/JavaSSOP/
SSOP_Help/tle_def.html)

.. [3] Broucke, R. A. & Cefola, P. J. 1972, Celestial Mechanics, Volume 5, Issue 3, pp. 303-310

"""

coordinates: CoordinateSystem = Property(
default=CoordinateSystem.CARTESIAN,
doc="The parameterisation used on initiation. Acceptable values "
"are 'CARTESIAN' (default), 'KEPLERIAN', 'TLE', or 'EQUINOCTIAL'. "
"All other inputs will return errors. Will accept string inputs."
)

grav_parameter: float = Property(
default=3.986004418e14,
doc=r"Standard gravitational parameter :math:\mu = G M. The default "
r"is :math:3.986004418 \times 10^{14} \, "
r":math:\mathrm{m}^3 \mathrm{s}^{-2}.")

)

def __init__(self, state_vector, *args, **kwargs):
""""""
if 'coordinates' in kwargs:
coordinates = CoordinateSystem(kwargs['coordinates'])
else:
coordinates = CoordinateSystem.CARTESIAN

# Check to see if the initialisation is via metadata
if coordinates.name == 'TLE' and (state_vector is None or len(state_vector) == 0):

# Resolve the timestamp
year = 2000 + int(line1[17:20])
day = line1[20:23]

hour = float(line1[23:32]) * 24
fhour = int(np.floor(hour))

minu = (hour - fhour) * 60
fminu = int(np.floor(minu))

seco = (minu - fminu) * 60
fseco = int(np.floor(seco))

mics = (seco - fseco) * 1e6
fmics = int(np.round(mics))

timestamp = datetime.datetime.strptime(
str(year) + " " + str(day) + " " + str(fhour) + " " + str(fminu) + " " +
str(fseco) + " " + str(
fmics), "%Y %j %H %M %S %f")

state_vector = StateVector([float(line2[8:16]) * np.pi / 180,
float(line2[17:25]) * np.pi / 180,
float('.' + line2[26:33]),
float(line2[34:42]) * np.pi / 180,
float(line2[43:51]) * np.pi / 180,
float(line2[52:63]) * 2 * np.pi / 86400])

kwargs['timestamp'] = timestamp
super().__init__(state_vector, *args, **kwargs)

else:
raise TypeError("State vector and metadata cannot both be empty")

# Otherwise check that the state vector is the right size
elif len(state_vector) != 6:
raise ValueError(
"State vector shape should be 6x1 : got {}".format(state_vector.shape))

# Coordinate type checks
if coordinates.name == 'CARTESIAN':
super().__init__(state_vector, *args, **kwargs)

elif coordinates.name == 'KEPLERIAN':

if np.less(state_vector[0], 0.0) | np.greater(state_vector[0], 1.0):
raise ValueError("Eccentricity should be between 0 and 1: got {}"
.format(state_vector[0]))

# And go ahead and initialise as previously
super().__init__(state_vector, *args, **kwargs)

# Convert Keplerian elements to Cartesian

# First enforce the correct type
state_vector[2] = Inclination(state_vector[2])
state_vector[3] = EclipticLongitude(state_vector[3])
state_vector[4] = EclipticLongitude(state_vector[4])
state_vector[5] = EclipticLongitude(state_vector[5])

self.state_vector = keplerian_to_rv(state_vector, grav_parameter=self.grav_parameter)

elif coordinates.name == 'TLE':

if np.less(state_vector[2], 0.0) | np.greater(state_vector[2], 1.0):
raise ValueError("Eccentricity should be between 0 and 1: got {}"
.format(state_vector[0]))

super().__init__(state_vector, *args, **kwargs)

# First enforce the correct type
state_vector[0] = Inclination(state_vector[0])
state_vector[1] = EclipticLongitude(state_vector[1])
state_vector[3] = EclipticLongitude(state_vector[3])
state_vector[4] = EclipticLongitude(state_vector[4])

# Get the semi-major axis from the mean motion
semimajor_axis = np.cbrt(self.grav_parameter / state_vector[5] ** 2)

# True anomaly from mean anomaly
tru_anom = tru_anom_from_mean_anom(state_vector[4], state_vector[2])

# Use given and derived quantities to convert from Keplarian to
# Cartesian
self.state_vector = keplerian_to_rv(StateVector([state_vector[2], semimajor_axis,
state_vector[0], state_vector[1],
state_vector[3], tru_anom]),
grav_parameter=self.grav_parameter)

elif coordinates.name == 'EQUINOCTIAL':

if np.less(state_vector[1], -1.0) | np.greater(state_vector[1], 1.0):
raise ValueError("Horizontal Eccentricity should be between -1 "
"and 1: got {}".format(state_vector[1]))
if np.less(state_vector[2], -1.0) | np.greater(state_vector[2], 1.0):
raise ValueError("Vertical Eccentricity should be between -1 and "
"1: got {}".format(state_vector[2]))

super().__init__(state_vector, *args, **kwargs)

# First enforce the correct type for mean longitude
state_vector[5] = EclipticLongitude(state_vector[5])

# Calculate the Keplarian element quantities
semimajor_axis = state_vector[0]
raan = np.arctan2(state_vector[3], state_vector[4])
inclination = 2 * np.arctan(state_vector[3] / np.sin(raan))
arg_per = np.arctan2(state_vector[1], state_vector[2]) - raan
mean_anomaly = state_vector[5] - arg_per - raan
eccentricity = state_vector[1] / (np.sin(arg_per + raan))

# True anomaly from mean anomaly
tru_anom = tru_anom_from_mean_anom(mean_anomaly, eccentricity)

# Convert from Keplarian to Cartesian
self.state_vector = keplerian_to_rv(StateVector([eccentricity, semimajor_axis,
inclination, raan, arg_per,
tru_anom]),
grav_parameter=self.grav_parameter)

# Some vector quantities
@property
def _nodeline(self):
"""The vector node line (defines the longitude of the ascending node in the
Primary-centered inertial frame)"""

k = np.array([0, 0, 1])
boldh = self.specific_angular_momentum

boldn = np.cross(k, boldh, axis=0)
n = np.sqrt(np.dot(boldn.T, boldn).item())

# If inclination is 0, the node line is [0] and has 0 magnitude. By
# convention in these situations, we set the node line as a unit vector
# pointing along x. Note that the magnitude of the vector is not
# consistent with that produced by the cross product. (But we assume
# that the node line is only used for directional information.
if n < np.finfo(n).eps:
return np.array([1, 0, 0])
else:
return boldn

@property
def _eccentricity_vector(self):
r""" The eccentricity vector :math:\mathbf{e}"""

rang = self.range
speed = self.speed
self.state_vector[3:6]).item()/rang

return (1/self.grav_parameter) * ((speed**2 - self.grav_parameter/rang)
* self.state_vector[0:3] - rang *
self.state_vector[3:6])

@property
def specific_angular_momentum(self):
r"""The specific angular momentum, :math:\mathbf{h}"""
return np.cross(self.state_vector[0:3], self.state_vector[3:6], axis=0)

@property
def cartesian_state_vector(self):
r"""The state vector :math:X_{t_0} = [r_x, r_y, r_z, \dot{r}_x, \dot{r}_y, \dot{r}_z]^{T}
in 'Primary-Centred' Inertial coordinates, equivalent to ECI in the case of the Earth
"""
return StateVector(self.state_vector)

# Some scalar quantities
@property
def epoch(self):
"""The epoch, or state timestamp"""
return self.timestamp

@property
def range(self):
"""The distance to object (from gravitational centre of primary)"""
return np.sqrt(np.dot(self.state_vector[0:3].T,
self.state_vector[0:3])).item()

@property
def speed(self):
"""The current instantaneous speed (scalar)"""
return np.sqrt(np.dot(self.state_vector[3:6].T,
self.state_vector[3:6]).item())

@property
def eccentricity(self):
r"""The orbital eccentricity, :math:e \; (0 \le e \le 1)

Note
----
This version of the calculation uses a form dependent only on scalars

"""
# TODO Check to see which of the following is quicker/better

# Either
# return np.sqrt(np.dot(self._eccentricity_vector.T,
# self._eccentricity_vector).item())
# or
return np.sqrt(1 + (self.mag_specific_angular_momentum ** 2 /
self.grav_parameter ** 2) *
(self.speed ** 2 - 2 * self.grav_parameter /
self.range))

@property
def semimajor_axis(self):
"""The orbital semi-major axis"""
return (self.mag_specific_angular_momentum**2 / self.grav_parameter) *\
(1 / (1 - self.eccentricity**2))

# Used to be this
# return 1/((2/self.range) - (self.speed**2)/self.grav_parameter)

@property
def inclination(self):
r"""Orbital inclination, :math:i \; (0 \le i < \pi), [rad]"""
boldh = self.specific_angular_momentum
h = self.mag_specific_angular_momentum

return Inclination(np.arccos(np.clip(boldh[2].item()/h, -1, 1)))

@property
def longitude_ascending_node(self):
r"""The longitude (or right ascension) of ascending node, :math:\Omega \; (0 \leq \Omega <
2\pi)"""
boldn = self._nodeline
n = np.sqrt(np.dot(boldn.T, boldn).item())

if boldn[1].item() >= 0:
return EclipticLongitude(np.arccos(np.clip(boldn[0].item()/n, -1, 1)))
else:  # boldn[1].item() < 0:
return EclipticLongitude(2*np.pi - np.arccos(np.clip(boldn[0].item()/n, -1, 1)))

@property
def argument_periapsis(self):
r"""The argument of periapsis, :math:\omega \; (0 \le \omega < 2\pi) in radians"""
boldn = self._nodeline
n = np.sqrt(np.dot(boldn.T, boldn).item())
bolde = self._eccentricity_vector

# If eccentricity is 0 then there's no unambiguous longitude of
# periapsis. In these situations we set the argument of periapsis to
# 0.
if self.eccentricity < np.finfo(self.eccentricity).eps:
return 0

# Quadrant ambiguity. The clip function is required to mitigate against
# the occasional floating-point errors which push the ratio outside the
# -1,1 region.
if bolde[2].item() >= 0:
return EclipticLongitude(np.arccos(np.clip(np.dot(boldn.T, bolde).item() /
(n * self.eccentricity), -1, 1)))
else:  # bolde[2].item() < 0:
return EclipticLongitude(2*np.pi - np.arccos(np.clip(np.dot(boldn.T, bolde).item() /
(n * self.eccentricity), -1, 1)))

@property
def true_anomaly(self):
r"""The true anomaly, :math:\theta \; (0 \le \theta < 2\pi) in radians"""
# Resolve the quadrant ambiguity.The clip function is required to
# mitigate against floating-point errors which push the ratio outside
# the -1,1 region.
self.state_vector[3:6]).item() / self.speed

return EclipticLongitude(np.arccos(np.clip(
np.dot(self._eccentricity_vector.T / self.eccentricity, self.state_vector[0:3] /
self.range).item(), -1, 1)))
return EclipticLongitude(2*np.pi - np.arccos(np.clip(
np.dot(self._eccentricity_vector.T / self.eccentricity, self.state_vector[0:3] /
self.range).item(), -1, 1)))

@property
def eccentric_anomaly(self):
r"""The eccentric anomaly, :math:E \; (0 \le E < 2\pi) in radians

Note
----
This computes the quantity exactly via the Keplerian
eccentricity and true anomaly rather than via the mean
anomaly using an iterative procedure.

"""
return EclipticLongitude(np.remainder(2 * np.arctan(np.sqrt((1 - self.eccentricity) /
(1 + self.eccentricity)) *
np.tan(self.true_anomaly / 2)),
2*np.pi))

@property
def mean_anomaly(self):
r"""Mean anomaly, :math:M \; (0 \le M < 2\pi), in radians

Note
----
Uses the eccentric anomaly and Kepler's equation to get
mean anomaly from true anomaly and eccentricity.

"""

return EclipticLongitude(self.eccentric_anomaly - self.eccentricity *
np.sin(self.eccentric_anomaly))  # Kepler's equation

@property
def period(self):
"""Orbital period, :math:T ([time])"""
return ((2 * np.pi) / np.sqrt(self.grav_parameter)) * \
np.power(self.semimajor_axis, 3 / 2)

@property
def mean_motion(self):
r"""The mean motion, :math:\frac{2 \pi}{T}, where :math:T is the period, (rad / [time])
"""
return 2 * np.pi / self.period

@property
def mag_specific_angular_momentum(self):
"""The magnitude of the specific angular momentum, :math:h"""
boldh = self.specific_angular_momentum
return np.sqrt(np.dot(boldh.T, boldh).item())

# Alternative via scalars
# return np.sqrt(self.grav_parameter * self.semimajor_axis *
# (1 - self.eccentricity ** 2))

@property
def specific_orbital_energy(self):
r"""Specific orbital energy (:math:\frac{-GM}{2a})"""
return -self.grav_parameter / (2 * self.semimajor_axis)

@property
def equinoctial_h(self):
r"""The horizontal component of the eccentricity in equinoctial coordinates is
:math:h = e \sin (\omega + \Omega)"""

return self.eccentricity * np.sin(self.argument_periapsis +
self.longitude_ascending_node)

@property
def equinoctial_k(self):
r"""The vertical component of the eccentricity in equinoctial coordinates is
:math:k = e \cos (\omega + \Omega)"""

return self.eccentricity * np.cos(self.argument_periapsis +
self.longitude_ascending_node)

@property
def equinoctial_p(self):
r"""The horizontal component of the inclination in equinoctial coordinates is
:math:p = \tan (i/2) \sin \Omega"""
return np.tan(self.inclination/2) * \
np.sin(self.longitude_ascending_node)

@property
def equinoctial_q(self):
r"""The vertical component of the inclination in equinoctial coordinates is
:math:q = \tan (i/2) \cos \Omega"""
return np.tan(self.inclination / 2) * np.cos(self.longitude_ascending_node)

@property
def mean_longitude(self):
r"""The mean longitude, defined as :math:\lambda = M_0 + \omega + \Omega (rad)"""
return EclipticLongitude(self.mean_anomaly + self.argument_periapsis +
self.longitude_ascending_node)

# The following return vectors of complete sets of elements
@property
def keplerian_elements(self):
r"""The vector of Keplerian elements :math:X = [e, a, i, \Omega, \omega, \theta]^{T}
where :math:e is the orbital eccentricity (unitless), :math:a the semi-major axis
([length]), :math:i the inclination (radian), :math:\Omega is the longitude of the
ascending node (radian), :math:\omega the argument of periapsis (radian), and
:math:\theta the true anomaly (radian)"""
return StateVector(np.array([[self.eccentricity],
[self.semimajor_axis],
[self.inclination],
[self.longitude_ascending_node],
[self.argument_periapsis],
[self.true_anomaly]]))

@property
def two_line_element(self):
r"""The Two-Line Element vector :math:X = [i, \Omega, e, \omega, M_0, n]^{T} where
:math:i the inclination (radian) :math:\Omega is the longitude of the ascending node
(radian), :math:e is the orbital eccentricity (unitless), :math:\omega the argument of
periapsis (radian), :math:M_0 the mean anomaly (radian) :math:n the mean motion
return StateVector(np.array([[self.inclination],
[self.longitude_ascending_node],
[self.eccentricity],
[self.argument_periapsis],
[self.mean_anomaly],
[self.mean_motion]]))

@property
def equinoctial_elements(self):
r"""The equinoctial elements, :math:X = [a, h, k, p, q, \lambda]^{T} where :math:a the
semi-major axis ([length]), :math:h and :math:k are the horizontal and vertical
components of the eccentricity respectively (unitless), :math:p and :math:q are the
horizontal and vertical components of the inclination respectively (radian) and
:math:\lambda is the mean longitude (radian) [3]_
"""
return StateVector(np.array([[self.semimajor_axis],
[self.equinoctial_h],
[self.equinoctial_k],
[self.equinoctial_p],
[self.equinoctial_q],
[self.mean_longitude]]))

[docs]class GaussianOrbitalState(GaussianState, OrbitalState):
"""An Orbital state for use in Kalman filters (and perhaps elsewhere). Inherits from
GaussianState so has covariance matrix. As no checks on the validity of the covariance
matrix are made, care should be exercised in its use. The propagator will generally require
a particular coordinate reference which must be understood.

All methods provided by OrbitalState are available.

"""