approx in transit

exoplanet/light_curve.py

@@ -39,7 +39,8 @@ def __init__(self, u, r_star=1.0, model=None):
         self.c = get_cl(u_ext)
         self.c_norm = self.c / (np.pi * (self.c[0] + 2 * self.c[1] / 3))
 
-    def get_light_curve(self, r, orbit, t, texp=None, oversample=7, order=2):
+    def get_light_curve(self, r, orbit, t, texp=None, oversample=7, order=2,
+                        use_approx_in_transit=True, duration_factor=3):
         """Get the light curve for an orbit at a set of times
 
         Args:
@@ -71,11 +72,21 @@ def get_light_curve(self, r, orbit, t, texp=None, oversample=7, order=2):
                 centered Riemann sum (equivalent to the "resampling" procedure
                 suggested by Kipping 2010), ``1`` for the trapezoid rule, or
                 ``2`` for Simpson's rule.
+            use_approx_in_transit (Optional[bool]): If ``True``, the model will
+                only be evaluated for the data points expected to be in transit
+                as computed using the ``approx_in_transit`` method on
+                ``orbit``.
 
         """
         r = tt.as_tensor_variable(r)
         t = tt.as_tensor_variable(t)
 
+        if use_approx_in_transit:
+            transit_model = tt.zeros_like(t)
+            inds = orbit.approx_in_transit(t, r=r,
+                                           duration_factor=duration_factor)
+            t = t[inds]
+
         if texp is None:
             tgrid = t
             rgrid = tt.shape_padleft(r, tgrid.ndim) \
@@ -121,8 +132,13 @@ def get_light_curve(self, r, orbit, t, texp=None, oversample=7, order=2):
             stencil = tt.shape_padright(tt.shape_padleft(stencil, t.ndim),
                                         r.ndim)
             lc = tt.sum(stencil * lc, axis=t.ndim)
+        lc = tt.sum(lc, axis=-1)
 
