aws-lambda-photometric-pipeline.py

import astropy.units as u
from astropy.io import fits
from astropy.coordinates import SkyCoord
from astropy.modeling.fitting import LevMarLSQFitter
from astropy.stats import gaussian_sigma_to_fwhm
from astropy.table import Table
import boto3
import io
from photutils.background import MADStdBackgroundRMS, MMMBackground
from photutils.detection import IRAFStarFinder
from photutils.psf import (
        DAOGroup, FittableImageModel, IterativelySubtractedPSFPhotometry
)
import roman_datamodels as rdm
import urllib.parse

verbose = True


def load_s3_object(bucket_name, key_name):
    s3_client = boto3.client('s3')
    response = s3_client.get_object(
        Bucket=bucket_name,
        Key=key_name
    )

    content = io.BytesIO(response["Body"].read())

    return content

def run_photutils_pipeline(l2_image, epsf_image, sigma_psf=1.0, init_guesses=None):
    """Photutils pipeline, based heavily on the guide at
    https://photutils.readthedocs.io/en/stable/psf.html#basic-usage
    If initial guesses are provided then hopefully this runs much
    more quickly.
    sigma_psf is set to 1.0, which is close enough for F158.
    Note that l2_image cannot have units attached, at least for now.
    """

    background_rms = MADStdBackgroundRMS()
    std = background_rms(l2_image)

    finder = IRAFStarFinder(
        threshold=50*std,
        fwhm=sigma_psf * gaussian_sigma_to_fwhm,
        minsep_fwhm=0.1,
        roundhi=2.0,
        roundlo=-2.0,
    )

    daogroup = DAOGroup(2.0 * sigma_psf * gaussian_sigma_to_fwhm)

    mmm_bkg = MMMBackground()

    fitter = LevMarLSQFitter()

    psf_model = FittableImageModel(epsf_image)

    photometry = IterativelySubtractedPSFPhotometry(
        finder=finder,
        group_maker=daogroup,
        bkg_estimator=mmm_bkg,
        psf_model=psf_model,
        fitter=LevMarLSQFitter(),
        niters=1,
        fitshape=(5, 5)
    )

    result_table = photometry(image=l2_image, init_guesses=init_guesses)

    return result_table

def initialize_guesses(gsc, wcs, flux_threshold=None):
    """Use guide star catalog for initial guesses on photometry step. Need to
    convert catalog to x and y coordinates in the l2 image frame.
    """

    if flux_threshold is not None:
        gsc = gsc[gsc["f814w_acswfc_flux"] > flux_threshold].copy()

    in_image = wcs.in_image(gsc["ra"], gsc["dec"])
    gsc = gsc[in_image]

    x0, y0 = wcs.invert(gsc["ra"], gsc["dec"])

    init_guesses = Table([x0, y0], names=["x_0", "y_0"])

    return init_guesses

def solve_offsets(phot, gsc, wcs, max_sep=1.0*u.arcsec, verbose=True):
    """Now uses guide star catalog (i.e. init_guesses) to
    crossmatch the ePSF photometry table. The wcs is again
    needed to convert between detector and sky coordinates.
    Returns the RMS offsets and number of offsets.
    """

    phot_coords = SkyCoord(*wcs(phot["x_0"], phot["y_0"]), unit="deg")
    gsc_coords = SkyCoord(*wcs(gsc["x_0"], gsc["y_0"]), unit="deg")

    idx, sep, _ = gsc_coords.match_to_catalog_sky(phot_coords)

    # phot_x_gsc_coords = phot_coords[sep < max_sep]
    # gsc_x_phot_coords = gsc_coords[idx[sep < max_sep]]

    print(f"Matches found for {sum(sep < max_sep)} of {len(sep)} sources.")

    return {"offset": (sep**2).mean()**0.5, "number": len(sep)}


def handler(event, context):

    # sca = "01"

    bucket_name = event['Records'][0]['s3']['bucket']['name']
    key_name = urllib.parse.unquote_plus(
        event['Records'][0]['s3']['object']['key'],
        encoding='utf-8'
    )

    sca = key_name[35:37]

    if verbose:
        print(f"New L2 image at {key_name}.")


    # load L2 image
    l2_image_content = load_s3_object(
        bucket_name=bucket_name,
        key_name=key_name,
    )
    af = rdm.open(l2_image_content)
    l2_image = af.data.value
    wcs = af.meta.wcs

    if verbose:
        print(f"Roman WFI L2 image (SCA {sca}) loaded from S3.")

    # load ePSF
    epsf_image_content = load_s3_object(
        bucket_name="wfi-epsf-library",
        key_name=f"F158/SCA{sca}_psflibipc.fits"
    )
    epsf_image = fits.getdata(epsf_image_content)

    if verbose:
        print(f"ePSF image loaded (SCA {sca}).")

    # load gsc and initialize guesses
    gsc_content = load_s3_object(
        bucket_name="roman-gsc",
        key_name="gaia_catalog.fits"
    )
    gsc = Table.read(gsc_content)

    if verbose:
        print("Guide star catalog loaded.")

    init_guesses = initialize_guesses(gsc, wcs)

    if verbose:
        print("Using guide stars to initialize photometry.")
    photometric_table = run_photutils_pipeline(l2_image, epsf_image, init_guesses=init_guesses)

    #photometric_table.write(f"WFI{sca}_photometry.ecsv", overwrite=True)
    #if verbose:
    #    print(f"Photometric table written to `WFI{sca}_photometry.ecsv`")

    # astrometry step
    astrom_info = solve_offsets(photometric_table, init_guesses, wcs, verbose=verbose)
    offset = astrom_info["offset"]
    n_sources = astrom_info["number"]

    # remove units -- imply arcsec -- can't be serialized
    offset = offset.to(u.arcsec).value

    if verbose:
        print(f"RMS offset is {offset:.4f} arcsec.")

    return {
        "body": {
            "bucket_name": bucket_name,
            "file_name": key_name,
            "astrometric_offset": float(offset),
            "number_of_sources": n_sources
        },
        "statusCode": 200,
    }