{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example of using PyImfit with Markov-Chain Monte Carlo code \"emcee\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a Jupyter notebook demonstrating how to use PyImfit with the MCMC code [emcee](https://github.com/dfm/emcee).\n",
    "\n",
    "If you are seeing this as part of the readthedocs.org HTML documentation, you can retrieve the original .ipynb file\n",
    "[here](https://github.com/perwin/pyimfit/blob/master/docs/pyimfit_emcee.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some initial setup for nice-looking plots:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/IPython/core/magics/pylab.py:160: UserWarning: pylab import has clobbered these variables: ['mean']\n",
      "`%matplotlib` prevents importing * from pylab and numpy\n",
      "  \"\\n`%matplotlib` prevents importing * from pylab and numpy\"\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "\n",
    "matplotlib.rcParams['figure.figsize'] = (8,6)\n",
    "matplotlib.rcParams['xtick.labelsize'] = 16\n",
    "matplotlib.rcParams['ytick.labelsize'] = 16\n",
    "matplotlib.rcParams['axes.labelsize'] = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create image-fitting model using PyImfit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the pymfit package; also load astropy.io.fits so we can read FITS files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyimfit\n",
    "from astropy.io import fits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load data image (in this case, a small cutout of an SDSS image showing a faint star):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "imageFile = \"./pyimfit_emcee_files/faintstar.fits\"\n",
    "image_faintstar = fits.getdata(imageFile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a ModelDescription instance based on an imfit configuration file (which specifies a single elliptical Gaussian model):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "configFile = \"./pyimfit_emcee_files/config_imfit_faintstar.dat\"\n",
    "model_desc = pyimfit.ModelDescription.load(configFile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternately, you can create the ModelDescription programmatically from within Python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a SimpleModelDescription instance (one function block); specify the x0,y0 center for the function block.\n",
    "model_desc = pyimfit.SimpleModelDescription()\n",
    "# define the X0,Y0 initial guess and limits\n",
    "model_desc.x0.setValue(5.0, [3.0,8.0])\n",
    "model_desc.y0.setValue(6.0, [3.0,8.0])\n",
    "# create a Gaussian image function for the star and set its parameters' initial guesses and limits\n",
    "star_function = pyimfit.make_imfit_function(\"Gaussian\")\n",
    "star_function.PA.setValue(155, [140,170])\n",
    "star_function.ell.setValue(0.1, [0,0.3])\n",
    "star_function.I_0.setValue(250, [220,320])\n",
    "star_function.sigma.setValue(1.0, [0.1,1.5])\n",
    "# now add the image function to the model\n",
    "model_desc.addFunction(star_function)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an Imfit instance containing the model, and add the image data and image-description info:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "imfit_fitter = pyimfit.Imfit(model_desc)\n",
    "imfit_fitter.loadData(image_faintstar, gain=4.72, read_noise=1.15, original_sky=124.94)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit the model to the data (using the default Levenberg-Marquardt solver) and extract the best-fitting parameter values\n",
    "(X0, Y0, PA, ellipticity, I_0, sigma):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best-fitting parameter values:\n",
      "5.64339, 6.18794, 155.354, 0.0950157, 268.92, 1.00772\n",
      "\n"
     ]
    }
   ],
   "source": [
    "results = imfit_fitter.doFit(getSummary=True)\n",
    "p_bestfit = results.params\n",
    "\n",
    "print(\"Best-fitting parameter values:\")\n",
    "for i in range(len(p_bestfit) - 1):\n",
    "    print(\"{0:g}\".format(p_bestfit[i]), end=\", \")\n",
    "print(\"{0:g}\\n\".format(p_bestfit[-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define log-probability functions for use with emcee"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Emcee requires a function which calculates and returns the log of the posterior probability (using the likelihood and the prior probability).\n",
    "\n",
    "We'll create a general function for this which takes as input the current model parameters, an Imfit instance which can compute the fit statistic for those parameters (= $-2 \\: \\times$ log likelihood) and a user-supplied function for computing the prior; this will return the sum of the log likelihood and the log of the prior:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lnPosterior_for_emcee( params, imfitter, lnPrior_func ):\n",
    "    \"\"\"\n",
    "    Returns log of posterior probability (which is calculated as the\n",
    "    product of the specified prior and the likelihood computed by the\n",
    "    Imfit object using the specified parameter values).\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    params : 1D numpy ndarray of float\n",
    "        vector of current parameter values\n",
    "    \n",
    "    imfitter : pyimfit.Imfit instance\n",
    "    \n",
    "    lnPrior_func : function or other callable\n",
    "        Should compute and return log of prior probability\n",
    "        signature = lnPrior_func(parameter_vector, Imfit_instance)\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    logPosterior : float\n",
    "    \"\"\"\n",
    "    lnPrior = lnPrior_func(params, imfitter)\n",
    "    if not np.isfinite(lnPrior):\n",
    "        return -np.inf\n",
    "    # note that Imfit.computeFitStatistic returns -2 log(likelihood)\n",
    "    lnLikelihood = -0.5 * imfitter.computeFitStatistic(params)\n",
    "    return lnPrior + lnLikelihood"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll create a prior-probability function.\n",
    "\n",
    "For simplicity, we'll use the case of constant priors within parameter limits, with the parameter limits obtained from a user-supplied Imfit instance. (But you can make the prior-probability function as complicated as you like.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lnPrior_limits( params, imfitter ):\n",
    "    \"\"\"\n",
    "    Defines prior-probability distributions as flat within specified limits.\n",
    "    If any parameter is outside the limits, we return -np.inf; otherwise, we\n",
    "    return ln(1.0) = 0 (not strictly speaking a correct probability, but it\n",
    "    works for this case).\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    params : 1D numpy ndarray of float\n",
    "    \n",
    "    imfitter : pyimfit.Imfit instance\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    logPrior : float\n",
    "    \"\"\"\n",
    "    parameterLimits = imfitter.getParameterLimits()\n",
    "    if None in parameterLimits:\n",
    "        raise ValueError(\"All parameters must have lower and upper limits.\")\n",
    "    nParams = len(params)\n",
    "    for i in range(nParams):\n",
    "        if params[i] < parameterLimits[i][0] or params[i] > parameterLimits[i][1]:\n",
    "            return -np.inf\n",
    "    return 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up and run Markov-Chain Monte Carlo using emcee"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import [emcee](https://emcee.readthedocs.io/en/latest/), and also [corner](https://corner.readthedocs.io/en/latest/) (so we can make a nice plot of the results):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import emcee\n",
    "import corner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Specify the number of dimensions (= number of parameters in the model) and a large number of walkers, then instantiate\n",
    "a standard emcee sampler, using our previously defined posterior function (the Imfit instance containing the data and model and the simple prior function are provided as extra arguments):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ndims, nwalkers = 6, 100\n",
    "\n",
    "sampler = emcee.EnsembleSampler(nwalkers, ndims, lnPosterior_for_emcee, args=(imfit_fitter, lnPrior_limits))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define some initial starting values -- 0.1% Gaussian perturbations around the previously determined best-fit parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_pos = [p_bestfit * (1 + 0.001*np.random.randn(ndims)) for i in range(nwalkers)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the sampler for 500 steps (reset it first, in case we're running this again, to ensure we start anew):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sampler.reset()\n",
    "final_state = sampler.run_mcmc(initial_pos, 500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot values from all the walkers versus step number to get an idea of where convergence might happend\n",
    "(here, we just plot the ellipticity and I_0 values):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 125040 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def PlotAllWalkers( sample_chain, parameterIndex, yAxisLabel ):\n",
    "    nWalkers = sample_chain.shape[0]\n",
    "    for i in range(nWalkers):\n",
    "        plot(sample_chain[i,:,parameterIndex], color='0.5')\n",
    "    xlabel('Step number')\n",
    "    ylabel(yAxisLabel)\n",
    "\n",
    "PlotAllWalkers(sampler.chain, 3, 'ellipticity')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 105408 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "PlotAllWalkers(sampler.chain, 4, 'I_0')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the \"converged\" subset of the chains as step numbers $\\ge 200$, and merge all the individual walkers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of samples in \"converged\" chain = 30000\n"
     ]
    }
   ],
   "source": [
    "converged_samples = sampler.chain[:, 200:, :].reshape((-1, ndims))\n",
    "print(\"Number of samples in \\\"converged\\\" chain = {0}\".format(len(converged_samples)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Corner plot of converged MCMC samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define some nice labels and parameter ranges for the corner plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "cornerLabels = [r\"$X_{0}$\", r\"$Y_{0}$\", \"PA\", \"ell\", r\"$I_{0}$\", r\"$\\sigma$\"] \n",
    "\n",
    "x0_range = (5.55, 5.73)\n",
    "y0_range = (6.09, 6.29)\n",
    "pa_range = (138,173)\n",
    "ell_range = (0, 0.2)\n",
    "i0_range = (240,300)\n",
    "sigma_range = (0.92, 1.1)\n",
    "ranges = [x0_range, y0_range, pa_range, ell_range, i0_range, sigma_range]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make a corner plot; the thin blue lines/points indicate best-fit values from above. [Note that we have to explicitly capture the Figure instance returned by corner.corner, otherwise we'll get a duplicate display of the plot]:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 500148 characters]\n",
      "text/plain": [
       "<Figure size 1000.8x1000.8 with 36 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = corner.corner(converged_samples, labels=cornerLabels, range=ranges, truths=p_bestfit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One thing to notice is that the PA values are running up against our (rather narrow) limits for that parameter, so a next step might be to re-run this with larger PA limits."
   ]
  }
 ],
 "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.4"
  },
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "99px",
    "width": "252px"
   },
   "navigate_menu": true,
   "number_sections": false,
   "sideBar": true,
   "threshold": "3",
   "toc_cell": false,
   "toc_section_display": "none",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
