docstrings

exoplanet/orbits/__init__.py

@@ -1,5 +1,9 @@
 # -*- coding: utf-8 -*-
 
-__all__ = ["KeplerianOrbit", "get_true_anomaly"]
+__all__ = [
+    "KeplerianOrbit", "get_true_anomaly",
+    "TTVOrbit",
+]
 
 from .keplerian import KeplerianOrbit, get_true_anomaly
+from .ttv import TTVOrbit

exoplanet/orbits/keplerian.py

@@ -21,18 +21,53 @@
 class KeplerianOrbit(object):
     """A system of bodies on Keplerian orbits around a common central
 
+    Given the input parameters, the values of all other parameters will be
+    computed so a ``KeplerianOrbit`` instance will always have attributes for
+    each argument. Note that the units of the computed attributes will all be
+    in the standard units of this class (``R_sun``, ``M_sun``, and ``days``)
+    except for ``rho_star`` which will be in ``g / cm^3``.
+
+    There are only specific combinations of input parameters that can be used:
+
+    1. First, either ``period`` or ``a`` must be given. If values are given
+       for both parameters, then neither ``m_star`` or ``rho_star`` can be
+       defined because the stellar density implied by each planet will be
+       computed in ``rho_star``.
+    2. Only one of ``incl`` and ``b`` can be given.
+    3. If a value is given for ``ecc`` then ``omega`` must also be given.
+    4. If no stellar parameters are given, the central body is assumed to be
+       the sun. If only ``rho_star`` is defined, the radius of the central is
+       assumed to be ``1 * R_sun``. Otherwise, at most two of ``m_star``,
+       ``r_star``, and ``rho_star`` can be defined.
+
+    Args:
+        period: The orbital periods of the bodies in days.
+        a: The semimajor axes of the orbits in ``R_sun``.
+        t0: The time of a reference transit for each orbits in days.
+        incl: The inclinations of the orbits in radians.
+        b: The impact parameters of the orbits.
+        ecc: The eccentricities of the orbits. Must be ``0 <= ecc < 1``.
+        omega: The arguments of periastron for the orbits in radians.
+        m_planet: The masses of the planets in units of ``m_planet_units``.
+        m_star: The mass of the star in ``M_sun``.
+        r_star: The radius of the star in ``R_sun``.
+        rho_star: The density of the star in units of ``rho_star_units``.
+        m_planet_units: An ``astropy.units`` compatible unit object giving the
+            units of the planet masses. If not given, the default is ``M_sun``.
+        rho_star_units: An ``astropy.units`` compatible unit object giving the
+            units of the stellar density. If not given, the default is
+            ``g / cm^3``.
+
     """
 
     __citations__ = ("astropy", )
 
     def __init__(self,
-                 period=None, a=None, rho_star=None,
-                 t0=0.0, incl=None, b=None,
-                 m_star=None, r_star=None,
-                 ecc=None, omega=None,
-                 m_planet=0.0, model=None,
-                 m_planet_units=None,
-                 rho_units=None,
+                 period=None, a=None, t0=0.0, incl=None, b=None,
+                 ecc=None, omega=None, m_planet=0.0,
+                 m_star=None, r_star=None, rho_star=None,
+                 m_planet_units=None, rho_star_units=None,
+                 model=None,
                  **kwargs):
         add_citations_to_model(self.__citations__, model=model)
 
@@ -47,7 +82,7 @@ def __init__(self,
 
         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,
-                                        rho_units)
+                                        rho_star_units)
         self.m_total = self.m_star + self.m_planet
 
         self.n = 2 * np.pi / self.period
@@ -62,6 +97,8 @@ def __init__(self,
                 self.b = tt.as_tensor_variable(b)
                 self.incl = tt.arccos(self.b / self.a_planet)
         else:
+            if b is not None:
+                raise ValueError("only one of 'incl' and 'b' can be given")
             self.incl = tt.as_tensor_variable(incl)
             self.b = self.a_planet * tt.cos(self.incl)
 
@@ -89,7 +126,7 @@ def __init__(self,
             self.K0 /= tt.sqrt(1 - self.ecc**2)
 
     def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star,
-                               rho_units):
+                               rho_star_units):
         if a is None and period is None:
             raise ValueError("values must be provided for at least one of a "
                              "and period")
@@ -108,7 +145,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
+            rho_star_units = None
 
         # Make sure that the right combination of stellar parameters are given
         if r_star is None and m_star is None:
@@ -121,8 +158,10 @@ def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star,
 
         if rho_star is not None:
             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 rho_star_units is not None:
+                rho_star *= (1 * rho_star_units).to(u.M_sun / u.R_sun**3).value
+            else:
+                rho_star /= gcc_to_sun
         if r_star is not None:
             r_star = tt.as_tensor_variable(r_star)
         if m_star is not None:
@@ -153,8 +192,11 @@ def _rotate_vector(self, x, y):
             b = self.sin_omega * x + self.cos_omega * y
         return (a, self.cos_incl * b, self.sin_incl * b)
 
+    def _warp_times(self, t):
+        return tt.shape_padright(t)
+
     def _get_true_anomaly(self, t):
-        M = (tt.shape_padright(t) - self.tref) * self.n
+        M = (self._warp_times(t) - self.tref) * self.n
         if self.ecc is None:
             return M
         _, f = self.kepler_op(M, self.ecc + tt.zeros_like(M))
