curve_fit_spectrograph.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SPECTRUM FITTING"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Available on GitHub at https://github.com/spyphyz/curve_fit\n",
    "\n",
    "The following codes approach the problem in a number of different ways. Scipy is used for continuum fitting, using a power law. Emission line fitting is performed using Gaussian (and trapezoid, if required) models in astropy. Both codes allow you to choose a range of wavelengths that you want to investigate (by changing 'lower_limit', and 'upper_limit'). You could no doubt flip this around, and perform either task using either python module, but I haven't had time to explore this further.\n",
    "\n",
    "I've had a go at using sherpa as well - Sherpa may be better at analysing uncertainties, although I've got a bit stuck with this. Powerlaw is yet another module I've played with at the end. \n",
    "\n",
    "An alternative way of doing these fits may be to invoke R code... embed in the Python code\n",
    "\n",
    "These codes import a csv file into a pandas dataframe; could utilise astropy.io to import FITS files if required - can also convert into a dataframe if required.\n",
    "\n",
    "Future possibilities: scikit_learn to automate curve fits, using list of emission line rest wavelengths to specify each gaussian mean, sherpa to analyse continuum, pipeline to automate upload of FITS files from a relational database (e.g. SDSS), blackbody_fit of continuum?\n",
    "\n",
    "Alternative: PySpecKit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scipy - Continuum Fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html for scipy.optimize.curve_fit details\n",
    "\n",
    "Here I present a power law fit - can also perform polynomial fits using curve_fit.\n",
    "\n",
    "Also see https://github.com/jeffalstott/powerlaw/blob/master/manuscript/Manuscript_Code.ipynb and https://arxiv.org/abs/1305.0215 for an alternative python powerlaw module, that could possibly automate the wavelength ranges fitted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "import pandas as pd\n",
    "from pandas import DataFrame\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import curve_fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input data into a pandas dataframe\n",
    "path = 'Dropbox/OpenUniversity/S382_Astrophysics/SDSS_Project/sdss_abell/extra_stuff/curve_fits'\n",
    "file = 'spectrum_z03_z18_auto_model.csv'\n",
    "df1 = pd.read_csv(os.path.join(path, file))\n",
    "df1.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose the range of wavelength you want to apply the continuum fit to\n",
    "lower_limit = 1320\n",
    "upper_limit = 6000\n",
    "df2 = df1[(df1['wavelength'] >= lower_limit) & (df1['wavelength'] <= upper_limit)]\n",
    "\n",
    "# plot this section of spectrum (spectral flux density is in arbitrary units)\n",
    "df2.plot(x='wavelength', y='flux', lw=0.4, color='green', label='data')\n",
    "plt.title('Continuum Section', fontsize=24, color='black')\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Spectral Flux Density / AU', fontsize=11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xData = df2['wavelength']\n",
    "yData = df2['flux']\n",
    "\n",
    "def f(x, N, a):\n",
    "    return N * x ** a\n",
    "\n",
    "# optimize\n",
    "popt, pcov = curve_fit(f, xData, yData)\n",
    "\n",
    "# the output here are the values for k and a, respectively, where y = kx^a\n",
    "popt[0], popt[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1SD of the parameters in popt, calculated from the covariance pcov array\n",
    "perr = np.sqrt(np.diag(pcov))\n",
    "perr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot data with power model\n",
    "plt.plot(xData, yData, 'g', label='data') # original spectrum data in green\n",
    "plt.plot(xData, popt[0]*xData**popt[1], 'r', label='power model') # power model data in red\n",
    "\n",
    "# plot standard deviation curves, based on perr\n",
    "plt.plot(xData, (popt[0]+perr[1])*xData**(popt[1]+perr[1]), '-', lw=0.2, color='purple', label='+1SD')\n",
    "plt.plot(xData, (popt[0]-perr[1])*xData**(popt[1]-perr[1]), '-', lw=0.2, color='purple', label='- 1SD')\n",
    "\n",
    "plt.title('Continuum Fit', fontsize=24, color='black')\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Spectral Flux Density / AU')\n",
    "plt.legend()\n",
    "plt.savefig(os.path.join(path, 'spectrum_continuum.png'), dpi = 1080)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Astropy - Emission Line Fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "See http://docs.astropy.org/en/stable/modeling/index.html for astropy model details\n",
    "This code works by inputting a csv file of data which has already had the baseline continuum model subtracted. You could also perform this subtraction easily within the pandas dataframe if so desired, and even input the power law model automatically from the Scipy code above if you were really clever. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from pandas import DataFrame\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from astropy.modeling import models, fitting\n",
    "from astropy.modeling.models import Gaussian1D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using the Levenberg-Marquardt algorithm and least squares statistic - LevMarLSQFitter()\n",
    "# Fitting Gaussian and Trapezoid models to a curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input data into a pandas dataframe\n",
    "path = 'Dropbox/OpenUniversity/S382_Astrophysics/SDSS_Project/sdss_abell/extra_stuff/curve_fits'\n",
    "file = 'spectrum_z03_z18_auto_model.csv'\n",
    "df1 = pd.read_csv(os.path.join(path, file))\n",
    "df1.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lower_limit = 4800\n",
    "upper_limit = 4950\n",
    "df2 = df1[(df1['wavelength'] >= lower_limit) & (df1['wavelength'] <= upper_limit)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df2.plot(x='wavelength', y='flux_minus_top_minus15', lw=0.6, color='blue', label='data') #lw = line width\n",
    "plt.title('Emission Peak Isolated', fontsize=14, color='black')\n",
    "plt.legend()\n",
    "# plt.grid(True)\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Continuum-Subtracted Spectral Flux Density / AU', fontsize=9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = df2['wavelength']\n",
    "y = df2['flux_minus_top_minus15']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fit the data using a box model\n",
    "# Insert best guesses for model parameters\n",
    "#t_init = models.Trapezoid1D(amplitude=58, x_0=2797., width=200, slope=0.5)\n",
    "#fit_t = fitting.LevMarLSQFitter()\n",
    "#t = fit_t(t_init, x, y)\n",
    "\n",
    "# Fit the data using a Gaussian\n",
    "# set mean as the rest emission wavelength for species (e.g. Mg II)\n",
    "g_init = models.Gaussian1D(amplitude=40, mean=4861, stddev=25)\n",
    "fit_g = fitting.LevMarLSQFitter()\n",
    "g = fit_g(g_init, x, y)\n",
    "\n",
    "# this prints the model output parameters after fitting\n",
    "print(g)\n",
    "#Prints FWHM (units are Angstroms)...\n",
    "g.fwhm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the data with the best-fit model\n",
    "plt.figure(figsize=(8,5))\n",
    "plt.plot(x, y, '.', label='data points')\n",
    "plt.plot(x, g(x), label='gaussian fit', color='red')\n",
    "\n",
    "plt.axvline(x=g.mean.value, label='mean', color='yellow', lw=0.5)\n",
    "plt.vlines(x=g.mean.value+0.5*g.fwhm, ymin=0, ymax=g(g.mean.value+0.5*g.fwhm), label='fwhm', color='purple', lw=0.5)\n",
    "plt.vlines(x=g.mean.value-0.5*g.fwhm, ymin=0, ymax=g(g.mean.value-0.5*g.fwhm), color='purple', lw=0.5)\n",
    "\n",
    "#plt.plot(x, t(x), label='Trapezoid')\n",
    "\n",
    "plt.title('Emission Peak Gaussian-Fitted (Astropy)', fontsize=14, color='black')\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Continuum-Subtracted Spectral Flux Density / AU')\n",
    "plt.ylim(0,)\n",
    "plt.legend()\n",
    "plt.savefig(os.path.join(path, 'species_curve_fit.png'), dpi = 1080)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this prints the model output parameters after fitting\n",
    "# print(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sherpa Fits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install sherpa and saba as per instructions: https://github.com/astropy/saba/blob/d729b34262f70d0679cdcb112ab2832b73791098/docs/index.rst#id17\n",
    "\n",
    "Further install details: https://github.com/sherpa/sherpa/\n",
    "\n",
    "Documentation: http://cxc.harvard.edu/sherpa4.10/, http://conference.scipy.org/proceedings/SciPy2009/paper_8/full_text.pdf, http://conference.scipy.org/proceedings/scipy2011/pdfs/brefsdal.pdf, http://adsabs.harvard.edu/abs/2007ASPC..376..543D, http://cxc.cfa.harvard.edu/contrib/sherpa/, and https://saba.readthedocs.io/en/latest/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from pandas import DataFrame\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from astropy.modeling.fitting import SherpaFitter\n",
    "from astropy.modeling.models import Gaussian1D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input data into a pandas dataframe\n",
    "path = 'Dropbox/OpenUniversity/S382_Astrophysics/SDSS_Project/sdss_abell/extra_stuff/curve_fits'\n",
    "file = 'spectrum_z03_z18_auto_model.csv'\n",
    "df1 = pd.read_csv(os.path.join(path, file))\n",
    "df1.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lower_limit = 4800\n",
    "upper_limit = 4950\n",
    "df2 = df1[(df1['wavelength'] >= lower_limit) & (df1['wavelength'] <= upper_limit)]\n",
    "\n",
    "df2.plot(x='wavelength', y='flux_minus_top_minus15', lw=0.6, color='blue', label='data') #lw = line width\n",
    "plt.title('Emission Peak Isolated', fontsize=14, color='black')\n",
    "plt.legend()\n",
    "# plt.grid(True)\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Continuum-Subtracted Spectral Flux Density / AU', fontsize=9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sfit = SherpaFitter(optimizer='levmar', estmethod='confidence') # defaults to statistic='leastsq', I think!\n",
    "sfit = SherpaFitter(statistic='chi2', optimizer='levmar', estmethod='confidence') # use if have error data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "true = Gaussian1D(amplitude=30, mean=4860, stddev=5, name='Mg II')\n",
    "err = 0.8\n",
    "step = 0.2\n",
    "x = df2['wavelength']\n",
    "y = df2['flux_minus_top_minus15']\n",
    "y_errs = df2['error']\n",
    "binsize = step * np.ones(x.shape) # binsize is the width of the bin!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_model = true.copy() # offset fit model from true\n",
    "fit_model.amplitude = 30\n",
    "fit_model.mean = 4860\n",
    "fit_model.stddev = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fitted_model = sfit(fit_model, x, y)\n",
    "fitted_model = sfit(fit_model, x, y, xbinsize=binsize, err=y_errs) # if error data available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitted_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitted_model.fwhm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitted_model.mean.value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the data with the best-fit model\n",
    "# plt.figure(figsize=(8,5))\n",
    "plt.plot(x, y, '.', label='data points')\n",
    "#plt.plot(x, t(x), label='Trapezoid')\n",
    "plt.plot(x, fitted_model(x), label='gaussian fit', color='red')\n",
    "plt.axvline(x=fitted_model.mean.value, label='mean', color='yellow', lw=0.5)\n",
    "plt.vlines(x=fitted_model.mean.value+0.5*fitted_model.fwhm, ymin=0, ymax=fitted_model(fitted_model.mean.value+0.5*fitted_model.fwhm), label='fwhm', color='purple', lw=0.5)\n",
    "plt.vlines(x=fitted_model.mean.value-0.5*fitted_model.fwhm, ymin=0, ymax=fitted_model(fitted_model.mean.value-0.5*fitted_model.fwhm), color='purple', lw=0.5)\n",
    "\n",
    "plt.title('Emission Peak Gaussian-Fitted (Sherpa)', fontsize=14, color='black')\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Continuum-Subtracted Spectral Flux Density / AU')\n",
    "plt.ylim(0,)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sfit.fit_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use this if statistic='chi2'\n",
    "param_errors = sfit.est_errors(sigma=3)  # Note that sigma can be an input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_model = fitted_model.copy()\n",
    "max_model = fitted_model.copy()\n",
    "\n",
    "for pname, pval, pmin, pmax in zip(*param_errors):\n",
    "        print(pname, pval, pmin, pmax)\n",
    "        getattr(min_model, pname).value = pval + pmin\n",
    "        getattr(max_model, pname).value = pval + pmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitted_model = sfit(fit_model, x, y, xbinsize=binwidth, err=yerrs)\n",
    "plt.plot(x, true(x), label=\"True\")\n",
    "plt.errorbar(x, y, xerr=binwidth, yerr=yerrs, ls=\"\", label=\"Data\")\n",
    "plt.plot(x, fit_model(x), label=\"Starting fit model\")\n",
    "plt.plot(x, fitted_model(x), label=\"Fitted model\")\n",
    "plt.legend(loc=2, frameon=False)\n",
    "plt.xlim((-3, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Powerlaw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "http://pythonhosted.org/powerlaw/; https://pypi.org/project/powerlaw/; http://arxiv.org/abs/1305.0215; http://arxiv.org/abs/0706.1062; http://www.plosone.org/article/info%3Adoi%2F10.1371%2Fjournal.pone.0019779; http://www.plosone.org/article/info%3Adoi%2F10.1371%2Fjournal.pone.0085777"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I don't yet understand how this works! This example code http://nbviewer.jupyter.org/github/jeffalstott/powerlaw/blob/master/manuscript/Manuscript_Code.ipynb will probably be the most useful to understand how to use powerlaw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import mpmath\n",
    "import scipy\n",
    "from pandas import DataFrame\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import powerlaw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input data into a pandas dataframe\n",
    "path = 'Dropbox/OpenUniversity/S382_Astrophysics/SDSS_Project/sdss_abell/extra_stuff/curve_fits'\n",
    "file = 'spectrum_z03_z18_auto_model.csv'\n",
    "df1 = pd.read_csv(os.path.join(path, file))\n",
    "df1.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose the range of wavelength you want to apply the continuum fit to\n",
    "lower_limit = 1320\n",
    "upper_limit = 2000\n",
    "df2 = df1[(df1['wavelength'] >= lower_limit) & (df1['wavelength'] <= upper_limit)]\n",
    "\n",
    "# plot this section of spectrum (spectral flux density is in arbitrary units)\n",
    "df2.plot(x='wavelength', y='flux', lw=0.4, color='green', label='data')\n",
    "plt.title('Continuum Section', fontsize=24, color='black')\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Spectral Flux Density / AU', fontsize=11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = df2['flux'] # data can be list or numpy array\n",
    "results = powerlaw.Fit(data)\n",
    "print(results.power_law.alpha)\n",
    "print(results.power_law.xmin)\n",
    "R, p = results.distribution_compare('power_law', 'lognormal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "R, p = results.distribution_compare('power_law', 'lognormal_positive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = df2['wavelength']\n",
    "y = df2['flux']\n",
    "df2['power_law']=results.power_law.xmin*df2['wavelength']**-results.power_law.alpha #don't understand this yet!!!\n",
    "\n",
    "ax1 = df2.plot(x='wavelength', y='flux', lw=0.6, color='blue', label='data') #lw = line width\n",
    "ax2 = df2.plot(x='wavelength', y='power_law', ax = ax1) # need to sort this line out to plot the power law\n",
    "\n",
    "plt.title('Emission Peak Isolated', fontsize=14, color='black')\n",
    "plt.legend()\n",
    "# plt.grid(True)\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Continuum-Subtracted Spectral Flux Density / AU', fontsize=9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.power_law.plot_pdf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.lognormal.plot_pdf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = results.power_law.pdf()\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pyspeckit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PySpecKit at http://pyspeckit.readthedocs.io/en/latest/index.html; https://github.com/pyspeckit/pyspeckit; https://keflavich.github.io/pyspeckit-paper/main_compressed.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy\n",
    "from pandas import DataFrame\n",
    "import matplotlib.pyplot as plt # don't necessarily need this as sp.plotter works\n",
    "import os\n",
    "\n",
    "from astropy import units as u\n",
    "import pyspeckit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input data into a pandas dataframe\n",
    "path = 'Dropbox/OpenUniversity/S382_Astrophysics/SDSS_Project/sdss_abell/extra_stuff/curve_fits'\n",
    "file = 'spectrum_z03_z18_auto_model.csv'\n",
    "df1 = pd.read_csv(os.path.join(path, file))\n",
    "df2.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose the range of wavelength you want to apply the continuum fit to\n",
    "lower_limit = 1320\n",
    "upper_limit = 2000\n",
    "df2 = df1[(df1['wavelength'] >= lower_limit) & (df1['wavelength'] <= upper_limit)]\n",
    "\n",
    "# plot this section of spectrum (spectral flux density is in arbitrary units)\n",
    "df2.plot(x='wavelength', y='flux', lw=0.4, color='green', label='data')\n",
    "plt.title('Continuum Section', fontsize=24, color='black')\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Spectral Flux Density / AU', fontsize=11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#xaxis = df2['wavelength'] #*u.km/u.s \n",
    "\n",
    "xaxis = np.linspace(-25, 25)*u.km/u.s\n",
    "sigma = 3.0*u.km/u.s\n",
    "data = 5*np.exp(-xaxis**2 / (2*sigma**2))*u.Jy\n",
    "error = np.ones_like(data) * 0.2\n",
    "sp = pyspeckit.Spectrum(xarr=xaxis, data=data, error=error)\n",
    "sp.plotter()\n",
    "#sp.plotter(errstyle='fill') # fill doesn't work. The errors should be shown with the ‘fill’ style and represent symmetric 1 − SD Gaussian errors\n",
    "#sp.plotter.savefig(\"example_fig_1.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the follwoing code from the pyspeckit website is for fitting an sdss seyfert spectrum\n",
    "\n",
    "import pyspeckit\n",
    "import numpy as np\n",
    "\n",
    "# Rest wavelengths of the lines we are fitting - use as initial guesses\n",
    "NIIa = 6549.86\n",
    "NIIb = 6585.27\n",
    "Halpha = 6564.614\n",
    "SIIa = 6718.29\n",
    "SIIb = 6732.68"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize spectrum object and plot region surrounding Halpha-[NII] complex\n",
    "# spec = pyspeckit.Spectrum('sample_sdss.txt', errorcol=2)\n",
    "xaxis = df2['wavelength']\n",
    "data = df2['flux']\n",
    "error = df2['error']\n",
    "spec = pyspeckit.Spectrum(xarr = xaxis, data = data, error = error)\n",
    "spec.plotter(xmin = 1320, xmax = 2000, ymin = 0, ymax = 150)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize spectrum object and plot region surrounding Halpha-[NII] complex\n",
    "spec = pyspeckit.Spectrum('sample_sdss.txt', errorcol=2)\n",
    "spec.plotter(xmin = 6450, xmax = 6775, ymin = 0, ymax = 150)\n",
    "\n",
    "# We fit the [NII] and [SII] doublets, and allow two components for Halpha.\n",
    "# The widths of all narrow lines are tied to the widths of [SII].\n",
    "guesses = [50, NIIa, 5, 100, Halpha, 5, 50, Halpha, 50, 50, NIIb, 5, 20, SIIa, 5, 20, SIIb, 5]\n",
    "tied = ['', '', 'p[17]', '', '', 'p[17]', '', 'p[4]', '', '3 * p[0]', '', 'p[17]', '', '', 'p[17]', '', '', '']\n",
    "\n",
    "# Actually do the fit.\n",
    "spec.specfit(guesses = guesses, tied = tied, annotate = False)\n",
    "spec.plotter.refresh()\n",
    "\n",
    "# Let's use the measurements class to derive information about the emission\n",
    "# lines.  The galaxy's redshift and the flux normalization of the spectrum\n",
    "# must be supplied to convert measured fluxes to line luminosities.  If the\n",
    "# spectrum we loaded in FITS format, 'BUNITS' would be read and we would not\n",
    "# need to supply 'fluxnorm'.\n",
    "spec.measure(z = 0.05, fluxnorm = 1e-17)\n",
    "\n",
    "# Now overplot positions of lines and annotate\n",
    "\n",
    "y = spec.plotter.ymax * 0.85    # Location of annotations in y\n",
    "\n",
    "for i, line in enumerate(spec.measurements.lines.keys()):\n",
    "\n",
    "    # If this line is not in our database of lines, don't try to annotate it\n",
    "    if line not in spec.speclines.optical.lines.keys(): continue\n",
    "\n",
    "    x = spec.measurements.lines[line]['modelpars'][1]   # Location of the emission line\n",
    "    spec.plotter.axis.plot([x]*2, [spec.plotter.ymin, spec.plotter.ymax], ls = '--', color = 'k')   # Draw dashed line to mark its position\n",
    "    spec.plotter.axis.annotate(spec.speclines.optical.lines[line][-1], (x, y), rotation = 90, ha = 'right', va = 'center')  # Label it\n",
    "\n",
    "# Make some nice axis labels\n",
    "spec.plotter.axis.set_xlabel(r'Wavelength / $\\AA$')\n",
    "spec.plotter.axis.set_ylabel(r'Flux $(10^{-17} \\mathrm{erg/s/cm^2/\\AA})$')\n",
    "spec.plotter.refresh()\n",
    "\n",
    "# Print out spectral line information\n",
    "print \"Line   Flux (erg/s/cm^2)     Amplitude (erg/s/cm^2)    FWHM (Angstrom)   Luminosity (erg/s)\"\n",
    "for line in spec.measurements.lines.keys():\n",
    "    print line, spec.measurements.lines[line]['flux'], spec.measurements.lines[line]['amp'], spec.measurements.lines[line]['fwhm'], \\\n",
    "        spec.measurements.lines[line]['lum']\n",
    "\n",
    "# Had we not supplied the objects redshift (or distance), the line\n",
    "# luminosities would not have been measured, but integrated fluxes would\n",
    "# still be derived.  Also, the measurements class separates the broad and\n",
    "# narrow H-alpha components, and identifies which lines are which. How nice!\n",
    "\n",
    "spec.specfit.plot_fit()\n",
    "\n",
    "# Save the figure\n",
    "# spec.plotter.figure.savefig(\"sdss_fit_example.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df2['wavelength']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ASTROPY POWERLAW FIT OF CONTINUUM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Licensed under a 3-clause BSD style license - see LICENSE.rst\n",
    "# starting point - \n",
    "# http://nullege.com/codes/show/src%40g%40a%40gammapy-HEAD%40docs%40spectrum%40fit_crab_spectrum.py/5/astropy.modeling.fitting.NonLinearLSQFitter/python\n",
    "\n",
    "%matplotlib inline\n",
    "import os\n",
    "import pandas as pd\n",
    "from pandas import DataFrame\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from astropy.modeling.models import PowerLaw1D\n",
    "from astropy.modeling import models, fitting\n",
    "from astropy.modeling.fitting import LevMarLSQFitter # a non-linear fitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input data into a pandas dataframe\n",
    "path = 'Dropbox/OpenUniversity/S382_Astrophysics/SDSS_Project/sdss_abell/extra_stuff/curve_fits'\n",
    "file = 'spectrum_z03_z18_auto_model.csv'\n",
    "df1 = pd.read_csv(os.path.join(path, file))\n",
    "df1.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose the range of wavelength you want to apply the continuum fit to\n",
    "lower_limit = 1320\n",
    "upper_limit = 6000\n",
    "df2 = df1[(df1['wavelength'] >= lower_limit) & (df1['wavelength'] <= upper_limit)]\n",
    "\n",
    "# plot this section of spectrum (spectral flux density is in arbitrary units)\n",
    "df2.plot(x='wavelength', y='flux', lw=0.4, color='green', label='data')\n",
    "plt.title('Continuum Section', fontsize=24, color='black')\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Spectral Flux Density / AU', fontsize=11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data = df2\n",
    "wavelength, flux, flux_err  = df2['wavelength'], df2['flux'], df2['error']\n",
    "\n",
    "# equation looking to fit is y = A(x/x_0)^alpha - in this case, x = wavelength, y will be flux\n",
    "\n",
    "model = PowerLaw1D(3e7, 1, 1.5) # best starting guesses for A, x_0 and alpha\n",
    "model.x_0.fixed = False\n",
    "fitter = LevMarLSQFitter()\n",
    "  \n",
    "model = fitter(model, wavelength, flux, weights=(1. / flux_err))\n",
    "print(model)\n",
    "  \n",
    "plt.errorbar(wavelength, wavelength ** 2 * flux, wavelength ** 2 * flux_err, fmt='.', label='raw data')\n",
    "plt.errorbar(wavelength, wavelength ** 2 * model(wavelength), fmt='r-', label='powerlaw fit')\n",
    "plt.loglog()\n",
    "plt.title('Continuum Section', fontsize=24, color='black')\n",
    "plt.xlabel(r'Emitted Wavelength / $\\AA$',  fontsize=13)\n",
    "plt.ylabel('Spectral Flux Density / AU', fontsize=11)\n",
    "plt.legend()\n",
    "# plt.savefig('astropy_model_curvefit_continuum.png')"
   ]
  },
  {
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}