more documentation, plus the ellipsoidal and doppler curves

docs/index.md

@@ -3,13 +3,53 @@
    You can adapt this file completely to your liking, but it should at least
    contain the root `toctree` directive. -->
 
-Welcome to squishyplanet's documentation!
-=========================================
+squishyplanet
+=============
 
-lorem ipsum
+``squishyplanet`` is a relatively lightweight package for creating transit light curves
+and phase curves of non-spherical exoplanets. It can generate the model, but then
+leaves the choice of inference framework up to you.
 
+We recommend that potential users start with the [geometry visualizations](geometry.md)
+to get a sense of the coordinate system and how the planet is defined. 
 
-# Test section
+## A summary:
+
+<span style="font-size:larger;">Transits</span>
+
+The transit portion can handle arbitrary
+order polynomial limb darkening by following most of the algorithm presented in
+[Agol, Luger, and Foreman-Mackey 2020](https://ui.adsabs.harvard.edu/abs/2020AJ....159..123A/abstract).
+However, where that publication pushes through to derive analytic expressions (which is
+possible but challenging when working with spherical planets), we don't even attempt 
+such wizardry for our triaxial planets and instead numerically integrate the final
+flux-blocking step.
+
+Even with this reliance on numerical solutions though, thanks to `JAX` and its ability 
+to just-in-time compile functions, these transit integrals are still relatively fast. 
+Users computing transits only can expect speeds slightly slower than but comparable to 
+[jaxoplanet](https://jax.exoplanet.codes/en/latest/). Also, in the limiting case where 
+the planet is forced to be spherical, `squishyplanet` is designed to be just as accurate
+as `jaxoplanet`. See the [Compare with jaxo/exoplanet notebook](tutorials/lightcurve_compare.ipynb)
+for more details.
+
+From the [quickstart](quickstart.ipynb) guide:
+
+![oblate_vs_spherical](oblate_vs_spherical.png)
+
+<span style="font-size:larger;">Phase Curves</span>
+
+`squishyplanet` can also compute reflected and emitted phase curves from the planet, as
+well as simple ellipsoidal and doppler variations from the star. Admittedly, the
+implementation of these features is much more crude than the transit portion: where that
+involves a lot of math/optimization, any component involving the planet is calculated 
+via brute-force Monte Carlo integration. So, while the transit portion is relatively fast,
+users should expect phase curve evaluations to be much slower, on the order of 100s of
+ms per evaluation.
+
+## Attribution
+
+\[insert JOSS citation/bibtex here someday\]
 
 ```{toctree}
 :maxdepth: 1
@@ -27,6 +67,7 @@ geometry
 :caption: Tutorials/Demos
 
 tutorials/create_a_lightcurve.ipynb
+tutorials/illustrations.ipynb
 tutorials/lightcurve_compare.ipynb
 tutorials/create_a_phase_curve.ipynb
 tutorials/fit_a_transit.ipynb
@@ -40,13 +81,3 @@ tutorials/fit_a_transit.ipynb
 
 api
 ```
-
-
-
-<!-- 
-Indices and tables
-==================
-
-* :ref:`genindex`
-* :ref:`modindex`
-* :ref:`search` -->

docs/tutorials/illustrations.ipynb

@@ -0,0 +1,25 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Illustrations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/tutorials/lightcurve_compare.ipynb

@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Compare with Jaxo/exoplanet\n",
+    "# Compare with jaxo/exoplanet\n",
     "\n",
     "While ``squishyplanet`` takes a more numerical approach to generating transit lightcurves\n",
     "than elegant packages such as ``exoplanet`` and ``starry``, it was designed to trade off\n",

squishyplanet/OblateSystem.py

@@ -19,6 +19,8 @@
     reflected_phase_curve,
     emission_phase_curve,
     phase_curve,
+    stellar_ellipsoidal_variations,
+    stellar_doppler_variations,
 )
 
 import copy
@@ -145,6 +147,16 @@ class OblateSystem:
         compute_emitted_phase_curve (bool, default=False):
             Whether to include flux emitted by the planet when calling
             :func:`lightcurve`.
+        compute_stellar_ellipsoidal_variations (bool, default=False):
+            Whether to include stellar ellipsoidal variations in the light curve. This
+            is the effect of the star's shape changing due to the gravitational pull of
+            the planet, and here is modeled as a simple sinusoidal variation with 4
+            peaks per orbit.
+        compute_stellar_doppler_variations (bool, default=False):
+            Whether to include stellar doppler variations in the light curve. This 
+            captures the effects of the star's radial velocity changing and boosting the
+            total flux/pushing some flux into/out of the bandpass of the observation.
+            Here, it is modeled as a simple sinusoidal variation with 2 peaks per orbit.
         phase_curve_nsamples (int, default=50_000):
             The number of random samples of the planet's surface to draw when performing
             Monte Carlo estimates of the emitted/reflected flux. A larger number will
@@ -186,11 +198,15 @@ def __init__(
         hotspot_concentration=0.2,
         reflected_albedo=1.0,
         emitted_scale=1e-6,
+        stellar_ellipsoidal_alpha=1e-6,
+        stellar_doppler_alpha=1e-6,
         systematic_trend_coeffs=jnp.array([0.0, 0.0]),
         log_jitter=-10,
         tidally_locked=True,
         compute_reflected_phase_curve=False,
         compute_emitted_phase_curve=False,
+        compute_stellar_ellipsoidal_variations=False,
+        compute_stellar_doppler_variations=False,
         phase_curve_nsamples=50_000,
         random_seed=0,
         data=jnp.array([1.0]),
@@ -217,11 +233,15 @@ def __init__(
             "hotspot_concentration",
             "reflected_albedo",
             "emitted_scale",
+            "stellar_ellipsoidal_alpha",
+            "stellar_doppler_alpha",
             "systematic_trend_coeffs",
             "log_jitter",
             "tidally_locked",
             "compute_reflected_phase_curve",
             "compute_emitted_phase_curve",
+            "compute_stellar_ellipsoidal_variations",
+            "compute_stellar_doppler_variations",
             "phase_curve_nsamples",
             "random_seed",
             "data",
@@ -488,6 +508,12 @@ def lightcurve(self, params={}):
         return _lightcurve(
             compute_reflected_phase_curve=self._state["compute_reflected_phase_curve"],
             compute_emitted_phase_curve=self._state["compute_emitted_phase_curve"],
+            compute_stellar_ellipsoidal_variations=self._state[
+                "compute_stellar_ellipsoidal_variations"
+            ],
+            compute_stellar_doppler_variations=self._state[
+                "compute_stellar_doppler_variations"
+            ],
             random_seed=self._state["random_seed"],
             phase_curve_nsamples=self._state["phase_curve_nsamples"],
             state=self._state,
@@ -518,6 +544,12 @@ def loglike(self, params={}):
         return _loglike(
             compute_reflected_phase_curve=self._state["compute_reflected_phase_curve"],
             compute_emitted_phase_curve=self._state["compute_emitted_phase_curve"],
+            compute_stellar_ellipsoidal_variations=self._state[
+                "compute_stellar_ellipsoidal_variations"
+            ],
+            compute_stellar_doppler_variations=self._state[
+                "compute_stellar_doppler_variations"
+            ],
             random_seed=self._state["random_seed"],
             phase_curve_nsamples=self._state["phase_curve_nsamples"],
             state=self._state,
@@ -532,100 +564,96 @@ def loglike(self, params={}):
         1,
         2,
         3,
+        4,
+        5,
     ),
 )
 def _lightcurve(
     compute_reflected_phase_curve,
     compute_emitted_phase_curve,
+    compute_stellar_ellipsoidal_variations,
+    compute_stellar_doppler_variations,
     random_seed,
     phase_curve_nsamples,
     state,
     params,
 ):
-
-    # if all you want is the transit, don't do any of the phase curve calculations
-    if (not compute_reflected_phase_curve) & (not compute_emitted_phase_curve):
-        for key in params.keys():
+    # always compute the primary transit and trend
+    for key in params.keys():
             state[key] = params[key]
-        trend = jnp.polyval(state["systematic_trend_coeffs"], state["times"])
-        return lightcurve(state) + trend
+    transit = lightcurve(state)
+    trend = jnp.polyval(state["systematic_trend_coeffs"], state["times"])
+
+    # if you don't want any phase curve stuff, you're done
+    if (not compute_reflected_phase_curve) & (not compute_emitted_phase_curve) and \
+        (not compute_stellar_doppler_variations) & (not compute_stellar_ellipsoidal_variations):
+        return transit + trend
 
-    # if you do want a phase curve, generate the radii and thetas that you'll reuse
-    # at each timestep
+    ######################################################
+    # compute the planet's contribution to the phase curve
+    ######################################################
+
+    # generate the radii and thetas that you'll reuse at each timestep
     sample_radii, sample_thetas = generate_sample_radii_thetas(
         jax.random.key(random_seed),
         jnp.arange(phase_curve_nsamples),
     )
+
+    # technically these are all calculated in "transit", but since phase
+    # curves are a) rare and b) expensive, we'll just do it again to keep
+    # the transit section of the code more self-contained
+    three = planet_3d_coeffs(**state)
+    two = planet_2d_coeffs(**three)
+    positions = skypos(**state)
+    x_c = positions[0, :]
+    y_c = positions[1, :]
+    z_c = positions[2, :]
 
     # just the reflected component
     if compute_reflected_phase_curve & (not compute_emitted_phase_curve):
-        for key in params.keys():
-            state[key] = params[key]
-
-        transit = lightcurve(state)
-
-        # technically these are all calculated in "transit", but since phase
-        # curves are a) rare and b) expensive, we'll just do it again to keep
-        # the transit section of the code more self-contained
-        three = planet_3d_coeffs(**state)
-        two = planet_2d_coeffs(**three)
-        positions = skypos(**state)
-        x_c = positions[0, :]
-        y_c = positions[1, :]
-        z_c = positions[2, :]
         reflected = reflected_phase_curve(
             sample_radii, sample_thetas, two, three, state, x_c, y_c, z_c
         )
+        emitted = 0.0
         # it really didn't make a difference in speed here
         # reflected = jax.vmap(
         #     reflected_phase_curve, in_axes=(0, 0, None, None, None, None, None, None)
         # )(sample_radii, sample_thetas, two, three, state, x_c, y_c, z_c)
         # reflected = jnp.mean(reflected, axis=0)
 
-        trend = jnp.polyval(state["systematic_trend_coeffs"], state["times"])
-        return transit + reflected + trend
-
     # just the emitted component
     elif (not compute_reflected_phase_curve) & compute_emitted_phase_curve:
-        for key in params.keys():
-            state[key] = params[key]
-        transit = lightcurve(state)
-
-        # technically these are all calculated in "transit", but since phase
-        # curves are a) rare and b) expensive, we'll just do it again to keep
-        # the transit section of the code more self-contained
-        three = planet_3d_coeffs(**state)
-        two = planet_2d_coeffs(**three)
-        positions = skypos(**state)
-        x_c = positions[0, :]
-        y_c = positions[1, :]
-        z_c = positions[2, :]
+        reflected = 0.0
         emitted = emission_phase_curve(sample_radii, sample_thetas, two, three, state)
 
-        trend = jnp.polyval(state["systematic_trend_coeffs"], state["times"])
-        return transit + emitted + trend
-
-    # both reflected and emitted components
+    # both reflected and emitted components. this function shares some of the
+    # computation between the two, so it's a bit faster than running them separately
     elif compute_reflected_phase_curve & compute_emitted_phase_curve:
-        for key in params.keys():
-            state[key] = params[key]
-        transit = lightcurve(state)
-
-        # technically these are all calculated in "transit", but since phase
-        # curves are a) rare and b) expensive, we'll just do it again to keep
-        # the transit section of the code more self-contained
-        three = planet_3d_coeffs(**state)
-        two = planet_2d_coeffs(**three)
-        positions = skypos(**state)
-        x_c = positions[0, :]
-        y_c = positions[1, :]
-        z_c = positions[2, :]
         reflected, emitted = phase_curve(
             sample_radii, sample_thetas, two, three, state, x_c, y_c, z_c
         )
 
-        trend = jnp.polyval(state["systematic_trend_coeffs"], state["times"])
-        return transit + reflected + emitted + trend
+    ####################################################
+    # compute the star's contribution to the phase curve
+    ####################################################
+
+    if compute_stellar_ellipsoidal_variations | compute_stellar_doppler_variations:
+        time_deltas = state["times"] - state["t_peri"]
+        mean_anomalies = 2 * jnp.pi * time_deltas / state["period"]
+        true_anomalies = kepler(mean_anomalies, state["e"])
+
+    if compute_stellar_ellipsoidal_variations:
+        ellipsoidal = stellar_ellipsoidal_variations(true_anomalies, state["stellar_ellipsoidal_alpha"], state["period"])
+    else:
+        ellipsoidal = 0.0
+
+    if compute_stellar_doppler_variations:
+        doppler = stellar_doppler_variations(true_anomalies, state["stellar_doppler_alpha"], state["period"])
+    else:
+        doppler = 0.0
+
+    # put it all together
+    return transit + trend + reflected + emitted + ellipsoidal + doppler
 
 
 @partial(
@@ -635,11 +663,15 @@ def _lightcurve(
         1,
         2,
         3,
+        4,
+        5,
     ),
 )
 def _loglike(
     compute_reflected_phase_curve,
     compute_emitted_phase_curve,
+    compute_stellar_ellipsoidal_variations,
+    compute_stellar_doppler_variations,
     random_seed,
     phase_curve_nsamples,
     state,