Quickstart: Accessing OpenUniverse2024 Data

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

import numpy as np
import s3fs
from matplotlib import pyplot as plt
import pyarrow.fs
import pyarrow.parquet as pq
import hpgeom
import json
from astroquery.ipac.irsa import Irsa
from astropy.coordinates import SkyCoord
from astropy import units as u
from astropy.io import fits

# Create an anonymous (public read-only) connection to the NASA IRSA S3 bucket.
s3 = s3fs.S3FileSystem(anon=True)

# Top-level path components
BUCKET_NAME = "nasa-irsa-simulations"
OU_PREFIX = "openuniverse2024"
ROMAN_TDS_PREFIX = "roman/full/RomanTDS/images/simple_model"

# Pick one band to explore
BAND = "J129"
band_directory = f"{BUCKET_NAME}/{OU_PREFIX}/{ROMAN_TDS_PREFIX}/{BAND}"

# List all pointings available for the chosen band
all_pointings = [p.split("/")[-1] for p in s3.ls(band_directory)]
print(f"Found {len(all_pointings)} pointings in band {BAND}:")
print(all_pointings[:10], "...")

# Select one pointing and list the files it contains
POINTING = "10190"
image_directory = f"{band_directory}/{POINTING}"

files = [f"s3://{f}" for f in s3.ls(image_directory)]
print(f"Found {len(files)} files in pointing {POINTING}")

# Open one FITS file and inspect its extensions
fname = files[0]
with fits.open(fname, use_fsspec=True, fsspec_kwargs={"anon": True}, memmap=False) as hdul:
    print(f"File: {fname}")
    print(f"Number of extensions: {len(hdul)}\n")
    hdul.info()

def show_gallery(files, max_images=9):
    """
    Display a gallery of FITS images.

    Parameters
    ----------
    files : list of str
        List of S3 URIs to FITS files.
    max_images : int, optional
        Maximum number of images to display (default: 9).
    """
    n_images = min(len(files), max_images)
    ncols = n_images if n_images < 4 else 3
    nrows = (n_images + ncols - 1) // ncols

    fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 4 * nrows))
    axes = np.atleast_1d(axes).ravel()

    for i, f in enumerate(files[:n_images]):
        with fits.open(f, fsspec_kwargs={"anon": True}, memmap=False) as hdul:
            data = hdul[1].data
            vmin, vmax = np.nanpercentile(data, [5, 99])
            axes[i].imshow(data, origin="lower", cmap="gray", vmin=vmin, vmax=vmax)
            axes[i].set_title(f.split("/")[-1], fontsize=8)
            axes[i].axis("off")

    for j in range(i + 1, len(axes)):
        axes[j].axis("off")

    plt.tight_layout()
    plt.show()

# Display up to 3 images from the selected directory.
show_gallery(files, max_images=3)

# The Roman Time-Domain Survey is centered near the LSST ELAIS-S1 Deep Drilling Field.
ra = 9.45
dec = -44.02

# Convert sky coordinates to a HEALPix region index (nside=32, RING ordering)
nside = 32
region = hpgeom.angle_to_pixel(nside, ra, dec, lonlat=True, nest=False)
print(f"HEALPix region for RA={ra}, Dec={dec}: {region}")

# Build the S3 paths for this region's catalog files
CATALOG_NAME = "roman_rubin_cats_v1.1.2_faint"
catalog_prefix = f"{BUCKET_NAME}/{OU_PREFIX}/roman/full/{CATALOG_NAME}"

snana_path    = f"{catalog_prefix}/snana_{region}.parquet"
galaxy_path   = f"{catalog_prefix}/galaxy_{region}.parquet"
gal_flux_path = f"{catalog_prefix}/galaxy_flux_{region}.parquet"

print("SNANA file:       ", snana_path)
print("Galaxy info file: ", galaxy_path)
print("Galaxy flux file: ", gal_flux_path)

def inspect_parquet_files(s3_path, *, region='us-east-1'):
    """
    Print the structure of a Parquet file on S3 without reading its data.

    Reads only the Parquet metadata footer (row count, column names and types),
    which is fast regardless of file size.

    Parameters
    ----------
    s3_path : str
        S3 path to the Parquet file (without the s3:// prefix).
    """
    fs = pyarrow.fs.S3FileSystem(region=region, anonymous=True)
    meta = pq.read_metadata(s3_path, filesystem=fs)
    schema = pq.read_schema(s3_path, filesystem=fs)

    print(f"Rows: {meta.num_rows}  |  Columns: {len(schema.names)}")
    print("\nColumn names:")
    for name in schema.names:
        print("  ", name)

