{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5d636239",
   "metadata": {},
   "source": [
    "(eb)=\n",
    "\n",
    "# Fitting a detached eclipsing binary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d610cf-71d2-4e9d-bcea-d21f18bc2c5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import exoplanet\n",
    "\n",
    "exoplanet.utils.docs_setup()\n",
    "print(f\"exoplanet.__version__ = '{exoplanet.__version__}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a8c9a97",
   "metadata": {},
   "source": [
    "In this case study, we'll go through the steps required to fit the light curve and radial velocity measurements for the detached eclipsing binary system HD 23642.\n",
    "This is a bright system that has been fit by many authors ([1](https://arxiv.org/abs/astro-ph/0403444), [2](https://arxiv.org/abs/astro-ph/0409507), [3](https://ui.adsabs.harvard.edu/abs/2007A%26A...463..579G/abstract), [4](https://arxiv.org/abs/1602.01901), and [5](https://arxiv.org/abs/1603.08484) to name a few) so this is a good benchmark for our demonstration.\n",
    "\n",
    "The light curve that we'll use is from K2 and we'll use the same radial velocity measurements as [David+ (2016)](https://arxiv.org/abs/1602.01901) compiled from [here](https://arxiv.org/abs/astro-ph/0403444) and [here](https://ui.adsabs.harvard.edu/abs/2007A%26A...463..579G/abstract).\n",
    "We'll use a somewhat simplified model for the eclipses that treats the stars as spherical and ignores the phase curve (we'll model it using a Gaussian process instead of a more physically motivated model).\n",
    "But, as you'll see, these simplifying assumptions are sufficient for this case of a detached and well behaved system.\n",
    "Unlike some previous studies, we will fit an eccentric orbit instead of fixing the eccentricity to zero.\n",
    "This probably isn't really necessary here, but it's useful to demonstrate how you would fit a more eccentric system.\n",
    "Finally, we model the phase curve and other trends in both the light curve and radial velocities using Gaussian processes.\n",
    "This will account for unmodeled stellar variability and residual systematics, drifts, and other effects left over from the data reduction procedure.\n",
    "\n",
    "## Data access\n",
    "\n",
    "First, let's define some values from the literature that will be useful below.\n",
    "Here we're taking the period and eclipse time from [David+ (2016)](https://arxiv.org/abs/1602.01901) as initial guesses for these parameters in our fit.\n",
    "We'll also include the same prior on the flux ratio of the two stars that was computed for the Kepler bandpass by [David+ (2016)](https://arxiv.org/abs/1602.01901)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd02d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "lit_period = 2.46113408\n",
    "lit_t0 = 119.522070 + 2457000 - 2454833\n",
    "\n",
    "# Prior on the flux ratio for Kepler\n",
    "lit_flux_ratio = (0.354, 0.035)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21d405ba",
   "metadata": {},
   "source": [
    "Then we'll download the Kepler data.\n",
    "In this case, the pipeline aperture photometry isn't very good (because this star is so bright!) so we'll just download the target pixel file and co-add all the pixels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cebe2058",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import lightkurve as lk\n",
    "\n",
    "tpf = lk.search_targetpixelfile(\"EPIC 211082420\").download()\n",
    "lc = tpf.to_lightcurve(aperture_mask=\"all\")\n",
    "lc = lc.remove_nans().normalize()\n",
    "\n",
    "hdr = tpf.hdu[1].header\n",
    "texp = hdr[\"FRAMETIM\"] * hdr[\"NUM_FRM\"]\n",
    "texp /= 60.0 * 60.0 * 24.0\n",
    "\n",
    "# To keep things fast for this example, we're only going to use a third of the data\n",
    "np.random.seed(68594)\n",
    "m = np.random.rand(len(lc.time)) < 1.0 / 3\n",
    "\n",
    "x = np.ascontiguousarray(lc.time.value[m], dtype=np.float64)\n",
    "y = np.ascontiguousarray(lc.flux[m], dtype=np.float64)\n",
    "mu = np.median(y)\n",
    "y = (y / mu - 1) * 1e3\n",
    "\n",
    "plt.plot(\n",
    "    (x - lit_t0 + 0.5 * lit_period) % lit_period - 0.5 * lit_period, y, \".k\"\n",
    ")\n",
    "plt.xlim(-0.5 * lit_period, 0.5 * lit_period)\n",
    "plt.xlabel(\"time since primary eclipse [days]\")\n",
    "_ = plt.ylabel(\"relative flux [ppt]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a76d0d2d",
   "metadata": {},
   "source": [
    "Then we'll enter the radial velocity data.\n",
    "I couldn't find these data online anywhere so I manually transcribed the data from the referenced papers (typos are my own!)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37554f0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ref1 = 2453000\n",
    "ref2 = 2400000\n",
    "rvs = np.array(\n",
    "    [\n",
    "        # https://arxiv.org/abs/astro-ph/0403444\n",
    "        (39.41273 + ref1, -85.0, 134.5),\n",
    "        (39.45356 + ref1, -88.0, 139.0),\n",
    "        (39.50548 + ref1, -91.0, 143.0),\n",
    "        (43.25049 + ref1, 105.5, -136.0),\n",
    "        (46.25318 + ref1, 29.5, -24.5),\n",
    "        # https://ui.adsabs.harvard.edu/abs/2007A%26A...463..579G/abstract\n",
    "        (52629.6190 + ref2, 88.8, -127.0),\n",
    "        (52630.6098 + ref2, -48.0, 68.0),\n",
    "        (52631.6089 + ref2, -9.5, 13.1),\n",
    "        (52632.6024 + ref2, 63.6, -90.9),\n",
    "        (52633.6162 + ref2, -94.5, 135.0),\n",
    "        (52636.6055 + ref2, 10.3, -13.9),\n",
    "        (52983.6570 + ref2, 18.1, -25.1),\n",
    "        (52987.6453 + ref2, -80.6, 114.5),\n",
    "        (52993.6322 + ref2, 49.0, -70.7),\n",
    "        (53224.9338 + ref2, 39.0, -55.7),\n",
    "        (53229.9384 + ref2, 57.2, -82.0),\n",
    "    ]\n",
    ")\n",
    "rvs[:, 0] -= 2454833\n",
    "rvs = rvs[np.argsort(rvs[:, 0])]\n",
    "\n",
    "x_rv = np.ascontiguousarray(rvs[:, 0], dtype=np.float64)\n",
    "y1_rv = np.ascontiguousarray(rvs[:, 1], dtype=np.float64)\n",
    "y2_rv = np.ascontiguousarray(rvs[:, 2], dtype=np.float64)\n",
    "\n",
    "fold = (rvs[:, 0] - lit_t0 + 0.5 * lit_period) % lit_period - 0.5 * lit_period\n",
    "plt.plot(fold, rvs[:, 1], \".\", label=\"primary\")\n",
    "plt.plot(fold, rvs[:, 2], \".\", label=\"secondary\")\n",
    "plt.legend(fontsize=10)\n",
    "plt.xlim(-0.5 * lit_period, 0.5 * lit_period)\n",
    "plt.ylabel(\"radial velocity [km / s]\")\n",
    "_ = plt.xlabel(\"time since primary eclipse [days]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8bbd5b4",
   "metadata": {},
   "source": [
    "## Probabilistic model\n",
    "\n",
    "Then we define the probabilistic model using PyMC3 and exoplanet.\n",
    "This is similar to the other tutorials and case studies, but here we're using a [SecondaryEclipseLightCurve](https://docs.exoplanet.codes/en/stable/user/api/#exoplanet.SecondaryEclipseLightCurve) to generate the model light curve and we're modeling the radial velocity trends using a Gaussian process instead of a polynomial.\n",
    "Otherwise, things should look pretty familiar!\n",
    "\n",
    "After defining the model, we iteratively clip outliers in the light curve using sigma clipping and then estimate the maximum a posteriori parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faf7d53e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc3 as pm\n",
    "import aesara_theano_fallback.tensor as tt\n",
    "\n",
    "import exoplanet as xo\n",
    "import pymc3_ext as pmx\n",
    "from celerite2.theano import terms, GaussianProcess\n",
    "\n",
    "\n",
    "def build_model(mask):\n",
    "    with pm.Model() as model:\n",
    "        # Systemic parameters\n",
    "        mean_lc = pm.Normal(\"mean_lc\", mu=0.0, sd=5.0)\n",
    "        mean_rv = pm.Normal(\"mean_rv\", mu=0.0, sd=50.0)\n",
    "        u1 = xo.QuadLimbDark(\"u1\")\n",
    "        u2 = xo.QuadLimbDark(\"u2\")\n",
    "\n",
    "        # Parameters describing the primary\n",
    "        log_M1 = pm.Normal(\"log_M1\", mu=0.0, sigma=10.0)\n",
    "        log_R1 = pm.Normal(\"log_R1\", mu=0.0, sigma=10.0)\n",
    "        M1 = pm.Deterministic(\"M1\", tt.exp(log_M1))\n",
    "        R1 = pm.Deterministic(\"R1\", tt.exp(log_R1))\n",
    "\n",
    "        # Secondary ratios\n",
    "        log_k = pm.Normal(\"log_k\", mu=0.0, sigma=10.0)  # radius ratio\n",
    "        log_q = pm.Normal(\"log_q\", mu=0.0, sigma=10.0)  # mass ratio\n",
    "        log_s = pm.Normal(\n",
    "            \"log_s\", mu=np.log(0.5), sigma=10.0\n",
    "        )  # surface brightness ratio\n",
    "        pm.Deterministic(\"k\", tt.exp(log_k))\n",
    "        pm.Deterministic(\"q\", tt.exp(log_q))\n",
    "        pm.Deterministic(\"s\", tt.exp(log_s))\n",
    "\n",
    "        # Prior on flux ratio\n",
    "        pm.Normal(\n",
    "            \"flux_prior\",\n",
    "            mu=lit_flux_ratio[0],\n",
    "            sigma=lit_flux_ratio[1],\n",
    "            observed=tt.exp(2 * log_k + log_s),\n",
    "        )\n",
    "\n",
    "        # Parameters describing the orbit\n",
    "        b = xo.ImpactParameter(\"b\", ror=tt.exp(log_k), testval=1.5)\n",
    "        log_period = pm.Normal(\"log_period\", mu=np.log(lit_period), sigma=1.0)\n",
    "        period = pm.Deterministic(\"period\", tt.exp(log_period))\n",
    "        t0 = pm.Normal(\"t0\", mu=lit_t0, sigma=1.0)\n",
    "\n",
    "        # Parameters describing the eccentricity: ecs = [e * cos(w), e * sin(w)]\n",
    "        ecs = pmx.UnitDisk(\"ecs\", testval=np.array([1e-5, 0.0]))\n",
    "        ecc = pm.Deterministic(\"ecc\", tt.sqrt(tt.sum(ecs**2)))\n",
    "        omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n",
    "\n",
    "        # Build the orbit\n",
    "        R2 = pm.Deterministic(\"R2\", tt.exp(log_k + log_R1))\n",
    "        M2 = pm.Deterministic(\"M2\", tt.exp(log_q + log_M1))\n",
    "        orbit = xo.orbits.KeplerianOrbit(\n",
    "            period=period,\n",
    "            t0=t0,\n",
    "            ecc=ecc,\n",
    "            omega=omega,\n",
    "            b=b,\n",
    "            r_star=R1,\n",
    "            m_star=M1,\n",
    "            m_planet=M2,\n",
    "        )\n",
    "\n",
    "        # Track some other orbital elements\n",
    "        pm.Deterministic(\"incl\", orbit.incl)\n",
    "        pm.Deterministic(\"a\", orbit.a)\n",
    "\n",
    "        # Noise model for the light curve\n",
    "        sigma_lc = pm.InverseGamma(\n",
    "            \"sigma_lc\",\n",
    "            testval=1.0,\n",
    "            **pmx.estimate_inverse_gamma_parameters(0.1, 2.0),\n",
    "        )\n",
    "        sigma_gp = pm.InverseGamma(\n",
    "            \"sigma_gp\",\n",
    "            testval=0.5,\n",
    "            **pmx.estimate_inverse_gamma_parameters(1.0, 5.0),\n",
    "        )\n",
    "        rho_gp = pm.InverseGamma(\n",
    "            \"rho_gp\",\n",
    "            testval=5.0,\n",
    "            **pmx.estimate_inverse_gamma_parameters(1.0, 5.0),\n",
    "        )\n",
    "        kernel_lc = terms.SHOTerm(sigma=sigma_gp, rho=rho_gp, Q=1.0 / 3)\n",
    "\n",
    "        # Noise model for the radial velocities\n",
    "        sigma_rv1 = pm.InverseGamma(\n",
    "            \"sigma_rv1\",\n",
    "            testval=1.0,\n",
    "            **pmx.estimate_inverse_gamma_parameters(0.5, 5.0),\n",
    "        )\n",
    "        sigma_rv2 = pm.InverseGamma(\n",
    "            \"sigma_rv2\",\n",
    "            testval=1.0,\n",
    "            **pmx.estimate_inverse_gamma_parameters(0.5, 5.0),\n",
    "        )\n",
    "        sigma_rv_gp = pm.InverseGamma(\n",
    "            \"sigma_rv_gp\",\n",
    "            testval=1.5,\n",
    "            **pmx.estimate_inverse_gamma_parameters(1.0, 5.0),\n",
    "        )\n",
    "        rho_rv_gp = pm.InverseGamma(\n",
    "            \"rho_rv_gp\",\n",
    "            testval=2.0,\n",
    "            **pmx.estimate_inverse_gamma_parameters(1.0, 25.0),\n",
    "        )\n",
    "        kernel_rv = terms.SHOTerm(sigma=sigma_rv_gp, w0=rho_rv_gp, Q=1.0 / 3)\n",
    "\n",
    "        # Set up the light curve model\n",
    "        lc = xo.SecondaryEclipseLightCurve(u1, u2, tt.exp(log_s))\n",
    "\n",
    "        def model_lc(t):\n",
    "            return (\n",
    "                mean_lc\n",
    "                + 1e3\n",
    "                * lc.get_light_curve(orbit=orbit, r=R2, t=t, texp=texp)[:, 0]\n",
    "            )\n",
    "\n",
    "        # Condition the light curve model on the data\n",
    "        gp_lc = GaussianProcess(kernel_lc, t=x[mask], yerr=sigma_lc)\n",
    "        gp_lc.marginal(\"obs_lc\", observed=y[mask] - model_lc(x[mask]))\n",
    "\n",
    "        # Set up the radial velocity model\n",
    "        def model_rv1(t):\n",
    "            return mean_rv + 1e-3 * orbit.get_radial_velocity(t)\n",
    "\n",
    "        def model_rv2(t):\n",
    "            return mean_rv - 1e-3 * orbit.get_radial_velocity(t) * tt.exp(\n",
    "                -log_q\n",
    "            )\n",
    "\n",
    "        # Condition the radial velocity model on the data\n",
    "        gp_rv1 = GaussianProcess(kernel_rv, t=x_rv, yerr=sigma_rv1)\n",
    "        gp_rv1.marginal(\"obs_rv1\", observed=y1_rv - model_rv1(x_rv))\n",
    "        gp_rv2 = GaussianProcess(kernel_rv, t=x_rv, yerr=sigma_rv2)\n",
    "        gp_rv2.marginal(\"obs_rv2\", observed=y2_rv - model_rv2(x_rv))\n",
    "\n",
    "        # Optimize the logp\n",
    "        map_soln = model.test_point\n",
    "\n",
    "        # First the RV parameters\n",
    "        map_soln = pmx.optimize(map_soln, [mean_rv, log_q])\n",
    "        map_soln = pmx.optimize(\n",
    "            map_soln, [mean_rv, sigma_rv1, sigma_rv2, sigma_rv_gp, rho_rv_gp]\n",
    "        )\n",
    "\n",
    "        # Then the LC parameters\n",
    "        map_soln = pmx.optimize(map_soln, [mean_lc, log_R1, log_k, log_s, b])\n",
    "        map_soln = pmx.optimize(\n",
    "            map_soln, [mean_lc, log_R1, log_k, log_s, b, u1, u2]\n",
    "        )\n",
    "        map_soln = pmx.optimize(\n",
    "            map_soln, [mean_lc, sigma_lc, sigma_gp, rho_gp]\n",
    "        )\n",
    "        map_soln = pmx.optimize(map_soln, [t0, log_period])\n",
    "\n",
    "        # Then all the parameters together\n",
    "        map_soln = pmx.optimize(map_soln, [mean_rv, log_q, ecs])\n",
    "        map_soln = pmx.optimize(map_soln)\n",
    "\n",
    "        extras = dict(\n",
    "            x=x[mask],\n",
    "            y=y[mask],\n",
    "            model_lc=model_lc,\n",
    "            model_rv1=model_rv1,\n",
    "            model_rv2=model_rv2,\n",
    "            gp_lc_pred=gp_lc.predict(y[mask] - model_lc(x[mask])),\n",
    "        )\n",
    "\n",
    "    return model, map_soln, extras\n",
    "\n",
    "\n",
    "def sigma_clip():\n",
    "    mask = np.ones(len(x), dtype=bool)\n",
    "    num = len(mask)\n",
    "\n",
    "    for i in range(10):\n",
    "        model, map_soln, extras = build_model(mask)\n",
    "\n",
    "        with model:\n",
    "            mdl = pmx.eval_in_model(\n",
    "                extras[\"model_lc\"](extras[\"x\"]) + extras[\"gp_lc_pred\"],\n",
    "                map_soln,\n",
    "            )\n",
    "\n",
    "        resid = y[mask] - mdl\n",
    "        sigma = np.sqrt(np.median((resid - np.median(resid)) ** 2))\n",
    "        mask[mask] = np.abs(resid - np.median(resid)) < 7 * sigma\n",
    "        print(\"Sigma clipped {0} light curve points\".format(num - mask.sum()))\n",
    "        if num - mask.sum() < 10:\n",
    "            break\n",
    "        num = mask.sum()\n",
    "\n",
    "    return model, map_soln, extras\n",
    "\n",
    "\n",
    "model, map_soln, extras = sigma_clip()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fbaca43",
   "metadata": {},
   "source": [
    "At these maximum a posteriori parameters, let's make some plots of the model predictions compared to the observations to make sure that things look reasonable.\n",
    "First the phase-folded radial velocities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a1d2f52",
   "metadata": {},
   "outputs": [],
   "source": [
    "period = map_soln[\"period\"]\n",
    "t0 = map_soln[\"t0\"]\n",
    "mean = map_soln[\"mean_rv\"]\n",
    "\n",
    "x_fold = (x_rv - t0 + 0.5 * period) % period - 0.5 * period\n",
    "plt.plot(fold, y1_rv - mean, \".\", label=\"primary\")\n",
    "plt.plot(fold, y2_rv - mean, \".\", label=\"secondary\")\n",
    "\n",
    "x_phase = np.linspace(-0.5 * period, 0.5 * period, 500)\n",
    "with model:\n",
    "    y1_mod, y2_mod = pmx.eval_in_model(\n",
    "        [extras[\"model_rv1\"](x_phase + t0), extras[\"model_rv2\"](x_phase + t0)],\n",
    "        map_soln,\n",
    "    )\n",
    "plt.plot(x_phase, y1_mod - mean, \"C0\")\n",
    "plt.plot(x_phase, y2_mod - mean, \"C1\")\n",
    "\n",
    "plt.legend(fontsize=10)\n",
    "plt.xlim(-0.5 * period, 0.5 * period)\n",
    "plt.ylabel(\"radial velocity [km / s]\")\n",
    "plt.xlabel(\"time since primary eclipse [days]\")\n",
    "_ = plt.title(\"HD 23642; map model\", fontsize=14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4a87106",
   "metadata": {},
   "source": [
    "And then the light curve.\n",
    "In the top panel, we show the Gaussian process model for the phase curve.\n",
    "It's clear that there's a lot of information there that we could take advantage of, but that's a topic for another day.\n",
    "In the bottom panel, we're plotting the phase folded light curve and we can see the ridiculous signal to noise that we're getting on the eclipses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2be079e",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_lc_pred = np.linspace(x.min(), x.max(), 3000)\n",
    "with model:\n",
    "    gp_pred = (\n",
    "        pmx.eval_in_model(extras[\"gp_lc_pred\"], map_soln) + map_soln[\"mean_lc\"]\n",
    "    )\n",
    "    lc = (\n",
    "        pmx.eval_in_model(extras[\"model_lc\"](t_lc_pred), map_soln)\n",
    "        - map_soln[\"mean_lc\"]\n",
    "    )\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(2, sharex=True, figsize=(12, 7))\n",
    "\n",
    "ax1.plot(extras[\"x\"], extras[\"y\"], \"k.\", alpha=0.2)\n",
    "ax1.plot(extras[\"x\"], gp_pred, color=\"C1\", lw=1)\n",
    "\n",
    "ax2.plot(extras[\"x\"], extras[\"y\"] - gp_pred, \"k.\", alpha=0.2)\n",
    "ax2.plot(t_lc_pred, lc, color=\"C2\", lw=1)\n",
    "ax2.set_xlim(extras[\"x\"].min(), extras[\"x\"].max())\n",
    "\n",
    "ax1.set_ylabel(\"raw flux [ppt]\")\n",
    "ax2.set_ylabel(\"de-trended flux [ppt]\")\n",
    "ax2.set_xlabel(\"time [KBJD]\")\n",
    "ax1.set_title(\"HD 23642; map model\", fontsize=14)\n",
    "\n",
    "fig.subplots_adjust(hspace=0.05)\n",
    "\n",
    "fig, ax1 = plt.subplots(1, figsize=(12, 3.5))\n",
    "\n",
    "x_fold = (\n",
    "    (extras[\"x\"] - map_soln[\"t0\"]) % map_soln[\"period\"] / map_soln[\"period\"]\n",
    ")\n",
    "inds = np.argsort(x_fold)\n",
    "\n",
    "ax1.plot(x_fold[inds], extras[\"y\"][inds] - gp_pred[inds], \"k.\", alpha=0.2)\n",
    "ax1.plot(x_fold[inds] - 1, extras[\"y\"][inds] - gp_pred[inds], \"k.\", alpha=0.2)\n",
    "ax2.plot(\n",
    "    x_fold[inds],\n",
    "    extras[\"y\"][inds] - gp_pred[inds],\n",
    "    \"k.\",\n",
    "    alpha=0.2,\n",
    "    label=\"data!\",\n",
    ")\n",
    "ax2.plot(x_fold[inds] - 1, extras[\"y\"][inds] - gp_pred, \"k.\", alpha=0.2)\n",
    "\n",
    "yval = extras[\"y\"][inds] - gp_pred\n",
    "bins = np.linspace(0, 1, 75)\n",
    "num, _ = np.histogram(x_fold[inds], bins, weights=yval)\n",
    "denom, _ = np.histogram(x_fold[inds], bins)\n",
    "ax2.plot(0.5 * (bins[:-1] + bins[1:]) - 1, num / denom, \".w\")\n",
    "\n",
    "args = dict(lw=1)\n",
    "\n",
    "x_fold = (t_lc_pred - map_soln[\"t0\"]) % map_soln[\"period\"] / map_soln[\"period\"]\n",
    "inds = np.argsort(x_fold)\n",
    "ax1.plot(x_fold[inds], lc[inds], \"C2\", **args)\n",
    "ax1.plot(x_fold[inds] - 1, lc[inds], \"C2\", **args)\n",
    "\n",
    "ax1.set_xlim(-1, 1)\n",
    "ax1.set_ylabel(\"de-trended flux [ppt]\")\n",
    "ax1.set_xlabel(\"phase\")\n",
    "_ = ax1.set_title(\"HD 23642; map model\", fontsize=14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "566f67a1",
   "metadata": {},
   "source": [
    "## Sampling\n",
    "\n",
    "Finally we can run the MCMC:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93fa4b3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(23642)\n",
    "with model:\n",
    "    trace = pmx.sample(\n",
    "        tune=1500,\n",
    "        draws=1000,\n",
    "        start=map_soln,\n",
    "        cores=2,\n",
    "        chains=2,\n",
    "        target_accept=0.95,\n",
    "        return_inferencedata=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "686a0bf7",
   "metadata": {},
   "source": [
    "As usual, we can check the convergence diagnostics for some of the key parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f2d4129",
   "metadata": {},
   "outputs": [],
   "source": [
    "import arviz as az\n",
    "\n",
    "az.summary(trace, var_names=[\"M1\", \"M2\", \"R1\", \"R2\", \"ecs\", \"incl\", \"s\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08e95ad5",
   "metadata": {},
   "source": [
    "## Results\n",
    "\n",
    "It can be useful to take a look at some diagnostic corner plots to see how the sampling went.\n",
    "First, let's look at some observables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f68aee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import corner\n",
    "\n",
    "_ = corner.corner(\n",
    "    trace,\n",
    "    var_names=[\"k\", \"q\", \"ecs\"],\n",
    "    labels=[\n",
    "        \"$k = R_2 / R_1$\",\n",
    "        \"$q = M_2 / M_1$\",\n",
    "        \"$e\\,\\cos\\omega$\",\n",
    "        \"$e\\,\\sin\\omega$\",\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fd825ab",
   "metadata": {},
   "source": [
    "And then we can look at the physical properties of the stars in the system.\n",
    "In this figure, we're comparing to the results from [David+ (2016)](https://arxiv.org/abs/1602.01901) (shown as blue crosshairs).\n",
    "The orange contours in this figure show the results transformed to a uniform prior on eccentricity as discussed below.\n",
    "These contours are provided to demonstrate (qualitatively) that these inferences are not sensitive to the choice of prior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d8a10cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "flat_samps = trace.posterior.stack(sample=(\"chain\", \"draw\"))\n",
    "weights = 1.0 / flat_samps[\"ecc\"].values\n",
    "weights *= len(weights) / np.sum(weights)\n",
    "fig = corner.corner(\n",
    "    trace,\n",
    "    var_names=[\"R1\", \"R2\", \"M1\", \"M2\"],\n",
    "    weights=weights,\n",
    "    plot_datapoints=False,\n",
    "    color=\"C1\",\n",
    "    smooth=1,\n",
    "    smooth1d=1,\n",
    ")\n",
    "_ = corner.corner(\n",
    "    trace,\n",
    "    var_names=[\"R1\", \"R2\", \"M1\", \"M2\"],\n",
    "    truths=[1.727, 1.503, 2.203, 1.5488],\n",
    "    fig=fig,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4c62c5c",
   "metadata": {},
   "source": [
    "## A note about eccentricities\n",
    "\n",
    "If you looked closely at the model defined above, you might have noticed that we chose a slightly odd eccentricity prior: $p(e) \\propto e$.\n",
    "This is implied by sampling with $e\\,\\cos\\omega$ and $e\\,\\sin\\omega$ as the parameters, as has been discussed many times in the literature.\n",
    "There are many options for correcting for this prior and instead assuming a uniform prior on eccentricity (for example, sampling with $\\sqrt{e}\\,\\cos\\omega$ and $\\sqrt{e}\\,\\sin\\omega$ as the parameters), but you'll find much worse sampling performance for this problem if you try any of these options (trust us, we tried!) because the geometry of the posterior surface becomes much less suitable for the sampling algorithm in PyMC3.\n",
    "Instead, we can re-weight the samples after running the MCMC to see how the results change under the new prior.\n",
    "Most of the parameter inferences are unaffected by this change (because the data are very constraining!), but the inferred eccentricity (and especially $e\\,\\sin\\omega$) will depend on this choice.\n",
    "The following plots show how these parameter inferences are affected.\n",
    "Note, especially, how the shape of the $e\\,\\sin\\omega$ density changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45dfe752",
   "metadata": {},
   "outputs": [],
   "source": [
    "esinw = flat_samps[\"ecc\"].values * np.sin(flat_samps[\"omega\"].values)\n",
    "plt.hist(\n",
    "    esinw,\n",
    "    50,\n",
    "    density=True,\n",
    "    histtype=\"step\",\n",
    "    label=\"$p(e) = 2\\,e$\",\n",
    ")\n",
    "plt.hist(\n",
    "    esinw,\n",
    "    50,\n",
    "    density=True,\n",
    "    histtype=\"step\",\n",
    "    weights=1.0 / flat_samps[\"ecc\"].values,\n",
    "    label=\"$p(e) = 1$\",\n",
    ")\n",
    "plt.xlabel(\"$e\\,\\sin(\\omega)$\")\n",
    "plt.ylabel(\"$p(e\\,\\sin\\omega\\,|\\,\\mathrm{data})$\")\n",
    "plt.yticks([])\n",
    "plt.legend(fontsize=12)\n",
    "\n",
    "plt.figure()\n",
    "plt.hist(\n",
    "    flat_samps[\"ecc\"].values,\n",
    "    50,\n",
    "    density=True,\n",
    "    histtype=\"step\",\n",
    "    label=\"$p(e) = 2\\,e$\",\n",
    ")\n",
    "plt.hist(\n",
    "    flat_samps[\"ecc\"].values,\n",
    "    50,\n",
    "    density=True,\n",
    "    histtype=\"step\",\n",
    "    weights=1.0 / flat_samps[\"ecc\"].values,\n",
    "    label=\"$p(e) = 1$\",\n",
    ")\n",
    "plt.xlabel(\"$e$\")\n",
    "plt.ylabel(\"$p(e\\,|\\,\\mathrm{data})$\")\n",
    "plt.yticks([])\n",
    "plt.xlim(0, 0.015)\n",
    "_ = plt.legend(fontsize=12)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c646fce",
   "metadata": {},
   "source": [
    "We can then use the `corner.quantile` function to compute summary statistics of the weighted samples as follows.\n",
    "For example, here how to compute the 90% posterior upper limit for the eccentricity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa3e52d",
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = 1.0 / flat_samps[\"ecc\"].values\n",
    "print(\n",
    "    \"for p(e) = 2*e: p(e < x) = 0.9 -> x = {0:.5f}\".format(\n",
    "        corner.quantile(flat_samps[\"ecc\"].values, [0.9])[0]\n",
    "    )\n",
    ")\n",
    "print(\n",
    "    \"for p(e) = 1:   p(e < x) = 0.9 -> x = {0:.5f}\".format(\n",
    "        corner.quantile(flat_samps[\"ecc\"].values, [0.9], weights=weights)[0]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a71cd0dd",
   "metadata": {},
   "source": [
    "Or, the posterior mean and variance for the radius of the primary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af0964b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = flat_samps[\"R1\"].values\n",
    "\n",
    "print(\n",
    "    \"for p(e) = 2*e: R1 = {0:.3f} ± {1:.3f}\".format(\n",
    "        np.mean(samples), np.std(samples)\n",
    "    )\n",
    ")\n",
    "\n",
    "mean = np.sum(weights * samples) / np.sum(weights)\n",
    "sigma = np.sqrt(np.sum(weights * (samples - mean) ** 2) / np.sum(weights))\n",
    "print(\"for p(e) = 1:   R1 = {0:.3f} ± {1:.3f}\".format(mean, sigma))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c595de6e",
   "metadata": {},
   "source": [
    "As you can see (and as one would hope) this choice of prior does not significantly change our inference of the primary radius."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ca0a3d",
   "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": "2f34e160",
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    txt, bib = xo.citations.get_citations_for_model()\n",
    "print(txt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f82b81fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(bib.split(\"\\n\\n\")[0] + \"\\n\\n...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7ed353a-c229-4809-b079-9d2ca202a06b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "encoding": "# -*- coding: utf-8 -*-"
  },
  "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
}