{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ef809242-d29f-44c2-8425-4c5c6f8584b4",
   "metadata": {},
   "source": [
    "(rv-multi)=\n",
    "\n",
    "# RVs with multiple instruments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "343551b8-b377-4674-8058-3964e5953077",
   "metadata": {},
   "outputs": [],
   "source": [
    "import exoplanet\n",
    "\n",
    "exoplanet.utils.docs_setup()\n",
    "print(f\"exoplanet.__version__ = '{exoplanet.__version__}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7728bae8-f878-4ed8-8a14-c0dfd78b4b33",
   "metadata": {},
   "source": [
    "In this case study, we will look at how we can use exoplanet and PyMC3 to combine datasets from different RV instruments to fit the orbit of an exoplanet system.\n",
    "Before getting started, I want to emphasize that the exoplanet code doesn't have strong opinions about how your data are collected, it only provides extensions that allow PyMC3 to evaluate some astronomy-specific functions.\n",
    "This means that you can build any kind of observation model that PyMC3 supports, and support for multiple instruments isn't really a *feature* of exoplanet, even though it is easy to implement.\n",
    "\n",
    "For the example, we'll use public observations of Pi Mensae which hosts two planets, but we'll ignore the inner planet because the significance of the RV signal is small enough that it won't affect our results.\n",
    "The datasets that we'll use are from the Anglo-Australian Planet Search (AAT) and the HARPS archive.\n",
    "As is commonly done, we will treat the HARPS observations as two independent datasets split in June 2015 when the HARPS hardware was upgraded.\n",
    "Therefore, we'll consider three datasets that we will allow to have different instrumental parameters (RV offset and jitter), but shared orbital parameters and stellar variability.\n",
    "In some cases you might also want to have a different astrophyscial variability model for each instrument (if, for example, the observations are made in very different bands), but we'll keep things simple for this example.\n",
    "\n",
    "The AAT data are available from [The Exoplanet Archive](https://exoplanetarchive.ipac.caltech.edu/) and the HARPS observations can be downloaded from the [ESO Archive](http://archive.eso.org/wdb/wdb/adp/phase3_spectral/form).\n",
    "For the sake of simplicity, we have extracted the HARPS RVs from the archive in advance using [Megan Bedell's harps_tools library](https://github.com/megbedell/harps_tools).\n",
    "\n",
    "To start, download the data and plot them with a (very!) rough zero point correction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc26e0c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from astropy.io import ascii\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "aat = ascii.read(\n",
    "    \"https://exoplanetarchive.ipac.caltech.edu/data/ExoData/0026/0026394/data/UID_0026394_RVC_001.tbl\"\n",
    ")\n",
    "harps = pd.read_csv(\n",
    "    \"https://raw.githubusercontent.com/exoplanet-dev/case-studies/main/data/pi_men_harps_rvs.csv\",\n",
    "    skiprows=1,\n",
    ")\n",
    "harps = harps.rename(lambda x: x.strip().strip(\"#\"), axis=1)\n",
    "harps_post = np.array(harps.date > \"2015-07-01\", dtype=int)\n",
    "\n",
    "t = np.concatenate((aat[\"JD\"], harps[\"bjd\"]))\n",
    "rv = np.concatenate((aat[\"Radial_Velocity\"], harps[\"rv\"]))\n",
    "rv_err = np.concatenate((aat[\"Radial_Velocity_Uncertainty\"], harps[\"e_rv\"]))\n",
    "inst_id = np.concatenate((np.zeros(len(aat), dtype=int), harps_post + 1))\n",
    "\n",
    "inds = np.argsort(t)\n",
    "t = np.ascontiguousarray(t[inds], dtype=float)\n",
    "rv = np.ascontiguousarray(rv[inds], dtype=float)\n",
    "rv_err = np.ascontiguousarray(rv_err[inds], dtype=float)\n",
    "inst_id = np.ascontiguousarray(inst_id[inds], dtype=int)\n",
    "\n",
    "inst_names = [\"aat\", \"harps_pre\", \"harps_post\"]\n",
    "num_inst = len(inst_names)\n",
    "\n",
    "for i, name in enumerate(inst_names):\n",
    "    m = inst_id == i\n",
    "    plt.errorbar(\n",
    "        t[m], rv[m] - np.min(rv[m]), yerr=rv_err[m], fmt=\".\", label=name\n",
    "    )\n",
    "\n",
    "plt.legend(fontsize=10)\n",
    "plt.xlabel(\"BJD\")\n",
    "_ = plt.ylabel(\"radial velocity [m/s]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c9ccb08",
   "metadata": {},
   "source": [
    "Then set up the probabilistic model.\n",
    "Most of this is similar to the model in the [Radial velocity fitting](https://gallery.exoplanet.codes/tutorials/rv/) tutorial, but there are a few changes to highlight:\n",
    "\n",
    "1. Instead of a polynomial model for trends, stellar variability, and inner planets, we're using a Gaussian process here. This won't have a big effect here, but more careful consideration should be performed when studying lower signal-to-noise systems.\n",
    "2. There are three radial velocity offsets and three jitter parameters (one for each instrument) that will be treated independently. This is the key addition made by this case study."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0aacd4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc3 as pm\n",
    "import exoplanet as xo\n",
    "import aesara_theano_fallback.tensor as tt\n",
    "\n",
    "import pymc3_ext as pmx\n",
    "from celerite2.theano import terms, GaussianProcess\n",
    "\n",
    "t_phase = np.linspace(-0.5, 0.5, 5000)\n",
    "\n",
    "with pm.Model() as model:\n",
    "    # Parameters describing the orbit\n",
    "    log_K = pm.Normal(\"log_K\", mu=np.log(300), sigma=10)\n",
    "    log_P = pm.Normal(\"log_P\", mu=np.log(2093.07), sigma=10)\n",
    "    K = pm.Deterministic(\"K\", tt.exp(log_K))\n",
    "    P = pm.Deterministic(\"P\", tt.exp(log_P))\n",
    "\n",
    "    ecs = pmx.UnitDisk(\"ecs\", testval=np.array([0.7, -0.3]))\n",
    "    ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2))\n",
    "    omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n",
    "    phase = pmx.UnitUniform(\"phase\")\n",
    "    tp = pm.Deterministic(\"tp\", 0.5 * (t.min() + t.max()) + phase * P)\n",
    "\n",
    "    orbit = xo.orbits.KeplerianOrbit(\n",
    "        period=P, t_periastron=tp, ecc=ecc, omega=omega\n",
    "    )\n",
    "\n",
    "    # Noise model parameters\n",
    "    log_sigma_gp = pm.Normal(\"log_sigma_gp\", mu=np.log(10), sigma=50)\n",
    "    log_rho_gp = pm.Normal(\"log_rho_gp\", mu=np.log(50), sigma=50)\n",
    "\n",
    "    # Per instrument parameters\n",
    "    means = pm.Normal(\n",
    "        \"means\",\n",
    "        mu=np.array([np.median(rv[inst_id == i]) for i in range(num_inst)]),\n",
    "        sigma=200,\n",
    "        shape=num_inst,\n",
    "    )\n",
    "    sigmas = pm.HalfNormal(\"sigmas\", sigma=10, shape=num_inst)\n",
    "\n",
    "    # Compute the RV offset and jitter for each data point depending on its instrument\n",
    "    mean = tt.zeros(len(t))\n",
    "    diag = tt.zeros(len(t))\n",
    "    for i in range(len(inst_names)):\n",
    "        mean += means[i] * (inst_id == i)\n",
    "        diag += (rv_err**2 + sigmas[i] ** 2) * (inst_id == i)\n",
    "    pm.Deterministic(\"mean\", mean)\n",
    "    pm.Deterministic(\"diag\", diag)\n",
    "    resid = rv - mean\n",
    "\n",
    "    def rv_model(x):\n",
    "        return orbit.get_radial_velocity(x, K=K)\n",
    "\n",
    "    kernel = terms.SHOTerm(\n",
    "        sigma=tt.exp(log_sigma_gp), rho=tt.exp(log_rho_gp), Q=1.0 / 3\n",
    "    )\n",
    "    gp = GaussianProcess(kernel, t=t, diag=diag, mean=rv_model)\n",
    "    gp.marginal(\"obs\", observed=resid)\n",
    "    pm.Deterministic(\"gp_pred\", gp.predict(resid, include_mean=False))\n",
    "    pm.Deterministic(\"rv_phase\", rv_model(P * t_phase + tp))\n",
    "\n",
    "    map_soln = model.test_point\n",
    "    map_soln = pmx.optimize(map_soln, [means])\n",
    "    map_soln = pmx.optimize(map_soln, [means, phase])\n",
    "    map_soln = pmx.optimize(map_soln, [means, phase, log_K])\n",
    "    map_soln = pmx.optimize(map_soln, [means, tp, K, log_P, ecs])\n",
    "    map_soln = pmx.optimize(map_soln, [sigmas, log_sigma_gp, log_rho_gp])\n",
    "    map_soln = pmx.optimize(map_soln)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a995b50",
   "metadata": {},
   "source": [
    "After fitting for the parameters that maximize the posterior probability, we can plot this model to make sure that things are looking reasonable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bb6cb4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_pred = np.linspace(t.min() - 400, t.max() + 400, 5000)\n",
    "with model:\n",
    "    plt.plot(\n",
    "        t_pred, pmx.eval_in_model(rv_model(t_pred), map_soln), \"k\", lw=0.5\n",
    "    )\n",
    "\n",
    "detrended = rv - map_soln[\"mean\"] - map_soln[\"gp_pred\"]\n",
    "plt.errorbar(t, detrended, yerr=rv_err, fmt=\",k\")\n",
    "plt.scatter(\n",
    "    t, detrended, c=inst_id, s=8, zorder=100, cmap=\"tab10\", vmin=0, vmax=10\n",
    ")\n",
    "plt.xlim(t_pred.min(), t_pred.max())\n",
    "plt.xlabel(\"BJD\")\n",
    "plt.ylabel(\"radial velocity [m/s]\")\n",
    "_ = plt.title(\"map model\", fontsize=14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39de44f9",
   "metadata": {},
   "source": [
    "That looks fine, so now we can run the MCMC sampler:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e5f9237",
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    trace = pmx.sample(\n",
    "        tune=1000,\n",
    "        draws=1000,\n",
    "        start=map_soln,\n",
    "        chains=2,\n",
    "        cores=2,\n",
    "        return_inferencedata=True,\n",
    "        random_seed=[39091, 39095],\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77cf0cae",
   "metadata": {},
   "source": [
    "Then we can look at some summaries of the trace and the constraints on some of the key parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dcf9f2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import corner\n",
    "import arviz as az\n",
    "\n",
    "corner.corner(trace, var_names=[\"P\", \"K\", \"tp\", \"ecc\", \"omega\"])\n",
    "\n",
    "az.summary(\n",
    "    trace, var_names=[\"P\", \"K\", \"tp\", \"ecc\", \"omega\", \"means\", \"sigmas\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57ee64aa",
   "metadata": {},
   "source": [
    "And finally we can plot the phased RV curve and overplot our posterior inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ec10ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "flat_samps = trace.posterior.stack(sample=(\"chain\", \"draw\"))\n",
    "\n",
    "mu = np.mean(flat_samps[\"mean\"].values + flat_samps[\"gp_pred\"].values, axis=-1)\n",
    "mu_var = np.var(flat_samps[\"mean\"], axis=-1)\n",
    "jitter_var = np.median(flat_samps[\"diag\"], axis=-1)\n",
    "period = np.median(flat_samps[\"P\"])\n",
    "tp = np.median(flat_samps[\"tp\"])\n",
    "\n",
    "detrended = rv - mu\n",
    "folded = ((t - tp + 0.5 * period) % period) / period\n",
    "plt.errorbar(folded, detrended, yerr=np.sqrt(mu_var + jitter_var), fmt=\",k\")\n",
    "plt.scatter(\n",
    "    folded,\n",
    "    detrended,\n",
    "    c=inst_id,\n",
    "    s=8,\n",
    "    zorder=100,\n",
    "    cmap=\"tab10\",\n",
    "    vmin=0,\n",
    "    vmax=10,\n",
    ")\n",
    "plt.errorbar(\n",
    "    folded + 1, detrended, yerr=np.sqrt(mu_var + jitter_var), fmt=\",k\"\n",
    ")\n",
    "plt.scatter(\n",
    "    folded + 1,\n",
    "    detrended,\n",
    "    c=inst_id,\n",
    "    s=8,\n",
    "    zorder=100,\n",
    "    cmap=\"tab10\",\n",
    "    vmin=0,\n",
    "    vmax=10,\n",
    ")\n",
    "\n",
    "x = t_phase + 0.5\n",
    "y = np.mean(flat_samps[\"rv_phase\"], axis=-1)\n",
    "plt.plot(x, y, \"k\", lw=0.5, alpha=0.5)\n",
    "plt.plot(x + 1, y, \"k\", lw=0.5, alpha=0.5)\n",
    "\n",
    "plt.axvline(1, color=\"k\", lw=0.5)\n",
    "plt.xlim(0, 2)\n",
    "plt.xlabel(\"phase\")\n",
    "plt.ylabel(\"radial velocity [m/s]\")\n",
    "_ = plt.title(\"posterior inference\", fontsize=14)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93d6d8ce",
   "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": "e5defff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    txt, bib = xo.citations.get_citations_for_model()\n",
    "print(txt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "471a1034",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(bib.split(\"\\n\\n\")[0] + \"\\n\\n...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56b7c69c-31db-45c2-b12e-92ea681d5d70",
   "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
}