{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run notebook_setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have not already read it, you may want to start with the first tutorial: [Getting started with The Joker](1-Getting-started.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sampling over long-term velocity trend parameters\n",
    "\n",
    "In addition to the default linear parameters (see [Tutorial 1](1-Getting-started.ipynb), or the documentation for ``JokerSamples.default()``), *The Joker* allows adding linear parameters to include a long-term polynomial velocity trend. While *The Joker* formally only works for the two-body problem, if the system is a triple, and the outer body is in a significantly longer period orbit, you may be able to represent its perturbation to the primary as a smooth polynomial (in time) over the observation window you have. Below, we will show this with an example\n",
    "\n",
    "First, some imports we will need later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import astropy.table as at\n",
    "import astropy.units as u\n",
    "from astropy.visualization.units import quantity_support\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "\n",
    "import thejoker as tj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up a random number generator to ensure reproducibility\n",
    "rnd = np.random.default_rng(seed=42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by loading some pre-generated data meant to represent radial velocity observations of a single luminous source with two faint companions: one on a shorter period orbit, the other with a period ~ 10x longer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = tj.RVData.guess_from_table(at.QTable.read('data-triple.ecsv'))\n",
    "data = data[rnd.choice(len(data), size=16, replace=False)]  # downsample data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = data.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first pretend that we don't know it is a triple, and try generating orbit samples assuming a binary with no polynomial velocity trend. We will set up the default prior with some reasonable parameters that we have used in previous tutorials, and generate a big cache of prior samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior = tj.JokerPrior.default(\n",
    "    P_min=2*u.day, P_max=1e3*u.day,\n",
    "    sigma_K0=30*u.km/u.s,\n",
    "    sigma_v=100*u.km/u.s)\n",
    "prior_samples = prior.sample(size=250_000, random_state=rnd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can run *The Joker* to generate posterior samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joker = tj.TheJoker(prior, random_state=rnd)\n",
    "samples = joker.rejection_sample(data, prior_samples, \n",
    "                                 max_posterior_samples=128)\n",
    "samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = tj.plot_rv_curves(samples, data=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Only one sample was returned, and it's not a very good fit to the data (see the plot above). This is because the data were generated from a hierarchical triple system, but fit as a two-body system. Let's now try generating Keplerian orbit samples for the inner binary, while including a polynomial trend in velocity to capture the long-term trend from the outer companion. To do this, we specify the number of polynomial trend coefficients to sample over: 1 is constant, 2 is linear, 3 is quadratic, etc. We also have to specify the standard deviations of the Gaussian priors on the trend parameters, so the units have to be set accordingly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_trend = tj.JokerPrior.default(\n",
    "    P_min=2*u.day, P_max=1e3*u.day,\n",
    "    sigma_K0=30*u.km/u.s,\n",
    "    sigma_v=[100*u.km/u.s, \n",
    "             0.5*u.km/u.s/u.day, \n",
    "             1e-2*u.km/u.s/u.day**2],\n",
    "    poly_trend=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice the additional parameters `v1`, `v2` in the prior: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_trend"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now set up to generate prior samples and run *The Joker* including the new linear trend parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_samples_trend = prior_trend.sample(size=250_000,\n",
    "                                         random_state=rnd)\n",
    "joker_trend = tj.TheJoker(prior_trend, random_state=rnd)\n",
    "samples_trend = joker_trend.rejection_sample(data, prior_samples_trend, \n",
    "                                             max_posterior_samples=128)\n",
    "samples_trend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = tj.plot_rv_curves(samples_trend, data=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Those orbit samples look much better at matching the data! In a real-world situation with these data and results, given that the samples look like they all share a similar period, at this point I would start standard MCMC to continue generating samples. But, that is covered in [Tutorial 4](4-Continue-sampling-mcmc.ipynb), so for now, we will proceed with only the samples returned from The Joker.\n",
    "\n",
    "So how do the sample values compare to the truth? I cached the true orbital parameter values for the inner binary of this system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('true-orbit-triple.pkl', 'rb') as f:\n",
    "    truth = pickle.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Truth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "truth['P'], truth['e'], truth['K']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assuming binary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples['P'], samples['e'], samples['K']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assuming binary + quadratic velocity trend:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_trend.mean()['P'], samples_trend.mean()['e'], samples_trend.mean()['K']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The period and semi-amplitude values we infer from including the velocity trend is less biased. Hurrah!"
   ]
  }
 ],
 "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.7.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}