Following a simulated supernova through the Roman Time Domain Survey

# Uncomment the next line to install dependencies if needed.
# !pip install astropy matplotlib numpy pandas pyarrow s3fs scipy astroquery hpgeom ipython

# Import modules
import warnings
import json

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import hpgeom
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.nddata import Cutout2D
from astropy.nddata.utils import NoOverlapError
from astropy.wcs import WCS, FITSFixedWarning
from astroquery.ipac.irsa import Irsa
from matplotlib import animation
from scipy.ndimage import rotate
from IPython.display import HTML

# Needed to access data in the cloud
import s3fs
s3 = s3fs.S3FileSystem(anon=True)  # create an S3 client

# Filter out the FITSFixedWarning, which is consequenceless and gets thrown every time you deal with a WCS
# in a Roman openuniverse simulated image using astropy.
warnings.simplefilter('ignore', category=FITSFixedWarning)

# Point the astroquery IRSA client at the simulated-data services, which are
# separate from the ones serving IRSA's observed data.
Irsa.sia_url = "https://irsa.ipac.caltech.edu/simulated/SIA"
Irsa.tap_url = "https://irsa.ipac.caltech.edu/simulated/TAP"

OU_ROMAN_SIA_COLLECTION = 'simulated_roman_openuniverse2024'

# Read in the (simulated) Observation Sequence File.

BUCKET_NAME = 'nasa-irsa-simulations'
ROMAN_PREFIX = 'openuniverse2024/roman/full'

ROMAN_TDS_PATH = f'{ROMAN_PREFIX}/RomanTDS'
FILENAME = 'Roman_TDS_obseq_11_6_23.fits'
OBSEQ_PATH = f's3://{BUCKET_NAME}/{ROMAN_TDS_PATH}/{FILENAME}'

obseq_hdu = fits.open(OBSEQ_PATH, fsspec_kwargs={"anon": True})
obseq = pd.DataFrame(obseq_hdu[1].data)

print(obseq)

# Find the ranges of RA, Dec, and date listed in the observation sequence file.

ra_min, dec_min = obseq[['ra','dec']].min()
ra_max, dec_max = obseq[['ra','dec']].max()
mjd_min = obseq['date'].min()
mjd_max = obseq['date'].max()

print("ra_min, ra_max:", ra_min, ra_max)
print("mjd_min, mjd_max:", mjd_min, mjd_max)

# The Roman Time-Domain Survey is centered near the LSST ELAIS-S1 Deep Drilling Field.
region = hpgeom.angle_to_pixel(32, 9.45, -44.02, lonlat=True, nest=False)

parquet_file = f's3://{BUCKET_NAME}/{ROMAN_PREFIX}/roman_rubin_cats_v1.1.2_faint/snana_{region}.parquet'
transients = pd.read_parquet(parquet_file, filesystem=s3)

#List the unique models in the SNANA file.
unique_models = pd.Series(transients['model_name']).drop_duplicates().tolist()
unique_models

# Most of the models are non SNIa (NON1ASED).
# Choose only the SNIa
sn1a = transients[transients['model_name'] == 'SALT3.NIR_WAVEEXT'] # SNe Ia only.
print('Number of SN1a in SNANA file: ', len(sn1a))

# Choose the SNIa that overlap with the spatial and temporal extent of the survey.
ra_mask = np.logical_and(sn1a['ra'] > ra_min, sn1a['ra'] < ra_max)
dec_mask = np.logical_and(sn1a['dec'] > dec_min, sn1a['dec'] < dec_max)
mjd_mask = np.logical_and(sn1a['start_mjd'] > mjd_min, sn1a['end_mjd'] < mjd_max)
all_mask = np.logical_and.reduce((ra_mask,dec_mask,mjd_mask))
covered_sn1a = sn1a[all_mask]
print('Number of SNIa within the survey:', len(covered_sn1a))

# Choose the SNIa that are nearby, at redshifts less than 0.5, so they are bright enough
# to stand out clearly against their host galaxy.
nearby_sn1a = covered_sn1a[covered_sn1a['z_CMB'] < 0.5]
print('Number of nearby SNIa:', len(nearby_sn1a))

