{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining a Milky Way potential model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Third-party dependencies\n",
    "from astropy.io import ascii\n",
    "import astropy.units as u\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import leastsq\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": [
    "## Introduction\n",
    "\n",
    "`gala` provides a simple and easy way to access and integrate orbits in an\n",
    "approximate mass model for the Milky Way. The parameters of the mass model are\n",
    "determined by least-squares fitting the enclosed mass profile of a pre-defined\n",
    "potential form to recent measurements compiled from the literature. These\n",
    "measurements are provided with the documentation of `gala` and are shown below.\n",
    "The radius units are kpc, and mass units are solar masses:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "tbl = ascii.read('data/MW_mass_enclosed.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "tbl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now plot the above data and uncertainties:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(4,4))\n",
    "\n",
    "ax.errorbar(tbl['r'], tbl['Menc'], yerr=(tbl['Menc_err_neg'], tbl['Menc_err_pos']), \n",
    "            marker='o', markersize=2, color='k', alpha=1., ecolor='#aaaaaa', \n",
    "            capthick=0, linestyle='none', elinewidth=1.)\n",
    "\n",
    "ax.set_xlim(1E-3, 10**2.6)\n",
    "ax.set_ylim(7E6, 10**12.25)\n",
    "\n",
    "ax.set_xlabel('$r$ [kpc]')\n",
    "ax.set_ylabel('$M(<r)$ [M$_\\odot$]')\n",
    "\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now need to assume some form for the potential. For simplicity and within reason, we'll use a four component potential model consisting of a Hernquist ([1990](https://ui.adsabs.harvard.edu/#abs/1990ApJ...356..359H/abstract)) bulge and nucleus, a Miyamoto-Nagai ([1975](https://ui.adsabs.harvard.edu/#abs/1975PASJ...27..533M/abstract)) disk, and an NFW ([1997](https://ui.adsabs.harvard.edu/#abs/1997ApJ...490..493N/abstract)) halo. We'll fix the parameters of the disk and bulge to be consistent with previous work ([Bovy 2015](https://ui.adsabs.harvard.edu/#abs/2015ApJS..216...29B/abstract) - please cite that paper if you use this potential model) and vary the scale mass and scale radius of the nucleus and halo, respectively. We'll fit for these parameters in log-space, so we'll first define a function that returns a `gala.potential.CCompositePotential` object given these four parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_potential(log_M_h, log_r_s, log_M_n, log_a):\n",
    "    mw_potential = gp.CCompositePotential()\n",
    "    mw_potential['bulge'] = gp.HernquistPotential(m=5E9, c=1., units=galactic)\n",
    "    mw_potential['disk'] = gp.MiyamotoNagaiPotential(m=6.8E10*u.Msun, a=3*u.kpc, b=280*u.pc,\n",
    "                                                     units=galactic)\n",
    "    mw_potential['nucl'] = gp.HernquistPotential(m=np.exp(log_M_n), c=np.exp(log_a)*u.pc,\n",
    "                                                 units=galactic)\n",
    "    mw_potential['halo'] = gp.NFWPotential(m=np.exp(log_M_h), r_s=np.exp(log_r_s), units=galactic)\n",
    "\n",
    "    return mw_potential"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now need to specify an initial guess for the parameters - let's do that (by making them up), and then plot the initial guess potential over the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Initial guess for the parameters- units are:\n",
    "#     [Msun, kpc, Msun, pc]\n",
    "x0 = [np.log(6E11), np.log(20.), np.log(2E9), np.log(100.)] \n",
    "init_potential = get_potential(*x0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "xyz = np.zeros((3, 256))\n",
    "xyz[0] = np.logspace(-3, 3, 256)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(4,4))\n",
    "\n",
    "ax.errorbar(tbl['r'], tbl['Menc'], yerr=(tbl['Menc_err_neg'], tbl['Menc_err_pos']), \n",
    "            marker='o', markersize=2, color='k', alpha=1., ecolor='#aaaaaa', \n",
    "            capthick=0, linestyle='none', elinewidth=1.)\n",
    "\n",
    "fit_menc = init_potential.mass_enclosed(xyz*u.kpc)\n",
    "ax.loglog(xyz[0], fit_menc.value, marker='', color=\"#3182bd\",\n",
    "          linewidth=2, alpha=0.7)\n",
    "\n",
    "ax.set_xlim(1E-3, 10**2.6)\n",
    "ax.set_ylim(7E6, 10**12.25)\n",
    "\n",
    "ax.set_xlabel('$r$ [kpc]')\n",
    "ax.set_ylabel('$M(<r)$ [M$_\\odot$]')\n",
    "\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks pretty good already! But let's now use least-squares fitting to optimize our nucleus and halo parameters. We first need to define an error function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def err_func(p, r, Menc, Menc_err):\n",
    "    pot = get_potential(*p)\n",
    "    xyz = np.zeros((3,len(r)))\n",
    "    xyz[0] = r\n",
    "    model_menc = pot.mass_enclosed(xyz).to(u.Msun).value\n",
    "    return (model_menc - Menc) / Menc_err"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because the uncertainties are all approximately but not exactly symmetric, we'll take the maximum of the upper and lower uncertainty values and assume that the uncertainties in the mass measurements are Gaussian (a bad but simple assumption):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "err = np.max([tbl['Menc_err_pos'], tbl['Menc_err_neg']], axis=0)\n",
    "p_opt, ier = leastsq(err_func, x0=x0, args=(tbl['r'], tbl['Menc'], err))\n",
    "assert ier in range(1,4+1), \"least-squares fit failed!\"\n",
    "fit_potential = get_potential(*p_opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a best-fit potential! Let's plot the enclosed mass of the fit potential over the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "xyz = np.zeros((3, 256))\n",
    "xyz[0] = np.logspace(-3, 3, 256)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(4,4))\n",
    "\n",
    "ax.errorbar(tbl['r'], tbl['Menc'], yerr=(tbl['Menc_err_neg'], tbl['Menc_err_pos']), \n",
    "            marker='o', markersize=2, color='k', alpha=1., ecolor='#aaaaaa', \n",
    "            capthick=0, linestyle='none', elinewidth=1.)\n",
    "\n",
    "fit_menc = fit_potential.mass_enclosed(xyz*u.kpc)\n",
    "ax.loglog(xyz[0], fit_menc.value, marker='', color=\"#3182bd\",\n",
    "          linewidth=2, alpha=0.7)\n",
    "\n",
    "ax.set_xlim(1E-3, 10**2.6)\n",
    "ax.set_ylim(7E6, 10**12.25)\n",
    "\n",
    "ax.set_xlabel('$r$ [kpc]')\n",
    "ax.set_ylabel('$M(<r)$ [M$_\\odot$]')\n",
    "\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This potential is already implemented in `gala` in `gala.potential.special`, and we can import it with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from gala.potential import MilkyWayPotential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "potential = MilkyWayPotential()\n",
    "potential"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}