inspect_parquet_files(snana_path)

# Read just the model_name column to see what transient types are in this region
fs = pyarrow.fs.S3FileSystem(region='us-east-1', anonymous=True)
model_names = pq.read_table(snana_path, filesystem=fs, columns=["model_name"]).to_pandas()
model_names["model_name"].unique()

inspect_parquet_files(galaxy_path)

inspect_parquet_files(gal_flux_path)

fs = pyarrow.fs.S3FileSystem(region='us-east-1', anonymous=True)

# Read the full SNANA catalog
df_snana = pq.read_table(snana_path, filesystem=fs).to_pandas()

# Pick one transient — here we grab the first row as an example
example_transient = df_snana.iloc[0]
print("Example transient:")
print(example_transient[["model_name", "host_id", "start_mjd", "end_mjd"]])

# Look up its host galaxy by matching host_id (SNANA) to galaxy_id (galaxy info file)
host = pq.read_table(
    galaxy_path,
    filesystem=fs,
    filters=[("galaxy_id", "==", example_transient["host_id"])]
).to_pandas()

print("\nHost galaxy info:")
host

# Point the astroquery IRSA client to the correct locations.
Irsa.sia_url = "https://irsa.ipac.caltech.edu/simulated/SIA"
Irsa.tap_url = "https://irsa.ipac.caltech.edu/simulated/TAP"

# List all available simulated image collections
Irsa.list_collections(servicetype='SIA')

# Collection names for OpenUniverse2024
OU_ROMAN_SIA_COLLECTION = 'simulated_roman_openuniverse2024'
OU_RUBIN_SIA_COLLECTION = 'simulated_rubin_openuniverse2024'

def get_s3_fpath(cloud_access):
    """Extract the S3 URI from the cloud_access JSON string in an SIA result."""
    cloud_info = json.loads(cloud_access)
    bucket_name = cloud_info['aws']['bucket_name']
    key = cloud_info['aws']['key']
    return f's3://{bucket_name}/{key}'

# Use the host galaxy position from Section 2 (or set any RA/Dec you want to query).
host_ra  = float(host.iloc[0]["ra"])
host_dec = float(host.iloc[0]["dec"])
search_radius = 1 * u.arcsec  # small radius: we just need images that contain this point

#convert ra, dec to SkyCoords for ease of use
coords = SkyCoord(host_ra, host_dec, unit='deg')

# Query Roman TDS images in the J129 band
sia_results = Irsa.query_sia(pos=(coords, search_radius.to(u.deg)),
                             collection=OU_ROMAN_SIA_COLLECTION)

# We first choose to look at the J129 band and the simple_model images
bandname = "J129"
roman_images = sia_results[
    ['TDS_simple_model' in r['obs_id'] and bandname in r['energy_bandpassname']
     for r in sia_results]
]
roman_images['s3_uri'] = [get_s3_fpath(r['cloud_access']) for r in roman_images]

print(f"Found {len(roman_images)} Roman {bandname} images at RA={host_ra:.4f}, Dec={host_dec:.4f}")
roman_images['obs_id', 't_min', 't_max', 's3_uri']

# The same search works for Rubin images — just swap the collection name and band filter.
# Unlike the Roman collection, the Rubin collection contains only one image type (calexp),
# so no obs_id filter is needed beyond selecting the desired band.
rubin_band = "r"
rubin_results = Irsa.query_sia(pos=(coords, search_radius.to(u.deg)),
                               collection=OU_RUBIN_SIA_COLLECTION)

rubin_images = rubin_results[
    [rubin_band in r['energy_bandpassname'] for r in rubin_results]
]
rubin_images['s3_uri'] = [get_s3_fpath(r['cloud_access']) for r in rubin_images]

print(f"Found {len(rubin_images)} Rubin {rubin_band}-band images at RA={host_ra:.4f}, Dec={host_dec:.4f}")
rubin_images['obs_id', 't_min', 't_max', 's3_uri']