# Let's choose SN 20131477, which peaks while its field is being visited regularly.
oid = 20131477
chosen_object = nearby_sn1a[nearby_sn1a['id'] == oid].iloc[0]

ra, dec = chosen_object['ra'], chosen_object['dec']
peak_mjd = chosen_object['peak_mjd']
coord = SkyCoord(ra*u.deg, dec*u.deg)

print(f"SN {oid}: RA={ra:.6f}, Dec={dec:.6f}, z={chosen_object['z_CMB']:.3f}, peaks at MJD {peak_mjd:.1f}")

sia_results = Irsa.query_sia(pos=(coord, 1 * u.arcsec),
                             collection=OU_ROMAN_SIA_COLLECTION)

print(f"Images covering this position: {len(sia_results)}")

band = 'R062'

is_tds = np.array(['TDS_simple_model' in str(obs_id) for obs_id in sia_results['obs_id']])
is_band = np.char.strip(np.array(sia_results['energy_bandpassname'], dtype=str)) == band

instances = sia_results[is_tds & is_band]
instances.sort('t_min')

print(f"{band} images covering SN {oid}: {len(instances)}")

epoch_mjd = np.asarray(instances['t_min'], dtype=float)
in_window = (epoch_mjd >= peak_mjd - 40) & (epoch_mjd <= peak_mjd + 80)

instances = instances[in_window]
epoch_mjd = epoch_mjd[in_window]

print(f"Epochs to animate: {len(instances)}, "
      f"spanning MJD {epoch_mjd.min():.1f} to {epoch_mjd.max():.1f}")

image_paths = [get_s3_fpath(row['cloud_access']) for row in instances]
image_paths[:3]

#Make the cutouts; this will take a few minutes.
stamps = []
mjd = []
for imgpath, epoch in zip(image_paths, epoch_mjd):
    print(imgpath)
    with fits.open(imgpath, fsspec_kwargs={"anon": True}) as hdu:
        img = hdu[1].data
        header = hdu[0].header
        wcs = WCS(header)
        x, y = wcs.world_to_pixel(coord)

        # Manually rotate the images so they are all aligned.
        CDmat = np.array([header['CD1_1'], header['CD1_2'],
                          header['CD2_1'], header['CD2_2']]).reshape(2,2)

        orientation = header['ORIENTAT']

        # These chips are "flipped".
        if header['SCA_NUM'] % 3 == 0:
            orientation += 180

        # Build rotation matrix.
        CD1_1_rot = np.cos(-orientation*np.pi/180)
        CD1_2_rot = -np.sin(-orientation*np.pi/180)
        CD2_1_rot = np.sin(-orientation*np.pi/180)
        CD2_2_rot = np.cos(-orientation*np.pi/180)

        RotMat = np.array([CD1_1_rot, CD1_2_rot,
                          CD2_1_rot, CD2_2_rot]).reshape(2,2)

        RotMat_inv = np.array([CD1_1_rot, -CD1_2_rot,
                              -CD2_1_rot, CD2_2_rot]).reshape(2,2)

        # Apply rotation to the CDi_j header keywords.
        CDmat_rot = np.dot(CDmat,RotMat_inv)

        # Update header.
        header['CD1_1'], header['CD1_2'] = CDmat_rot[0]
        header['CD2_1'], header['CD2_2'] = CDmat_rot[1]
        header['ORIENTAT'] -= orientation

        # Rotate the image.
        rot_img = rotate(img,angle=orientation,reshape=False,cval=np.nan)

        rot_wcs = WCS(header)

        try:
            # Make cutout around SN Ia location.
            cutout = Cutout2D(rot_img,coord,100,wcs=rot_wcs,mode='partial')
            stamps.append(cutout.data)
            mjd.append(epoch)
        except NoOverlapError:
            pass

print(f"Collected {len(stamps)} cutouts")

savepath = f'SN{oid}.gif'
savepath
anim = animate_stamps(stamps, savepath, labels=[f'MJD {m:.1f}' for m in mjd])

HTML(anim.to_jshtml(default_mode='once'))