-        return lc
+        if use_approx_in_transit:
+            transit_model = tt.set_subtensor(transit_model[inds], lc)
+            return transit_model
+        else:
+            return lc
 
     def compute_light_curve(self, b, r, los=None):
         """Compute the light curve for a set of impact parameters and radii

exoplanet/light_curve_test.py

@@ -10,6 +10,7 @@
 
 import starry
 
+from .orbits import KeplerianOrbit
 from .light_curve import StarryLightCurve
 
 
@@ -41,3 +42,32 @@ def test_light_curve_grad():
 
     lc = lambda u, b, r: StarryLightCurve(u).compute_light_curve(b, r)  # NOQA
     utt.verify_grad(lc, [u_val, b_val, r_val])
+
+
+def test_approx_in_transit():
+    t = np.linspace(-20, 20, 1000)
+    m_planet = np.array([0.3, 0.5])
+    m_star = 1.45
+    orbit = KeplerianOrbit(
+        m_star=m_star,
+        r_star=1.5,
+        t0=np.array([0.5, 17.4]),
+        period=np.array([10.0, 5.3]),
+        ecc=np.array([0.1, 0.8]),
+        omega=np.array([0.5, 1.3]),
+        m_planet=m_planet,
+    )
+    u = np.array([0.2, 0.3, 0.1, 0.5])
+    r = np.array([0.1, 0.01])
+
+    lc = StarryLightCurve(u)
+    model1 = lc.get_light_curve(r, orbit, t)
+    model2 = lc.get_light_curve(r, orbit, t, use_approx_in_transit=False)
+    vals = theano.function([], [model1, model2])()
+    utt.assert_allclose(*vals)
+
+    model1 = lc.get_light_curve(r, orbit, t, texp=0.1)
+    model2 = lc.get_light_curve(r, orbit, t, texp=0.1,
+                                use_approx_in_transit=False)
+    vals = theano.function([], [model1, model2])()
+    utt.assert_allclose(*vals)

exoplanet/orbits/keplerian.py

@@ -30,7 +30,10 @@ def __init__(self,
                  t0=0.0, incl=None, b=None,
                  m_star=None, r_star=None,
                  ecc=None, omega=None,
-                 m_planet=0.0, model=None, **kwargs):
+                 m_planet=0.0, model=None,
+                 m_planet_units=None,
+                 rho_units=None,
+                 **kwargs):
         add_citations_to_model(self.__citations__, model=model)
 
         self.kepler_op = KeplerOp(**kwargs)
@@ -39,9 +42,12 @@ def __init__(self,
         self.period = tt.as_tensor_variable(period)
         self.t0 = tt.as_tensor_variable(t0)
         self.m_planet = tt.as_tensor_variable(m_planet)
+        if m_planet_units is not None:
+            self.m_planet *= (1 * m_planet_units).to(u.M_sun).value
 
         self.a, self.period, self.rho_star, self.r_star, self.m_star = \
-            self._get_consistent_inputs(a, period, rho_star, r_star, m_star)
+            self._get_consistent_inputs(a, period, rho_star, r_star, m_star,
+                                        rho_units)
         self.m_total = self.m_star + self.m_planet
 
         self.n = 2 * np.pi / self.period
@@ -82,7 +88,8 @@ def __init__(self,
 
             self.K0 /= tt.sqrt(1 - self.ecc**2)
 
-    def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star):
+    def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star,
+                               rho_units):
         if a is None and period is None:
             raise ValueError("values must be provided for at least one of a "
                              "and period")
@@ -101,6 +108,7 @@ def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star):
                 r_star = 1.0
             rho_star = 3*np.pi*(a / r_star)**3 / (G_grav*period**2)
             rho_star -= 3*self.m_planet/(4*np.pi*r_star**3)
+            rho_units = None
 
         # Make sure that the right combination of stellar parameters are given
         if r_star is None and m_star is None:
@@ -112,8 +120,9 @@ def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star):
                              "rho_star, m_star, and r_star")
 
         if rho_star is not None:
-            # Convert density to M_sun / R_sun^3
-            rho_star = tt.as_tensor_variable(rho_star) / gcc_to_sun
+            rho_star = tt.as_tensor_variable(rho_star)
+            if rho_units is not None:
+                rho_star *= (1 * rho_units).to(u.M_sun / u.R_sun**3).value
         if r_star is not None:
             r_star = tt.as_tensor_variable(r_star)
         if m_star is not None:
@@ -161,17 +170,44 @@ def _get_position(self, a, t):
         return self._rotate_vector(r * cosf, r * tt.sin(f))
 
     def get_planet_position(self, t):
-        """The planets' positions in Solar radii"""
+        """The planets' positions in the barycentric frame
+
+        Args:
+            t: The times where the position should be evaluated.
+
+        Returns:
+            x, y, z: The components of the position vector at ``t`` in units
+                of ``R_sun``.
+
+        """
         return tuple(tt.squeeze(x)
                      for x in self._get_position(self.a_planet, t))
 
     def get_star_position(self, t):
-        """The star's position in Solar radii"""
+        """The star's position in the barycentric frame
+
+        Args:
+            t: The times where the position should be evaluated.
+
+        Returns:
+            x, y, z: The components of the position vector at ``t`` in units
+                of ``R_sun``.
+
+        """
         return tuple(tt.squeeze(x)
                      for x in self._get_position(self.a_star, t))
 
     def get_relative_position(self, t):
