fixing approx_in_transit with texp included

docs/user/api.rst

@@ -11,9 +11,25 @@ Orbits
 .. autoclass:: exoplanet.orbits.KeplerianOrbit
    :inherited-members:
 
+
 Light curve models
 ------------------
 
 .. autoclass:: exoplanet.light_curve.StarryLightCurve
    :inherited-members:
 
+
+Estimators
+----------
+
+.. autofunction:: exoplanet.estimators.estimate_minimum_mass
+
+
+Distributions
+-------------
+
+.. autoclass:: exoplanet.distributions.UnitVector
+.. autoclass:: exoplanet.distributions.Angle
+.. autoclass:: exoplanet.distributions.Triangle
+.. autoclass:: exoplanet.distributions.RadiusImpactParameter
+

exoplanet/__init__.py

@@ -9,7 +9,8 @@
 
 if not __EXOPLANET_SETUP__:
     __all__ = ["distributions", "gp", "orbits", "sampling", "utils",
+               "estimators",
                "StarryLightCurve"]
 
-    from . import distributions, gp, orbits, sampling, utils
+    from . import distributions, gp, orbits, sampling, utils, estimators
     from .light_curve import StarryLightCurve

exoplanet/estimators.py

@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import division, print_function
+
+__all__ = ["estimate_minimum_mass"]
+
+import numpy as np
+
+import astropy.units as u
+
+
+def _get_design_matrix(periods, t0s, x):
+    if t0s is not None:
+        return np.vstack([
+            np.cos(2*np.pi*(x - (t0s[i] - 0.25*periods[i])) / periods[i])
+            for i in range(len(periods))
+        ] + [np.ones(len(x))]).T
+    return np.concatenate([
+        (np.sin(2*np.pi*x / periods[i]),
+         np.cos(2*np.pi*x / periods[i]))
+        for i in range(len(periods))
+    ] + [np.ones((1, len(x)))], axis=0).T
+
+
+def estimate_minimum_mass(periods, x, y, yerr=None, t0s=None, m_star=1):
+    """Estimate the minimum mass(es) for planets in an RV series
+
+    Args:
+        periods: The periods of the planets. Assumed to be in ``days`` if not
+            an AstroPy Quantity.
+        x: The observation times. Assumed to be in ``days`` if not an AstroPy
+            Quantity.
+        y: The radial velocities. Assumed to be in ``m/s`` if not an AstroPy
+            Quantity.
+        yerr (Optional): The uncertainty on ``y``.
+        t0s (Optional): The time of a reference transit for each planet, if
+            known.
+        m_star (Optional): The mass of the star. Assumed to be in ``M_sun``
+            if not an AstroPy Quantity.
+
+    Returns:
+        msini: An estimate of the minimum mass of each planet as an AstroPy
+            Quantity with units of ``M_jupiter``.
+
+    """
+    if yerr is None:
+        ivar = np.ones_like(y)
+    else:
+        ivar = 1.0 / yerr**2
+
+    m_star = u.Quantity(m_star, unit=u.M_sun)
+    periods = u.Quantity(np.atleast_1d(periods), unit=u.day)
+    if t0s is not None:
+        t0s = u.Quantity(np.atleast_1d(t0s), unit=u.day).value
+    x = u.Quantity(np.atleast_1d(x), unit=u.day)
+    y = u.Quantity(np.atleast_1d(y), unit=u.m / u.s)
+    ivar = u.Quantity(np.atleast_1d(ivar), unit=(u.s / u.m) ** 2)
+
+    D = _get_design_matrix(periods.value, t0s, x.value)
+    w = np.linalg.solve(np.dot(D.T, D*ivar.value[:, None]),
+                        np.dot(D.T, y.value*ivar.value))
+    if t0s is not None:
+        K = w[:-1]
+    else:
+        w = w[:-1]
+        K = np.sqrt(w[::2]**2 + w[1::2]**2)
+
+    m_J = K / 28.4329 * m_star.value**(2./3)
+    m_J *= (periods.to(u.year)).value**(1./3)
+
+    return m_J * u.M_jupiter

exoplanet/light_curve.py

@@ -83,7 +83,7 @@ def get_light_curve(self, r, orbit, t, texp=None, oversample=7, order=2,
 
         if use_approx_in_transit:
             transit_model = tt.zeros_like(t)
-            inds = orbit.approx_in_transit(t, r=r,
+            inds = orbit.approx_in_transit(t, r=r, texp=texp,
                                            duration_factor=duration_factor)
             t = t[inds]
 

exoplanet/orbits/keplerian.py

@@ -267,7 +267,7 @@ def get_radial_velocity(self, t, output_units=None):
         v = self.get_star_velocity(t)
         return conv * v[2]
 
-    def approx_in_transit(self, t, r=0.0, duration_factor=3):
+    def approx_in_transit(self, t, r=0.0, duration_factor=3, texp=None):
         """Get a list of timestamps that are expected to be in transit
 
         Args:
@@ -277,6 +277,7 @@ def approx_in_transit(self, t, r=0.0, duration_factor=3):
                 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.
+            texp (Optional[float]): The exposure time.
 
         Returns:
             inds (vector): The indices of the timestamps that are expected to
@@ -296,6 +297,9 @@ def approx_in_transit(self, t, r=0.0, duration_factor=3):
 
         # 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)
+        tol = 0.5 * duration_factor * max_dur
+        if texp is not None:
+            tol += 0.5 * texp
+        mask = tt.any(tt.abs_(dt) < tol, axis=-1)
 
         return tt.arange(t.size)[mask]