New Feature: ASDF Cutout Functionality

astrocut/asdf_cutouts.py

@@ -0,0 +1,57 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+
+"""This module implements cutout functionality similar to fitscut, but for the ASDF file format."""
+
+import asdf
+import astropy
+
+
+def get_center_pixel(file, ra, dec):
+    """ Get the center pixel from a roman 2d science image """
+
+    # Get the roman 2D science image
+    with asdf.open(file) as f:
+        # Get the WCS
+        wcs = f['roman']['meta']['wcs']
+
+        # Get the WCS header
+        header = wcs.to_fits_sip()
+
+    # Update WCS header with some keywords that it's missing.
+    # Otherwise, it won't work with astropy.wcs tools (TODO: Figure out why. What are these keywords for?)
+    for k in ['cpdis1', 'cpdis2', 'det2im1', 'det2im2', 'sip']:
+        header[k] = 'na'
+
+    # New WCS object with updated header
+    wcs_updated = astropy.wcs.WCS(header)
+
+    # Turn input RA, Dec into a SkyCoord object
+    coordinates = astropy.coordinates.SkyCoord(ra, dec, unit='deg')
+
+    # Map the coordinates to a pixel's location on the Roman 2d array (row, col)
+    row, col = astropy.wcs.utils.skycoord_to_pixel(coords=coordinates, wcs=wcs_updated)
+
+    return (row, col), wcs_updated
+
+
+def get_cutout(file, coords, wcs, size=20, outfile="example_roman_cutout.fits"):
+    """ Get a roman image cutout """
+
+    # Get the 2D science image
+    with asdf.open(file) as f:
+        data = f['roman']['data']
+
+        cutout = astropy.nddata.Cutout2D(data, position=coords, wcs=wcs, size=(size, size))
+
+        astropy.io.fits.writeto(outfile, data=cutout.data, header=cutout.wcs.to_header(), overwrite=True)
+
+
+def asdf_cut(input_file, ra, dec, *, cutout_size=20, output_file="example_roman_cutout.fits"):
+    """ Preliminary proof-of-concept functionality.
+    Takes a single ASDF input file (``input_file``) and generates a cutout of designated size ``cutout_size``
+    around the given coordinates (``coordinates``).
+    """
+
+    pixel_coordinates, wcs = get_center_pixel(input_file, ra, dec)
+
+    get_cutout(input_file, pixel_coordinates, wcs, size=cutout_size, outfile=output_file)

astrocut/tests/test_asdf_cut.py

@@ -0,0 +1,122 @@
+
+import numpy as np
+import pytest
+
+import asdf
+from astropy.modeling import models
+from astropy import coordinates as coord
+from astropy import units as u
+from astropy.io import fits
+from astropy.wcs import WCS
+from gwcs import wcs
+from gwcs import coordinate_frames as cf
+from astrocut.asdf_cutouts import get_center_pixel, get_cutout, asdf_cut
+
+
+def make_wcs(xsize, ysize, ra=30., dec=45.):
+    """ create a fake gwcs object """
+    # todo - refine this to better reflect roman wcs
+
+    # create transformations
+    pixelshift = models.Shift(-xsize) & models.Shift(-ysize)
+    pixelscale = models.Scale(0.1 / 3600.) & models.Scale(0.1 / 3600.)  # 0.1 arcsec/pixel
+    tangent_projection = models.Pix2Sky_TAN()
+    celestial_rotation = models.RotateNative2Celestial(ra, dec, 180.)
+
+    # transform pixels to sky
+    det2sky = pixelshift | pixelscale | tangent_projection | celestial_rotation
+
+    # define the wcs object
+    detector_frame = cf.Frame2D(name="detector", axes_names=("x", "y"), unit=(u.pix, u.pix))
+    sky_frame = cf.CelestialFrame(reference_frame=coord.ICRS(), name='icrs', unit=(u.deg, u.deg))
+    return wcs.WCS([(detector_frame, det2sky), (sky_frame, None)])
+
+
+@pytest.fixture()
+def fakedata():
+    """ fixture to create fake data and wcs """
+    # set up initial parameters
+    nx = 100
+    ny = 100
+    size = nx * ny
+    ra = 30.
+    dec = 45.
+
+    # create the wcs
+    wcsobj = make_wcs(nx/2, ny/2, ra=ra, dec=dec)
+    wcsobj.bounding_box = ((0, nx), (0, ny))
+
+    # create the data
+    data = np.arange(size).reshape(nx, ny)
+
+    yield data, wcsobj
+
+
+@pytest.fixture()
+def make_file(tmp_path, fakedata):
+    """ fixture to create a fake dataset """
+    # get the fake data
+    data, wcsobj = fakedata
+
+    # create meta
+    meta = {'wcs': wcsobj}
+
+    # create and write the asdf file
+    tree = {'roman': {'data': data, 'meta': meta}}
+    af = asdf.AsdfFile(tree)
+
+    path = tmp_path / "roman"
+    path.mkdir(exist_ok=True)
+    filename = path / "test_roman.asdf"
+    af.write_to(filename)
+
+    yield filename
+
+
+def test_get_center_pixel(make_file):
+    """ test we can get the correct center pixel """
+    pixel_coordinates, wcs = get_center_pixel(make_file, 30., 45.)
+    assert np.allclose(pixel_coordinates, (np.array(50.), np.array(50.)))
+    assert np.allclose(wcs.celestial.wcs.crval, np.array([30., 45.]))
+
+
+@pytest.fixture()
+def output_file(tmp_path):
+    """ fixture to create the output path """
+    # create output fits path
+    out = tmp_path / "roman"
+    out.mkdir(exist_ok=True, parents=True)
+    output_file = out / "test_output_cutout.fits"
+    yield output_file
+
+
+def test_get_cutout(make_file, output_file, fakedata):
+    """ test we can create a cutout """
+
+    # get the input wcs
+    __, ww = fakedata
+    skycoord = ww(25, 25, with_units=True)
+    wcs = WCS(ww.to_fits_sip())
+
+    # create cutout
+    get_cutout(make_file, skycoord, wcs, size=10, outfile=output_file)
+
+    # test output
+    with fits.open(output_file) as hdulist:
+        data = hdulist[0].data
+        assert data.shape == (10, 10)
+        assert data[5, 5] == 2525
+
+
+def test_asdf_cutout(make_file, output_file):
+    """ test we can make a cutout """
+    # make cutout
+    ra, dec = (29.99901792, 44.99930555)
+    asdf_cut(make_file, ra, dec, cutout_size=10, output_file=output_file)
+
+    # test output
+    with fits.open(output_file) as hdulist:
+        data = hdulist[0].data
+        assert data.shape == (10, 10)
+        assert data[5, 5] == 2526
+

setup.cfg

@@ -16,9 +16,11 @@ packages = find:
 python_requires = >=3.8
 setup_requires = setuptools_scm
 install_requires =
+    asdf>=2.15.0 # for ASDF file format
     astropy>=5.2 # astropy with s3fs support
     fsspec[http]>=2022.8.2  # for remote cutouts
     s3fs>=2022.8.2  # for remote cutouts
+    roman_datamodels>=0.17.0 # for roman file support
     scipy
     Pillow