Merge 60c83dc into bcb794f

src/exoplanet/orbits/constants.py

@@ -1,6 +1,12 @@
 # -*- coding: utf-8 -*-
 
-__all__ = ["G_grav", "gcc_per_sun", "au_per_R_sun", "day_per_yr_over_2pi"]
+__all__ = [
+    "G_grav",
+    "gcc_per_sun",
+    "au_per_R_sun",
+    "day_per_yr_over_2pi",
+    "c_light",
+]
 
 import warnings
 
@@ -12,6 +18,7 @@
     G_grav = c.G.to(u.R_sun ** 3 / u.M_sun / u.day ** 2).value
     gcc_per_sun = (u.M_sun / u.R_sun ** 3).to(u.g / u.cm ** 3)
     au_per_R_sun = u.R_sun.to(u.au)
+    c_light = c.c.to(u.R_sun / u.day).value
 
     day_per_yr_over_2pi = (
         (1.0 * u.au) ** (3 / 2)

src/exoplanet/orbits/keplerian.py

@@ -13,7 +13,7 @@
 from ..theano_ops.contact import ContactPointsOp
 from ..theano_ops.kepler import KeplerOp
 from ..units import has_unit, to_unit, with_unit
-from .constants import G_grav, au_per_R_sun, gcc_per_sun
+from .constants import G_grav, au_per_R_sun, gcc_per_sun, c_light
 
 
 class KeplerianOrbit:
@@ -343,23 +343,97 @@ def _get_position(self, a, t, parallax=None):
 
         return self._rotate_vector(r * cosf, r * sinf)
 
-    def get_planet_position(self, t, parallax=None):
+    def _get_retarded_position(self, a, t, parallax=None, z0=0.0):
+        """Get the retarded position of a body, accounting for light delay.
+
+        Args:
+            a: the semi-major axis of the orbit.
+            t: the time (or tensor of times) to calculate the position.
+            parallax: (arcseconds) if provided, return the position in
+                units of arcseconds.
+            z0: the reference point along the z axis whose light travel time
+                delay is taken to be zero. Default is the origin.
+        Returns:
+            The position of the body in the observer frame. Default is in units
+            of R_sun, but if parallax is provided, then in units of arcseconds.
+
+        """
+        sinf, cosf = self._get_true_anomaly(t)
+
+        # Compute the orbital radius and the component of the velocity
+        # in the z direction
+        angvel = 2 * np.pi / self.period
+        if self.ecc is None:
+            r = a
+            vamp = angvel * a
+            vz = vamp * tt.sin(self.incl) * cosf
+        else:
+            r = a * (1.0 - self.ecc ** 2) / (1 + self.ecc * cosf)
+            vamp = angvel * a / tt.sqrt(1 - self.ecc ** 2)
+            cwf = tt.cos(self.omega) * cosf - tt.sin(self.omega) * sinf
+            vz = (
+                vamp
+                * tt.sin(self.incl)
+                * (self.ecc * tt.cos(self.omega) + cwf)
+            )
+
+        # True position of the body
+        x, y, z = self._rotate_vector(r * cosf, r * sinf)
+
+        # Component of the acceleration in the z direction
+        az = -angvel ** 2 * (a / r) ** 3 * z
+
+        # Compute the time delay at the **retarded** position, accounting
+        # for the instantaneous velocity and acceleration of the body.
+        # See the derivation at https://github.com/rodluger/starry/issues/66
+        delay = tt.switch(
+            tt.lt(tt.abs_(az), 1.0e-10),
+            (z0 - z) / (c_light + vz),
+            (c_light / az)
+            * (
+                (1 + vz / c_light)
+                - tt.sqrt(
+                    (1 + vz / c_light) * (1 + vz / c_light)
+                    - 2 * az * (z0 - z) / c_light ** 2
+                )
+            ),
+        )
+
+        # Re-compute Kepler's equation, this time at the **retarded** position
+        return tuple(
+            tt.squeeze(tt.diag(x.dimshuffle(1, 2, 0)))
+            for x in self._get_position(
+                a, tt.shape_padright(t) - delay, parallax
+            )
+        )
+
+    def get_planet_position(self, t, parallax=None, light_delay=False):
         """The planets' positions in the barycentric frame
 
         Args:
             t: The times where the position should be evaluated.
+            light_delay: account for the light travel time delay? Default is
+                False.
 
         Returns:
             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, parallax)
-        )
+        if light_delay is False:
+            return tuple(
+                tt.squeeze(x)
+                for x in self._get_position(self.a_planet, t, parallax)
+            )
+        else:
+            return tuple(
+                tt.squeeze(x)
+                for x in self._get_retarded_position(
+                    self.a_planet, t, parallax
+                )
+            )
 
-    def get_star_position(self, t, parallax=None):
+    def get_star_position(self, t, parallax=None, light_delay=False):
         """The star's position in the barycentric frame
 
         .. note:: If there are multiple planets in the system, this will
@@ -369,17 +443,26 @@ def get_star_position(self, t, parallax=None):
 
         Args:
             t: The times where the position should be evaluated.
+            light_delay: account for the light travel time delay? Default is
+                False.
 
         Returns:
             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, parallax)
