IRIS Imager: 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, make_point_source_image
from liger_iris_sim.expose import compute_iris_tput, expose_imager
from liger_iris_sim.utils import add_meta

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 = 'J'
size = (4096, 4096)
read_noise = 9 # e- RMS
dark_current = 0.025 # e-/sec
scale = 0.004 # arcsec / pixel
itime = 300 # 5 min
n_frames = 1 # 1 coadd
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)

# Create field of stars
n_stars = 10
mag_range = (24, 26)
source_rate = np.zeros(size, 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=}, {y=}")
    flux = filter_data['zpphot'] * 10**(-0.4 * mag) # photons / sec / m^2
    make_point_source_image(
        y, x, psf,
        flux=flux,
        out=source_rate,
    )

# Total sky
sky_emission_rate = filter_data['zpphot'] * 10**(-0.4 * filter_data['backmag']) # photons / sec / arcsec^2 / m^2
sky_emission_rate *= scale**2  # Integrate over 2D pixel (photons / sec / m^2)

# Expose
sim = expose_imager(
    source_rate,
    itime=itime, n_frames=n_frames, collarea=collarea,
    sky_emission_rate=sky_emission_rate,
    tput=tput, read_noise=read_noise, dark_current=dark_current,
)

# Initialize meta dict to match DataModel interface
meta = {
    'target.name': 'Stars',
    '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' : 'IMG1',
    '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)