{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting a Moffat function to a series of FITS images from a focus sweep.\n",
    "\n",
    "Import a number of libraries, numpy, fits, pyplot, Axes3D, ndimage, models, fitting, LogNorm, warnings.\n",
    "\n",
    "Creat an empty numpy array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from astropy.io import fits\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from scipy import ndimage\n",
    "from astropy.modeling import models, fitting\n",
    "from matplotlib.colors import LogNorm\n",
    "import warnings\n",
    "\n",
    "a = np.array([])\n",
    "print(type(a))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open one of our focus sweep FITS files, print its shape, and display using a LogNorm gray color map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hdul = fits.open(\"./gwaves-test2/62_300.fits\")\n",
    "data = hdul[0].data\n",
    "print(data.shape)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data, cmap='gray', norm=LogNorm())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, I'm trying to find the center of the star using center_of_mass.\n",
    "\n",
    "Then I cut out a 30 x 30 piece of the image data centered on this center_of_mass.\n",
    "\n",
    "Plot the fit data using gray scale LogNorm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "position = ndimage.measurements.center_of_mass(data[25:125,25:125])\n",
    "posx = int(round(position[0])) + 25\n",
    "posy = int(round(position[1])) + 25\n",
    "position = ndimage.measurements.center_of_mass(data[posx-50:posx+50,\n",
    "                                                    posy-50:posy+50])\n",
    "posx = int(round(position[0])) + 25\n",
    "posy = int(round(position[1])) + 25\n",
    "myfitdata = data[posx-15:posx+15,\n",
    "                 posy-15:posy+15]\n",
    "#print(myfitdata)\n",
    "plt.figure()\n",
    "plt.imshow(myfitdata, cmap='gray', norm=LogNorm())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, I want to remove the background as best I can so the fitting function will work well.\n",
    "\n",
    "I determine the indexes into my data array where the data is in the lower 98% of values.\n",
    "\n",
    "I average these pixels to get a background value and subtract that value from the data.\n",
    "\n",
    "Now the background should be near zero and I display the resulting image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thresh = np.percentile(data, 98.0)\n",
    "idx = data[:,:] < thresh\n",
    "bval = np.average(data[idx])\n",
    "plt.yscale('log', nonposy='clip')\n",
    "print(int(bval))\n",
    "myfitdata = myfitdata - bval\n",
    "histogram = plt.hist(myfitdata.flatten(), 100)\n",
    "plt.figure()\n",
    "plt.imshow(myfitdata, cmap=\"gray\", vmin = 0, vmax = 7000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, plot the image larger, using the default color map, and use bicubic interpolation to smooth it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 1, figsize=(12,12))\n",
    "axs.matshow(myfitdata, interpolation=\"bicubic\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Moffat Function\n",
    "$$f(x, y) = A \\left(1 + \\frac{\\left(x - x_{0}\\right)^{2} +\n",
    "\\left(y - y_{0}\\right)^{2}}{\\gamma^{2}}\\right)^{- \\alpha}$$\n",
    "alpha is power index of Moffat model\n",
    "\n",
    "gamma is core width of Moffat model\n",
    "\n",
    "A is amplitude\n",
    "\n",
    "For the astropy modeling fitter LevMarLSQFitter\n",
    "Arguments are:\n",
    "\n",
    "    model\n",
    "    \n",
    "    x coords\n",
    "    \n",
    "    y coords\n",
    "    \n",
    "    data\n",
    "    \n",
    "We use the Levenberg Marquardt Least Squares fitter to fit the Moffet function to the data.\n",
    "\n",
    "Along the way, we use numpy mgrid to make two arrays of coordinates needed by the fitter.\n",
    "\n",
    "I print one of these out to show what they look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xmod = models\n",
    "p_init = xmod.Moffat2D(x_0=15, y_0=15, alpha=0.0)\n",
    "        \n",
    "        \n",
    "xfit = fitting\n",
    "fit_p = xfit.LevMarLSQFitter()\n",
    "\n",
    "Xin, Yin = np.mgrid[0:30, 0:30]\n",
    "print(Xin)\n",
    "#print(Yin)\n",
    "\n",
    "with warnings.catch_warnings():\n",
    "    # Ignore model linearity warning from the fitter\n",
    "    warnings.simplefilter('ignore')\n",
    "    p = fit_p(p_init, Xin, Yin, myfitdata)\n",
    "        \n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now plot the image again and overplot a contour plot of the calculated Moffat function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 1, figsize=(12,12))\n",
    "plt.figure()\n",
    "#plt.imshow(myfitdata, cmap='gray', norm=LogNorm())\n",
    "axs.matshow(myfitdata, interpolation=\"bicubic\")\n",
    "\n",
    "\n",
    "axs.contour(p(Xin, Yin), colors='white', alpha=.9, linewidths=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print a few of the parameters of the fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (p.x_0)\n",
    "print (p.y_0)\n",
    "print (p.fwhm)\n",
    "print()\n",
    "print(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, use a 3D projection to make a surface plot of the fitted Moffat function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12,12))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "x = range(30)\n",
    "y = range(30)\n",
    "X,Y = np.meshgrid(x,y)\n",
    "Z = np.array(p(Xin, Yin))\n",
    "\n",
    "ax.plot_surface(X,Y,Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, putting it all together, fit Moffat functions to stars in all 9 images, plot the result.\n",
    "\n",
    "Along the way, we are storing the FWHM of the fitted Moffat functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = open(\"./gwaves-test2/focuslist\")\n",
    "line = f.readline()\n",
    "#print(os.path.basename(line))\n",
    "figsrow = 0\n",
    "figscol = 0\n",
    "fig, axs = plt.subplots(3, 3, figsize=(15,15))\n",
    "while line:\n",
    "    hdul = fits.open(line.strip())\n",
    "    data = hdul[0].data\n",
    "    position = ndimage.measurements.center_of_mass(data[25:125,25:125])\n",
    "    posx = int(round(position[0])) + 25\n",
    "    posy = int(round(position[1])) + 25\n",
    "    position = ndimage.measurements.center_of_mass(data[posx-50:posx+50,\n",
    "                                                    posy-50:posy+50])\n",
    "    posx = int(round(position[0])) + 25\n",
    "    posy = int(round(position[1])) + 25\n",
    "    myfitdata = data[posx-15:posx+15,\n",
    "                 posy-15:posy+15]\n",
    "\n",
    "    thresh = np.percentile(data, 98.0)\n",
    "    idx = data[:,:] < thresh\n",
    "    bval = np.average(data[idx])\n",
    "    #print(bval)\n",
    "    myfitdata = myfitdata - bval\n",
    "    \n",
    "    axs[figsrow, figscol].matshow(myfitdata, interpolation=\"bicubic\")\n",
    "        \n",
    "    xmod = models\n",
    "    p_init = xmod.Moffat2D(x_0=15, y_0=15, alpha=0.0)\n",
    "        \n",
    "        \n",
    "    xfit = fitting\n",
    "    fit_p = xfit.LevMarLSQFitter()\n",
    "\n",
    "    Xin, Yin = np.mgrid[0:30, 0:30]\n",
    "        \n",
    "    with warnings.catch_warnings():\n",
    "        # Ignore model linearity warning from the fitter\n",
    "        warnings.simplefilter('ignore')\n",
    "        p = fit_p(p_init, Xin, Yin, myfitdata)\n",
    "            \n",
    "    #print (p)\n",
    "    axs[figsrow, figscol].contour(p(Xin, Yin), colors='white', alpha=.5, linewidths=1)\n",
    "    \n",
    "    figscol = figscol + 1\n",
    "    if(figscol > 2):\n",
    "        figscol = 0\n",
    "        figsrow = figsrow + 1\n",
    "        \n",
    "    a = np.append(a, p.fwhm)\n",
    "    print(round(bval), round(100.*p.fwhm)/100., os.path.basename(line))\n",
    "    line = f.readline()\n",
    "    \n",
    "            \n",
    "f.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, plot the 9 FWHM values and fit a parabola to the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(a, linestyle='None', marker='o')\n",
    "x = np.arange(len(a))\n",
    "z = np.polyfit(x, a, 2)\n",
    "x = np.arange(len(a), step=.01)\n",
    "plt.plot(x, np.polyval(z,x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}