# Source code for orbitize.kepler

"""
This module solves for the orbit of the planet given Keplerian parameters.
"""
import numpy as np
import astropy.units as u
import astropy.constants as consts

try:
from . import _kepler
cext = True
except ImportError:
print("WARNING: KEPLER: Unable to import C-based Kepler's \
equation solver. Falling back to the slower NumPy implementation.")
cext = False

[docs]def calc_orbit(epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, mass_for_Kamp=None, tau_ref_epoch=0, tolerance=1e-9, max_iter=100):
"""
Returns the separation and radial velocity of the body given array of
orbital parameters (size n_orbs) at given epochs (array of size n_dates)

Based on orbit solvers from James Graham and Rob De Rosa. Adapted by Jason Wang and Henry Ngo.

Args:
epochs (np.array): MJD times for which we want the positions of the planet
sma (np.array): semi-major axis of orbit [au]
ecc (np.array): eccentricity of the orbit [0,1]
aop (np.array): argument of periastron [radians]
pan (np.array): longitude of the ascending node [radians]
tau (np.array): epoch of periastron passage in fraction of orbital period past MJD=0 [0,1]
plx (np.array): parallax [mas]
mtot (np.array): total mass of the two-body orbit (M_* + M_planet) [Solar masses]
mass_for_Kamp (np.array, optional): mass of the body that causes the RV signal.
For example, if you want to return the stellar RV, this is the planet mass.
If you want to return the planetary RV, this is the stellar mass. [Solar masses].
For planet mass ~ 0, mass_for_Kamp ~ M_tot, and function returns planetary RV (default).
tau_ref_epoch (float, optional): reference date that tau is defined with respect to (i.e., tau=0)
tolerance (float, optional): absolute tolerance of iterative computation. Defaults to 1e-9.
max_iter (int, optional): maximum number of iterations before switching. Defaults to 100.

Return:
3-tuple:

raoff (np.array): array-like (n_dates x n_orbs) of RA offsets between the bodies
(origin is at the other body) [mas]

deoff (np.array): array-like (n_dates x n_orbs) of Dec offsets between the bodies [mas]

vz (np.array): array-like (n_dates x n_orbs) of radial velocity of one of the bodies
(see mass_for_Kamp description)  [km/s]

Written: Jason Wang, Henry Ngo, 2018
"""

n_orbs  = np.size(sma)  # num sets of input orbital parameters
n_dates = np.size(epochs) # number of dates to compute offsets and vz

# return planetary RV if mass_for_Kamp is not defined
if mass_for_Kamp is None:
mass_for_Kamp = mtot

# Necessary for _calc_ecc_anom, for now
if np.isscalar(epochs): # just in case epochs is given as a scalar
epochs = np.array([epochs])
ecc_arr = np.tile(ecc, (n_dates, 1))

# Compute period (from Kepler's third law) and mean motion
period = np.sqrt(4*np.pi**2.0*(sma*u.AU)**3/(consts.G*(mtot*u.Msun)))
period = period.to(u.day).value
mean_motion = 2*np.pi/(period) # in rad/day

# # compute mean anomaly (size: n_orbs x n_dates)
manom = (mean_motion*(epochs[:, None] - tau_ref_epoch) - 2*np.pi*tau) % (2.0*np.pi)

# compute eccentric anomalies (size: n_orbs x n_dates)
eanom = _calc_ecc_anom(manom, ecc_arr, tolerance=tolerance, max_iter=max_iter)

# compute the true anomalies (size: n_orbs x n_dates)
# Note: matrix multiplication makes the shapes work out here and below
tanom = 2.*np.arctan(np.sqrt( (1.0 + ecc)/(1.0 - ecc))*np.tan(0.5*eanom) )
# compute 3-D orbital radius of second body (size: n_orbs x n_dates)
radius = sma * (1.0 - ecc * np.cos(eanom))

