import numpy as np

# Try to import Kepler's equation solver written in C
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 rv_drive(t, orbel, use_c_kepler_solver=cext):
"""RV Drive

Args:
t (array of floats): times of observations
orbel (array of floats): [per, tp, e, om, K].\
Omega is expected to be\
use_c_kepler_solver (bool): (default: True) If \
True use the Kepler solver written in C, else \
use the Python/NumPy version.

Returns:
rv: (array of floats): radial velocity model

"""

# unpack array of parameters
per, tp, e, om, k = orbel

# Performance boost for circular orbits
if e == 0.0:
m = 2 * np.pi * (((t - tp) / per) - np.floor((t - tp) / per))
return k * np.cos(m + om)

if per < 0:
per = 1e-4
if e < 0:
e = 0
if e > 0.99:
e = 0.99

# Calculate the approximate eccentric anomaly, E1, via the mean anomaly  M.
if use_c_kepler_solver:
rv = _kepler.rv_drive_array(t, per, tp, e, om, k)
else:
nu = radvel.orbit.true_anomaly(t, tp, per, e)
rv = k * (np.cos(nu + om) + e * np.cos(om))

return rv

[docs]def kepler(inbigM, inecc):
"""Solve Kepler's Equation

Args:
inbigM (array): input Mean anomaly
inecc (array): eccentricity

Returns:
eccentric anomaly: array

"""

Marr = inbigM  # protect inputs; necessary?
eccarr = inecc
conv = 1.0e-12  # convergence criterion
k = 0.85

Earr = Marr + np.sign(np.sin(Marr)) * k * eccarr  # first guess at E
# fiarr should go to zero when converges
fiarr = ( Earr - eccarr * np.sin(Earr) - Marr)
convd = np.where(np.abs(fiarr) > conv)  # which indices have not converged
nd = len(convd)  # number of unconverged elements
count = 0

while nd > 0:  # while unconverged elements exist
count += 1

M = Marr[convd]  # just the unconverged elements ...
ecc = eccarr[convd]
E = Earr[convd]

fi = fiarr[convd]  # fi = E - e*np.sin(E)-M    ; should go to 0
fip = 1 - ecc * np.cos(E)  # d/dE(fi) ;i.e.,  fi^(prime)
fipp = ecc * np.sin(E)  # d/dE(d/dE(fi)) ;i.e.,  fi^(\prime\prime)
fippp = 1 - fip  # d/dE(d/dE(d/dE(fi))) ;i.e.,  fi^(\prime\prime\prime)

# first, second, and third order corrections to E
d1 = -fi / fip
d2 = -fi / (fip + d1 * fipp / 2.0)
d3 = -fi / (fip + d2 * fipp / 2.0 + d2 * d2 * fippp / 6.0)
E = E + d3
Earr[convd] = E
fiarr = ( Earr - eccarr * np.sin( Earr ) - Marr) # how well did we do?
convd = np.abs(fiarr) > conv  # test for convergence
nd = np.sum(convd is True)

if Earr.size > 1:
return Earr
else:
return Earr

def profile():
# Profile and compare C-based Kepler solver with
# Python/Numpy implementation

import timeit

ecc = 0.20
numloops = 5000

for size in [10, 30, 100, 300, 1000]:

setup = """\
import numpy as np

gc.enable()

ecc = %f
orbel = [32.468, 2456000, ecc, np.pi/2, 10.0]

t = np.linspace(2455000, 2457000, %d)
""" % (ecc, size)

print("\nProfiling pure C code for an RV time series with {} "
"observations".format(size))
tc = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=True)',
setup=setup, number=numloops)
print("Ran %d model calculations in %5.3f seconds" % (numloops, tc))

print("Profiling Python code for an RV time series with {} "
"observations".format(size))
tp = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=False)',
setup=setup, number=numloops)
print("Ran %d model calculations in %5.3f seconds" % (numloops, tp))
print("The C version runs %5.2f times faster" % (tp/tc))