{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Integrate an orbit and compute uncertainties in Milky Way potential model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`gala` provides a simple mass model for the Milky Way based on recent measurements of the enclosed mass compiled from the literature. See the [Defining a Milky Way potential model](../potential/define-milky-way-model.ipynb) documentation for more information about how this model was defined.\n",
    "\n",
    "In this example, we'll use the position and velocity and uncertainties of the Milky Way satellite galaxy \"Draco\" to integrate orbits in a Milky Way mass model starting from samples from the error distribution over initial conditions defined by its observed kinematics. We'll then compute distributions of orbital properties like orbital period, pericenter, and eccentricity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some imports we'll need later:\n",
    "\n",
    "# Third-party\n",
    "import astropy.units as u\n",
    "import astropy.coordinates as coord\n",
    "from astropy.io import ascii\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Gala\n",
    "from gala.mpl_style import mpl_style\n",
    "plt.style.use(mpl_style)\n",
    "import gala.dynamics as gd\n",
    "import gala.integrate as gi\n",
    "import gala.potential as gp\n",
    "from gala.units import galactic\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will also set the default Astropy Galactocentric frame parameters to the values adopted in Astropy v4.0:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coord.galactocentric_frame_defaults.set('v4.0')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the Milky Way model, we'll use the built-in potential class in `gala` (see above for definition):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "potential = gp.MilkyWayPotential()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the sky position and distance of Draco, we'll use measurements from [Bonanos et al. 2004](https://arxiv.org/abs/astro-ph/0310477). For proper motion components, we'll use the recent HSTPROMO measurements ([Sohn et al. 2017](https://arxiv.org/abs/1707.02593)) and the line-of-sight velocity from [Walker et al. 2007](https://arxiv.org/abs/0708.0010)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "icrs = coord.SkyCoord(ra=coord.Angle('17h 20m 12.4s'), \n",
    "                      dec=coord.Angle('+57° 54′ 55″'),\n",
    "                      distance=76*u.kpc,\n",
    "                      pm_ra_cosdec=0.0569*u.mas/u.yr,\n",
    "                      pm_dec=-0.1673*u.mas/u.yr,\n",
    "                      radial_velocity=-291*u.km/u.s)\n",
    "\n",
    "icrs_err = coord.SkyCoord(ra=0*u.deg, dec=0*u.deg, distance=6*u.kpc,\n",
    "                          pm_ra_cosdec=0.009*u.mas/u.yr,\n",
    "                          pm_dec=0.009*u.mas/u.yr,\n",
    "                          radial_velocity=0.1*u.km/u.s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by transforming the measured values to a Galactocentric reference frame so we can integrate an orbit in our Milky Way model. We'll do this using the velocity transformation support in [`astropy.coordinates`](http://docs.astropy.org/en/stable/coordinates/velocities.html) (new to Astropy v2.0). We first have to define the position and motion of the sun relative to the Galactocentric frame, and create an [`astropy.coordinates.Galactocentric`](http://docs.astropy.org/en/stable/api/astropy.coordinates.Galactocentric.html#astropy.coordinates.Galactocentric) object with these parameters. We could specify these things explicitly, but instead we will use the default values that were recently updated in Astropy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gc_frame = coord.Galactocentric()\n",
    "gc_frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To transform the mean observed kinematics to this frame, we simply do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "galcen = icrs.transform_to(gc_frame)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! Now we have to turn the resulting `Galactocentric` object into orbital initial conditions, and integrate the orbit in our Milky Way model. We'll use a timestep of 0.5 Myr and integrate the orbit backwards for 10000 steps (5 Gyr):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = gd.PhaseSpacePosition(galcen.data)\n",
    "orbit = potential.integrate_orbit(w0, dt=-0.5*u.Myr, n_steps=10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize the orbit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = orbit.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the `orbit` object, we can easily compute quantities like the pericenter, apocenter, or eccentricity of the orbit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orbit.pericenter(), orbit.apocenter(), orbit.eccentricity()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use these functions to get the time of each pericenter or apocenter - let's plot the time of pericenter, and time of apocenter over the time series of the Galactocentric radius of the orbit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(orbit.t, orbit.spherical.distance, marker='None')\n",
    "\n",
    "per, per_times = orbit.pericenter(return_times=True, func=None)\n",
    "apo, apo_times = orbit.apocenter(return_times=True, func=None)\n",
    "\n",
    "for t in per_times:\n",
    "    plt.axvline(t.value, color='#67a9cf')\n",
    "    \n",
    "for t in apo_times:\n",
    "    plt.axvline(t.value, color='#ef8a62')\n",
    "    \n",
    "plt.xlabel('$t$ [{0}]'.format(orbit.t.unit.to_string('latex')))\n",
    "plt.ylabel('$r$ [{0}]'.format(orbit.x.unit.to_string('latex')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll sample from the error distribution over the distance, proper motions, and radial velocity, compute orbits, and plot distributions of mean pericenter and apocenter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_samples = 32\n",
    "dist = np.random.normal(icrs.distance.value, icrs_err.distance.value, \n",
    "                        n_samples) * icrs.distance.unit\n",
    "\n",
    "pm_ra_cosdec = np.random.normal(icrs.pm_ra_cosdec.value, \n",
    "                                icrs_err.pm_ra_cosdec.value, \n",
    "                                n_samples) * icrs.pm_ra_cosdec.unit\n",
    "\n",
    "pm_dec = np.random.normal(icrs.pm_dec.value, \n",
    "                          icrs_err.pm_dec.value, \n",
    "                          n_samples) * icrs.pm_dec.unit\n",
    "\n",
    "rv = np.random.normal(icrs.radial_velocity.value, icrs_err.radial_velocity.value, \n",
    "                      n_samples) * icrs.radial_velocity.unit\n",
    "\n",
    "ra = np.full(n_samples, icrs.ra.degree) * u.degree\n",
    "dec = np.full(n_samples, icrs.dec.degree) * u.degree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "icrs_samples = coord.SkyCoord(ra=ra, dec=dec, distance=dist,\n",
    "                              pm_ra_cosdec=pm_ra_cosdec,\n",
    "                              pm_dec=pm_dec, radial_velocity=rv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "icrs_samples.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "galcen_samples = icrs_samples.transform_to(gc_frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w0_samples = gd.PhaseSpacePosition(galcen_samples.data)\n",
    "orbit_samples = potential.integrate_orbit(w0_samples, dt=-1*u.Myr, n_steps=4000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orbit_samples.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pers = orbit_samples.pericenter(approximate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "apos = orbit_samples.apocenter(approximate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eccs = orbit_samples.eccentricity(approximate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(12, 4))\n",
    "\n",
    "axes[0].hist(pers.to_value(u.kpc), bins='auto')\n",
    "axes[0].set_xlabel('pericenter [kpc]')\n",
    "\n",
    "axes[1].hist(apos.to_value(u.kpc), bins='auto')\n",
    "axes[1].set_xlabel('apocenter [kpc]')\n",
    "\n",
    "axes[2].hist(eccs.value, bins='auto')\n",
    "axes[2].set_xlabel('eccentricity');"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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
}