{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1072f391-e4b4-476c-bec3-211fda0ef8ed",
   "metadata": {},
   "source": [
    "(ttv)=\n",
    "\n",
    "# Fitting transit times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12785329-3284-4b9e-b546-5ab4458561b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import exoplanet\n",
    "\n",
    "exoplanet.utils.docs_setup()\n",
    "print(f\"exoplanet.__version__ = '{exoplanet.__version__}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd479fba-73ba-40f1-a45e-82ea34822187",
   "metadata": {},
   "source": [
    "Fitting for or marginalizing over the transit times or transit timing variations (TTVs) can be useful for several reasons, and it is a compelling use case for *exoplanet* becuase the number of parameters in the model increases significantly because there will be a new parameter for each transit.\n",
    "The performance of the NUTS sampler used by *exoplanet* scales well with the number of parameters, so a TTV model should be substantially faster to run to convergence with *exoplanet* than with other tools.\n",
    "There are a few definitions and subtleties that should be considered before jumping in.\n",
    "\n",
    "In this tutorial, we will be using a \"descriptive\" model [orbits.TTVOrbit](https://docs.exoplanet.codes/en/stable/user/api/#exoplanet.orbits.TTVOrbit) to fit the light curve where the underlying motion is still Keplerian, but the time coordinate is warped to make `t0` a function of time.\n",
    "All of the other orbital elements besides `t0` are shared across all orbits, but the `t0` for each transit will be a parameter.\n",
    "This means that other variations (like transit duration variations) are not currently supported, but it would be possible to include more general effects.\n",
    "`exoplanet` also supports photodynamics modeling using the [orbits.ReboundOrbit](https://docs.exoplanet.codes/en/stable/user/api/#exoplanet.orbits.ReboundOrbit) for more detailed analysis, but that is a topic for a future tutorial.\n",
    "\n",
    "It is also important to note that \"transit time\" within *exoplanet* (and most other transit fitting software) is defined as the time of conjunction (called `t0` in the code): the time when the true anomaly is $\\pi/2 - \\omega$.\n",
    "Section 18 of [the EXOFASTv2 paper](https://arxiv.org/abs/1907.09480) includes an excellent discussion of some of the commonly used definitions of \"transit time\" in the literature.\n",
    "\n",
    "Finally, there is a subtlety in the definition of the \"period\" of an orbit with TTVs.\n",
    "Two possible definitions are: (1) the average time between transits, or (2) the slope of a least squares fit to the transit times as a function of transit number.\n",
    "In *exoplanet*, we use the latter definition and call this parameter the `ttv_period` to distinguish it from the `period` of the underlying Keplerian motion which sets the shape and duration of the transit.\n",
    "By default, these two periods are constrained to be equal, but it can be useful to fit for both parameters since the shape of the transit might not be perfectly described by the same period.\n",
    "That being said, if you fit for both periods, make sure that you constrain `ttv_period` and `period` to be similar or things can get a bit ugly.\n",
    "\n",
    "To get started, let's generate some simulated transit times.\n",
    "We'll use the [orbits.ttv.compute_expected_transit_times](https://docs.exoplanet.codes/en/stable/user/api/#exoplanet.orbits.ttv.compute_expected_transit_times) function to get the expected transit times for a linear ephemeris within some observation baseline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "512156b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import exoplanet as xo\n",
    "\n",
    "np.random.seed(3948)\n",
    "true_periods = np.random.uniform(8, 12, 2)\n",
    "true_t0s = true_periods * np.random.rand(2)\n",
    "t = np.arange(0, 80, 0.01)\n",
    "texp = 0.01\n",
    "yerr = 5e-4\n",
    "\n",
    "# Compute the transit times for a linear ephemeris\n",
    "true_transit_times = xo.orbits.ttv.compute_expected_transit_times(\n",
    "    t.min(), t.max(), true_periods, true_t0s\n",
    ")\n",
    "\n",
    "# Simulate transit timing variations using a simple model\n",
    "true_ttvs = [\n",
    "    (0.05 - (i % 2) * 0.1) * np.sin(2 * np.pi * tt / 23.7)\n",
    "    for i, tt in enumerate(true_transit_times)\n",
    "]\n",
    "true_transit_times = [tt + v for tt, v in zip(true_transit_times, true_ttvs)]\n",
    "\n",
    "# Plot the true TTV model\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 5), sharex=True)\n",
    "ax1.plot(true_transit_times[0], true_ttvs[0], \".k\")\n",
    "ax1.axhline(0, color=\"k\", lw=0.5)\n",
    "ax1.set_ylim(np.max(np.abs(ax1.get_ylim())) * np.array([-1, 1]))\n",
    "ax1.set_ylabel(\"$O-C$ [days]\")\n",
    "\n",
    "ax2.plot(true_transit_times[1], true_ttvs[1], \".k\")\n",
    "ax2.axhline(0, color=\"k\", lw=0.5)\n",
    "ax2.set_ylim(np.max(np.abs(ax2.get_ylim())) * np.array([-1, 1]))\n",
    "ax2.set_ylabel(\"$O-C$ [days]\")\n",
    "\n",
    "ax2.set_xlabel(\"transit time [days]\")\n",
    "_ = ax1.set_title(\"true TTVs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86f7e561",
   "metadata": {},
   "source": [
    "Now, like in the [Transit fitting](https://docs.exoplanet.codes/en/stable/tutorials/transit/) tutorial, we'll set up the the model using *PyMC3* and *exoplanet*, and then simulate a data set from that model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dafa187",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc3 as pm\n",
    "import pymc3_ext as pmx\n",
    "import aesara_theano_fallback.tensor as tt\n",
    "\n",
    "np.random.seed(9485023)\n",
    "\n",
    "with pm.Model() as model:\n",
    "    # This part of the model is similar to the model in the `transit` tutorial\n",
    "    mean = pm.Normal(\"mean\", mu=0.0, sd=1.0)\n",
    "    u = xo.QuadLimbDark(\"u\", testval=np.array([0.3, 0.2]))\n",
    "    logr = pm.Uniform(\n",
    "        \"logr\",\n",
    "        lower=np.log(0.01),\n",
    "        upper=np.log(0.1),\n",
    "        shape=2,\n",
    "        testval=np.log([0.04, 0.06]),\n",
    "    )\n",
    "    r = pm.Deterministic(\"r\", tt.exp(logr))\n",
    "    b = xo.ImpactParameter(\n",
    "        \"b\", ror=r, shape=2, testval=0.5 * np.random.rand(2)\n",
    "    )\n",
    "\n",
    "    # Now we have a parameter for each transit time for each planet:\n",
    "    transit_times = []\n",
    "    for i in range(2):\n",
    "        transit_times.append(\n",
    "            pm.Normal(\n",
    "                \"tts_{0}\".format(i),\n",
    "                mu=true_transit_times[i],\n",
    "                sd=1.0,\n",
    "                shape=len(true_transit_times[i]),\n",
    "            )\n",
    "        )\n",
    "\n",
    "    # Set up an orbit for the planets\n",
    "    orbit = xo.orbits.TTVOrbit(b=b, transit_times=transit_times)\n",
    "\n",
    "    # It will be useful later to track some parameters of the orbit\n",
    "    pm.Deterministic(\"t0\", orbit.t0)\n",
    "    pm.Deterministic(\"period\", orbit.period)\n",
    "    for i in range(2):\n",
    "        pm.Deterministic(\"ttvs_{0}\".format(i), orbit.ttvs[i])\n",
    "\n",
    "    # The rest of this block follows the transit fitting tutorial\n",
    "    light_curves = xo.LimbDarkLightCurve(u).get_light_curve(\n",
    "        orbit=orbit, r=r, t=t, texp=texp\n",
    "    )\n",
    "    light_curve = pm.math.sum(light_curves, axis=-1) + mean\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",
    "    pm.Normal(\"obs\", mu=light_curve, sd=yerr, observed=y)\n",
    "\n",
    "    map_soln = model.test_point\n",
    "    map_soln = pmx.optimize(start=map_soln, vars=transit_times)\n",
    "    map_soln = pmx.optimize(start=map_soln, vars=[r, b])\n",
    "    map_soln = pmx.optimize(start=map_soln, vars=transit_times)\n",
    "    map_soln = pmx.optimize(start=map_soln)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fd9aaa5",
   "metadata": {},
   "source": [
    "Here's our simulated light curve and the initial model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4134703",
   "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": "feb3b63a",
   "metadata": {},
   "source": [
    "This looks similar to the light curve from the [Transit fitting](https://docs.exoplanet.codes/en/stable/tutorials/transit/) tutorial, but if we try plotting the folded transits, we can see that something isn't right: these transits look pretty smeared out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db223b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for n, letter in enumerate(\"bc\"):\n",
    "    plt.figure()\n",
    "\n",
    "    # Get the posterior median orbital parameters\n",
    "    p = map_soln[\"period\"][n]\n",
    "    t0 = map_soln[\"t0\"][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",
    "    other = map_soln[\"light_curves\"][:, (n + 1) % 2]\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",
    "    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": "42c90578",
   "metadata": {},
   "source": [
    "Instead, we can correct for the transit times by subtracting these estimates and plot that instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7085e09b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    t_warp = pmx.eval_in_model(orbit._warp_times(t), map_soln)\n",
    "\n",
    "for n, letter in enumerate(\"bc\"):\n",
    "    plt.figure()\n",
    "\n",
    "    p = map_soln[\"period\"][n]\n",
    "    other = map_soln[\"light_curves\"][:, (n + 1) % 2]\n",
    "\n",
    "    # NOTE: 't0' has already been subtracted!\n",
    "    x_fold = (t_warp[:, n] + 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",
    "    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": "957750a3",
   "metadata": {},
   "source": [
    "That looks better!\n",
    "\n",
    "## Sampling\n",
    "\n",
    "Now let's run some MCMC as usual:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2369a6e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(230948)\n",
    "with model:\n",
    "    trace = pmx.sample(\n",
    "        tune=1000,\n",
    "        draws=1000,\n",
    "        start=map_soln,\n",
    "        cores=2,\n",
    "        chains=2,\n",
    "        return_inferencedata=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64f52752",
   "metadata": {},
   "source": [
    "Then check the convergence diagnostics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7180bcc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import arviz as az\n",
    "\n",
    "az.summary(trace, var_names=[\"mean\", \"u\", \"logr\", \"b\", \"tts_0\", \"tts_1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57ad7a4a",
   "metadata": {},
   "source": [
    "And plot the corner plot of the physical parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ded78309",
   "metadata": {},
   "outputs": [],
   "source": [
    "import corner\n",
    "\n",
    "names = [\"period\", \"r\", \"b\"]\n",
    "\n",
    "with model:\n",
    "    truths = dict(zip(names, pmx.eval_in_model([orbit.period, r, b])))\n",
    "\n",
    "_ = corner.corner(\n",
    "    trace,\n",
    "    truths=truths,\n",
    "    var_names=names,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d386c48",
   "metadata": {},
   "source": [
    "We could also plot corner plots of the transit times, but they're not terribly enlightening in this case so let's skip it.\n",
    "\n",
    "Finally, let's plot the posterior estimates of the the transit times in an O-C diagram:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ebe6afd",
   "metadata": {},
   "outputs": [],
   "source": [
    "flat_samps = trace.posterior.stack(sample=(\"chain\", \"draw\"))\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 5), sharex=True)\n",
    "\n",
    "q = np.percentile(flat_samps[\"ttvs_0\"], [16, 50, 84], axis=-1)\n",
    "ax1.fill_between(\n",
    "    np.mean(flat_samps[\"tts_0\"], axis=-1),\n",
    "    q[0],\n",
    "    q[2],\n",
    "    color=\"C0\",\n",
    "    alpha=0.4,\n",
    "    edgecolor=\"none\",\n",
    ")\n",
    "ref = np.polyval(\n",
    "    np.polyfit(true_transit_times[0], true_ttvs[0], 1), true_transit_times[0]\n",
    ")\n",
    "ax1.plot(true_transit_times[0], true_ttvs[0] - ref, \".k\")\n",
    "ax1.axhline(0, color=\"k\", lw=0.5)\n",
    "ax1.set_ylim(np.max(np.abs(ax1.get_ylim())) * np.array([-1, 1]))\n",
    "\n",
    "ax1.set_ylabel(\"$O-C$ [days]\")\n",
    "\n",
    "q = np.percentile(flat_samps[\"ttvs_1\"], [16, 50, 84], axis=-1)\n",
    "ax2.fill_between(\n",
    "    np.mean(flat_samps[\"tts_1\"], axis=-1),\n",
    "    q[0],\n",
    "    q[2],\n",
    "    color=\"C1\",\n",
    "    alpha=0.4,\n",
    "    edgecolor=\"none\",\n",
    ")\n",
    "ref = np.polyval(\n",
    "    np.polyfit(true_transit_times[1], true_ttvs[1], 1), true_transit_times[1]\n",
    ")\n",
    "ax2.plot(true_transit_times[1], true_ttvs[1] - ref, \".k\", label=\"truth\")\n",
    "ax2.axhline(0, color=\"k\", lw=0.5)\n",
    "ax2.set_ylim(np.max(np.abs(ax2.get_ylim())) * np.array([-1, 1]))\n",
    "\n",
    "ax2.legend(fontsize=10)\n",
    "ax2.set_ylabel(\"$O-C$ [days]\")\n",
    "ax2.set_xlabel(\"transit time [days]\")\n",
    "_ = ax1.set_title(\"posterior inference\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "230caf4f",
   "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": "57fb744d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    txt, bib = xo.citations.get_citations_for_model()\n",
    "print(txt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1c4f6fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(bib.split(\"\\n\\n\")[0] + \"\\n\\n...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e3dc28-baba-45de-b934-ef868945a821",
   "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
}