Embed ASDF data into FITS extension

Author: snbianco

astrocut/asdf_cutout.py

@@ -17,6 +17,7 @@
 from astropy.utils.decorators import deprecated_renamed_argument
 from astropy.wcs import WCS
 from s3path import S3Path
+from stdatamodels import asdf_in_fits
 
 from . import log, __version__
 from .image_cutout import ImageCutout
@@ -106,17 +107,21 @@ def fits_cutouts(self) -> List[fits.HDUList]:
         """
         if not self._fits_cutouts:
             fits_cutouts = []
-            for file, cutouts in self.cutouts_by_file.items():
-                # TODO: Create a FITS object with ASDF extension
-                # Create a primary FITS header to hold data and WCS
+            for i, (file, cutouts) in enumerate(self.cutouts_by_file.items()):                
                 cutout = cutouts[0]
-                primary_hdu = fits.PrimaryHDU(data=cutout.data, header=cutout.wcs.to_header(relax=True))
+                if self._lite:
+                    tree = self._get_lite_tree(file, cutout, self._gwcs_objects[i])
+                else:
+                    tree = self._asdf_trees[i]
 
-                # Add original file to header
-                primary_hdu.header['ORIG_FLE'] = str(file)
+                # Create a primary FITS header to hold data and WCS
+                primary_hdu = fits.PrimaryHDU(data=cutout.data, header=cutout.wcs.to_header(relax=True))
+                primary_hdu.header['ORIG_FLE'] = file  # Add original file to header
+                hdul = fits.HDUList([primary_hdu])
 
-                # Write to HDUList
-                fits_cutouts.append(fits.HDUList([primary_hdu]))
+                # Embed ASDF into FITS
+                hdul_embed = asdf_in_fits.to_hdulist(tree, hdul)
+                fits_cutouts.append(hdul_embed)
             self._fits_cutouts = fits_cutouts
         return self._fits_cutouts
 
@@ -130,13 +135,7 @@ def asdf_cutouts(self) -> List[asdf.AsdfFile]:
             for i, (file, cutouts) in enumerate(self.cutouts_by_file.items()):
                 cutout = cutouts[0]
                 if self._lite:
-                    tree = {
-                        self._mission_kwd: {
-                            'meta': {'wcs': self._slice_gwcs(cutout, self._gwcs_objects[i]),
-                                     'orig_file': str(file)},
-                            'data': cutout.data
-                        }
-                    }
+                    tree = self._get_lite_tree(file, cutout, self._gwcs_objects[i])
                 else:
                     tree = self._asdf_trees[i]
 
@@ -156,6 +155,32 @@ def asdf_cutouts(self) -> List[asdf.AsdfFile]:
 
             self._asdf_cutouts = asdf_cutouts
         return self._asdf_cutouts
+    
+    def _get_lite_tree(self, file: str, cutout: Cutout2D, gwcs: gwcs.wcs.WCS) -> dict:
+        """
+        Helper function to create an ASDF tree in lite mode.
+
+        Parameters
+        ----------
+        file : str
+            The input filename.
+        cutout : `~astropy.nddata.Cutout2D`
+            The cutout object.
+        gwcs : gwcs.wcs.WCS
+            The original GWCS object.
+
+        Returns
+        -------
+        tree : dict
+            The ASDF tree in lite mode. The tree contains only the cutout data and the sliced GWCS.
+        """
+        return {
+            self._mission_kwd: {
+                'meta': {'wcs': self._slice_gwcs(cutout, gwcs),
+                        'orig_file': file},
+                'data': cutout.data
+            }
+        }
 
     def _get_cloud_http(self, input_file: Union[str, S3Path]) -> str:
         """ 
@@ -208,7 +233,7 @@ def _load_file_data(self, input_file: Union[str, Path, S3Path]) -> dict:
 
         return tree
 
