Merge pull request #148 from christinahedges/simple

src/exoplanet/orbits/simple.py

@@ -2,6 +2,7 @@
 
 __all__ = ["SimpleTransitOrbit"]
 
+import numpy as np
 import theano.tensor as tt
 
 from ..utils import as_tensor_variable
@@ -22,16 +23,17 @@ class SimpleTransitOrbit:
 
     """
 
-    def __init__(self, period=None, t0=0.0, b=0.0, duration=None, r_star=1.0):
+    def __init__(self, period, duration, t0=0.0, b=0.0, r_star=1.0, ror=0):
         self.period = as_tensor_variable(period)
         self.t0 = as_tensor_variable(t0)
         self.b = as_tensor_variable(b)
         self.duration = as_tensor_variable(duration)
         self.r_star = as_tensor_variable(r_star)
 
         self._b_norm = self.b * self.r_star
-        x2 = self.r_star ** 2 - self._b_norm ** 2
-        self.speed = 2 * tt.sqrt(x2) / self.duration
+        x2 = r_star ** 2 * ((1 + ror) ** 2 - b ** 2)
+        self.speed = 2 * np.sqrt(x2) / duration
+
         self._half_period = 0.5 * self.period
         self._ref_time = self.t0 - self._half_period
 

tests/orbits/simple_test.py

@@ -5,7 +5,8 @@
 import theano
 
 from exoplanet.light_curves import LimbDarkLightCurve
-from exoplanet.orbits.simple import SimpleTransitOrbit
+from exoplanet.orbits import SimpleTransitOrbit
+from exoplanet.orbits import KeplerianOrbit
 
 
 def test_simple():
@@ -21,7 +22,9 @@ def test_simple():
         < 0.5 * duration
     )
 
-    orbit = SimpleTransitOrbit(period, t0, b, duration, r_star)
+    orbit = SimpleTransitOrbit(
+        period=period, duration=duration, t0=t0, b=b, r_star=r_star
+    )
 
     x, y, z = theano.function([], orbit.get_planet_position(t))()
     b_val = np.sqrt(x ** 2 + y ** 2)
@@ -48,30 +51,34 @@ def test_simple():
         orbit.get_radial_velocity(t)
 
 
-def test_simple_light_curve():
-    period = 3.456
-    t0 = 1.45
-    b = 0.5
-    duration = 0.12
-    r_star = 1.345
+def test_simple_light_curve_compare_kepler():
+    t = np.linspace(0.0, 1, 1000)
+    # We use a long period, because at short periods there is a big difference
+    # between a circular orbit and an object moving on a straight line.
+    period = 1000
+    t0 = 0.5
+    r = 0.01
+    r_star = 1
+    b = 1 - r / r_star * 3
 
-    t = t0 + np.linspace(-2 * period, 2 * period, 5000)
-    m0 = (
-        np.abs((t - t0 + 0.5 * period) % period - 0.5 * period)
-        < 0.5 * duration
+    star = LimbDarkLightCurve([0])
+    orbit_keplerian = KeplerianOrbit(
+        period=period, t0=t0, b=b, r_star=r_star, m_star=1
     )
-    orbit = SimpleTransitOrbit(period, t0, b, duration, r_star)
-
-    star = LimbDarkLightCurve([0.2, 0.3])
-    lc1 = star.get_light_curve(
-        orbit=orbit, r=0.01, t=t, use_in_transit=False
+    duration = (period / np.pi) * np.arcsin(
+        ((r_star + r) ** 2 - (b * r_star) ** 2) ** 0.5 / orbit_keplerian.a
     ).eval()
-    lc2 = star.get_light_curve(orbit=orbit, r=0.01, t=t).eval()
-    assert np.allclose(lc1, lc2)
-    assert np.allclose(lc2[~m0], 0.0)
 
-    lc1 = star.get_light_curve(
-        orbit=orbit, r=0.01, t=t, texp=0.01, use_in_transit=False
-    ).eval()
-    lc2 = star.get_light_curve(orbit=orbit, r=0.01, t=t, texp=0.01).eval()
-    assert np.allclose(lc1, lc2)
+    lc_keplerian = star.get_light_curve(orbit=orbit_keplerian, r=r, t=t)
+    orbit_simple1 = SimpleTransitOrbit(
+        period=period,
+        t0=t0,
+        b=b,
+        duration=duration,
+        r_star=r_star,
+        ror=r / r_star,
+    )
+    lc_simple1 = star.get_light_curve(orbit=orbit_simple1, r=r, t=t)
+
+    # Should look similar to Keplerian orbit
+    assert np.allclose(lc_keplerian.eval(), lc_simple1.eval(), rtol=0.001)