<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Extracting spectra from KCWI datacubes
Step1:  Using SpectralCube
Read https
Step2:  From this spectrum, it looks like there are some bad wavelengths near 5500 A. I will mask this out before proceeding.
Masking
Further inspection reveals the bad wavelengths are between 5575 A and 5582 A. To mask out a chunk of the cube, one must set the boolean values corresponding to that section to False.
Step3:  Create a whitelight image
kcwi.get_img can produce a projection along the spectral axis. This projection can be simply between two wavelength limits or using a more complicated transmission curve (i.e. to produce synthetic photometry). Here, we only show the white light image.
Step4:  Extract spectrum using a spatial mask
One can use a 2D mask to extract a spectrum within that region. An intuitive way to do this would be to use PyRegion and create a mask from a ds9.reg file.
Step5:  It is crucial that the masks are always boolean. Using integer or float masks will make the process very slow.
Step6:  Get spectra from elliptical masks
This is a thin wrapper around the spec_from_mask function because ellipses are so ubiquitous.
Step7:  Find sources in an image
This is a wrapper for SExtractor's extract function. I'm passing a ds9 region file to exclude the edges from analysis.
Step8:  Extract spectra and create a MARZ input file
The final step is to get spectra of the identified sources and to create a MARZ input file.
<ASSISTANT_TASK:>
Python Code:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import astropy.units as u
from astropy.table import Table
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.wcs import WCS, utils as wcsutils
from astropy.stats import sigma_clipped_stats
import astropy.visualization as vis
from photutils.aperture import EllipticalAperture, SkyEllipticalAperture
import os, sys, glob, pdb, sep
from importlib import reload
import pyregion as pyreg
from spectral_cube import SpectralCube
from frb.frb import FRB
from frb.analysis import kcwi 
plt.rcParams['font.size'] = 17
flam = r"$F_\lambda$ $\rm (10^{-20}~erg~s^{-1}~cm^{-2}~\AA^{-1}$"
flam2 = r"$\sigma(F_\lambda)^2$ $\rm (10^{-40}~erg^2~s^{-2}~cm^{-4}~\AA^{-2}$"
def imshow_astro(img, wcsinfo = None, figsize = (10,10), colorbar =True,
               cblabel="", cbfrac = 0.035, norm = None,
               stretch = vis.LinearStretch(), cmap = "hot",
               vrange = (None, None)):
    
    _, med, std = sigma_clipped_stats(img.data)
    if not norm:
        norm = vis.ImageNormalize(stretch = stretch)
    fig =  plt.figure(figsize = figsize)
    if wcsinfo:
        ax = plt.subplot(projection = wcsinfo)
    else:
        ax = plt.subplot()
    
    vmin, vmax = vrange
    
    if not vmin:
        vmin = med
    if not vmax:
        vmax = med + 6*std
    
    im = ax.imshow(img, vmin = vmin, vmax = vmax, norm = norm, cmap = cmap)
    if colorbar:
        cb = plt.colorbar(im, label = cblabel, fraction =cbfrac)
    return fig, ax