-        )
+        if light_delay is False:
+            return tuple(
+                tt.squeeze(x)
+                for x in self._get_position(self.a_star, t, parallax)
+            )
+        else:
+            return tuple(
+                tt.squeeze(x)
+                for x in self._get_retarded_position(self.a_star, t, parallax)
+            )
 
-    def get_relative_position(self, t, parallax=None):
+    def get_relative_position(self, t, parallax=None, light_delay=False):
         """The planets' positions relative to the star in the X,Y,Z frame.
 
         .. note:: This treats each planet independently and does not take the
@@ -391,17 +474,25 @@ def get_relative_position(self, t, parallax=None):
 
         Args:
             t: The times where the position should be evaluated.
+            light_delay: account for the light travel time delay? Default is
+                False.
 
         Returns:
             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, t, parallax)
-        )
+        if light_delay is False:
+            return tuple(
+                tt.squeeze(x) for x in self._get_position(-self.a, t, parallax)
+            )
+        else:
+            return tuple(
+                tt.squeeze(x)
+                for x in self._get_retarded_position(-self.a, t, parallax)
+            )
 
-    def get_relative_angles(self, t, parallax=None):
+    def get_relative_angles(self, t, parallax=None, light_delay=False):
         """The planets' relative position to the star in the sky plane, in
         separation, position angle coordinates.
 
@@ -411,14 +502,18 @@ def get_relative_angles(self, t, parallax=None):
 
         Args:
             t: The times where the position should be evaluated.
