{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "72b2f829",
   "metadata": {},
   "source": [
    "======================== Import Packages =========================="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4769e901",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os, pdb, glob\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from astropy.io import fits\n",
    "from astropy.table import Table, join\n",
    "from astropy.coordinates import SkyCoord\n",
    "from astropy import units as u\n",
    "from astroquery.vizier import Vizier\n",
    "import warnings\n",
    "from astropy.logger import AstropyWarning\n",
    "warnings.filterwarnings('ignore', category=AstropyWarning)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c002da4",
   "metadata": {},
   "source": [
    "===================== Define Functions ==================="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1c2d8757",
   "metadata": {},
   "outputs": [],
   "source": [
    "def readfits(file):\n",
    "\n",
    "    \"\"\"\n",
    "    PURPOSE:    Read in FITS file and header info\n",
    "\n",
    "    INPUT:      Path to FITS file (str)\n",
    "\n",
    "    OUTPUT:     img = image (float arr)\n",
    "                xcen, ycen = image center coordinates in pixels (float)\n",
    "                xpix, ypix = image pixel width in deg/pix units (float)\n",
    "                bmaj, bmin, bpa = beam major axis, minor axis, position angle (float)\n",
    "                xcen_ra, ycen_ra = image center coordinates in deg units (float)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    ### READ IN FITS FILE\n",
    "    hdulist = fits.open(file)\n",
    "    data = hdulist[0].data[0, 0, :, :]\n",
    "    head = hdulist[0].header\n",
    "    hdulist.close()\n",
    "\n",
    "    ### GET HEADER INFO\n",
    "    xcen = head['CRPIX1']\n",
    "    ycen = head['CRPIX2']\n",
    "    xpix = head['CDELT1']\n",
    "    ypix = head['CDELT2']\n",
    "    xcen_ra = head['CRVAL1']\n",
    "    xcen_de = head['CRVAL2']\n",
    "    bmaj = head['bmaj']\n",
    "    bmin = head['bmin']\n",
    "    bpa  = head['bpa']\n",
    "\n",
    "    return(data, xcen, ycen, xpix, ypix, bmaj, bmin, bpa, xcen_ra, xcen_de)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "20008f42",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_fits(img, line, i, xpix_all, ypix_all, bmaj_all, bmin_all, bpa_all):\n",
    "\n",
    "    \"\"\"\n",
    "    PURPOSE:    Write FITS file and header info\n",
    "\n",
    "    INPUT:      Image to write (array)\n",
    "                Line that was stacked (str)\n",
    "                Pixel width in deg/pix units for all images in stack (array)\n",
    "                Beam major axis, minor axis, position angle for all images in stack (array)\n",
    "\n",
    "    OUTPUT:     Stacked image (FITS file)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    os.system('rm ../output/stack_nd_'+line+'_'+str(i)+'.fits')\n",
    "    hdu = fits.PrimaryHDU()\n",
    "    hdu.data = img\n",
    "\n",
    "    hdu.header['CRPIX1'] = hdu.header['NAXIS1']/2\n",
    "    hdu.header['CRPIX2'] = hdu.header['NAXIS2']/2\n",
    "    hdu.header['bmaj'] = bmaj_all.mean()\n",
    "    hdu.header['bmin'] = bmin_all.mean()\n",
    "    hdu.header['bpa'] = bpa_all.mean()\n",
    "    hdu.header['cdelt1'] = xpix_all.mean()\n",
    "    hdu.header['cdelt2'] = ypix_all.mean()\n",
    "\n",
    "    hdu.writeto('../output/stack_nd_'+line+'_'+str(i)+'.fits')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "44e7bf16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop_img(file_img, hw_as, c_obj):\n",
    "\n",
    "    \"\"\"\n",
    "    PURPOSE:    Crop & center an image\n",
    "\n",
    "    INPUT:      file_img = path to FITS file (str)\n",
    "                hw_as = half-width of desired cropped image size in arcsec (float)\n",
    "                c_obj = coordinates of object (AstroPy SkyCoords)\n",
    "\n",
    "    OUTPUT:     img = cropped & centered image (float arr)\n",
    "                width_pix = half-width of cropped image in pixels (int)\n",
    "                xpix_img, ypix_img = image pixel width in deg/pix units (float)\n",
    "                Beam major axis, minor axis, position angle (float)\n",
    "                bmaj, bmin, bpa = beam major axis, minor axis, position angle (float)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    ### LOAD IMAGE AND GET CENTER COORDINATES\n",
    "    img, xcen_img, ycen_img, xpix_img, ypix_img, bmaj_img, bmin_img, bpa_img, xcen_ra_img, ycen_de_img = readfits(file_img)\n",
    "    c_img = SkyCoord(xcen_ra_img, ycen_de_img, frame='icrs', unit='deg')\n",
    "    \n",
    "    ### CENTER IMAGE ON OBJECT LOCATION \n",
    "    dra, ddec = c_img.spherical_offsets_to(c_obj)\n",
    "    width_pix = int(round(hw_as / (ypix_img * 3600.0)))\n",
    "    xctr = xcen_img + dra.value / xpix_img\n",
    "    yctr = ycen_img + ddec.value / ypix_img\n",
    "\n",
    "    ### CROP IMAGE\n",
    "    img = img[int(round(yctr - width_pix)):int(round(yctr + width_pix)),\n",
    "              int(round(xctr - width_pix)):int(round(xctr + width_pix))]\n",
    "    \n",
    "    return img, width_pix, xpix_img, ypix_img, bmaj_img, bmin_img, bpa_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "06595dee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stackme(t, line):\n",
    "\n",
    "    \"\"\"\n",
    "    PURPOSE:    Stack image\n",
    "\n",
    "    INPUT:      Table of sources to be stacked (AstroPy Table)\n",
    "                Line name (str; must be 'cont', '13CO', C18O')\n",
    "\n",
    "    OUTPUT:     Stacked image (array)\n",
    "                Pixel width in deg/pix units for all images in stack (array)\n",
    "                Beam major axis, minor axis, position angle for all images in stack (array)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    xpix_all, ypix_all = np.empty(len(t)), np.empty(len(t))\n",
    "    bmaj_all, bmin_all, bpa_all = np.empty(len(t)), np.empty(len(t)), np.empty(len(t))\n",
    "    \n",
    "    for i, val in enumerate(t['Name']):\n",
    "\n",
    "        ### GET IMAGE FILE NAME & CHECK FILE EXISTS\n",
    "        if (line == 'C18O'): suffix = '_c18o32.mom0_cropped.fits'\n",
    "        if (line == '13CO'): suffix = '_13co32.mom0_cropped.fits'\n",
    "        if (line == 'cont'):   suffix = '_cont_cropped.fits'\n",
    "        file_img = '../data/FITS/' + val + suffix\n",
    "        file_img = file_img.replace(' ', '_')\n",
    "        if os.path.isfile(file_img) is False:\n",
    "            print('missing FITS file for ' + val, line)\n",
    "            pdb.set_trace()\n",
    "\n",
    "        ### GET COORDINATES OF OBJECT FROM PAPER TABLE\n",
    "        de_obj = str(t['DEJ2000'][i].split(' ')[0][0]) + str(t['DEJ2000'][i].split(' ')[0][1:]) + 'd' + str(t['DEJ2000'][i].split(' ')[1]) + 'm' + str(t['DEJ2000'][i].split(' ')[2]) + 's'\n",
    "        ra_obj = str(t['RAJ2000'][i].split(' ')[0]) + 'h' + str(t['RAJ2000'][i].split(' ')[1]) + 'm' + str(t['RAJ2000'][i].split(' ')[2]) + 's'\n",
    "        c_obj = SkyCoord(ra_obj, de_obj, frame='icrs')\n",
    "\n",
    "        ### CROP IMAGE\n",
    "        img_cont, width_pix, xpix_img, ypix_img, bmaj_img, bmin_img, bpa_img = crop_img(file_img, 8.0, c_obj)\n",
    "        xpix_all[i], ypix_all[i] = xpix_img, ypix_img\n",
    "        bmaj_all[i], bmin_all[i], bpa_all[i] = bmaj_img, bmin_img, bpa_img\n",
    "\n",
    "        ### SCALE WITH DISTANCE AND PUT INTO MJY UNITS\n",
    "        img_cont = 1e3 * img_cont * ((t['Dist'][i] / 200.)**2)\n",
    "\n",
    "        ### ADD IMAGE TO STACK ARRAY\n",
    "        if (i==0):\n",
    "            img_all = np.zeros([2 * width_pix, 2 * width_pix, 1])\n",
    "            temp = img_cont.reshape((2 * width_pix, 2 * width_pix, 1))\n",
    "            img_all = temp\n",
    "            \n",
    "        else:\n",
    "            temp = img_cont.reshape((2 * width_pix, 2 * width_pix, 1))\n",
    "            img_all = np.append(img_all, temp, axis=2)\n",
    "\n",
    "    ### COLLAPSE IMAGE STACK ARRAY\n",
    "    stacked = np.sum(img_all, 2) / len(t)\n",
    "\n",
    "    return stacked, xpix_all, ypix_all, bmaj_all, bmin_all, bpa_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b43ea6fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(catalog, join_key='Name'):\n",
    "\n",
    "    \"\"\"\n",
    "    PURPOSE:    Get data from literature with Vizier\n",
    "\n",
    "    INPUT:      catalog = ctalog name on Vizier (str)\n",
    "                join_key = column header to join tables, if multiple (str; optional)\n",
    "\n",
    "    OUTPUT:     t = data table (AstroPy Table)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    ### GET FULL CATALOG (ALL COLUMNS, ALL ROWS)\n",
    "    viz = Vizier(catalog=catalog, columns=['**'])\n",
    "    viz.ROW_LIMIT = -1\n",
    "    tv = viz.get_catalogs(catalog)\n",
    "\n",
    "    ### IF MULTIPLE TABLES, JOIN THEN\n",
    "    for i, val in enumerate(tv.keys()):\n",
    "        if i == 0:\n",
    "            t = tv[val]\n",
    "        else:\n",
    "            tt = tv[val]\n",
    "            if join_key in tt.columns:\n",
    "                t = join(t, tt, join_type='inner', keys=join_key)\n",
    "\n",
    "    return t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea3a396f",
   "metadata": {},
   "source": [
    "========================== Code =========================="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e114e5ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "### GET LUPUS DATA\n",
    "T = get_data(\"J/ApJ/828/46\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "956527f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "### REMOVE SZ 82\n",
    "T.remove_row(np.where(T['Name'] == 'Sz 82')[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3e6ba1f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "### INDEX DUST NON-DETECTIONS\n",
    "ind_dust_nd = T['F890'] / T['e_F890'] < 3.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f6731392",
   "metadata": {},
   "outputs": [],
   "source": [
    "### INDEX dust detections, 13CO non-detections, C18O non-detections\n",
    "ind_gas_nd = ((T['F890'] / T['e_F890'] > 3.0)  & (T['l_F13CO'] == '<') & (T['l_F18CO'] == '<') )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ae9a0327",
   "metadata": {},
   "outputs": [],
   "source": [
    "### INDEX dust detections, 13CO detections, 18CO non-detections\n",
    "ind_C18O_nd = ((T['F890'] / T['e_F890'] > 3.0)  & (T['l_F13CO'] != '<') & (T['l_F18CO'] == '<') )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "893ac6cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "### STACK IMAGES\n",
    "lines = ['13CO', 'C18O']\n",
    "for n, nval in enumerate(lines):\n",
    "    for i, val in enumerate([T[ind_dust_nd], T[ind_gas_nd], T[ind_C18O_nd]]):\n",
    "        stacked, xpix_all, ypix_all, bmaj_all, bmin_all, bpa_all = stackme(val, lines[n])\n",
    "        write_fits(stacked, lines[n], i, xpix_all, ypix_all, bmaj_all, bmin_all, bpa_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3f74e218",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.close('all')"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}