# Source code for irfpy.util.ballistic

"""Ballistic trajectory related module

- :mod:irfpy.util.keplernumeric
"""
import numpy as np

[docs]class Trajectory:

def __init__(self, r0, th0, vr0, om0, mass, g=6.674e-11, m=1):
r""" Initialize the trajectory.

:param r0: Initial position, :math:r_0.
:param th0: Initial phase, :math:\theta_0. Radians.
:param vr0: Initial radial velocity, :math:v_0=dr/dt(t=0).
:param om0: Initial angular veloicty, :math:\omega_0=d\theta/dt(t=0).
:param mass: Mass of the central body (or effective mass of the center of gravity).
:keyword g: Gravity of constant, :math:G.  Default value is for MKSA unit value.
:keyword m: Mass of the moving object.  It DOES NOT concern the trajectory,
but some physical quantities do.

*Theory*: From the initial conditions, the trajectory is expressed in the parametric equation
:math:r = LL / (1 + e \cos (\theta + \Phi). Here, r is the distance from the
central body, e is eccentricity, Phi represents the phase shift.
"""
# From the initial condition, the trajectory is expressed in the parametric
# elipse equation: r = LL / (1 + e cos (_theta + Phi) ).
# LL, e, and Phi are calculated analytically below.

self._r0 = r0
self._th0 = th0
self._vr0 = vr0
self._om0 = om0
self._M = mass
self._GM = mass * g
self._m = m

self._L = m * r0 * r0 * om0    # Angular moment. Constant over the trajectory.

self._LL = ((self._L / self._m) ** 2) / self._GM  # Parametric ellipse equation: LL/(1+e cos theta)

# The phase at the initial time
self._Phi = np.arctan2(self._vr0 * self._LL / (self._r0 ** 2) / self._om0, self._LL / self._r0 - 1) - self._th0

phi0 = self._Phi + self._th0
sin = np.sin(phi0)
cos = np.cos(phi0)

self._es = self._vr0 * self._LL / (self._r0 * self._r0 * self._om0 * sin)
self._ec = ((self._LL / self._r0) - 1) / cos
if np.abs(sin) > np.abs(cos):
### Use sin conterpart of formulation
self._e = self._es
else:
### Use cos conterpart of formulation
self._e = self._ec

[docs]    def print_values(self):
print('r0:', self._r0)
print('th0:', self._th0)
print('v0:', self._vr0)
print('om0:', self._om0)
print('M:', self._M)
print('GM:', self._GM)
print('m:', self._m)
print('L:', self._L)
print('LL:', self._LL)
print('Phi:', self._Phi)
print('es:', self._es)
print('ec:', self._ec)
print('e:', self._e)

[docs]    def angular_momentum(self):
r""" Return the angular momentum.

The angular momentum is constant.
It is defined by the initial condition: L=m r0^2 omega0.

.. math::

L = m r_0^2 \omega_0
"""
return self._L

[docs]    def eccentricity(self):
return self._e

[docs]    def r(self, th):
r""" Return the distance r as a function of th

:param th: The phase angle in radians.

*Theory*: From LL, e, and Phi, the radial distance is

:math:LL / (1 + e \cos(\theta+\Phi)).
"""
r = self._LL / (1 + self._e * np.cos(th + self._Phi))
return r

[docs]    def omega(self, th):
r""" Return :math:omega as a function of th.

:param th: The phase angle in radians.
:returns: The angular speed, omega.

*Theory*: The angular momentum (L see :meth:angular_momentum) is conserved, so that the
angular speed omega can be retrieved via the equation

.. math::

L = m r(\theta)^2 \omega(\theta).

The distance r is retrieved from :meth:r
"""
r = self.r(th)
L = self._L
m = self._m

return L / m / r / r

[docs]    def v(self, th):
""" Return the radial velocity as a function of th.

Deprecated. Use :meth:vr.
"""
return self.vr(th)

[docs]    def vr(self, th):
r""" Return the radial velocity as a function of th

:param th: The phase angle in radians.

*Theory:* We derived the vr as follows.

.. math::

r = \frac{LL}{1+e\cos(\theta+\Phi)}.

By taking time derivative, we obtain

.. math::

v_r = \frac{dr}{dt} = e \sin(\theta+\Phi)\frac{LL}{(1+e\cos(\theta+\Phi))^2}\frac{d\theta}{dt}

By simplifying the equation, we obtain

.. math::

v_r = \frac{e l}{m LL}\sin{\theta + \Phi}.
"""
L = self._L
m = self._m
e = self._e
LL = self._LL
phi = self._Phi

v = L / m * e * np.sin(th + phi) / LL
return v