+            light_delay: account for the light travel time delay? Default is
+                False.
 
         Returns:
             The separation (arcseconds) and position angle (radians,
             measured east of north) of the planet relative to the star.
 
         """
-
-        X, Y, Z = self._get_position(-self.a, t, parallax)
+        if light_delay is False:
+            X, Y, Z = self._get_position(-self.a, t, parallax)
+        else:
+            X, Y, Z = self._get_retarded_position(-self.a, t, parallax)
 
         # calculate rho and theta
         rho = tt.squeeze(tt.sqrt(X ** 2 + Y ** 2))  # arcsec

src/exoplanet/orbits/keplerian_test.py

@@ -1,11 +1,13 @@
 # -*- coding: utf-8 -*-
 
 import astropy.units as u
+from astropy.constants import c
 import numpy as np
 import pytest
 import theano
 import theano.tensor as tt
 from theano.tests import unittest_tools as utt
+from scipy.optimize import minimize
 
 from ..units import with_unit
 from .keplerian import KeplerianOrbit, _get_consistent_inputs
@@ -407,3 +409,130 @@ def test_get_consistent_inputs():
 
     with pytest.raises(ValueError):
         _get_consistent_inputs(a3, None, rho_star3, r_star3, m_star3, None)
+
+
+def test_light_delay():
+
+    # Instantiate the orbit
+    m_star = tt.scalar()
+    period = tt.scalar()
+    ecc = tt.scalar()
+    omega = tt.scalar()
+    Omega = tt.scalar()
+    incl = tt.scalar()
+    m_planet = tt.scalar()
+    t = tt.scalar()
+    orbit = KeplerianOrbit(
+        m_star=m_star,
+        r_star=1.0,
+        t0=0.0,
+        period=period,
+        ecc=ecc,
+        omega=omega,
+        Omega=Omega,
+        incl=incl,
+        m_planet=m_planet,
+    )
+
+    # True position
+    get_position = theano.function(
+        [t, m_star, period, ecc, omega, Omega, incl, m_planet],
+        orbit.get_planet_position([t], light_delay=False),
+    )
+
+    # Retarded position
+    get_retarded_position = theano.function(
+        [t, m_star, period, ecc, omega, Omega, incl, m_planet],
+        orbit.get_planet_position([t], light_delay=True),
+    )
+
+    # Retarded position (numerical)
+    def get_exact_retarded_position(t, *args):
+        def loss(params):
+            ti, = params
+            xr, yr, zr = get_position(ti, *args)
+            delay = (zr * u.Rsun / c).to(u.day).value
+            return (ti - delay - t) ** 2
+
+        tr = minimize(loss, t).x[0]
+        return get_position(tr, *args)
+
+    # Compare for 100 different orbits
+    np.random.seed(13)
+    for i in range(100):
+        m_star = 0.1 + np.random.random() * 1.9
+        period = np.random.random() * 500
+        ecc = np.random.random()
+        omega = np.random.random() * 2 * np.pi
+        Omega = np.random.random() * 2 * np.pi
+        incl = np.random.random() * 0.5 * np.pi
+        m_planet = np.random.random()
+        t = np.random.random() * period
+        args = (m_star, period, ecc, omega, Omega, incl, m_planet)
+        assert np.allclose(
+            np.reshape(get_retarded_position(t, *args), (-1,)),
+            np.reshape(get_exact_retarded_position(t, *args), (-1,)),
+        )
+
+
+def test_light_delay_shape_two_planets_vector_t():
+    orbit = KeplerianOrbit(period=[1.0, 2.0])
+    x, y, z = orbit.get_planet_position([1.0, 2.0], light_delay=False)
+    xr, yr, zr = orbit.get_planet_position([1.0, 2.0], light_delay=True)
+    assert np.array_equal(x.shape.eval(), xr.shape.eval())
+
+
+@pytest.mark.xfail
+def test_light_delay_shape_scalar_t():
+    # TODO: Fix shapes
+    # `light_delay=True` only works if `t` is an array;
+    # throws `ValueError` if it is a scalar.
+    orbit = KeplerianOrbit(period=1.0)
+    x, y, z = orbit.get_planet_position(1.0, light_delay=False)
+    xr, yr, zr = orbit.get_planet_position(1.0, light_delay=True)
+    assert np.array_equal(x.shape.eval(), xr.shape.eval())
+
+
+@pytest.mark.xfail
+def test_light_delay_shape_single_t():
+    # TODO: Fix shapes
+    # `light_delay=False` returns a scalar
+    # `light_delay=True` returns a (1, 1) matrix
+    orbit = KeplerianOrbit(period=1.0)
+    x, y, z = orbit.get_planet_position([1.0], light_delay=False)
+    xr, yr, zr = orbit.get_planet_position([1.0], light_delay=True)
+    assert np.array_equal(x.shape.eval(), xr.shape.eval())
+
+
+@pytest.mark.xfail
+def test_light_delay_shape_vector_t():
+    # TODO: Fix shapes
+    # `light_delay=False` returns a vector
+    # `light_delay=True` returns a (1, 1) matrix
+    orbit = KeplerianOrbit(period=1.0)
+    x, y, z = orbit.get_planet_position([1.0, 2.0], light_delay=False)
+    xr, yr, zr = orbit.get_planet_position([1.0, 2.0], light_delay=True)
+    assert np.array_equal(x.shape.eval(), xr.shape.eval())
+
+
+@pytest.mark.xfail
+def test_light_delay_shape_two_planets_scalar_t():
+    # TODO: Fix shapes
+    # `light_delay=True` only works if `t` is an array;
+    # throws `ValueError` if it is a scalar.
+    orbit = KeplerianOrbit(period=[1.0, 2.0])
+    x, y, z = orbit.get_planet_position(1.0, light_delay=False)
+    xr, yr, zr = orbit.get_planet_position(1.0, light_delay=True)
+    assert np.array_equal(x.shape.eval(), xr.shape.eval())
+
+
+@pytest.mark.xfail
+def test_light_delay_shape_two_planets_single_t():
+    # TODO: Fix shapes
+    # `light_delay=False` returns a vector
+    # `light_delay=True` returns a (1, 2) matrix
+    orbit = KeplerianOrbit(period=[1.0, 2.0])
+    x, y, z = orbit.get_planet_position([1.0], light_delay=False)
+    xr, yr, zr = orbit.get_planet_position([1.0], light_delay=True)
+    assert np.array_equal(x.shape.eval(), xr.shape.eval())
+