IRIS IFU: z~1 spiral galaxy in H-alpha emission

# Imports
import liger_iris_sim
from liger_iris_sim.loaders import load_filter_data, load_iris_psf
from liger_iris_sim.sources import bin_image, resample_convolve_cube
from liger_iris_sim.expose import compute_iris_tput, expose_ifu
from liger_iris_sim.utils import add_meta
from liger_iris_sim.sky import get_maunakea_spectral_sky_emission, get_maunakea_spectral_sky_transmission
from liger_iris_sim.utils import generate_wave_grid

import os
import pickle
import numpy as np
from scipy.constants import c, h
from astropy.io import fits

# Output directory
out_dir = os.getcwd() + '/results/'
os.makedirs(out_dir, exist_ok=True)

# Load Source
gal_scale = 0.025 # arcsec / pixel
with fits.open('id11978_Halpha_cube_wl.fits') as f:
    gal_wave = f[0].data # microns
with fits.open('id11978_Halpha_cube.fits') as f:
    # units = erg / s / cm^2 / pixel = erg / (s * cm^2 * pixel)
    Ephot = (h * c / (1.3 * 1E-6)) * 1E7 # erg / photon
    gal_cube = f[0].data / Ephot * 100**2 # photons / (s * m^2 * pixel)
    gal_cube = gal_cube[30:-30, 30:-30, :] # crop
    gal_cube = np.clip(gal_cube, 0, None)


# Atmosphere, instrument, and exposure params
mode = 'slicer' # img, lenslet, slicer
filt = 'JN3'
grating = 'J4000'
resolution = 4000
bandpass = 0.05
scale = 0.025 # arcsec / pixel
size = (
    int(gal_scale / scale * gal_cube.shape[0]),
    int(gal_scale / scale * gal_cube.shape[1])
)
read_noise = 9 # e- RMS
num_detector_pixels = 3 # How many detector pixels correspond to a single voxel
dark_current = 0.025 # e-/sec
itime = 900 # 15 min
n_frames = 8 # 2 hours
airmass = 1.4
tel_tput = 0.91 # Default
ao_tput = 0.8 # Default
filt_tput = 0.9 # Default
collarea = liger_iris_sim.tmt_collarea # m^2

# Load filter data
filter_data = load_filter_data(filt)

# Calculate total throughput
tput = compute_iris_tput(mode=mode, wavelength=filter_data['wavecenter'], tel=tel_tput, ao=ao_tput, filt=filt_tput)

# Load on axis PSF at this wavelength and bin to match pixel scale
psf, psf_info = load_iris_psf(mode, filter_data['wavecenter'], xs=0, ys=0)
psf = bin_image(psf, scale_in=psf_info['psf_sampling'], scale_out=scale, fix_even=True)

# NQ sampled wave grid for this filter
wave = generate_wave_grid(filter_data['wavecenter'], bandpass, resolution)

# Sky emission
sky_em = get_maunakea_spectral_sky_emission(
    wave,
    resolution=resolution,
    T_tel=275, T_atm=258, T_aos=243, T_zod=5800,
    Em_tel=0.09, Em_atm=0.2, Em_aos=0.01, Em_zod=1.47E-12,
)

# Integrate over sky pixel
sky_emission_rate = sky_em['sky_emission'] * scale**2 # Integrate over pixel: photons / (s * m^2 * wavebin)

# Sky transmission
sky_transmission = get_maunakea_spectral_sky_transmission(
    wave,
    resolution=resolution,
    airmass=1.4
)

# Resample onto IRIS grid
gal_cube_resampled = resample_convolve_cube(
    wave_in=gal_wave, cube_in=gal_cube, scale_in=gal_scale,
    wave_out=wave,
    scale_out=scale,
    size_out=size,
    resolution=resolution,
    psf=psf
)

# Expose
sim = expose_ifu(
    gal_cube_resampled,
    itime=itime, n_frames=n_frames, collarea=collarea,
    sky_emission_rate=sky_emission_rate,
    sky_transmission=sky_transmission, num_detector_pixels=3,
    tput=tput, read_noise=read_noise, dark_current=dark_current
)

snr_tot = np.nanquantile(np.sqrt(np.nansum(sim['snr']**2, axis=0)), 0.995)
print(f"99.5th percentile SNR: {snr_tot}")

# Target info
meta = {
    'model_type' : 'IFUImageModel',
    'target.name': 'Z1Galaxy',
    'target.ra' : 0,
    'target.dec' : 0,
    'target.airmass_start' : 1.4,
    'exposure.exposure_time' : itime,
    'exposure.nframes' : n_frames,
    'exposure.jd_start' : 2460577.5,
    'instrument.detector' : 'IFU',
    'instrument.name' : 'IRIS',
    'instrument.grating' : grating,
    'instrument.mode' : 'IFU',
    'instrument.ifumode' : 'slicer',
    'instrument.filter' : filt,
    'instrument.scale' : scale,
}
add_meta(meta, size)

# Save results
sim['sky_emission_rate'] = sky_emission_rate
sim['sky_transmission'] = sky_transmission
sim['psf'] = psf

# Save results into pickle
sim['psf_info'] = psf_info
filename = out_dir + f"SIM_{meta['instrument.name']}_{meta['instrument.detector']}_{meta['target.name']}_{meta['instrument.filter']}_{int(meta['instrument.scale']*1000)}mas.pkl"
print(f"Saving {filename}")
with open(filename, 'wb') as f:
    pickle.dump(dict(sim=sim, meta=meta), f)