05. Create and Fit Radial Profiles.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "from astropy.table import Table\n",
    "import numpy.ma as ma\n",
    "import numpy as np\n",
    "import emcee\n",
    "import corner\n",
    "import os\n",
    "import sys\n",
    "from astropy.io import fits\n",
    "import astropy.units as u\n",
    "from regions import PixCoord, EllipsePixelRegion, CircleAnnulusPixelRegion, CirclePixelRegion\n",
    "from math import sqrt\n",
    "\n",
    "# stuff that I wrote\n",
    "from utilities import parallel_process, compound_regions\n",
    "from model import beta_model, integ_beta_model, prob\n",
    "sys.path.append(f'{os.environ[\"HOME\"]}/Projects/planckClusters/catalogs')\n",
    "from load_catalogs import load_PSZcatalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility functions\n",
    "def clean_files(outpath):\n",
    "    print('Cleaning files...')\n",
    "    \n",
    "    os.system(f'find {outpath}/PSZ* -name \"*corner_*.png\" -type f -delete')\n",
    "    os.system(f'find {outpath}/PSZ* -name \"*_mcmcfits.txt\" -type f -delete')\n",
    "    os.system(f'find {outpath}/PSZ* -name \"*_radproffit_*.png\" -type f -delete')\n",
    "    os.system(f'find {outpath}/PSZ* -name \"*.radprof\" -type f -delete')\n",
    "    \n",
    "def cp_results(outpath):\n",
    "    print('Copying files...')\n",
    "    \n",
    "    os.system(f'find {outpath}/PSZ* -name \"*corner_*.png\" -exec cp -t {outpath}/pngs/fits/ ' '{} \\;')\n",
    "    os.system(f'find {outpath}/PSZ* -name \"*radproffit_*.png\" -exec cp -t {outpath}/pngs/fits/ ' '{} \\;')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_radprof(name, outpath):\n",
    "    pixscale = 6.548089E-04 * 3600\n",
    "\n",
    "    # model parameters\n",
    "    params = {}\n",
    "    params['n'] = 75\n",
    "    params['r0'] = 0\n",
    "    params['dr'] = 4\n",
    "\n",
    "    # check for files\n",
    "    # events image\n",
    "    if not os.path.isfile(f'{outpath}/{name}/{name}_img_50-200.fits'):\n",
    "        return\n",
    "    else:\n",
    "        evnts = f'{outpath}/{name}/{name}_img_50-200.fits'\n",
    "    # expo map\n",
    "    if not os.path.isfile(f'{outpath}/{name}/{name}_exp.fits'):\n",
    "        return\n",
    "    else:\n",
    "        expmap = f'{outpath}/{name}/{name}_exp.fits'\n",
    "    # detections\n",
    "    if not os.path.isfile(f'{outpath}/{name}/{name}_vtp.detect'):\n",
    "        return\n",
    "    else:\n",
    "        srcs = f'{outpath}/{name}/{name}_vtp.detect'\n",
    "\n",
    "    # load the data\n",
    "    evnt_data = fits.getdata(evnts)\n",
    "    expmap_data = fits.getdata(expmap)\n",
    "\n",
    "    # now we need to read the individual detections\n",
    "    detects = Table.read(srcs, hdu=1)\n",
    "\n",
    "    # we are going to compute profiles only for the sources within 6'\n",
    "    try:\n",
    "        detect_mask = detects['PSZ_dist'] <= 6\n",
    "    except KeyError:\n",
    "        print(name, outpath)\n",
    "        return\n",
    "\n",
    "    for i in np.arange(len(detects))[detect_mask]:\n",
    "        # mask out all sources except the source we are working with\n",
    "        if len(detects) < 2:\n",
    "            # here there is only one source in the field. Don't mask anything.\n",
    "            mask_sources = False\n",
    "        else:\n",
    "            mask_sources = True\n",
    "            src_mask = [j for j in range(len(detects)) if j != i]\n",
    "\n",
    "        if mask_sources:\n",
    "            regs = []\n",
    "            ### have to subtract 1 from the region positions, because FITS are 1-index'd and python\n",
    "            ### is 0-index'd.\n",
    "            for idx, xc1, yc1, rc, rotc in detects[src_mask][[\n",
    "                    'INDEX', 'X', 'Y', 'R', 'ROTANG'\n",
    "            ]]:\n",
    "                ellipse_region = EllipsePixelRegion(center=PixCoord(x=xc1 - 1,\n",
    "                                                                    y=yc1 - 1),\n",
    "                                                    width=2 * rc[0] * 1.5,\n",
    "                                                    height=2 * rc[1] * 1.5,\n",
    "                                                    angle=rotc * u.deg,\n",
    "                                                    meta={'ID': idx})\n",
    "                regs.append(ellipse_region)\n",
    "\n",
    "            # create the compound regions\n",
    "            cmp_reg = compound_regions(regs, evnt_data.shape[0])\n",
    "            # invert the comp region -- we want the places where there aren't regions\n",
    "            cmp_regmask = np.where(cmp_reg == 0, 1, 0)\n",
    "\n",
    "        else:\n",
    "            cmp_regmask = np.ones_like(evnt_data)\n",
    "\n",
    "        with open(f'{outpath}/{name}/{name}_vtp_{detects[\"INDEX\"][i]}.radprof',\n",
    "                  'w') as radprof:\n",
    "            radprof.write(\n",
    "                f\"# source number and position: {name}_{detects['INDEX'][i]} \"\n",
    "                f\"{detects['X'][i]:.3f} {detects['Y'][i]:.3f}\\n\")\n",
    "            radprof.write(\n",
    "                f\"# profile parameters: {params['n']} {params['r0']} {params['dr']}\\n\"\n",
    "            )\n",
    "            radprof.write(f\"# bin r1 r2 x y w sb sb_err psf area psfsb\\n\")\n",
    "\n",
    "            for irad in range(params['n']):\n",
    "                r1 = params['r0'] + irad * params['dr']\n",
    "                r2 = params['r0'] + (irad + 1) * params['dr']\n",
    "\n",
    "                center = PixCoord(x=detects['X'][i] - 1, y=detects['Y'][i] - 1)\n",
    "                if r1 == 0:\n",
    "                    ann_reg = CirclePixelRegion(center=center, radius=r2)\n",
    "                else:\n",
    "                    ann_reg = CircleAnnulusPixelRegion(center=center,\n",
    "                                                       inner_radius=r1,\n",
    "                                                       outer_radius=r2)\n",
    "                # make it into an image\n",
    "                ann_regmask = compound_regions([ann_reg], evnt_data.shape[0])\n",
    "\n",
    "                # combine the regions with the annulus.\n",
    "                final_regmask = (ann_regmask * cmp_regmask)\n",
    "                final_regmask_inv = np.where(\n",
    "                    final_regmask == 0, 1,\n",
    "                    0)  # this is the final mask for the data\n",
    "\n",
    "                evnt_data_masked = ma.masked_array(evnt_data,\n",
    "                                                   mask=final_regmask_inv)\n",
    "                expmap_data_masked = ma.masked_array(expmap_data,\n",
    "                                                     mask=final_regmask_inv)\n",
    "\n",
    "                x = evnt_data_masked.sum()\n",
    "                y = expmap_data_masked.mean()\n",
    "                w = expmap_data_masked.count()\n",
    "\n",
    "                sb = x / (y * w * pixscale**2 / 3600)  #  cts/s/arcmin^2\n",
    "                err_gehrels = sqrt(x + 0.75)\n",
    "                sbe = (1. + err_gehrels) / (y * w * pixscale**2 / 3600\n",
    "                                            )  #  cts/s/arcmin^2\n",
    "\n",
    "                # build the model of the psf\n",
    "                psfr1 = integ_beta_model(r1, rc=(5.6 / 2.36), beta=-1/3)  # cnts\n",
    "                psfr2 = integ_beta_model(r2, rc=(5.6 / 2.36), beta=-1/3)  # cnts\n",
    "                psf = psfr2 - psfr1  # total flux in the annulus\n",
    "\n",
    "                area = np.pi * r2**2 - np.pi * r1**2  # number of square pixels\n",
    "\n",
    "                psfsb = psf / (y * area * pixscale**2 / 3600)\n",
    "\n",
    "                radprof.write(\n",
    "                    f'{irad:3d} {r1:7d} {r2:7d} {x:9.1f} {y:9.1f} {w:9.1f}'\n",
    "                    f'{sb:12.4e} {sbe:12.4e} '\n",
    "                    f'{psf:9.1f} {area:9.1f} {psfsb:12.4e} \\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_radprof(name, outpath):\n",
    "\n",
    "    # swift pixelscale in degrees\n",
    "    pixsc = 6.548089E-04 * 3600\n",
    "\n",
    "    if os.path.isfile(f'{outpath}/{name}/{name}_vtp.detect'):\n",
    "        srcs = f'{outpath}/{name}/{name}_vtp.detect'\n",
    "        # now we need to read the individual detections\n",
    "        detects = Table.read(srcs, hdu=1)\n",
    "    else:\n",
    "        return\n",
    "\n",
    "    # we are going to compute profiles only for the sources within 6'\n",
    "    detects = detects[detects['PSZ_dist'] <= 6]\n",
    "\n",
    "    # open a file to write out the results\n",
    "    outfile = open(f'{outpath}/{name}/{name}_mcmcfits.txt', 'w')\n",
    "    outfile.write(\n",
    "        '# Field ID So_50 So_84 So_16 '\n",
    "        'rc_50 rc_84 rc_16 beta_50 beta_84 beta_16 bg_50 bg_84 bg_16\\n')\n",
    "\n",
    "    # loop over the sources -- only the first 20\n",
    "    for i in detects['INDEX']:\n",
    "\n",
    "        if os.path.isfile(f'{outpath}/{name}/{name}_vtp_{i}.radprof'):\n",
    "            data = Table.read(f'{outpath}/{name}/{name}_vtp_{i}.radprof',\n",
    "                              format='ascii',\n",
    "                              header_start=2)\n",
    "        else:\n",
    "            continue\n",
    "\n",
    "        # don't try to fit things if there isn't any data\n",
    "        if len(data) <= 10:\n",
    "            continue\n",
    "\n",
    "        # mask out any points where the sb is zero\n",
    "        mask = data['sb'] != 0\n",
    "        data = data[mask]\n",
    "\n",
    "        # x-axis, in arcminutes\n",
    "        x = (data['r1'] + data['r2']) / 2. / 60. * pixsc\n",
    "\n",
    "        # this is the parameter fitting\n",
    "        ndim = 4  # number of parameters in the model\n",
    "        nwalkers = 100  # number of MCMC walkers\n",
    "        nburn = 100  # \"burn-in\" period to let chains stabilize\n",
    "        nsteps = 500  # number of MCMC steps to take\n",
    "\n",
    "        pos = [\n",
    "            np.array([0.001, 1., 1., 0.001]) + 1e-4 * np.random.randn(ndim)\n",
    "            for i in range(nwalkers)\n",
    "        ]\n",
    "        sampler = emcee.EnsembleSampler(nwalkers,\n",
    "                                        ndim,\n",
    "                                        prob,\n",
    "                                        args=(x, data['sb'], data['sb_err']))\n",
    "        sampler.run_mcmc(pos, nsteps)\n",
    "        emcee_trace = sampler.get_chain(discard=nburn, thin=10, flat=True)\n",
    "\n",
    "        # plot the result\n",
    "        fignum = np.random.randint(1, 1000)\n",
    "        fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(6.5, 5), num=fignum)\n",
    "        # data\n",
    "        ax.errorbar(x, data['sb'], yerr=data['sb_err'], fmt='o', label='data')\n",
    "        for So, rc, beta, bg in emcee_trace[np.random.randint(len(emcee_trace),\n",
    "                                                              size=100)]:\n",
    "            ax.plot(x, beta_model(So, x, rc, beta) + bg, color='k', alpha=0.1)\n",
    "        ylims = ax.get_ylim()\n",
    "\n",
    "        # fit -- median (and error) values\n",
    "        So, rc, beta, bg = map(\n",
    "            lambda v: (v[1], v[2] - v[1], v[1] - v[0]),\n",
    "            zip(*np.percentile(emcee_trace, [16, 50, 84], axis=0)))\n",
    "        ax.plot(x,\n",
    "                beta_model(So[0], x, rc[0], beta[0]) + bg[0],\n",
    "                label=r'$\\beta$ + bg')\n",
    "        ax.plot(x, beta_model(So[0], x, rc[0], beta[0]), label=r'$\\beta$')\n",
    "\n",
    "        # add the psf\n",
    "        ax.plot(x, data['psfsb'] + bg[0], label='psf', color='orange')\n",
    "\n",
    "        # write out the fit parameters\n",
    "        outfile.write(f'{name} ')\n",
    "        outfile.write(f'{i} ')\n",
    "        outfile.write(f'{So[0]:.5f} {So[1]:.5f} {So[2]:.5f} ')\n",
    "        outfile.write(f'{rc[0]:.5f} {rc[1]:.5f} {rc[2]:.5f} ')\n",
    "        outfile.write(f'{beta[0]:.5f} {beta[1]:.5f} {beta[2]:.5f} ')\n",
    "        outfile.write(f'{bg[0]:.5f} {bg[1]:.5f} {bg[2]:.5f}\\n')\n",
    "\n",
    "        # rc and bg lines\n",
    "        ax.axhline(bg[0], label='bg', zorder=0, lw=1)\n",
    "        ax.axvline(rc[0], zorder=0, lw=1)\n",
    "\n",
    "        ax.semilogy()\n",
    "\n",
    "        if ylims[0] < 1e-5:\n",
    "            ax.text(rc[0] + 0.1,\n",
    "                    2e-5,\n",
    "                    f\"rc = {rc[0]:.2f}'\",\n",
    "                    rotation='vertical',\n",
    "                    ha='left')\n",
    "            ax.set_ylim(1e-5, ylims[1])\n",
    "        else:\n",
    "            ax.text(rc[0] + 0.1,\n",
    "                    ylims[0],\n",
    "                    f\"rc = {rc[0]:.2f}'\",\n",
    "                    rotation='vertical',\n",
    "                    ha='left')\n",
    "            ax.set_ylim(ylims)\n",
    "        ax.set_xlabel('Radius [arcmin]')\n",
    "        ax.set_ylabel('Flux $[cnts/s/arcmin^2]$')\n",
    "        ax.legend(loc='upper right')\n",
    "\n",
    "        fig.savefig(f'{outpath}/{name}/{name}_radproffit_{i}.png',\n",
    "                    bbox='tight',\n",
    "                    dpi=180)\n",
    "        plt.close(fig)\n",
    "\n",
    "        # add a corner plot of the fits.\n",
    "        # plot the bg and So in log to make them fit on the plot better\n",
    "        emcee_trace[:, 0] = np.log10(emcee_trace[:, 0])\n",
    "        emcee_trace[:, 3] = np.log10(emcee_trace[:, 3])\n",
    "        f = corner.corner(\n",
    "            emcee_trace,\n",
    "            labels=['log So', 'rc', r'$\\beta$', 'log bg'],\n",
    "            bins=[50, 50, 50, 50],\n",
    "            smooth=True,\n",
    "            fill_contours=True,\n",
    "            truths=[np.log10(So[0]), rc[0], beta[0],\n",
    "                    np.log10(bg[0])])\n",
    "\n",
    "        f.savefig(f'{outpath}/{name}/{name}_corner_{i}.png', bbox='tight')\n",
    "        plt.close(f)\n",
    "    outfile.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# get file data\n",
    "data = load_PSZcatalog()\n",
    "data = data.sort_index(axis=1)\n",
    "\n",
    "outpath = './data'\n",
    "\n",
    "arr = [{'name':n.replace(' ', '_'), 'outpath':outpath} for n in data['NAME']]\n",
    "\n",
    "###\n",
    "# This is the order you should call the functions. \n",
    "# There are other functions in this notebook, but these are the only ones \n",
    "# that should be called directly.\n",
    "###\n",
    "\n",
    "clean_files(outpath)\n",
    "\n",
    "parallel_process(arr, get_radprof, use_kwargs=True, n_jobs=6) # ~7 minutes\n",
    "parallel_process(arr, fit_radprof, use_kwargs=True, n_jobs=6) # ~20 minutes\n",
    "\n",
    "cp_results(outpath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### For testing ###\n",
    "\n",
    "outpath = './data'\n",
    "name = 'PSZ2_G078.36-47.22'\n",
    "# get_radprof(name, outpath)\n",
    "fit_radprof(name, outpath)"
   ]
  },
  {
   "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.7.1"
  },
  "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
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}