{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e3a9f10d",
   "metadata": {},
   "source": [
    "(transit)=\n",
    "\n",
    "# Transit fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a78ae748-869b-4a44-8733-b145c79fbef2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import exoplanet\n",
    "\n",
    "exoplanet.utils.docs_setup()\n",
    "print(f\"exoplanet.__version__ = '{exoplanet.__version__}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9bab082",
   "metadata": {},
   "source": [
    "*exoplanet* includes methods for computing the light curves transiting planets.\n",
    "In its simplest form this can be used to evaluate a light curve like you would do with [batman](https://astro.uchicago.edu/~kreidberg/batman/), for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a90311d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import exoplanet as xo\n",
    "\n",
    "# The light curve calculation requires an orbit\n",
    "orbit = xo.orbits.KeplerianOrbit(period=3.456)\n",
    "\n",
    "# Compute a limb-darkened light curve using starry\n",
    "t = np.linspace(-0.1, 0.1, 1000)\n",
    "u = [0.3, 0.2]\n",
    "light_curve = (\n",
    "    xo.LimbDarkLightCurve(*u)\n",
    "    .get_light_curve(orbit=orbit, r=0.1, t=t, texp=0.02)\n",
    "    .eval()\n",
    ")\n",
    "# Note: the `eval` is needed because this is using Theano in\n",
    "# the background\n",
    "\n",
    "plt.plot(t, light_curve, color=\"C0\", lw=2)\n",
    "plt.ylabel(\"relative flux\")\n",
    "plt.xlabel(\"time [days]\")\n",
    "_ = plt.xlim(t.min(), t.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9f84dd5",
   "metadata": {},
   "source": [
    "But the real power comes from the fact that this is defined as a [Aesara/Theano operation](https://aesara.readthedocs.io/en/latest/extending/index.html) so it can be combined with PyMC3 to do gradient-based inference.\n",
    "\n",
    "## The transit model in PyMC3\n",
    "\n",
    "In this section, we will construct a simple transit fit model using *PyMC3* and then we will fit a two planet model to simulated data.\n",
    "To start, let's randomly sample some periods and phases and then define the time sampling:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cb97f4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "periods = np.random.uniform(5, 20, 2)\n",
    "t0s = periods * np.random.rand(2)\n",
    "t = np.arange(0, 80, 0.02)\n",
    "yerr = 5e-4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bc61a15",
   "metadata": {},
   "source": [
    "Then, define the parameters.\n",
    "In this simple model, we'll just fit for the limb darkening parameters of the star, and the period, phase, impact parameter, and radius ratio of the planets (note: this is already 10 parameters and running MCMC to convergence using [emcee](https://emcee.readthedocs.io) would probably take at least an hour).\n",
    "For the limb darkening, we'll use a quadratic law as parameterized by [Kipping (2013)](https://arxiv.org/abs/1308.0009).\n",
    "This reparameterizations is implemented in *exoplanet* as custom *PyMC3* distribution :class:`exoplanet.distributions.QuadLimbDark`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b375c261",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc3 as pm\n",
    "import pymc3_ext as pmx\n",
    "\n",
    "with pm.Model() as model:\n",
    "    # The baseline flux\n",
    "    mean = pm.Normal(\"mean\", mu=0.0, sd=1.0)\n",
    "\n",
    "    # The time of a reference transit for each planet\n",
    "    t0 = pm.Normal(\"t0\", mu=t0s, sd=1.0, shape=2)\n",
    "\n",
    "    # The log period; also tracking the period itself\n",
    "    logP = pm.Normal(\"logP\", mu=np.log(periods), sd=0.1, shape=2)\n",
    "    period = pm.Deterministic(\"period\", pm.math.exp(logP))\n",
    "\n",
    "    # The Kipping (2013) parameterization for quadratic limb darkening paramters\n",
    "    u = xo.distributions.QuadLimbDark(\"u\", testval=np.array([0.3, 0.2]))\n",
    "\n",
    "    r = pm.Uniform(\n",
    "        \"r\", lower=0.01, upper=0.1, shape=2, testval=np.array([0.04, 0.06])\n",
    "    )\n",
    "    b = xo.distributions.ImpactParameter(\n",
    "        \"b\", ror=r, shape=2, testval=np.random.rand(2)\n",
    "    )\n",
    "\n",
    "    # Set up a Keplerian orbit for the planets\n",
    "    orbit = xo.orbits.KeplerianOrbit(period=period, t0=t0, b=b)\n",
    "\n",
    "    # Compute the model light curve using starry\n",
    "    light_curves = xo.LimbDarkLightCurve(u[0], u[1]).get_light_curve(\n",
    "        orbit=orbit, r=r, t=t\n",
    "    )\n",
    "    light_curve = pm.math.sum(light_curves, axis=-1) + mean\n",
    "\n",
    "    # Here we track the value of the model light curve for plotting\n",
    "    # purposes\n",
    "    pm.Deterministic(\"light_curves\", light_curves)\n",
    "\n",
    "    # ******************************************************************* #\n",
    "    # On the folowing lines, we simulate the dataset that we will fit     #\n",
    "    #                                                                     #\n",
    "    # NOTE: if you are fitting real data, you shouldn't include this line #\n",
    "    #       because you already have data!                                #\n",
    "    # ******************************************************************* #\n",
    "    y = pmx.eval_in_model(light_curve)\n",
    "    y += yerr * np.random.randn(len(y))\n",
    "    # ******************************************************************* #\n",
    "    # End of fake data creation; you want to include the following lines  #\n",
    "    # ******************************************************************* #\n",
    "\n",
    "    # The likelihood function assuming known Gaussian uncertainty\n",
    "    pm.Normal(\"obs\", mu=light_curve, sd=yerr, observed=y)\n",
    "\n",
    "    # Fit for the maximum a posteriori parameters given the simuated\n",
    "    # dataset\n",
    "    map_soln = pmx.optimize(start=model.test_point)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e286731",
   "metadata": {},
   "source": [
    "Now we can plot the simulated data and the maximum a posteriori model to make sure that our initialization looks ok."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36b6e0eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(t, y, \".k\", ms=4, label=\"data\")\n",
    "for i, l in enumerate(\"bc\"):\n",
    "    plt.plot(\n",
    "        t, map_soln[\"light_curves\"][:, i], lw=1, label=\"planet {0}\".format(l)\n",
    "    )\n",
    "plt.xlim(t.min(), t.max())\n",
    "plt.ylabel(\"relative flux\")\n",
    "plt.xlabel(\"time [days]\")\n",
    "plt.legend(fontsize=10)\n",
    "_ = plt.title(\"map model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e58ba4e",
   "metadata": {},
   "source": [
    "## Sampling\n",
    "\n",
    "Now, let's sample from the posterior defined by this model.\n",
    "As usual, there are strong covariances between some of the parameters so we'll use `pmx.sample` from [pymc3-ext](https://github.com/exoplanet-dev/pymc3-ext)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80d57b30",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "with model:\n",
    "    trace = pmx.sample(\n",
    "        tune=1000,\n",
    "        draws=1000,\n",
    "        start=map_soln,\n",
    "        cores=2,\n",
    "        chains=2,\n",
    "        target_accept=0.9,\n",
    "        return_inferencedata=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf521c51",
   "metadata": {},
   "source": [
    "After sampling, it's important that we assess convergence.\n",
    "We can do that using the `summary` function from `ArviZ`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4cdbb2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import arviz as az\n",
    "\n",
    "az.summary(trace, var_names=[\"period\", \"t0\", \"r\", \"b\", \"u\", \"mean\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "302a9c91",
   "metadata": {},
   "source": [
    "That looks pretty good!\n",
    "Fitting this without *exoplanet* would have taken a lot more patience.\n",
    "\n",
    "Now we can also look at the [corner plot](https://corner.readthedocs.io) of some of that parameters of interest:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d48da383",
   "metadata": {},
   "outputs": [],
   "source": [
    "import corner\n",
    "\n",
    "truth = dict(\n",
    "    zip(\n",
    "        [\"period\", \"r\"],\n",
    "        pmx.eval_in_model([period, r], model.test_point, model=model),\n",
    "    )\n",
    ")\n",
    "_ = corner.corner(\n",
    "    trace,\n",
    "    var_names=[\"period\", \"r\"],\n",
    "    truths=truth,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c08a166d",
   "metadata": {},
   "source": [
    "## Phase plots\n",
    "\n",
    "Like in the [Radial velocity fitting](./rv.ipynb) tutorial, we can make plots of the model predictions for each planet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cee95a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for n, letter in enumerate(\"bc\"):\n",
    "    plt.figure()\n",
    "\n",
    "    # Get the posterior median orbital parameters\n",
    "    period_trace = trace.posterior[\"period\"].values[:, :, n]\n",
    "    p = np.median(period_trace)\n",
    "    t0 = np.median(trace.posterior[\"t0\"].values[:, :, n])\n",
    "\n",
    "    # Compute the median of posterior estimate of the contribution from\n",
    "    # the other planet. Then we can remove this from the data to plot\n",
    "    # just the planet we care about.\n",
    "    lcs = trace.posterior[\"light_curves\"].values\n",
    "    other = np.median(lcs[:, :, :, (n + 1) % 2], axis=(0, 1))\n",
    "\n",
    "    # Plot the folded data\n",
    "    x_fold = (t - t0 + 0.5 * p) % p - 0.5 * p\n",
    "    plt.errorbar(\n",
    "        x_fold, y - other, yerr=yerr, fmt=\".k\", label=\"data\", zorder=-1000\n",
    "    )\n",
    "\n",
    "    # Plot the folded model\n",
    "    inds = np.argsort(x_fold)\n",
    "    inds = inds[np.abs(x_fold)[inds] < 0.3]\n",
    "    pred = lcs[:, :, inds, n] + trace.posterior[\"mean\"].values[:, :, None]\n",
    "    pred = np.median(pred, axis=(0, 1))\n",
    "    plt.plot(x_fold[inds], pred, color=\"C1\", label=\"model\")\n",
    "\n",
    "    # Annotate the plot with the planet's period\n",
    "    txt = \"period = {0:.4f} +/- {1:.4f} d\".format(\n",
    "        np.mean(period_trace), np.std(period_trace)\n",
    "    )\n",
    "    plt.annotate(\n",
    "        txt,\n",
    "        (0, 0),\n",
    "        xycoords=\"axes fraction\",\n",
    "        xytext=(5, 5),\n",
    "        textcoords=\"offset points\",\n",
    "        ha=\"left\",\n",
    "        va=\"bottom\",\n",
    "        fontsize=12,\n",
    "    )\n",
    "\n",
    "    plt.legend(fontsize=10, loc=4)\n",
    "    plt.xlim(-0.5 * p, 0.5 * p)\n",
    "    plt.xlabel(\"time since transit [days]\")\n",
    "    plt.ylabel(\"relative flux\")\n",
    "    plt.title(\"planet {0}\".format(letter))\n",
    "    plt.xlim(-0.3, 0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6c47be4",
   "metadata": {},
   "source": [
    "## Citations\n",
    "\n",
    "As described in the [citation tutorial](https://docs.exoplanet.codes/en/stable/tutorials/citation/), we can use [citations.get_citations_for_model](https://docs.exoplanet.codes/en/stable/user/api/#exoplanet.citations.get_citations_for_model) to construct an acknowledgement and BibTeX listing that includes the relevant citations for this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5461564",
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    txt, bib = xo.citations.get_citations_for_model()\n",
    "print(txt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43ecde5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(bib.split(\"\\n\\n\")[0] + \"\\n\\n...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afcce422-b61b-4c16-ab44-75179f7741e3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}