Explanation: Extracting spectra from KCWI datacubes
End of explanation
# Load a datacube and its variance
# Insert your favourite datacube here
cubefile = "/home/sunil/Desktop/FRB/190608/KCWI/data/HG190608_5_icubes.fits"
varfile = "/home/sunil/Desktop/FRB/190608/KCWI/data/HG190608_5_vQubes.fits"
cube = SpectralCube.read(cubefile)*10000 # Convert to FLAM20
varcube = (SpectralCube.read(varfile)*10000)**2 # The file actually has std.dev. So I'm convertting it to variance
cube
# Show wavelength info
cube.spectral_axis
# Show WCS
cube.wcs
# Celestial WCS
spat_wcs = cube.wcs.celestial
spat_wcs
# Spatial slice
testslice = cube[100,:, :]
fig, ax = imshow_astro(testslice.value, spat_wcs, cblabel = flam)
ax. set_title("Cube slice")
plt.show()
# Spectral ray
testray = cube[:, 50, 50]
plt.figure(figsize = (10,6))
plt.step(cube.spectral_axis, testray)
plt.xlabel(r"$\lambda~(\rm\AA)$")
plt.ylabel(flam)
#plt.xlim(5574, 5583)
plt.show()
Explanation: Using SpectralCube
Read https://spectral-cube.readthedocs.io/en/latest/index.html for the latest documentation. We present a very quick introduction to a limited set of relevant functions here.
End of explanation
# Mask out bad wavelengths
badwave = (cube.spectral_axis > 5574*u.AA)&(cube.spectral_axis <5583*u.AA)
# Tile wavelength mask to fill the cube. This step is required because
# SpectralCube can't yet figure out how to tile a 1D mask.
goodwavemask = np.tile(~badwave, (cube.shape[2],cube.shape[1],1)).T
# Mask cubes
cube = cube.with_mask(goodwavemask)
varcube = varcube.with_mask(goodwavemask)
# write to files
if not os.path.isdir("/home/sunil/Desktop/FRB/190608/KCWI/masked"):
    os.mkdir("/home/sunil/Desktop/FRB/190608/KCWI/masked")