[docs]    def integ_distance_num(self, end_phase, ndiv=36000):
r""" Integrate the distance flown from initial phase until the given end_phase.

The distance along the trajectory is integrated along the orbit.
Simple daikei-kousiki (?what is the English?) is used.

:param end_phase: The end of the phase (:math:\theta_1). Radians.
:param ndiv: Division used for the separation.
:return: The distance.

*Sample*

Let's see a lunar orbiter. The 100 km circular orbit (distance of 1838 km) is considered.
It is easier to use the helper function :func:speed_from_distances to get the
velocity for such orbiter.

>>> vper, vapo = speeds_from_distances(1838e3, 1838e3, 7.342e22)
>>> print('{:.1f} m/s'.format(vper))
1632.8 m/s

Create the :class:Trajectory object.

>>> orbiter = Trajectory(1838e3, 0, 0, 1632.8 / 1838e3, 7.342e22)
>>> print('{:.2f}'.format(orbiter._e))   # Eccentrycity. This is (almost) circle.
0.00

Calculate the full length of the trajectory. It is :math:2\pi R where R is
the radius of the orbiter (1838 km). Thus, 11548.5 km.

>>> trajectory_length = orbiter.integ_distance_num(2 * np.pi)
>>> print('{:.3e} m'.format(trajectory_length))
1.155e+07 m
"""
# separate the phase angle by the division
thetalist = np.linspace(self._th0, end_phase, ndiv)
thetalist_c = (thetalist[1:] + thetalist[:-1]) / 2.
d_thetalist = thetalist[1:] - thetalist[:-1]
rlist = self.r(thetalist_c)    # r(th)
distance_list = rlist * d_thetalist   # r dth
distance = distance_list.sum()
return distance

[docs]    def integ_time_num(self, end_phase, ndiv=36000):
r""" Calcuate the time from the initial time to the end at the end_phase.

The calculation uses numerical calculation

.. math::

t = \int_{\theta_0}^{\theta_1} \frac{d\theta}{\omega}

:param end_phase: The end of the phase (:math:\theta_1). Radians.
:param ndiv: Division used for the separation.
:return: The distance.

*Sample*

Let's see a lunar orbiter. The 100 km circular orbit (distance of 1838 km) is considered.
Create the :class:Trajectory object, in the same manner as :meth:integ_distance_num.

>>> orbiter = Trajectory(1838e3, 0, 0, 1632.8 / 1838e3, 7.342e22)
>>> print('{:.2f}'.format(orbiter._e))   # Eccentrycity. This is (almost) circle.
0.00

Calculate the circular period. In this case, the total distance flying over is
11548.5 km (see :meth:integ_distance_num).
The speed of the orbiter is 1632.8 m/s.
This resuts in 7072.8 s.

>>> period = orbiter.integ_time_num(2 * np.pi)
>>> print('{:.0f} s'.format(period))
7073 s
"""

# separate the phase angle by the division
thetalist = np.linspace(self._th0, end_phase, ndiv)
thetalist_c = (thetalist[1:] + thetalist[:-1]) / 2.
d_thetalist = thetalist[1:] - thetalist[:-1]
wlist = self.omega(thetalist_c)    # omega(th)
dt_list = d_thetalist / wlist
t = dt_list.sum()
return t

[docs]def speeds_from_distances(dperi, dapo, mass, g=6.674e-11):
r""" Calculate the speeds at peri- and apo-centers from the distances

It is common request to calulate the trajectory for, say, spacecraft with
specification of the heights.

:param dperi: Pericenter distance from the center
:param dapo: Apocenter distance from the center
:param mass: Mass of the central body
:keyword g: Constant of gravity, :math:G
:returns: (vper, vapo), where vper is the pericenter spped and vapo for apocenter.

The formulation is simple: Using the energy and angular momentum conversations.
Derivation is based on MKSA system.

.. math::

v_p = \frac{1}{d_p}\sqrt{\frac{2GM}{\frac{1}{d_a}+\frac{1}{d_p}}}    \\
v_a = \frac{1}{d_a}\sqrt{\frac{2GM}{\frac{1}{d_a}+\frac{1}{d_p}}}    \\

Example below is an orbiter around the Moon.  The pericenter altitude 30 km and
apocenter altitude of 270 km are considered.  Moon radius of 1738 km is taken.

>>> dp = (30 + 1738) * 1000    # in meter
>>> da = (270 + 1738) * 1000   # in meter
>>> from irfpy.util.planets import moon
>>> v_p, v_a = speeds_from_distances(dp, da, moon['mass'])
>>> print('{:.3f}'.format(v_p))
1717.547
>>> print('{:.3f}'.format(v_a))
1512.262

The results say the pericenter velocity of 1.718 km/s and apocenter velocity of 1.512 km/s

"""
sq = np.sqrt(2 * g * mass / (1 / dapo + 1 / dperi))
return sq / dperi, sq / dapo

