notebooks/SpotSolve.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Where's the spot?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An example of how to solve for the location, amplitude, and size of a star spot.\n",
    "\n",
    "As we discuss in [this notebook](StarSpots.ipynb), `starry` isn't really meant for modeling discrete features such as star spots; rather, `starry` employs spherical harmonics to model the surface brightness distribution as a smooth, continuous function. We generally recommend approaching the mapping problem in this fashion; see the Eclipsing Binary tutorials for more information on how to do this. However, if you really want to model the surface of a star with star spots, read on!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's begin by importing stuff as usual:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "%run notebook_setup.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import starry\n",
    "import exoplanet as xo\n",
    "import pymc3 as pm\n",
    "import pymc3_ext as pmx\n",
    "import matplotlib.pyplot as plt\n",
    "from corner import corner\n",
    "\n",
    "starry.config.quiet = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We discussed how to add star spots to a `starry` map in [this tutorial](StarSpots.ipynb). Here, we'll generate a synthetic light curve from a star with a single large spot with the following parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# True values\n",
    "truth = dict(contrast=0.25, radius=20, lat=30, lon=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Those are, respectively, the fractional amplitude of the spot, its standard deviation (recall that the spot is modeled as a Gaussian in $\\cos\\Delta\\theta$, its latitude and its longitude.\n",
    "\n",
    "To make things simple, we'll assume we know the inclination and period of the star exactly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Things we'll assume are known\n",
    "inc = 60.0\n",
    "P = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's instantiate a 15th degree map and give it those properties:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map = starry.Map(15)\n",
    "map.inc = inc\n",
    "map.spot(\n",
    "    contrast=truth[\"contrast\"],\n",
    "    radius=truth[\"radius\"],\n",
    "    lat=truth[\"lat\"],\n",
    "    lon=truth[\"lon\"],\n",
    ")\n",
    "map.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll generate a synthetic light curve with some noise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.linspace(0, 3.0, 500)\n",
    "flux0 = map.flux(theta=360.0 / P * t).eval()\n",
    "np.random.seed(0)\n",
    "flux_err = 2e-4\n",
    "flux = flux0 + flux_err * np.random.randn(len(t))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what that looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(t, flux)\n",
    "plt.xlabel(\"time [days]\")\n",
    "plt.ylabel(\"normalized flux\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we are going to derive posterior constraints on the spot properties. Let's define a `pymc3` model and within it, our priors and flux model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model() as model:\n",
    "\n",
    "    # Priors\n",
    "    contrast = pm.Uniform(\"contrast\", lower=0.0, upper=1.0, testval=0.5)\n",
    "    radius = pm.Uniform(\"radius\", lower=10.0, upper=35.0, testval=15.0)\n",
    "    lat = pm.Uniform(\"lat\", lower=-90.0, upper=90.0, testval=0.1)\n",
    "    lon = pm.Uniform(\"lon\", lower=-180.0, upper=180.0, testval=0.1)\n",
    "\n",
    "    # Instantiate the map and add the spot\n",
    "    map = starry.Map(ydeg=15)\n",
    "    map.inc = inc\n",
    "    map.spot(contrast=contrast, radius=radius, lat=lat, lon=lon)\n",
    "\n",
    "    # Compute the flux model\n",
    "    flux_model = map.flux(theta=360.0 / P * t)\n",
    "    pm.Deterministic(\"flux_model\", flux_model)\n",
    "\n",
    "    # Save our initial guess\n",
    "    flux_model_guess = pmx.eval_in_model(flux_model)\n",
    "\n",
    "    # The likelihood function assuming known Gaussian uncertainty\n",
    "    pm.Normal(\"obs\", mu=flux_model, sd=flux_err, observed=flux)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    ".. note::\n",
    "   It's important to define a nonzero ``testval`` for both the latitude and longitude of the spot.\n",
    "   If you don't do that, it's likely the optimizer will get stuck at ``lat, lon = 0, 0``, which is\n",
    "   a local maximum of the likelihood."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've placed some generous uniform priors on the four quantities we're solving for. Let's run a quick gradient descent to get a good starting position for the sampler:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    map_soln = pmx.optimize(start=model.test_point)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the data and our initial guess before and after the optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 5))\n",
    "plt.plot(t, flux, \"k.\", alpha=0.3, ms=2, label=\"data\")\n",
    "plt.plot(t, flux_model_guess, \"C1--\", lw=1, alpha=0.5, label=\"Initial\")\n",
    "plt.plot(\n",
    "    t, pmx.eval_in_model(flux_model, map_soln, model=model), \"C1-\", label=\"MAP\", lw=1\n",
    ")\n",
    "plt.legend(fontsize=10, numpoints=5)\n",
    "plt.xlabel(\"time [days]\", fontsize=24)\n",
    "plt.ylabel(\"relative flux\", fontsize=24);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here are the maximum a posteriori (MAP) values next to the true values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"{0:12s} {1:10s} {2:10s}\".format(\"\", \"truth\", \"map_soln\"))\n",
    "for key in truth.keys():\n",
    "    print(\"{0:10s} {1:10.5f} {2:10.5f}\".format(key, truth[key], map_soln[key]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not bad! Looks like we recovered the correct spot properties. But we're not done! Let's get posterior constraints on them by sampling with `pymc3`. Since this is such a simple problem, the following cell should run in about a minute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    trace = pmx.sample(tune=250, draws=500, start=map_soln, chains=4, target_accept=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot some diagnostics to assess convergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var_names = [\"contrast\", \"radius\", \"lat\", \"lon\"]\n",
    "display(pm.summary(trace, var_names=var_names))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, the corner plot showing the joint posteriors and the true values (in blue):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = pm.trace_to_dataframe(trace, varnames=var_names)\n",
    "corner(samples, truths=[truth[name] for name in var_names]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're done! It's easy to extend this to multiple spots, simply by calling `map.spot` once for each spot in the model, making sure you define new `pymc3` variables for the spot properties of each spot."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Raw Cell Format",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}