{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot astrometric orbits for RV samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we'll assume (1) we know the sky position and distance of the barycenter of the observed two-body system at some epoch, and (2) we later observe some radial velocity data for the primary star. We'll use the radial velocity measurements to empirically estimate the amplitude of the astrometric orbit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import astropy.coordinates as coord\n",
    "from astropy.time import Time\n",
    "import astropy.units as u\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import schwimmbad\n",
    "%matplotlib inline\n",
    "\n",
    "from twobody import KeplerOrbit, Barycenter, P_m_to_a\n",
    "from thejoker.data import RVData\n",
    "from thejoker.sampler import JokerParams, TheJoker\n",
    "from thejoker.plot import plot_rv_curves\n",
    "\n",
    "rnd = np.random.RandomState(seed=123)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To generate simulated data, we'll simulate observations of the exoplanet host GJ 876 (with parameters taken from [exoplanets.org](http://exoplanets.org/detail/GJ_876_b)). We'll assume the star's parallax, radial velocity, and sky position are that of the barycenter of the system, and that these are known."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "origin = coord.SkyCoord(ra='22 53 16.73352', dec='-14 15 49.3186', \n",
    "                        unit=(u.hourangle, u.deg),\n",
    "                        distance=4.689*u.pc, \n",
    "                        radial_velocity=-1.519*u.km/u.s)\n",
    "barycenter = Barycenter(origin=origin, t0=Time('J2000'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "m1 = 0.320*u.Msun\n",
    "m2 = 1.95*u.Mjup\n",
    "\n",
    "truth = dict()\n",
    "truth['P'] = 61.1166 * u.day\n",
    "truth['a'] = 0.2081*u.au * m2 / m1 # convert a2 to a1\n",
    "truth['e'] = 0.0324\n",
    "truth['M0'] = 0 * u.radian\n",
    "truth['t0'] = Time(2450546.80, format='jd', scale='utc')\n",
    "truth['omega'] = 50.3 * u.degree\n",
    "truth['barycenter'] = barycenter\n",
    "\n",
    "orbit = KeplerOrbit(**truth, \n",
    "                    Omega=0*u.deg, i=90*u.deg) # these angle don't matter for radial velocity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll generate 31 randomly spaced \"observations\" from some arbitrary epoch (MJD=58140.8) over ~1 year, with uncertainties of 25 m/s:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_data = 31\n",
    "t = Time(58140.8 + rnd.uniform(0, 350, n_data), \n",
    "         format='mjd', scale='tcb')\n",
    "t.sort()\n",
    "rv = orbit.radial_velocity(t)\n",
    "\n",
    "err = np.full(len(t), 25) * u.m/u.s\n",
    "rv = rv + rnd.normal(0, err.value)*err.unit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data = RVData(t=t, rv=rv, stddev=err)\n",
    "\n",
    "ax = data.plot()\n",
    "ax.set_xlabel(\"BMJD\")\n",
    "ax.set_ylabel(\"RV [km/s]\")\n",
    "ax.set_title('simulated data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll run The Joker with a fairly constrained period range to infer the orbital parameters consistent with these data. To start, we have to specify the hyperparameters (period prior bounds):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "params = JokerParams(P_min=8*u.day, P_max=256*u.day, anomaly_tol=1E-11)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll generate ~2,000,000 prior samples to rejection sample with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "with schwimmbad.MultiPool() as pool:\n",
    "    joker = TheJoker(params, pool=pool)\n",
    "#     samples = joker.rejection_sample(data, n_prior_samples=2**21)\n",
    "    samples = joker.rejection_sample(data, n_prior_samples=2**17)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the $2^{21}$ prior samples, only 12 samples were accepted. Let's now visualize the orbits defined by those samples, plotted over the \"data\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8,5)) \n",
    "t_grid = np.linspace(data.t.mjd.min()-10, data.t.mjd.max()+10, 1024)\n",
    "fig = plot_rv_curves(samples, t_grid, rv_unit=u.km/u.s, data=data, ax=ax,\n",
    "                     plot_kwargs=dict(color='#74a9cf', zorder=-100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks like all of the samples returned are within the same period mode. We'll take the mean of the samples to compress the posterior and use the mean values as our assumed orbit fit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mean_sample = samples.mean()\n",
    "mean_sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'd like to see what astrometric orbits are generated by these orbital parameters. But, the radial velocity measurements alone don't constrain two more key orbital elements needed to compute the astrometric orbits: the inclination of the orbit, $i$ or in code `i`, and the longitude of the ascending node, $\\Omega$ or in code `Omega`. We also need to specify a coordinate frame to transform the orbit to. For the former, we'll generate samples in the angles assuming the probability of the orbital orientation is isotropic in these angles: uniform in $\\cos i$ and uniform in $\\Omega$. For the latter, we'll use a spherical coordinate system centered on the position of the barycenter at the epoch we specified, with longitude and latitude that increase in the same sense as the ICRS frame. If you scroll up, this sky location is `ra='22 53 16.73352'`, `dec='-14 15 49.3186'`, and the epoch is J2000. We can use astropy to transform to this frame using the `astropy.coordinates.SkyOffsetFrame` with a specified origin:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "barycenter_frame = coord.SkyOffsetFrame(origin=barycenter.origin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(5, 5))\n",
    "\n",
    "for j in range(32):\n",
    "    a = P_m_to_a(mean_sample['P'], m1+m2)\n",
    "    orbit_sample = mean_sample.get_orbit(0, a=a, \n",
    "        i=np.arccos(np.random.uniform(0, 1)) * u.rad, \n",
    "        Omega=np.random.uniform(0, 2*np.pi) * u.rad,\n",
    "        barycenter=barycenter)\n",
    "    \n",
    "    t_grid = Time('J2014') + np.linspace(0, 5, 1000) * u.year\n",
    "    ref_plane = orbit_sample.reference_plane(t_grid)\n",
    "    icrs = orbit_sample.icrs(t_grid)\n",
    "    sky_offset = icrs.transform_to(barycenter_frame)\n",
    "\n",
    "    ax.plot(sky_offset.lon.wrap_at(180*u.deg).milliarcsecond, \n",
    "            sky_offset.lat.milliarcsecond, \n",
    "            alpha=0.25, linestyle='-')\n",
    "    \n",
    "ax.set_xlim(-50, 50)\n",
    "ax.set_ylim(-50, 50)\n",
    "\n",
    "ax.set_xlabel(r'$\\Delta \\alpha$ [{0:latex_inline}]'.format(u.mas))\n",
    "ax.set_ylabel(r'$\\Delta \\delta$ [{0:latex_inline}]'.format(u.mas))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "For this system, the expected astrometric deviations from a straight-line (due to the barycentric proper motion) are around ~40 milliarcseconds. If we know the proper motion of the barycenter, or if we assume that the proper motion of the host star is the same as the proper motion of the barycenter, we can predict the full astrometric orbit of the system. To do this, we have to define a `Barycenter` object with proper motion components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "origin_pm = coord.SkyCoord(ra='22 53 16.73352', dec='-14 15 49.3186', \n",
    "                           unit=(u.hourangle, u.deg),\n",
    "                           distance=4.689*u.pc, \n",
    "                           pm_ra_cosdec=959.84*u.mas/u.yr,\n",
    "                           pm_dec=-675.33*u.mas/u.yr,\n",
    "                           radial_velocity=-1.519*u.km/u.s)\n",
    "barycenter_pm = Barycenter(origin=origin_pm, t0=Time('J2000'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(5, 5))\n",
    "\n",
    "for j in range(8):\n",
    "    a = P_m_to_a(mean_sample['P'], m1+m2)\n",
    "    orbit_sample = mean_sample.get_orbit(0, a=a, \n",
    "        i=np.arccos(np.random.uniform(0, 1)) * u.rad, \n",
    "        Omega=np.random.uniform(0, 2*np.pi) * u.rad,\n",
    "        barycenter=barycenter_pm)\n",
    "    \n",
    "    t_grid = Time('J2014') + np.linspace(0, 5, 1000) * u.year\n",
    "    ref_plane = orbit_sample.reference_plane(t_grid)\n",
    "    icrs = orbit_sample.icrs(t_grid)\n",
    "    sky_offset = icrs.transform_to(barycenter_frame)\n",
    "\n",
    "    ax.plot(sky_offset.lon.wrap_at(180*u.deg).milliarcsecond, \n",
    "            sky_offset.lat.milliarcsecond, \n",
    "            alpha=0.25, linestyle='-')\n",
    "\n",
    "ax.set_xlabel(r'$\\Delta \\alpha$ [{0:latex_inline}]'.format(u.mas))\n",
    "ax.set_ylabel(r'$\\Delta \\delta$ [{0:latex_inline}]'.format(u.mas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}