-    def _make_cutout(self, array, position, wcs):
+    def _make_cutout(self, array: np.ndarray, position: tuple, wcs: WCS) -> Cutout2D:
         """
         Helper to generate a Cutout2D and return plain ndarray data.
 

astrocut/tests/test_asdf_cutout.py

@@ -14,6 +14,7 @@
 from astropy.io import fits
 from gwcs import wcs, coordinate_frames
 from PIL import Image
+from stdatamodels import asdf_in_fits
 
 from astrocut.asdf_cutout import ASDFCutout, asdf_cut, get_center_pixel
 from astrocut.exceptions import DataWarning, InvalidInputError, InvalidQueryError
@@ -158,24 +159,28 @@ def test_asdf_cutout(test_images, center_coord, cutout_size):
 
 
 def test_asdf_cutout_write_to_file(test_images, center_coord, cutout_size, tmpdir):
+    def check_asdf_metadata(af, original_file, cutout_data):
+        """Check that ASDF file contains correct metadata"""
+        assert 'roman' in af
+        assert 'meta' in af['roman']
+        # Check cutout data and metadata
+        for key in ['data', 'dq', 'err', 'context']:
+            assert key in af['roman']
+            assert np.all(af['roman'][key] == cutout_data)
+        meta = af['roman']['meta']
+        assert meta['wcs'].pixel_shape == (10, 10)
+        assert meta['product_type'] == 'l2'
+        assert meta['file_date'] == Time('2023-10-01T00:00:00', format='isot')
+        assert meta['origin'] == 'STSCI/SOC'
+        assert meta['orig_file'] == original_file.as_posix()
+    
     # Write cutouts to ASDF files on disk
     cutout = ASDFCutout(test_images, center_coord, cutout_size)
     asdf_files = cutout.write_as_asdf(output_dir=tmpdir)
     assert len(asdf_files) == 3
     for i, asdf_file in enumerate(asdf_files):
         with asdf.open(asdf_file) as af:
-            assert 'roman' in af
-            assert 'meta' in af['roman']
-            # Check cutout data and metadata
-            for key in ['data', 'dq', 'err', 'context']:
-                assert key in af['roman']
-                assert np.all(af['roman'][key] == cutout.cutouts[i].data)
-            meta = af['roman']['meta']
-            assert meta['wcs'].pixel_shape == (10, 10)
-            assert meta['product_type'] == 'l2'
-            assert meta['file_date'] == Time('2023-10-01T00:00:00', format='isot')
-            assert meta['origin'] == 'STSCI/SOC'
-            assert meta['orig_file'] == test_images[i].as_posix()
+            check_asdf_metadata(af, test_images[i], cutout.cutouts[i].data)
             # Check file size is smaller than original
             assert Path(asdf_file).stat().st_size < Path(test_images[i]).stat().st_size
 
@@ -191,6 +196,10 @@ def test_asdf_cutout_write_to_file(test_images, center_coord, cutout_size, tmpdi
             assert hdul[0].header['ORIG_FLE'] == test_images[i].as_posix()
             assert Path(fits_file).stat().st_size < Path(test_images[i]).stat().st_size
 
+        # Check ASDF extension contents
+        with asdf_in_fits.open(fits_file) as af:
+            check_asdf_metadata(af, test_images[i], cutout.cutouts[i].data)
+
 
 @pytest.mark.parametrize('output_format', ['.asdf', '.fits'])
 def test_asdf_cutout_write_to_zip(tmpdir, test_images, center_coord, cutout_size, output_format):
@@ -215,7 +224,7 @@ def test_asdf_cutout_write_to_zip(tmpdir, test_images, center_coord, cutout_size
         else:
             with fits.open(io.BytesIO(data)) as hdul:
                 assert isinstance(hdul, fits.HDUList)
-                assert len(hdul) == 1
+                assert len(hdul) == 2  # primary + embedded ASDF extension
                 assert hdul[0].data.shape == (cutout_size, cutout_size)
 
 
@@ -226,10 +235,9 @@ def test_asdf_cutout_write_to_zip_invalid_format(tmpdir, test_images, center_coo
         cutout.write_as_zip(output_dir=tmpdir, output_format='.invalid')
 
 
-def test_asdf_cutout_lite(test_images, center_coord, cutout_size, tmpdir):
-    # Write cutouts to ASDF objects in lite mode
-    cutout = ASDFCutout(test_images, center_coord, cutout_size, lite=True)
-    for af in cutout.asdf_cutouts:
+def test_asdf_cutout_lite(test_images, center_coord, cutout_size):
+    def check_lite_metadata(af):
+        """Check that ASDF file contains only lite metadata"""
         assert 'roman' in af
         assert 'data' in af['roman']
         assert 'meta' in af['roman']
@@ -238,11 +246,21 @@ def test_asdf_cutout_lite(test_images, center_coord, cutout_size, tmpdir):
         assert len(af['roman']) == 2  # only data and meta
         assert len(af['roman']['meta']) == 2  # only wcs and original filename
 
+    # Write cutouts to ASDF objects in lite mode
+    cutout = ASDFCutout(test_images, center_coord, cutout_size, lite=True)
+    for af in cutout.asdf_cutouts:
+        check_lite_metadata(af)
+
     # Write cutouts to HDUList objects in lite mode
     cutout = ASDFCutout(test_images, center_coord, cutout_size, lite=True)
     for hdul in cutout.fits_cutouts:
-        assert len(hdul) == 1  # primary HDU only
+        assert len(hdul) == 2  # primary HDU + embedded ASDF extension
         assert hdul[0].name == 'PRIMARY'
+        assert hdul[1].name == 'ASDF'
+
+        # Check ASDF extension contents
+        with asdf_in_fits.open(hdul) as af:
+            check_lite_metadata(af)
 
 
 def test_asdf_cutout_partial(test_images, center_coord, cutout_size):

setup.cfg

@@ -22,10 +22,10 @@ install_requires =
     fsspec[http]>=2022.8.2  # for remote cutouts
     s3fs>=2022.8.2  # for remote cutouts
     s3path>=0.5.7 # for remote file paths
-    roman_datamodels>=0.19.0 # for roman file support
     requests>=2.32.3 # for making HTTP requests
     spherical_geometry>=1.3.0
     gwcs>=0.21.0
+    stdatamodels>=4.1.0
     scipy
     Pillow