eb.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run notebook_setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Eclipsing binary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import exoplanet as xo\n",
    "import lightkurve as lk\n",
    "import pymc3 as pm\n",
    "import theano.tensor as tt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lit_period = 1.4408031\n",
    "lit_t0 = 894.37787 + 2456000 - 2454833"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lcf = lk.search_lightcurvefile(\"EPIC 203476597\").download(quality_bitmask=0)\n",
    "\n",
    "hdr = lcf.header(1)\n",
    "texp = hdr[\"FRAMETIM\"] * hdr[\"NUM_FRM\"]\n",
    "texp /= 60.0 * 60.0 * 24.0\n",
    "\n",
    "lc = lcf.PDCSAP_FLUX\n",
    "lc = lc.remove_nans()\n",
    "\n",
    "# lc = lc.flatten(window_length=401)\n",
    "\n",
    "# lc = lc.to_corrector(\"sff\").correct(windows=20)\n",
    "\n",
    "lc = lc.normalize()\n",
    "\n",
    "x = np.ascontiguousarray(lc.time, dtype=np.float64)\n",
    "y = np.ascontiguousarray(1e3 * (lc.flux - 1), dtype=np.float64)\n",
    "yerr = np.ascontiguousarray(1e3 * lc.flux_err, dtype=np.float64)\n",
    "\n",
    "# plt.plot(x, y)\n",
    "plt.scatter((x - lit_t0 + 0.5 * lit_period) % lit_period - 0.5 * lit_period, y, c=x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(mask):\n",
    "\n",
    "    with pm.Model() as model:\n",
    "\n",
    "        rho_star = pm.Bound(pm.Normal, lower=0)(\"rho_star\", mu=0.27, sigma=0.03)\n",
    "\n",
    "        mean = pm.Normal(\"mean\", mu=0.0, sd=5.0)\n",
    "        flux_ratio = pm.Lognormal(\"flux_ratio\", mu=np.log(50.0), sigma=10.0)\n",
    "        u1 = xo.QuadLimbDark(\"u1\")\n",
    "        u2 = xo.QuadLimbDark(\"u2\")\n",
    "\n",
    "        period = pm.Lognormal(\"period\", mu=np.log(lit_period), sigma=1.0)\n",
    "        t0 = pm.Normal(\"t0\", mu=lit_t0, sigma=1.0)\n",
    "        ror = pm.Lognormal(\"ror\", mu=np.log(0.07), sigma=5.0)\n",
    "        b = xo.ImpactParameter(\"b\", ror=ror, testval=0.5)\n",
    "                   \n",
    "        hk = xo.UnitDisk(\"hk\", testval=np.array([1e-4, 0.0]))\n",
    "        ecc = pm.Deterministic(\"ecc\", tt.sum(hk ** 2))\n",
    "        omega = pm.Deterministic(\"omega\", tt.arctan2(hk[0], hk[1]))\n",
    "\n",
    "        orbit = xo.orbits.KeplerianOrbit(period=period, t0=t0, ecc=ecc, omega=omega, b=b, rho_star=rho_star)\n",
    "\n",
    "        lc = xo.SecondaryEclipseLightCurve(u1, u2, 1e-3 * flux_ratio)\n",
    "\n",
    "        def lc_model(t):\n",
    "            return 1e3 * lc.get_light_curve(orbit=orbit, r=ror, t=t, texp=texp)[:, 0]\n",
    "\n",
    "        def full_model(t):\n",
    "            return mean + lc_model(t)\n",
    "\n",
    "        sigma = pm.InverseGamma(\"sigma\", testval=1.0, **xo.estimate_inverse_gamma_parameters(0.1, 2.0))\n",
    "        S_tot = pm.InverseGamma(\"S_tot\", testval=2.5, **xo.estimate_inverse_gamma_parameters(1.0, 5.0))\n",
    "        ell = pm.InverseGamma(\"ell\", testval=2.0, **xo.estimate_inverse_gamma_parameters(1.0, 5.0))\n",
    "        kernel = xo.gp.terms.SHOTerm(S_tot=S_tot, w0=2 * np.pi / ell, Q=1. / 3)\n",
    "        gp = xo.gp.GP(kernel, x[mask], yerr[mask] ** 2 + sigma ** 2, mean=full_model)\n",
    "\n",
    "        gp.marginal(\"obs\", observed=y[mask])\n",
    "        \n",
    "        map_soln = model.test_point\n",
    "        map_soln = xo.optimize(map_soln, [mean, u1, u2, ror, b, rho_star])\n",
    "        map_soln = xo.optimize(map_soln, [mean, sigma, S_tot, ell])\n",
    "        map_soln = xo.optimize(map_soln, [t0, period])\n",
    "        map_soln = xo.optimize(map_soln, [mean, flux_ratio])\n",
    "        map_soln = xo.optimize(map_soln, [hk])\n",
    "        map_soln = xo.optimize(map_soln)\n",
    "\n",
    "        model.gp = gp\n",
    "        model.lc_model = lc_model\n",
    "        model.full_model = full_model\n",
    "        model.x = x[mask]\n",
    "        model.y = y[mask]\n",
    "        model.yerr = yerr[mask]\n",
    "        \n",
    "    return model, map_soln\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 = build_model(mask)\n",
    "\n",
    "        with model:\n",
    "            mdl = xo.eval_in_model(model.full_model(x[mask]) + model.gp.predict(), map_soln)\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)) < 5 * sigma \n",
    "        print(num, mask.sum())\n",
    "        if num == mask.sum():\n",
    "            break\n",
    "        num = mask.sum()\n",
    "        \n",
    "    return model, map_soln"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, map_soln = sigma_clip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_soln"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:    \n",
    "    gp_pred = xo.eval_in_model(model.gp.predict(), map_soln) + map_soln[\"mean\"]\n",
    "    lc = xo.eval_in_model(model.lc_model(model.x), map_soln)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(2, sharex=True, figsize=(12, 7))\n",
    "\n",
    "ax1.plot(model.x, model.y, \"k.\", alpha=0.2)\n",
    "ax1.plot(model.x, gp_pred, color=\"C1\", lw=1)\n",
    "\n",
    "ax2.plot(model.x, model.y - gp_pred, \"k.\", alpha=0.2)\n",
    "ax2.plot(model.x, lc, color=\"C2\", lw=1)\n",
    "ax2.set_xlim(model.x.min(), model.x.max())\n",
    "\n",
    "ax1.set_title(\"EPIC 203476597\", fontsize=14)\n",
    "ax1.set_ylabel(\"raw flux [ppt]\")\n",
    "ax2.set_ylabel(\"de-trended flux [ppt]\")\n",
    "ax2.set_xlabel(\"time [KBJD]\")\n",
    "\n",
    "fig.subplots_adjust(hspace=0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax1 = plt.subplots(1, figsize=(12, 3.5))\n",
    "\n",
    "x_fold = (model.x - map_soln[\"t0\"]) % map_soln[\"period\"] / map_soln[\"period\"]\n",
    "inds = np.argsort(x_fold)\n",
    "\n",
    "ax1.plot(x_fold[inds], model.y[inds] - gp_pred[inds], \"k.\", alpha=0.2)\n",
    "ax1.plot(x_fold[inds] - 1, model.y[inds] - gp_pred[inds], \"k.\", alpha=0.2)\n",
    "ax2.plot(x_fold[inds], model.y[inds] - gp_pred[inds], \"k.\", alpha=0.2, label=\"data!\")\n",
    "ax2.plot(x_fold[inds] - 1, model.y[inds] - gp_pred, \"k.\", alpha=0.2)\n",
    "\n",
    "yval = model.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",
    "ax1.plot(x_fold[inds], lc[inds], \"C1\", **args)\n",
    "ax1.plot(x_fold[inds] - 1, lc[inds], \"C1\", **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(\"EPIC 203476597\", fontsize=14);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(203476597)\n",
    "with model:\n",
    "    trace = pm.sample(\n",
    "        tune=2000,\n",
    "        draws=2000,\n",
    "        start=map_soln,\n",
    "        cores=2,\n",
    "        chains=2,\n",
    "        step=xo.get_dense_nuts_step(start=map_soln, target_accept=0.9),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pm.traceplot(trace);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pm.summary(trace)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(trace[\"hk\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import corner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = pm.trace_to_dataframe(trace, varnames=[v.name for v in model.vars])\n",
    "corner.corner(samples);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = pm.trace_to_dataframe(trace, varnames=[\"hk\", \"ecc\", \"omega\"])\n",
    "corner.corner(samples, labels=[\"$\\sqrt{e}\\sin\\omega$\", \"$\\sqrt{e}\\cos\\omega$\", \"$e$\", \"$\\omega$\"]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = pm.trace_to_dataframe(trace, varnames=[\"hk\", \"ecc\", \"omega\"])\n",
    "samples[\"ses\"] = np.sqrt(samples[\"ecc\"]) * np.sin(samples[\"omega\"])\n",
    "samples[\"sec\"] = np.sqrt(samples[\"ecc\"]) * np.cos(samples[\"omega\"])\n",
    "weights = samples[\"ecc\"]\n",
    "corner.corner(samples, weights=weights);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = pm.trace_to_dataframe(trace, varnames=[\"hk_unitdisk__\", \"rho_star\", \"Ae\", \"Ab\", \"Ar\", \"lag\", \"flux_ratio\", \"b\", \"ror\", \"sigma\", \"S_tot\", \"ell\", \"mean\", \"t0\", \"period\"])\n",
    "# samples[\"esinw\"] = samples[\"ecc\"] * np.sin(samples[\"omega\"])\n",
    "# samples[\"ecosw\"] = samples[\"ecc\"] * np.cos(samples[\"omega\"])\n",
    "corner.corner(samples);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ecosw = trace[\"ecc\"] * np.cos(trace[\"omega\"])\n",
    "esinw = trace[\"ecc\"] * np.sin(trace[\"omega\"])\n",
    "\n",
    "delta = ecosw / (0.25 * np.pi)\n",
    "frac = (1 + esinw) / (1 - esinw)\n",
    "\n",
    "corner.corner(np.concatenate((esinw[:, None], ecosw[:, None], trace[\"hk\"]), axis=-1));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corner.corner(np.dot(np.array([[np.cos(np.pi / 2), np.sin(np.pi / 2)], [-np.sin(np.pi / 2), np.cos(np.pi / 2)]]), trace[\"hk_unitdisk__\"].T).T);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}