@@ -293,7 +335,7 @@ def approx_in_transit(self, t, r=0.0, duration_factor=3, texp=None):
 
         # 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
+        dt = tt.mod(self._warp_times(t) - self.t0 + hp, self.period) - hp
 
         # Estimate the data points that are within the maximum duration of the
         # transit

exoplanet/orbits/ttv.py

@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import division, print_function
+
+__all__ = ["TTVOrbit"]
+
+import theano.tensor as tt
+
+from .keplerian import KeplerianOrbit
+
+
+class TTVOrbit(KeplerianOrbit):
+    """A generalization of a Keplerian orbit with transit timing variations
+
+    Only one of the arguments ``ttvs`` or ``transit_times`` can be given and
+    the other will be computed from the one that was provided.
+
+    Args:
+        ttvs: A list (with on entry for each planet) of "O-C" vectors for each
+            transit of each planet in units of days. "O-C" means the
+            difference between the observed transit time and the transit time
+            expected for a regular periodic orbit.
+        transit_times: A list (with on entry for each planet) of transit times
+            for each transit of each planet in units of days. These times will
+            be used to compute the implied (least squares) ``period`` and
+            ``t0`` so these parameters cannot also be given.
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        ttvs = kwargs.pop("ttvs", None)
+        transit_times = kwargs.pop("transit_times", None)
+        if ttvs is None and transit_times is None:
+            raise ValueError("one of 'ttvs' or 'transit_times' must be "
+                             "defined")
+        if ttvs is not None:
+            self.ttvs = [tt.as_tensor_variable(ttv) for ttv in ttvs]
+        else:
+            if kwargs.pop("period", None) is not None:
+                raise ValueError("a period cannot be given if 'transit_times' "
+                                 "is defined")
+
+            self.transit_times = []
+            self.ttvs = []
+            period = []
+            t0 = []
+            for i, times in enumerate(transit_times):
+                times = tt.as_tensor_variable(times)
+
+                N = times.shape[0]
+                AT = tt.stack((tt.arange(N, dtype=times.dtype),
+                               tt.ones_like(times)), axis=0)
+                A = tt.transpose(AT)
+                ATA = tt.dot(AT, A)
+                ATy = tt.dot(AT, times)
+                w = tt.slinalg.solve_symmetric(ATA, ATy)
+                expect = tt.dot(w, AT)
+
+                period.append(w[0])
+                t0.append(w[1])
+                self.ttvs.append(times - expect)
+                self.transit_times.append(times)
+
+            kwargs["period"] = tt.stack(period)
+            kwargs["t0"] = tt.stack(t0)
+
+        super(TTVOrbit, self).__init__(*args, **kwargs)
+        self._base_time = 0.5 - self.t0 / self.period
+
+        if ttvs is not None:
+            self.transit_times = [
+                self.t0[i] + self.period[i] * tt.arange(ttv.shape[0]) + ttv
+                for i, ttv in enumerate(self.ttvs)]
+
+    def _warp_times(self, t):
+        delta = tt.shape_padleft(t) / tt.shape_padright(self.period, t.ndim)
+        delta += tt.shape_padright(self._base_time, t.ndim)
+        ind = tt.cast(tt.floor(delta), "int64")
+        dt = tt.stack([ttv[tt.clip(ind[i], 0, ttv.shape[0]-1)]
+                       for i, ttv in enumerate(self.ttvs)], -1)
+        return tt.shape_padright(t) + dt

exoplanet/sampling.py

@@ -113,6 +113,7 @@ def warmup(self, start=None, step_kwargs=None, **kwargs):
             step_kwargs = {}
         step = self.get_step_for_trace(**step_kwargs)
         self._extend(self.start, start=start, step=step, **kwargs)
+        self._current_step = self.get_step_for_trace(**step_kwargs)
         return self._current_trace
 
     def _get_start_and_step(self, start=None, step_kwargs=None, trace=None,

exoplanet/utils.py

@@ -11,6 +11,19 @@
 import pymc3 as pm
 
 
+def get_args_for_theano_function(point=None, model=None):
+    model = pm.modelcontext(model)
+    if point is None:
+        point = model.test_point
+    return [point[k.name] for k in model.vars]
+
+
+def get_theano_function_for_var(var, model=None, **kwargs):
+    model = pm.modelcontext(model)
+    kwargs["on_unused_input"] = kwargs.get("on_unused_input", "ignore")
+    return theano.function(model.vars, var, **kwargs)
+
+
 def eval_in_model(var, point=None, return_func=False, model=None, **kwargs):
     """Evaluate a Theano tensor or PyMC3 variable in a PyMC3 model
 
@@ -34,14 +47,8 @@ def eval_in_model(var, point=None, return_func=False, model=None, **kwargs):
         or this value, the Theano function, and the input arguments.
 
     """
-    model = pm.modelcontext(model)
-    if point is None:
-        point = model.test_point
-
-    kwargs["on_unused_input"] = kwargs.get("on_unused_input", "ignore")
-    func = theano.function(model.vars, var, **kwargs)
-    args = [point[k.name] for k in model.vars]
-
+    func = get_theano_function_for_var(var, model=model, **kwargs)
+    args = get_args_for_theano_function(point=point, model=model)
     if return_func:
         return func(*args), func, args
     return func(*args)