-        """The planets' positions relative to the star"""
+        """The planets' positions relative to the star
+
+        Args:
+            t: The times where the position should be evaluated.
+
+        Returns:
+            x, y, z: The components of the position vector at ``t`` in units
+                of ``R_sun``.
+
+        """
         star = self._get_position(self.a_star, t)
         planet = self._get_position(self.a_planet, t)
         return tuple(tt.squeeze(b-tt.shape_padright(tt.sum(a, axis=-1)))
@@ -185,11 +221,81 @@ def _get_velocity(self, m, t):
         return self._rotate_vector(-K*tt.sin(f), K*(tt.cos(f) + self.ecc))
 
     def get_planet_velocity(self, t):
-        """The planets' velocities in R_sun / day"""
+        """Get the planets' velocity vector
+
+        Args:
+            t: The times where the velocity should be evaluated.
+
+        Returns:
+            vx, vy, vz: The components of the velocity vector at ``t`` in units
+                of ``M_sun/day``.
+
+        """
         return tuple(tt.squeeze(x)
                      for x in self._get_velocity(-self.m_star, t))
 
     def get_star_velocity(self, t):
-        """The star's velocity in R_sun / day"""
+        """Get the star's velocity vector
+
+        Args:
+            t: The times where the velocity should be evaluated.
+
+        Returns:
+            vx, vy, vz: The components of the velocity vector at ``t`` in units
+                of ``M_sun/day``.
+
+        """
         return tuple(tt.squeeze(x)
                      for x in self._get_velocity(self.m_planet, t))
+
+    def get_radial_velocity(self, t, output_units=None):
+        """Get the radial velocity of the star
+
+        Args:
+            t: The times where the radial velocity should be evaluated.
+            output_units (Optional): An AstroPy velocity unit. If not given,
+                the output will be evaluated in ``m/s``.
+
+        Returns:
+            vrad: The radial velocity evaluated at ``t`` in units of
+                ``output_units``.
+
+        """
+        if output_units is None:
+            output_units = u.m / u.s
+        conv = (1 * u.R_sun / u.day).to(output_units).value
+        v = self.get_star_velocity(t)
+        return conv * v[2]
+
+    def approx_in_transit(self, t, r=0.0, duration_factor=3):
+        """Get a list of timestamps that are expected to be in transit
+
+        Args:
+            t (vector): A vector of timestamps to be evaluated.
+            r (Optional): The radii of the planets.
+            duration_factor (Optional[float]): The factor by which to multiply
+                the approximate duration when computing the in transit points.
+                Larger values will be more conservative and might be needed for
+                large planets or very eccentric orbits.
+
+        Returns:
+            inds (vector): The indices of the timestamps that are expected to
+                be in transit.
+
+        """
+        # Estimate the maximum duration of the transit using the equations
+        # from Winn (2010)
+        arg = (self.r_star + r) / (-self.a_planet)
+        max_dur = self.period * tt.arcsin(arg) / np.pi
+        if self.ecc is not None:
+            max_dur *= tt.sqrt(1-self.ecc**2) / (1+self.ecc*self.sin_omega)
+
+        # Wrap the times into time since transit
+        hp = 0.5 * self.period
+        dt = tt.mod(tt.shape_padright(t) - self.t0 + hp, self.period) - hp
+
+        # Estimate the data points that are within the maximum duration of the
+        # transit
+        mask = tt.any(tt.abs_(dt) < 0.5 * duration_factor * max_dur, axis=-1)
+
+        return tt.arange(t.size)[mask]

exoplanet/orbits/keplerian_test.py

@@ -105,3 +105,26 @@ def test_velocity():
         g = theano.grad(tt.sum(planet_pos[i]), t_tensor)
         planet_vel_expect[i] = theano.function([t_tensor], g)(t)
     utt.assert_allclose(planet_vel, planet_vel_expect)
+
+
+def test_approx_in_transit():
+    t = np.linspace(-20, 20, 1000)
+    m_planet = np.array([0.3, 0.5])
+    m_star = 1.45
+    orbit = KeplerianOrbit(
+        m_star=m_star,
+        r_star=1.5,
+        t0=np.array([0.5, 17.4]),
+        period=np.array([10.0, 5.3]),
+        ecc=np.array([0.1, 0.8]),
+        omega=np.array([0.5, 1.3]),
+        m_planet=m_planet,
+    )
+
+    coords = theano.function([], orbit.get_relative_position(t))()
+    r = np.sqrt(coords[0]**2 + coords[1]**2)
+    inds = theano.function([], orbit.approx_in_transit(t))()
+    m = np.isin(np.arange(len(t)), inds)
+    ok = (r[~m] > 2) | (coords[2][~m] > 0)
+
+    assert np.all(ok)

exoplanet/theano_ops/kepler/solver.py

@@ -70,6 +70,3 @@ def R_op(self, inputs, eval_points):
         if eval_points[0] is None:
             return eval_points
         return self.grad(inputs, eval_points)
-
-    def c_compile_args(self, compiler):
-        return ["-O3"]