IRIS IFU: Field of stars

# 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, convolve_spectrum, get_stellar_spectrum, make_point_source_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

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

# Atmosphere, instrument, and exposure params
mode = 'lenslet' # img, lenslet, slicer
filt = 'JN3'
resolution = 4000
bandpass = 0.05
size = (128, 128)
read_noise = 9 # e- RMS
num_detector_pixels = 3 # How many detector pixels correspond to a single voxel
dark_current = 0.025 # e-/sec
scale = 0.004 # arcsec / pixel
itime = 300 # 5 min
n_frames = 1 # 1 coadd
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)

# Rerieve phoenix stellar spectrum sampled on wave grid, normalize to 1.
template_spec = get_stellar_spectrum(
    T_eff=5500, log_g=4.5,
    wave_out=wave,
    flux=1.0
)

# Convolve with spectral LSF
template_spec = convolve_spectrum(wave, template_spec, resolution)

# Create field of stars
n_stars = 10
mag_range = (24, 26)
source_rate = np.zeros((size[0], size[1], len(wave)), dtype=np.float32)
for i in range(n_stars):
    mag = np.random.uniform(mag_range[0], mag_range[1])
    x = int(np.random.uniform(5, size[0] - 5))
    y = int(np.random.uniform(5, size[1] - 5))
    print(f"Star {i+1} / {n_stars}: {mag=}, x={x}, y={y}")
    flux = filter_data['zpphot'] * 10**(-0.4 * mag) # photons / sec / m^2
    make_point_source_cube(
        y, x,
        wave, template_spec,
        flux=flux,
        resolution=resolution,
        psf=psf, out=source_rate,
    )

# 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
)

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


# Initialize meta dict to match DataModel interface
meta = {
    'target.name': 'Stars',
    'target.airmass_start' : airmass,
    'target.ra' : 0.0,
    'target.dec' : 0.0,
    'exposure.exposure_type' : 'SCI',
    'exposure.exposure_time' : itime,
    'exposure.nframes' : n_frames,
    'exposure.jd_start' : 2460577.5,
    'instrument.detector' : 'IFU',
    'instrument.name' : 'IRIS',
    'instrument.mode' : 'lenslet',
    'instrument.filter' : filt,
    'instrument.scale' : scale,
}

# Populate remaining meta fields
add_meta(meta, size)

# 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)