Skip to content

Commit

Permalink
Merge 48d650e into 65bf031
Browse files Browse the repository at this point in the history
  • Loading branch information
@devincody
devincody committed Sep 16, 2019
2 parents 65bf031 + 48d650e commit b1405a4
Show file tree
Hide file tree
Showing 5 changed files with 196 additions and 73 deletions.
26 changes: 26 additions & 0 deletions orbitize/_kepler.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,11 @@ cdef extern from "kepler.c":
const int max_iter,
double eanom[])

cdef extern from "kepler.c":
void mikkola_array(const int n_elements,
const double manom[],
const double ecc[],
double eanom[])

def _c_newton_solver(np.ndarray[np.double_t,ndim=1] manom,
np.ndarray[np.double_t,ndim=1] ecc,
Expand Down Expand Up @@ -36,4 +41,25 @@ def _c_newton_solver(np.ndarray[np.double_t,ndim=1] manom,

newton_array(len(manom), <double*> manom.data, <double*> ecc.data, tolerance, max_iter, <double*> eanom.data)

return eanom

def _c_mikkola_solver(np.ndarray[np.double_t,ndim=1] manom,
np.ndarray[np.double_t,ndim=1] ecc):
"""
Wrapper function for C implementation of 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: Devin Cody, 2018
"""

# Initialize at E=M, E=pi is better at very high eccentricities
cdef np.ndarray eanom
eanom = np.zeros(len(manom))

mikkola_array(len(manom), <double*> manom.data, <double*> ecc.data, <double*> eanom.data)

return eanom
165 changes: 114 additions & 51 deletions orbitize/kepler.c
View file
Original file line number Diff line number Diff line change
@@ -1,69 +1,132 @@
#include <stdio.h>
#include <math.h>

# define M_PI 3.14159265358979323846 /* pi */
#ifndef M_PI
#define M_PI 3.14159265358979323846 /* pi */
#endif

void newton_array(const int n_elements,
const double manom[],
const double ecc[],
const double tol,
const int max_iter,
double eanom[]){
/*
Vectorized C++ Newton-Raphson solver for eccentric anomaly.
const double manom[],
const double ecc[],
const double tol,
const int max_iter,
double eanom[]){
/*
Vectorized C Newton-Raphson solver for eccentric anomaly.
Args:
manom (double[]): array of mean anomalies
ecc (double[]): array of eccentricities
eanom0 (double[]): array of first guess for eccentric anomaly, same shape as manom (optional)
Return:
None: eanom is changed by reference
None: eanom (double[]): is changed by reference
Written: Devin Cody, 2018
*/
register int i;
for (i = 0; i < n_elements; i ++){
double diff;
int niter = 0;
int half_max = max_iter/2.0; // divide max_iter by 2 using bit shift
// Let's do one iteration to start with
eanom[i] -= (eanom[i] - (ecc[i] * sin(eanom[i])) - manom[i]) / (1.0 - (ecc[i] * cos(eanom[i])));
diff = (eanom[i] - (ecc[i] * sin(eanom[i])) - manom[i]) / (1.0 - (ecc[i] * cos(eanom[i])));

while ((fabs(diff) > tol) && (niter <= max_iter)){
eanom[i] -= diff;

// If it hasn't converged after half the iterations are done, try starting from pi
if (niter == half_max) {
eanom[i] = M_PI;
}

diff = (eanom[i] - (ecc[i] * sin(eanom[i])) - manom[i]) / (1.0 - (ecc[i] * cos(eanom[i])));
niter += 1;
}

// If it has not converged, set eccentricity to -1 to signal that it needs to be
// solved using the analytical version. Note this behavior is a bit different from the
// numpy implementation
if (niter >= max_iter){
printf("%f %f %f %f >= %d iter\n", manom[i], eanom[i], diff, ecc[i], max_iter);
eanom[i] = -1;
}
}
*/
int i;
for (i = 0; i < n_elements; i ++){
double diff;
int niter = 0;
int half_max = max_iter/2.0; // divide max_iter by 2 using bit shift
// Let's do one iteration to start with
eanom[i] -= (eanom[i] - (ecc[i] * sin(eanom[i])) - manom[i]) / (1.0 - (ecc[i] * cos(eanom[i])));
diff = (eanom[i] - (ecc[i] * sin(eanom[i])) - manom[i]) / (1.0 - (ecc[i] * cos(eanom[i])));

while ((fabs(diff) > tol) && (niter <= max_iter)){
eanom[i] -= diff;

// If it hasn't converged after half the iterations are done, try starting from pi
if (niter == half_max) {
eanom[i] = M_PI;
}

diff = (eanom[i] - (ecc[i] * sin(eanom[i])) - manom[i]) / (1.0 - (ecc[i] * cos(eanom[i])));
niter += 1;
}

// If it has not converged, set eccentricity to -1 to signal that it needs to be
// solved using the analytical version. Note this behavior is a bit different from the
// numpy implementation
if (niter >= max_iter){
printf("%f %f %f %f >= %d iter\n", manom[i], eanom[i], diff, ecc[i], max_iter);
eanom[i] = -1;
}
}
}


void mikkola_array(const int n_elements, const double manom[], const double ecc[], double eanom[]){
/*
Vectorized C 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 (double[]): mean anomaly, must be between 0 and pi.
ecc (double[]): eccentricity
eanom0 (double[]): array for eccentric anomaly
Return:
None: eanom (double[]): is changed by reference
Written: Devin Cody, 2019
*/

int i;
double alpha, beta, aux, z, s0, s1, se0, ce0;
double f, f1, f2, f3, f4, u1, u2, u3;

for (i = 0; i < n_elements; i++){
alpha = (1.0 - ecc[i]) / ((4.0 * ecc[i]) + 0.5);
beta = (0.5 * manom[i]) / ((4.0 * ecc[i]) + 0.5);

aux = sqrt(beta*beta + alpha*alpha*alpha);
z = pow(fabs(beta + aux), (1.0/3.0));

s0 = z - (alpha/z);
s1 = s0 - (0.078*(pow(s0, 5))) / (1.0 + ecc[i]);
eanom[i] = manom[i] + (ecc[i] * (3.0*s1 - 4.0*(s1*s1*s1)));

se0=sin(eanom[i]);
ce0=cos(eanom[i]);

f = eanom[i]-ecc[i]*se0-manom[i];
f1 = 1.0-ecc[i]*ce0;
f2 = ecc[i]*se0;
f3 = ecc[i]*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);
eanom[i] += -f/(f1+0.5*f2*u3+(1.0/6.0)*f3*u3*u3+(1.0/24.0)*f4*(u3*u3*u3));
}
}


int main(void){
// Test functions with a small program
// Test functions with a small program

// Define variables for newton array method
double m[] = {.5, 1, 1.5};
double ecc[] = {.25, .75, .83};
double tol = 1e-9;
int mi = 100;
double eanom[] = {0, 0, 0};

// test newton_array
// Answer should be: [ 0.65161852, 1.73936894, 2.18046524])

newton_array(3, m, ecc, tol, mi, eanom);
int i;
for (i = 0; i < 3; i++){
printf("eanom[%d] = %f\n", i, eanom[i]);
eanom[i] = 0;
}

//Define variables for newton array method
double m[] = {.5, 1, 1.5};
double ecc[] = {.25, .75, .83};
double tol = 1e-9;
int mi = 100;
double eanom[] = {0, 0, 0};
// test mikkola_array
// Answer should be: [ 0.65161852, 1.73936894, 2.18046524])

//test newton_array
newton_array(3, m, ecc, tol, mi, eanom);
//Answer should be: [ 0.65161852, 1.73936894, 2.18046524])
mikkola_array(3, m, ecc, eanom);
for (i = 0; i < 3; i++){
printf("eanom[%d] = %f\n", i, eanom[i]);
}

return 0;
return 0;
}
15 changes: 9 additions & 6 deletions orbitize/kepler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ def calc_orbit(epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, mass_for_Kamp=No

return raoff, deoff, vz

def _calc_ecc_anom(manom, ecc, tolerance=1e-9, max_iter=100, use_cpp=False):
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)
Expand Down Expand Up @@ -152,7 +152,7 @@ def _calc_ecc_anom(manom, ecc, tolerance=1e-9, max_iter=100, use_cpp=False):

# Now low eccentricities
ind_low = np.where(~ecc_zero & ecc_low)
if cext and use_cpp:
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
Expand All @@ -163,7 +163,7 @@ def _calc_ecc_anom(manom, ecc, tolerance=1e-9, max_iter=100, use_cpp=False):
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])
if len(ind_high[0]) > 0: eanom[ind_high] = _mikkola_solver_wrapper(manom[ind_high], ecc[ind_high], use_c)

return np.squeeze(eanom)[()]

Expand Down Expand Up @@ -209,11 +209,11 @@ def _newton_solver(manom, ecc, tolerance=1e-9, max_iter=100, eanom0=None):

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]) # Send remaining orbits to the analytical version, this has not happened yet...
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):
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.
Expand All @@ -229,7 +229,10 @@ def _mikkola_solver_wrapper(manom, ecc):

ind_change = np.where(manom > np.pi)
manom[ind_change] = (2.0 * np.pi) - manom[ind_change]
eanom = _mikkola_solver(manom, ecc)
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
Expand Down
1 change: 1 addition & 0 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
try:
from Cython.Build import cythonize
except:
print("Error: Importing cython build environment failed")
USE_C_KEPLER_MODULE = 0


Expand Down

0 comments on commit b1405a4

Please sign in to comment.