[docs]class Trajectory3D:
r""" 3-D trajectory of ballistic motion.

*Theory*

The class represents the ballistic motion of a mass *m* under a huge gravity star of mass *M*.
(*M >> m* is assumed, but theoretically, if you use "equivalent mass" then one can extend this to more generic problems)

Consider the mass *M* at origin, and stationary.
The mass *m* moves around under Kepler motion.
Let's take the initial position :math:\vec{r_0} with the velocity :math:\vec{v_0}.
:math:r_0 = |\vec{r_0}| and :math:v_0 = |\vec{v_0}| should not be zero.
Also they should not be parallel, i.e. :math:\vec{r_0}\times\vec{v_0}\ne 0.

This is a 2-D problem, which can be represented by :class:Trajectory.
This, a conversion between the 3-D coordinates to the 2-D in :class:Trajectory is the key to solve.

Conversion:

Let's take the new axis as follows.

- :math:\hat{s}: Along the positon vector, i.e., :math:\hat{s} = \frac{\vec{r_0}}{r_0}.
- :math:\hat{t}: :math:\hat{s}-\hat{t} plane to include the velocity vector, with :math:v_t>0.
This means that :math:\hat{t} // (\vec{r_0}\times\vec{v_0})\times\vec{r_0}.

*Usage*
Let's try the following example.

- An orbiter around the Moon (mass=7.342e22 kg). In inertia frame
- The position of orbiter at t=0 is 1850 km at the equator, longitude at 45.

>>> r0 = 1850e3 * np.array([np.cos(np.pi / 4), np.sin(np.pi / 4), 0])
>>> print(r0)    # doctest: +NORMALIZE_WHITESPACE
[1308147.54519511  1308147.54519511        0.        ]

- The initial velocity of the orbiter at t=0 is (0, 0.01, 1.65) km/s.

>>> v0 = 1000 * np.array([0, 0.01, 1.65])
>>> print(v0)      # doctest: +NORMALIZE_WHITESPACE
[    0.    10.  1650.]

Now it is time to make the Trajectory3D object.

>>> traj = Trajectory3D(r0, v0, 7.342e22)
>>> print(traj)
Trajectory3D: s=(0.707,0.707,0) t=(-0.00303,0.00303,1) r0=1.85e+06 vs0=7.07 w0=0.000892

The initial position, r0, should be converted to the original position in xyz.

>>> print(traj.st2xyz(traj.s(0), traj.t(0)))  # doctest: +NORMALIZE_WHITESPACE
[1308147.54519511  1308147.54519511        0.        ]

The initial velocity, v0, should be converted to the original velocity in xyz.
>>> v0recal = traj.st2xyz(traj.vs(0), traj.vt(0))
>>> print('{v:.3f} {v:.3f} {v:.3f}'.format(v=v0recal))
0.000 10.000 1650.000

Later, the velocity can be retrieved by vxyz method.
>>> vxyz1 = traj.vxyz(1)     # Velocity at phase = 1 radian
>>> print('{v:.3f} {v:.3f} {v:.3f}'.format(v=vxyz1))  # It is in m/s
-952.900 -947.372 912.080

>>> traj.print_values()    # doctest: +SKIP
--- 2d ---
r0: 1850000.0
th0: 0.0
v0: 7.07106781187
om0: 0.00086487331077
...

The period of such orbit is ~7450 sec.
>>> print('{:.3f}'.format(traj.integ_time_num(np.pi * 2)))
7452.076

*Notes*

All the quantities are derived via the parameter :math:\theta.  This is frequently a difficulty, because
it is more natural to use the time *t* as a parameter. However, the theory of two-body collision is built
such that.  I do not want to dig into more in details now.
Of course practically, it could be worthwhile to get a function such like :math:\theta=\theta(t)
by inverting :math:t=t(\theta).
"""
def __init__(self, r0vec, v0vec, mass, g=6.674e-11, m=1):
r""" Initialize the trajectory.

:param r0vec: Initial position, :math:\vec{r_0}.
:param v0vec: Initial radial velocity, :math:\vec{v_0}.
:param mass: Mass of the central body (or effective mass of the center of gravity).
:keyword g: Gravity of constant, :math:G.  Default value is for MKSA unit value.
:keyword m: Mass of the moving object.  It DOES NOT concern the trajectory,
but some deriving physical quantities do.
"""
### New axis
self.r0vec = np.array(r0vec)
self.r0 = np.sqrt((self.r0vec ** 2).sum())
if self.r0vec.shape != (3,) or self.r0 == 0:
raise ValueError('Position should be non-zero 3-D vector, but {}'.format(str(self.r0vec)))

