nircam_parallel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Crowded field photometry --- 1b. NIRCam parallel imaging\n",
    "\n",
    "This notebook demonstrates how to create JWST data by using [MIRAGE](https://mirage-data-simulator.readthedocs.io/en/latest/).\n",
    "This is one of series of notebooks that focus on crwoded field photometry (see other notebooks in the directory).\n",
    "\n",
    "\n",
    "This notebook will demonstrate how to create NIRISS WFSS and direct image simulation. It will require the following input files;\n",
    "1. APT output files (.xml and .pointing) that specify the observation\n",
    "2. A source catalog that lists source properties (magnitude/morphology...), taken from HST.\n",
    "\n",
    "Users may also provide their catalog/apt files, to create own scenes, but may need to allocate columns for physical parameters.\n",
    "\n",
    "\n",
    "## To Do:\n",
    "- 1.Read input catalogs\n",
    "- 2.Generate two Mirage friendly source catalogs (one for point source, and the other for extended).\n",
    "- 3.Generate a set of yaml files, setup files for Mirage.\n",
    "- 4.Run simulation for NIRCAM parallel images.\n",
    "\n",
    "*Reduction of these generated raw data will be demonstrated in another notebook.\n",
    "\n",
    "### Requirement:\n",
    "- Mirage environment. [see here](https://mirage-data-simulator.readthedocs.io/en/latest/install.html)\n",
    "\n",
    "\n",
    "### It is recommended to run this notebook after 1a, as detailed explanations can be found there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set environment variables\n",
    "# It may be helpful to set these within your .bashrc or .cshrc file, so that CRDS will\n",
    "# know where to look for reference files during future runs of the JWST calibration\n",
    "# pipeline.\n",
    "\n",
    "import os\n",
    "\n",
    "# if you are in the institute network, read here;\n",
    "# https://mirage-data-simulator.readthedocs.io/en/latest/reference_files.html\n",
    "# Otherwise, you need to download to your local directory.\n",
    "\n",
    "path_mirage = '/path/into/which/files/are/downlaoded/'\n",
    "#path_mirage = '/Volumes/EHDD1/mirage_data/'\n",
    "\n",
    "if os.path.exists(path_mirage):\n",
    "    os.environ[\"MIRAGE_DATA\"] = path_mirage\n",
    "else:\n",
    "    from mirage.reference_files import downloader\n",
    "    download_path = path_mirage\n",
    "    downloader.download_reffiles(download_path, instrument='all', dark_type='linearized', skip_darks=False, skip_cosmic_rays=False, skip_psfs=False, skip_grism=False)\n",
    "    os.environ[\"MIRAGE_DATA\"] = path_mirage\n",
    "\n",
    "os.environ[\"CRDS_PATH\"] = os.path.join(os.path.expandvars('$HOME'), \"crds_cache\")\n",
    "os.environ[\"CDRS_SERVER_URL\"] = \"https://jwst-cdrs.stsci.edu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "import pkg_resources\n",
    "import yaml\n",
    "import zipfile\n",
    "import urllib.request\n",
    "\n",
    "from astropy.io import fits\n",
    "import astropy.units as u\n",
    "from astropy.visualization import simple_norm, imshow_norm\n",
    "import h5py\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from mirage import imaging_simulator\n",
    "from mirage import wfss_simulator\n",
    "from mirage.utils.constants import FLAMBDA_CGS_UNITS, FLAMBDA_MKS_UNITS, FNU_CGS_UNITS \n",
    "from mirage.yaml import yaml_generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mirage\n",
    "import astropy\n",
    "\n",
    "print('mirage ver:',mirage.__version__)\n",
    "print('astropy ver:',astropy.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The flow should be similar to example 1a."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "<a id='yaml_from_apt'></a>\n",
    "## Create a series of yaml files from [APT](https://jwst-docs.stsci.edu/display/JPP/JWST+Astronomers+Proposal+Tool+Overview)\n",
    "\n",
    "To set Mirage up for your observation designed in APT, you need to get .xml and .pointing files out of your APT file and provide them to Mirage.\n",
    "You can get these files from APT's tool bar ([File]->[Export]).\n",
    "\n",
    "As an example, observation here is taken from [GLASS-ERS](http://www.stsci.edu/jwst/observing-programs/approved-ers-programs/program-1324). Some modifications were added.\n",
    "\n",
    "Let's download input files first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "boxlink = 'https://data.science.stsci.edu/redirect/JWST/jwst-data_analysis_tools/NIRISS_lensing_cluster/files.zip'\n",
    "boxfile = './files.zip'\n",
    "DIR_DATA = './files/'\n",
    "\n",
    "if not os.path.exists(boxfile):\n",
    "    urllib.request.urlretrieve(boxlink, boxfile)\n",
    "    zf = zipfile.ZipFile(boxfile, 'r')\n",
    "    zf.extractall()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input files from APT\n",
    "# These can be obtained from a tab, [File]->[Export], in APT.\n",
    "# There is some modification:\n",
    "# 1.Deep direct imaging (3600s per exposure)\n",
    "# 2. Cut out NIRCam parallel imaging, and NIRSPEC observations, to run through the notebook here.\n",
    "\n",
    "xml_file = '%sA2744_example_nrcb_par.xml'%DIR_DATA\n",
    "pointing_file = '%sA2744_example_nrcb_par.pointing'%DIR_DATA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set reference file values. \n",
    "# Setting to 'crds_full_name' will search for and download needed\n",
    "# calibration reference files (commonly referred to as CRDS reference files) when\n",
    "# the yaml_generator is run. \n",
    "# \n",
    "# Setting to 'crds' will put placeholders in the yaml files and save the downloading\n",
    "# for when the simulated images are created.\n",
    "reffile_defaults = 'crds'\n",
    "\n",
    "# Optionally set the cosmic ray library and rate\n",
    "cosmic_rays = {'library': 'SUNMAX', 'scale': 1.0}\n",
    "\n",
    "# Optionally set the observation date to use for the data. Note that this information\n",
    "# is placed in the headers of the output files, but not used by Mirage in any way.\n",
    "dates = '2022-07-01'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optionally set the background signal rates to be used\n",
    "# As in Cami's note, this may not work.\n",
    "background = 'low' #'medium'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optionally set the telescope roll angle (PAV3) for the observations\n",
    "pav3 = 0 #12.5 + 45."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datatype = 'raw'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NIRCam parallel imaging;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create catalog for parallel field;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read the original catalog;\n",
    "PARFLD    = 'abell2744par'\n",
    "PARFLDID  = '11'\n",
    "magzp   = 25.0\n",
    "\n",
    "from astropy.io import ascii\n",
    "\n",
    "# This is for flux, where the magnitude zeropoint is mag0.\n",
    "t = ascii.read('%sinput_%s.cat'%(DIR_DATA,PARFLDID))\n",
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sectractor;\n",
    "pixscale = 0.06\n",
    "ts = ascii.read('%sf160w_%s.cat'%(DIR_DATA,PARFLDID))\n",
    "\n",
    "idsext = ts['NUMBER']\n",
    "radius = ts['FLUX_RADIUS']*pixscale # in arcsec\n",
    "elong  = ts['ELONGATION'] # = a/b\n",
    "pa     = 90. - ts['THETA_IMAGE'] # Theta_mirage = 90 - Theta_Sext\n",
    "pa += pav3 + pav3 # Doe to a conflict with Mirage.\n",
    "magauto= ts['MAG_AUTO']\n",
    "class_star = ts['CLASS_STAR']\n",
    "ellip  = 1. - 1/elong\n",
    "\n",
    "ts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since SExtractor does not provide Sersic index;\n",
    "# Set those bright, but extended object to n=4.\n",
    "ser = radius * 0 + 1.0\n",
    "con_ser = (magauto<20) & (class_star<0.9)\n",
    "ser[con_ser] = 4.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Primary infos;\n",
    "id  = t['id']\n",
    "ra  = ts['X_WORLD']\n",
    "dec = ts['Y_WORLD']\n",
    "#red = ez['z_m1']\n",
    "\n",
    "# Manage flux column;\n",
    "# This corresponds to photometric catalog;\n",
    "filt  = ['F105W','F125W','F140W','F160W','F435W','F606W','F814W']\n",
    "efilt = ['202','203','204','205','1','4','6'] # EAZY column number.\n",
    "\n",
    "\n",
    "flux = np.zeros((len(filt),len(id)),dtype='float32')\n",
    "eflux= np.zeros((len(filt),len(id)),dtype='float32')\n",
    "\n",
    "for ff in range(len(filt)):\n",
    "    flux[ff,:]  = t['F%s'%(efilt[ff])]\n",
    "    eflux[ff,:] = t['E%s'%(efilt[ff])]\n",
    "    \n",
    "mag160 = -2.5 * np.log10(flux[3,:]) + magzp\n",
    "mag160"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For the sake of test, only bright sources;\n",
    "mag_lim = 21 # But, remember; more source, more computing time\n",
    "filt_targ = 'F200W' # Targer filter of this simulation. \n",
    "\n",
    "# Just for index info.\n",
    "ii105 = np.where(np.asarray(filt[:]) == 'F105W')[0][0]\n",
    "ii140 = np.where(np.asarray(filt[:]) == 'F140W')[0][0]\n",
    "ii160 = np.where(np.asarray(filt[:]) == 'F160W')[0][0]\n",
    "ii814 = np.where(np.asarray(filt[:]) == 'F814W')[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write into two files;\n",
    "# 1. extended source catalog\n",
    "# 2. point source catalog\n",
    "file_extend_name = './sources_extend_%s.cat'%(PARFLDID)\n",
    "file_point_name  = './sources_point_%s.cat'%(PARFLDID)\n",
    "\n",
    "\n",
    "# Start writing;\n",
    "# Make sure you change niriss to nircam for header.\n",
    "fw = open(file_point_name,'w')\n",
    "fw.write('# position_RA_Dec\\n# abmag\\n#\\n#\\n')\n",
    "# Let's approximate WFC3IR F140W to NIRISS F150W\n",
    "# Let's approximate WFC3IR F160W to NIRISS F200W\n",
    "# Future work will extract magnitude from the SED fitting result the corresponding filters.\n",
    "fw.write('index   x_or_RA      y_or_Dec    niriss_f115w_magnitude    nircam_f150w_magnitude    niriss_f090w_magnitude    nircam_f200w_magnitude\\n')\n",
    "\n",
    "# Extended source catalog; \n",
    "# see https://mirage-data-simulator.readthedocs.io/en/latest/catalogs.html#extended-obj for more.\n",
    "fw2 = open(file_extend_name,'w')\n",
    "fw2.write('# position_RA_Dec\\n# abmag\\n#\\n#\\n')\n",
    "fw2.write('index   x_or_RA      y_or_Dec    radius    ellipticity    pos_angle       sersic_index     niriss_f115w_magnitude    nircam_f150w_magnitude    niriss_f090w_magnitude    nircam_f200w_magnitude\\n')\n",
    "\n",
    "# Adding a constant to the second catalog;\n",
    "nrand = 10000\n",
    "\n",
    "for ii in range(len(id)):\n",
    "    # Check you have positive flux;\n",
    "    if flux[ii140,ii]>0:\n",
    "        mag_140 = -2.5 * np.log10(flux[ii140,ii])+magzp\n",
    "    else:\n",
    "        mag_140 = 99\n",
    "\n",
    "    if flux[ii105,ii]>0:\n",
    "        mag_105 = -2.5 * np.log10(flux[ii105,ii])+magzp\n",
    "    else:\n",
    "        mag_105 = 99\n",
    "\n",
    "    if flux[ii160,ii]>0:\n",
    "        mag_160 = -2.5 * np.log10(flux[ii160,ii])+magzp\n",
    "    else:\n",
    "        mag_160 = 99\n",
    "\n",
    "    if flux[ii814,ii]>0:\n",
    "        mag_814 = -2.5 * np.log10(flux[ii814,ii])+magzp\n",
    "    else:\n",
    "        mag_814 = 99\n",
    "\n",
    "    # As mirage cannot handle faint objects...\n",
    "    if mag_160<mag_lim:\n",
    "        iisext = np.where(id[ii] == idsext[:])\n",
    "        if class_star[ii] > 0.90: # As point sources\n",
    "            fw.write('%d %.8f %.8f %.3f %.3f %.3f %.3f\\n'%(id[ii],ra[iisext],dec[iisext],mag_105,mag_140,mag_814,mag_160))\n",
    "        else: # Extended source\n",
    "            fw2.write('%d %.8f %.8f %.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\\n'%(id[ii]+nrand,ra[iisext],dec[iisext],radius[iisext],ellip[iisext],pa[iisext],ser[iisext],mag_105,mag_140,mag_814,mag_160))\n",
    "\n",
    "fw.close()\n",
    "fw2.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Need to create yaml file, as we want to update input source catalog; \n",
    "catalogs = {'point_source': file_point_name,'galaxy': file_extend_name}\n",
    "catalogs\n",
    "\n",
    "# Also, update yaml output directory;\n",
    "yaml_output_dir = './yaml_nircam/' #TEST_DATA_DIRECTORY\n",
    "if not os.path.exists(yaml_output_dir):\n",
    "    os.mkdir(yaml_output_dir)    \n",
    "simulations_output_dir = './output_nircam/' #TEST_DATA_DIRECTORY\n",
    "if not os.path.exists(simulations_output_dir):\n",
    "    os.mkdir(simulations_output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run the yaml generator\n",
    "yam = yaml_generator.SimInput(input_xml=xml_file, pointing_file=pointing_file,\n",
    "                              catalogs=catalogs, cosmic_rays=cosmic_rays,\n",
    "                              background=background,roll_angle=pav3,\n",
    "                              dates=dates, reffile_defaults=reffile_defaults,\n",
    "                              verbose=True, output_dir=yaml_output_dir,\n",
    "                              simdata_output_dir=simulations_output_dir,\n",
    "                              datatype=datatype)\n",
    "yam.create_inputs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yaml_files = glob(os.path.join(yaml_output_dir,\"jw*.yaml\"))\n",
    "\n",
    "yaml_NIRCAM_imaging_A_files = []\n",
    "yaml_NIRCAM_imaging_B_files = []\n",
    "\n",
    "for f in yaml_files:\n",
    "    my_dict = yaml.safe_load(open(f))\n",
    "    if my_dict[\"Inst\"][\"mode\"]==\"imaging\" and my_dict['Inst']['instrument']=='NIRCAM'\\\n",
    "    and my_dict['Readout']['filter']==filt_targ and my_dict['Readout']['array_name'][:4]=='NRCA':\n",
    "        yaml_NIRCAM_imaging_A_files.append(f)\n",
    "    elif my_dict[\"Inst\"][\"mode\"]==\"imaging\" and my_dict['Inst']['instrument']=='NIRCAM'\\\n",
    "    and my_dict['Readout']['filter']==filt_targ and my_dict['Readout']['array_name'][:4]=='NRCB':\n",
    "        yaml_NIRCAM_imaging_B_files.append(f)\n",
    "    \n",
    "print(\"Nircam Module B imaging files:\",len(yaml_NIRCAM_imaging_B_files))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running Mirage on NIRCam module B\n",
    "\n",
    "*Module A is not inclded here, as input sources are not distributed in the covered region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run all steps of the imaging simulator for yaml file #1\n",
    "for file in yaml_NIRCAM_imaging_B_files:\n",
    "    img_sim = imaging_simulator.ImgSim()\n",
    "    img_sim.paramfile = file\n",
    "    img_sim.create()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Take a look at one image;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare with the real F160W image, if you like;\n",
    "file_f160w = 'hlsp_frontier_hst_wfc3-60mas_abell2744-hffpar_f160w_v1.0-epoch2_drz.fits'\n",
    "if not os.path.exists(file_f160w):\n",
    "    link_hst = 'https://archive.stsci.edu/pub/hlsp/frontier/abell2744/images/hst/v1.0-epoch2/%s'%file_f160w\n",
    "    urllib.request.urlretrieve(link_hst, file_f160w)\n",
    "\n",
    "with fits.open(file_f160w) as hdulist:\n",
    "    data_real = hdulist[0].data\n",
    "    hdulist.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkid = yaml_NIRISS_imaging_files[0].split('/')[-1].replace('.yaml','') #'jw00042001001_01201_00001_nrcb1'\n",
    "\n",
    "final_file = os.path.join(simulations_output_dir, '%s_uncal.fits'%checkid)\n",
    "with fits.open(final_file) as hdulist:\n",
    "    data = hdulist['SCI'].data\n",
    "    hdulist.info()    \n",
    "    \n",
    "fig, ax = plt.subplots(figsize=(10, 10))\n",
    "norm = simple_norm(data[0, -1, :, :], stretch='log', min_cut=5000, max_cut=30000)\n",
    "cax = ax.imshow(data[0, -1, :, :], norm=norm)\n",
    "cbar = fig.colorbar(cax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import ndimage\n",
    "\n",
    "fig = plt.figure(figsize=(15.,10.))\n",
    "fig.subplots_adjust(top=0.98, bottom=0.16, left=0.1, right=0.99, hspace=0.15, wspace=0.25)\n",
    "ax1 = fig.add_subplot(121)\n",
    "ax2 = fig.add_subplot(122)\n",
    "\n",
    "norm = simple_norm(data[:,:], stretch='log', min_cut=0, max_cut=10)\n",
    "rotated_img = ndimage.rotate(data[:,:], -pav3)\n",
    "\n",
    "cax = ax1.imshow(rotated_img, norm=norm)\n",
    "\n",
    "norm = simple_norm(data_real[1200:3800, 1200:3800], stretch='log', min_cut=0, max_cut=10)\n",
    "cax = ax2.imshow(data_real[1200:3800, 1200:3800], norm=norm)\n",
    "cbar = fig.colorbar(cax)\n",
    "ax1.set_title('Mirage')\n",
    "ax2.set_title('HST F160W')\n",
    "plt.savefig('02_comparison.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next : Reduce uncal images with pipeline in a different notebook."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}