# compute ra/dec offsets (size: n_orbs x n_dates)
# math from James Graham. Lots of trig
c2i2 = np.cos(0.5*inc)**2
s2i2 = np.sin(0.5*inc)**2
arg1 = tanom + aop + pan
arg2 = tanom + aop - pan
c1 = np.cos(arg1)
c2 = np.cos(arg2)
s1 = np.sin(arg1)
s2 = np.sin(arg2)

# updated sign convention for Green Eq. 19.4-19.7
raoff = radius * (c2i2*s1 - s2i2*s2) * plx
deoff = radius * (c2i2*c1 + s2i2*c2) * plx

# compute the radial velocity (vz) of the body (size: n_orbs x n_dates)
# first comptue the RV semi-amplitude (size: n_orbs x n_dates)
Kv = np.sqrt(consts.G / (1.0 - ecc**2)) * (mass_for_Kamp * u.Msun * np.sin(inc)) / np.sqrt(mtot * u.Msun) / np.sqrt(sma * u.au)
# Convert to km/s
Kv = Kv.to(u.km/u.s)
# compute the vz
vz =  Kv.value * ( ecc*np.cos(aop) + np.cos(aop + tanom) )

# Squeeze out extra dimension (useful if n_orbs = 1, does nothing if n_orbs > 1)
vz = np.squeeze(vz)[()]

return raoff, deoff, vz

def _calc_ecc_anom(manom, ecc, tolerance=1e-9, max_iter=100, use_c=False):
"""
Computes the eccentric anomaly from the mean anomlay.
Code from Rob De Rosa's orbit solver (e < 0.95 use Newton, e >= 0.95 use Mikkola)

Args:
manom (float/np.array): mean anomaly, either a scalar or np.array of any shape
ecc (float/np.array): eccentricity, either a scalar or np.array of the same shape as manom
tolerance (float, optional): absolute tolerance of iterative computation. Defaults to 1e-9.
max_iter (int, optional): maximum number of iterations before switching. Defaults to 100.
Return:
eanom (float/np.array): eccentric anomalies, same shape as manom

Written: Jason Wang, 2018
"""

if np.isscalar(ecc) or (np.shape(manom) == np.shape(ecc)):
pass
else:
raise ValueError("ecc must be a scalar, or ecc.shape == manom.shape")

# If manom is a scalar, make it into a one-element array
if np.isscalar(manom):
manom = np.array((manom, ))

# If ecc is a scalar, make it the same shape as manom
if np.isscalar(ecc):
ecc = np.full(np.shape(manom), ecc)

# Initialize eanom array
eanom = np.full(np.shape(manom), np.nan)

# Save some boolean arrays
ecc_zero = ecc == 0.0
ecc_low = ecc < 0.95

# First deal with e == 0 elements
ind_zero = np.where(ecc_zero)
if len(ind_zero[0]) > 0: eanom[ind_zero] = manom[ind_zero]

# Now low eccentricities
ind_low = np.where(~ecc_zero & ecc_low)
if cext and use_c:
if len(ind_low[0]) > 0: eanom[ind_low] = _kepler._c_newton_solver(manom[ind_low], ecc[ind_low], tolerance=tolerance, max_iter=max_iter)

# the C solver returns eanom = -1 if it doesnt converge after max_iter iterations
m_one = eanom == -1
ind_high = np.where(~ecc_zero & ~ecc_low | m_one)
else:
if len(ind_low[0]) > 0: eanom[ind_low] = _newton_solver(manom[ind_low], ecc[ind_low], tolerance=tolerance, max_iter=max_iter)
ind_high = np.where(~ecc_zero & ~ecc_low)

# Now high eccentricities
if len(ind_high[0]) > 0: eanom[ind_high] = _mikkola_solver_wrapper(manom[ind_high], ecc[ind_high], use_c)

return np.squeeze(eanom)[()]

def _newton_solver(manom, ecc, tolerance=1e-9, max_iter=100, eanom0=None):
"""
Newton-Raphson solver for eccentric anomaly.
Args:
manom (np.array): array of mean anomalies
ecc (np.array): array of eccentricities
eanom0 (np.array): array of first guess for eccentric anomaly, same shape as manom (optional)
Return:
eanom (np.array): array of eccentric anomalies

Written: Rob De Rosa, 2018

"""
# Ensure manom and ecc are np.array (might get passed as astropy.Table Columns instead)
manom = np.array(manom)
ecc = np.array(ecc)