cube.write("/home/sunil/Desktop/FRB/190608/KCWI/masked/HG190608_5_icubes_masked.fits", overwrite=True)
varcube.write("/home/sunil/Desktop/FRB/190608/KCWI/masked/HG190608_5_vQubes_masked.fits", overwrite=True)
# Check masking:
testray = cube[:, 50, 50]
plt.figure(figsize = (10,6))
plt.step(cube.spectral_axis, testray)
plt.xlabel(r"$\lambda~(\rm\AA)$")
plt.ylabel(flam)
#plt.xlim(5574, 5583)
plt.show()
Explanation: From this spectrum, it looks like there are some bad wavelengths near 5500 A. I will mask this out before proceeding.
Masking
Further inspection reveals the bad wavelengths are between 5575 A and 5582 A. To mask out a chunk of the cube, one must set the boolean values corresponding to that section to False.
End of explanation
help(kcwi.get_img)
img = kcwi.get_img(cubefile)
fig, ax = imshow_astro(img.value, img.wcs, cblabel = flam2)
ax.set_title("White light image")
plt.show()
varimg = kcwi.get_img(varfile)
fig, ax = imshow_astro(varimg.value, varimg.wcs, cblabel=flam2)
ax.set_title("White light image variance")
plt.show()
Explanation: Create a whitelight image
kcwi.get_img can produce a projection along the spectral axis. This projection can be simply between two wavelength limits or using a more complicated transmission curve (i.e. to produce synthetic photometry). Here, we only show the white light image.
End of explanation
regfile = "/home/sunil/Desktop/FRB/190608/KCWI/whitelight/test.reg"
reg = pyreg.open(regfile).as_imagecoord(img.header)
mask = reg.get_filter().mask(img.data)
mask.shape
Explanation: Extract spectrum using a spatial mask
One can use a 2D mask to extract a spectrum within that region. An intuitive way to do this would be to use PyRegion and create a mask from a ds9.reg file.
End of explanation
mask.dtype
mask_vis = np.where(mask, 1, 0)
fig, ax = imshow_astro(img.value, img.wcs, colorbar = False)
im2 = ax.imshow(mask_vis, cmap ="Greys_r", alpha = 0.5)
plt.show()
# Get spectrum from the masked region
help(kcwi.spec_from_mask)
spec, varspec = kcwi.spec_from_mask(cube,mask,varcube, kind = "mean")
plt.figure(figsize = (10,6))
plt.plot(cube.spectral_axis, spec, label = "Specific intensity")
plt.plot(cube.spectral_axis, np.sqrt(varspec.data), label = "Std.dev.")
plt.legend()
plt.xlabel(r"$\lambda (\rm \AA)$")
plt.ylabel(r"$F_\lambda$ $\rm (10^{-20}\ erg\ s^{-1}\ cm^{-2}\ \AA^{-1})$")
plt.ylim(-5, 10)
plt.show()
Explanation: It is crucial that the masks are always boolean. Using integer or float masks will make the process very slow.
End of explanation
help(kcwi.spec_from_ellipse)
# Visualize elliptical aperture
aperpos_sky = SkyCoord("22h16m3.18s", "-7d53m35s")
a, b = 3*u.arcsec, 2*u.arcsec
rot = -60*u.deg
skyaper = SkyEllipticalAperture(aperpos_sky, a, b, rot)
aper = skyaper.to_pixel(img.wcs)
fig, ax = imshow_astro(img.value, img.wcs, colorbar = False)
aper.plot(ax, color="w", lw =2, ls="--")
plt.show()
# How is aper defined?
aper
# Get spectrum using aper
x0, y0 = aper.positions
spec, varspec =  kcwi.spec_from_ellipse(cube, varcube, x0, y0 , aper.a, aper.b, aper.theta)
plt.figure(figsize = (10,6))
plt.plot(cube.spectral_axis, spec.data, label = "Specific intensity")
plt.plot(cube.spectral_axis, np.sqrt(varspec.data), label = "Std.dev.")
plt.legend()
plt.xlabel(r"$\lambda (\rm \AA)$")
plt.ylabel(r"$F_\lambda$ $\rm (10^{-20}\ erg\ s^{-1}\ cm^{-2}\ \AA^{-1})$")
plt.ylim(-5, 10)
plt.show()
Explanation: Get spectra from elliptical masks
This is a thin wrapper around the spec_from_mask function because ellipses are so ubiquitous.
End of explanation
help(kcwi.find_sources)
imgfile = "/home/sunil/Desktop/FRB/190608/KCWI/whitelight/HG190608_5_white.fits"
regfile = "/home/sunil/Desktop/FRB/190608/KCWI/whitelight/ds9.reg"
img = kcwi.get_img(cubefile)
img.write(imgfile, overwrite=True)
# Only the light region is accepted for analysis
reg = pyreg.open(regfile).as_imagecoord(img.header)
mask = reg.get_filter().mask(img.data)
mask_vis = np.where(mask, 0, 1)
fig, ax = imshow_astro(img.value, img.wcs, colorbar = False)
im2 = ax.imshow(mask_vis, cmap ="Blues_r", alpha = 0.6)
plt.show()
objects, segmap = kcwi.find_sources(imgfile,regfile = regfile, nsig=1.5, minarea=5)
objects
fig,ax = imshow_astro(segmap, img.wcs, colorbar=False, vrange= (0,1))
ax.set_title("Segmentation map from sep.extract")
plt.show()
# Visualise the ellipses that extract has returned.
fig, ax = imshow_astro(img.value, img.wcs, colorbar = False)
for obj in objects:
    aper = EllipticalAperture((obj['x'], obj['y']), 2*obj['a'], 2*obj['b'], obj['theta'])
    aper.plot(ax, color="w", lw=2, ls="--")
plt.show()
Explanation: Find sources in an image
This is a wrapper for SExtractor's extract function. I'm passing a ds9 region file to exclude the edges from analysis.
End of explanation
reload(kcwi)
help(kcwi.get_source_spectra)
reload(kcwi)
cubefile = "/home/sunil/Desktop/FRB/190608/KCWI/masked/HG190608_5_icubes_masked.fits"
varcubefile = "/home/sunil/Desktop/FRB/190608/KCWI/masked/HG190608_5_vQubes_masked.fits"
outdir = "/home/sunil/Desktop/FRB/190608/KCWI/HG_5_spectra/"
marzfile = "HG_5_marz.fits"
speclist, varspeclist, wave = kcwi.get_source_spectra(cubefile, varfile, objects,outdir = outdir, marzfile=marzfile)
# All outputs
sorted(glob.glob(outdir+"*.fits"))
# Look at the contents of the MARZ file
marzhdus = fits.open("/home/sunil/Desktop/FRB/190608/KCWI/HG_5_spectra/HG_5_marz.fits")
marzhdus.info()
Explanation: Extract spectra and create a MARZ input file
The final step is to get spectra of the identified sources and to create a MARZ input file.
End of explanation
<END_TASK>