self.v0vec = np.array(v0vec)
self.v0 = np.linalg.norm(self.v0vec)
if self.v0vec.shape != (3,) or self.v0 == 0:
raise ValueError('Velocity should be non-zero 3-D vector, but {}'.format(str(self.v0vec)))

self.hats = self.r0vec / self.r0

self.hatt = np.cross(self.r0vec, self.v0vec)
if (self.hatt == 0).all():
raise ValueError('Position and Velocity should note be paralell, but {} and {}.'.format(str(self.r0vec), str(self.v0vec)))

self.hatt = np.cross(self.hatt, self.r0vec)
self.hatt = self.hatt / np.linalg.norm(self.hatt)

self.th0 = 0.     # Initial phase is always zero in s-t plane.

cosalpha = np.clip(np.dot(self.r0vec, self.v0vec) / self.r0 / self.v0, -1, 1)
self.alpha = np.arccos(cosalpha)
sinalpha = np.sqrt(1 - cosalpha ** 2)

self.vs0 = self.v0 * cosalpha
self.om0 = self.v0 * sinalpha / self.r0

self._2dtraj = Trajectory(self.r0, self.th0, self.vs0, self.om0, mass, g=g, m=m)

[docs]    def st2xyz(self, s, t):
"""
"""
return s * self.hats + t * self.hatt

[docs]    def r(self, th):
""" Return r from the phase angle as a parameter.
"""
return self._2dtraj.r(th)

[docs]    def omega(self, th):
""" Return the tangential angular velocity
"""
return self._2dtraj.omega(th)

[docs]    def s(self, th):
""" Return the s value
"""
return self.r(th) * np.cos(th)

[docs]    def t(self, th):
""" Return the t value
"""
return self.r(th) * np.sin(th)

[docs]    def v(self, th):
"""
return self._2dtraj.v(th)

[docs]    def vr(self, th):
""" Return the radial velociyt. Same as :meth:v
"""
return self.v(th)

[docs]    def vth(self, th):
""" Return the phase velocity."""
return self.omega(th) * self.r(th)

[docs]    def vs(self, th):
""" Return the s-direction velocity
"""
return self.v(th) * np.cos(th) - self.vth(th) * np.sin(th)

[docs]    def vt(self, th):
""" Return the tangential velocity
"""
return self.v(th) * np.sin(th) + self.vth(th) * np.cos(th)

[docs]    def xyz(self, th):
""" Return the (x, y, z) coordinates
"""
s = self.s(th)
t = self.t(th)
return self.st2xyz(s, t)

[docs]    def longitude(self, th):
"""Return the longitude in degrees at the phase angle |theta|.

:param th: Phase angle |theta|
:return: The longitude in degrees.
"""
xyz = self.xyz(th)
return lon

[docs]    def latitude(self, th):
""" Return the Latitude in degrees at the phase angle |theta|.

:param th: The phase angle |theta|
:return: The latitude in degrees
"""
r = self.r(th)
z = self.xyz(th)
cos = np.clip(z / r, -1, 1)

[docs]    def vxyz(self, th):
""" Return (vx, vy, vz)
"""
vs = self.vs(th)
vt = self.vt(th)
return self.st2xyz(vs, vt)

[docs]    def integ_time_num(self, end_phase, ndiv=36000):
return self._2dtraj.integ_time_num(end_phase, ndiv=ndiv)

[docs]    def integ_distance_num(self, end_phase, ndiv=36000):
return self._2dtraj.integ_distance_num(end_phase, ndiv=ndiv)

[docs]    def print_values(self):
print('--- 2d ---')
print('r0:', self._2dtraj._r0)
print('th0:', self._2dtraj._th0)
print('v0:', self._2dtraj._vr0)
print('om0:', self._2dtraj._om0)
print('M:', self._2dtraj._M)
print('GM:', self._2dtraj._GM)
print('m:', self._2dtraj._m)
print('L:', self._2dtraj._L)
print('LL:', self._2dtraj._LL)
print('Phi:', self._2dtraj._Phi)
print('es:', self._2dtraj._es)
print('ec:', self._2dtraj._ec)
print('e:', self._2dtraj._e)

def __str__(self):
s = 'Trajectory3D: s=({s:.3g},{s:.3g},{s:.3g}) t=({t:.3g},{t:.3g},{t:.3g}) '.format(s=self.hats, t=self.hatt)
s = s +  'r0={:.3g} vs0={:.3g} w0={:.3g}'.format( self.r0, self.vs0, self.om0)
return s