# Initialize at E=M, E=pi is better at very high eccentricities
if eanom0 is None:
eanom = np.copy(manom)
else:
eanom = np.copy(eanom0)

eanom -= (eanom - (ecc * np.sin(eanom)) - manom) / (1.0 - (ecc * np.cos(eanom)))

diff = (eanom - (ecc * np.sin(eanom)) - manom) / (1.0 - (ecc * np.cos(eanom)))
abs_diff = np.abs(diff)
ind = np.where(abs_diff > tolerance)
niter = 0
while ((ind[0].size > 0) and (niter <= max_iter)):
eanom[ind] -= diff[ind]
# If it hasn't converged after half the iterations are done, try starting from pi
if niter == (max_iter//2):
eanom[ind] = np.pi
diff[ind] = (eanom[ind] - (ecc[ind] * np.sin(eanom[ind])) - manom[ind]) / (1.0 - (ecc[ind] * np.cos(eanom[ind])))
abs_diff[ind] = np.abs(diff[ind])
ind = np.where(abs_diff > tolerance)
niter += 1

if niter >= max_iter:
print(manom[ind], eanom[ind], diff[ind], ecc[ind], '> {} iter.'.format(max_iter))
eanom[ind] = _mikkola_solver_wrapper(manom[ind], ecc[ind], use_c) # Send remaining orbits to the analytical version, this has not happened yet...

return eanom

def _mikkola_solver_wrapper(manom, ecc, use_c):
"""
Analtyical Mikkola solver (S. Mikkola. 1987. Celestial Mechanics, 40, 329-334.) for the eccentric anomaly.
Wrapper for the python implemenation of the IDL version. From Rob De Rosa.

Args:
manom (np.array): array of mean anomalies between 0 and 2pi
ecc (np.array): eccentricity
Return:
eanom (np.array): array of eccentric anomalies

Written: Jason Wang, 2018
"""

ind_change = np.where(manom > np.pi)
manom[ind_change] = (2.0 * np.pi) - manom[ind_change]
if cext and use_c:
eanom = _kepler._c_mikkola_solver(manom, ecc)
else:
eanom = _mikkola_solver(manom, ecc)
eanom[ind_change] = (2.0 * np.pi) - eanom[ind_change]

return eanom

def _mikkola_solver(manom, ecc):
"""
Analtyical Mikkola solver for the eccentric anomaly.
Adapted from IDL routine keplereq.pro by Rob De Rosa http://www.lpl.arizona.edu/~bjackson/idl_code/keplereq.pro

Args:
manom (float or np.array): mean anomaly, must be between 0 and pi.
ecc (float or np.array): eccentricity
Return:
eanom (np.array): array of eccentric anomalies

Written: Jason Wang, 2018
"""

alpha = (1.0 - ecc) / ((4.0 * ecc) + 0.5)
beta = (0.5 * manom) / ((4.0 * ecc) + 0.5)

aux = np.sqrt(beta**2.0 + alpha**3.0)
z = np.abs(beta + aux)**(1.0/3.0)

s0 = z - (alpha/z)
s1 = s0 - (0.078*(s0**5.0)) / (1.0 + ecc)
e0 = manom + (ecc * (3.0*s1 - 4.0*(s1**3.0)))

se0=np.sin(e0)
ce0=np.cos(e0)

f  = e0-ecc*se0-manom
f1 = 1.0-ecc*ce0
f2 = ecc*se0
f3 = ecc*ce0
f4 = -f2
u1 = -f/f1
u2 = -f/(f1+0.5*f2*u1)
u3 = -f/(f1+0.5*f2*u2+(1.0/6.0)*f3*u2*u2)
u4 = -f/(f1+0.5*f2*u3+(1.0/6.0)*f3*u3*u3+(1.0/24.0)*f4*(u3**3.0))

return (e0 + u4)