hacking on kepler solver again

exoplanet-dev · Nov 27, 2018 · a5b0407 · a5b0407
1 parent 96585d7
commit a5b0407
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 34 deletions.
diff --git a/exoplanet/theano_ops/kepler/solver.cc b/exoplanet/theano_ops/kepler/solver.cc
@@ -1,5 +1,9 @@
 #section support_code
 
+inline int sign(double x) {
+  return (x > 0) - (x < 0);
+}
+
 inline double wrap_into (double x, double period) {
     return x - period * floor(x / period);
 }
@@ -74,13 +78,10 @@ int APPLY_SPECIFIC(solver)(
   DTYPE_OUTPUT_0* E_out = (DTYPE_OUTPUT_0*)PyArray_DATA(*output0);
   DTYPE_OUTPUT_1* f_out = (DTYPE_OUTPUT_1*)PyArray_DATA(*output1);
 
-  T M_old = M_in[0], e_old = e_in[0];
-  T M, e, E;
-  T sinE, cosE, tanE2, gp, delta;
-  int flag = 0;
+  T M, e, E, delta, f, fp, fpp, fppp, ffp, fp2, arg, arg2, sinE, tanE2;
 
   for (npy_intp n = 0; n < N; ++n) {
-    M = wrap_into(M_in[n] + M_PI, 2 * M_PI) - M_PI;
+    M = M_in[n];
     e = e_in[n];
 
     if (e >= 1) {
@@ -92,39 +93,23 @@ int APPLY_SPECIFIC(solver)(
 
       // Special case for zero eccentricity
       E_out[n] = M;
-      f_out[n] = M;
-      flag = 0;
+      f_out[n] = wrap_into(M + M_PI, 2 * M_PI) - M_PI;
 
     } else {
 
-      // Estimate the initialization using the first order Taylor series
-      if (flag) {
-        delta = ((e - e_old) * sinE + (M - M_old)) / gp;
-        E += delta;
-      } else {
-        E = M;
-      }
-
-      // Do some iterations of Newton Raphson
+      E = M + e*sin(M);
       sinE = sin(E);
-      cosE = cos(E);
-      for (int i = 0; i < maxiter; ++i) {
-        gp = 1 - e * cosE;
-        delta = (E - e * sinE - M) / gp;
-        if (fabs(delta) <= T(tol)) break;
-        E -= delta;
+      for (int n = 0; n < maxiter; ++n) {
+        delta = e * sinE + M;
+        E -= (E - delta) / (1 - e * cos(E));
         sinE = sin(E);
-        cosE = cos(E);
+        if (fabs(E - e * sinE - M) <= tol) break;
       }
 
       // Save the result and compute the true anomaly
       E_out[n] = E;
-      tanE2 = sinE / (1 + cosE);  // tan(0.5*E)
+      tanE2 = sinE / (1 + cos(E));  // tan(0.5*E)
       f_out[n] = 2 * atan(sqrt((1+e)/(1-e))*tanE2);
-
-      flag = 1;
-      e_old = e;
-      M_old = M;
     }
   }
 

diff --git a/exoplanet/theano_ops/kepler/solver.py b/exoplanet/theano_ops/kepler/solver.py
@@ -18,7 +18,7 @@ class KeplerOp(gof.COp):
     func_file = "./solver.cc"
     func_name = "APPLY_SPECIFIC(solver)"
 
-    def __init__(self, tol=1e-8, maxiter=2000, **kwargs):
+    def __init__(self, tol=1e-12, maxiter=1000, **kwargs):
         self.tol = float(tol)
         self.maxiter = int(maxiter)
         super(KeplerOp, self).__init__(self.func_file, self.func_name)

diff --git a/exoplanet/theano_ops/kepler/solver_test.py b/exoplanet/theano_ops/kepler/solver_test.py
@@ -23,9 +23,6 @@ def _get_M_and_f(self, e, E):
         f = 2 * np.arctan(np.sqrt((1+e)/(1-e)) * np.tan(0.5*E))
         return M, f
 
-    def _wrap_E(self, E):
-        return (E + np.pi) % (2*np.pi) - np.pi
-
     def test_edge(self):
         E = np.array([0.0, 2*np.pi, -226.2])
         e = (1 - 1e-6) * np.ones_like(E)
@@ -37,7 +34,7 @@ def test_edge(self):
         E0, f0 = func(M, e)
 
         assert np.all(np.isfinite(E0))
-        utt.assert_allclose(self._wrap_E(E), E0)
+        utt.assert_allclose(E, E0)
         assert np.all(np.isfinite(f0))
         utt.assert_allclose(f, f0)
 
@@ -52,7 +49,12 @@ def test_solver(self):
         e_t = tt.matrix()
         func = theano.function([M_t, e_t], self.op(M_t, e_t))
         E0, f0 = func(M, e)
-        utt.assert_allclose(self._wrap_E(E), E0)
+
+        delta = np.abs(E - E0)
+        ind = np.unravel_index(np.argmax(delta), delta.shape)
+        print(M[ind], e[ind], E[ind], E0[ind])
+
+        utt.assert_allclose(E, E0)
         utt.assert_allclose(f, f0)
 
     def test_infer_shape(self):