More robust error handling, use ImageHDU for cutout data

Author: snbianco

CHANGES.rst

@@ -7,7 +7,7 @@ Unreleased
   writing multiple cutouts into a single ZIP archive. [#167]
 - Added ``get_tess_sectors`` function to return TESS sector information for sectors whose footprints overlap with 
   the given sky coordinates and cutout size. [#168]
-- Cutouts of ASDF data in FITS format now include embedded ASDF metadata and cutout data in an "ASDF" extension within the FITS file for 
+- Cutouts of ASDF data in FITS format now include embedded ASDF metadata in an "ASDF" extension within the FITS file for 
   Python versions greater than or equal to 3.11. [#170]
 
 Breaking Changes
@@ -32,6 +32,9 @@ Breaking Changes
     - New (no unit - pixels assumed): ``..._83.4063100_-62.4897710_64-x-64_astrocut.fits``
     - New (with units): ``..._83.4063100_-62.4897710_5arcmin-x-4arcmin_astrocut.fits``
 
+  - ASDF cutouts in FITS format now include cutout data in an ``ImageHDU`` extension called "CUTOUT". Code that reads ASDF cutouts from FITS files 
+    should be updated to access the "CUTOUT" extension for cutout data rather than the "PRIMARY" extension. [#170]
+
 
 1.1.0 (2025-09-15)
 ------------------

astrocut/asdf_cutout.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from time import monotonic
 from typing import List, Tuple, Union, Optional
+from datetime import date
 
 import asdf
 import gwcs
@@ -85,7 +86,7 @@ def __init__(self, input_files: List[Union[str, Path, S3Path]], coordinates: Uni
         super().__init__(input_files, coordinates, cutout_size, fill_value, verbose=verbose)
 
         # Must be using Python 3.11 or higher to support stdatamodels and ASDF-in-FITS embedding
-        self._supports_stdatamodels = sys.version_info >= (3, 11)
+        self._py311_or_higher = sys.version_info >= (3, 11)
 
         # Assign AWS credential attributes
         self._key = key
@@ -110,32 +111,62 @@ def fits_cutouts(self) -> List[fits.HDUList]:
         """
         if not self._fits_cutouts:
             # Try to import stdatamodels for ASDF-in-FITS embedding
-            if self._supports_stdatamodels:
+            if self._py311_or_higher:
                 try:
-                    from stdatamodels import asdf_in_fits
-                except ModuleNotFoundError:
-                    warnings.warn('The `stdatamodels` package is required for ASDF-in-FITS embedding. '
-                                  'Skipping embedding for these cutouts. To install astrocut with optional '
-                                  '`stdatamodels` support, run: pip install "astrocut[all]"', ModuleWarning)
-                    self._supports_stdatamodels = False
+                    # Check version of stdatamodels
+                    from stdatamodels import __version__ as stdata_version, asdf_in_fits
+                    if stdata_version < '4.1.0':
+                        warnings.warn(
+                            'The `stdatamodels` package is not available in the correct version (>=4.1.0); '
+                            'ASDF-in-FITS embedding will be skipped for these cutouts. Install the optional '
+                            'dependency with: pip install "astrocut[all]" or pip install stdatamodels>=4.1.0',
+                            ModuleWarning
+                        )
+                        self._can_embed_asdf_in_fits = False
+                    else:
+                        self._can_embed_asdf_in_fits = True
+                except ImportError:
+                    warnings.warn(
+                        'The `stdatamodels` package cannot be imported; ASDF-in-FITS embedding will be '
+                        'skipped for these cutouts. Install the optional dependency with: '
+                        'pip install "astrocut[all]" or pip install stdatamodels>=4.1.0',
+                        ModuleWarning
+                    )
+                    self._can_embed_asdf_in_fits = False
             else:
                 warnings.warn('ASDF-in-FITS embedding requires Python 3.11 or higher. '
                               'Skipping embedding for these cutouts.', ModuleWarning)
+                self._can_embed_asdf_in_fits = False
 
             fits_cutouts = []
             for i, (file, cutouts) in enumerate(self.cutouts_by_file.items()):                
                 cutout = cutouts[0]
                 if self._lite:
-                    tree = self._get_lite_tree(str(file), cutout, self._gwcs_objects[i])
+                    tree = {
+                        # Tree should only include sliced WCS and original filename
+                        self._mission_kwd: {
+                            'meta': {'wcs': self._slice_gwcs(cutout, self._gwcs_objects[i]),
+                                     'orig_file': str(file)}
+                        }
+                    }
                 else:
                     tree = self._asdf_trees[i]
-
-                # 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'] = str(file)  # Add original file to header
-                hdul = fits.HDUList([primary_hdu])
-
-                if self._supports_stdatamodels:
+                    # Tree should only include meta
+                    tree[self._mission_kwd] = {'meta': tree[self._mission_kwd]['meta']}
+
+                # Build the PrimaryHDU with keywords
+                primary_hdu = fits.PrimaryHDU()
+                primary_hdu.header.extend([('ORIGIN', 'STScI/MAST', 'institution responsible for creating this file'),
+                                           ('DATE', str(date.today()), 'file creation date'),
+                                           ('PROCVER', __version__, 'software version')])
+                
+                # Build ImageHDU with cutout data and WCS
+                image_hdu = fits.ImageHDU(data=cutout.data, header=cutout.wcs.to_header(relax=True))
+                image_hdu.header['ORIG_FLE'] = str(file)  # Add original file to header
+                image_hdu.header['EXTNAME'] = 'CUTOUT'
+                hdul = fits.HDUList([primary_hdu, image_hdu])
+
+                if self._can_embed_asdf_in_fits:
                     hdul_embed = asdf_in_fits.to_hdulist(tree, hdul)
                 else:
                     hdul_embed = hdul

astrocut/tests/test_asdf_cutout.py

@@ -4,6 +4,7 @@
 import zipfile
 import io
 import importlib.util
+from unittest.mock import MagicMock, patch
 
 import asdf
 from astropy import coordinates as coord
@@ -17,7 +18,7 @@
 from PIL import Image
 
 from astrocut.asdf_cutout import ASDFCutout, asdf_cut, get_center_pixel
-from astrocut.exceptions import DataWarning, InvalidInputError, InvalidQueryError
+from astrocut.exceptions import DataWarning, InvalidInputError, InvalidQueryError, ModuleWarning
 
 
 def make_wcs(xsize, ysize, ra=30., dec=45.):
@@ -159,20 +160,22 @@ 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):
+    def check_asdf_metadata(af, original_file, cutout_data, meta_only=False):
         """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()
+
+        if not meta_only:
+            # 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)
 
     # Write cutouts to ASDF files on disk
     cutout = ASDFCutout(test_images, center_coord, cutout_size)
@@ -190,17 +193,19 @@ def check_asdf_metadata(af, original_file, cutout_data):
     assert len(fits_files) == 3
     for i, fits_file in enumerate(fits_files):
         with fits.open(fits_file) as hdul:
-            assert np.all(hdul[0].data == cutout.cutouts[i].data)
-            assert hdul[0].header['NAXIS1'] == 10
-            assert hdul[0].header['NAXIS2'] == 10
-            assert hdul[0].header['ORIG_FLE'] == test_images[i].as_posix()
+            assert hdul[0].name == 'PRIMARY'
+            assert hdul[1].name == 'CUTOUT'
+            assert np.all(hdul[1].data == cutout.cutouts[i].data)
+            assert hdul[1].header['NAXIS1'] == 10
+            assert hdul[1].header['NAXIS2'] == 10
+            assert hdul[1].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 (stdatamodels optional)
         if importlib.util.find_spec('stdatamodels') is not None:
             from stdatamodels import asdf_in_fits
             with asdf_in_fits.open(fits_file) as af:
-                check_asdf_metadata(af, test_images[i], cutout.cutouts[i].data)
+                check_asdf_metadata(af, test_images[i], cutout.cutouts[i].data, meta_only=True)
 
 
 @pytest.mark.parametrize('output_format', ['.asdf', '.fits'])
@@ -226,8 +231,8 @@ 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) == 2 if importlib.util.find_spec('stdatamodels') is not None else 1
-                assert hdul[0].data.shape == (cutout_size, cutout_size)
+                assert len(hdul) == 3 if importlib.util.find_spec('stdatamodels') is not None else 2
+                assert hdul[1].data.shape == (cutout_size, cutout_size)
 
 
 def test_asdf_cutout_write_to_zip_invalid_format(tmpdir, test_images, center_coord, cutout_size):
@@ -238,15 +243,16 @@ def test_asdf_cutout_write_to_zip_invalid_format(tmpdir, test_images, center_coo
 
 
 def test_asdf_cutout_lite(test_images, center_coord, cutout_size):
-    def check_lite_metadata(af):
+    def check_lite_metadata(af, meta_only=False):
         """Check that ASDF file contains only lite metadata"""
         assert 'roman' in af
-        assert 'data' in af['roman']
         assert 'meta' in af['roman']
         assert 'wcs' in af['roman']['meta']
         assert 'orig_file' in af['roman']['meta']
-        assert len(af['roman']) == 2  # only data and meta
+        assert len(af['roman']) == (1 if meta_only else 2)
         assert len(af['roman']['meta']) == 2  # only wcs and original filename
+        if not meta_only:
+            assert 'data' in af['roman']
 
     # Write cutouts to ASDF objects in lite mode
     cutout = ASDFCutout(test_images, center_coord, cutout_size, lite=True)
@@ -257,15 +263,16 @@ def check_lite_metadata(af):
     cutout = ASDFCutout(test_images, center_coord, cutout_size, lite=True)
     for hdul in cutout.fits_cutouts:
         has_stdatamodels = importlib.util.find_spec('stdatamodels') is not None
-        assert len(hdul) == 2 if has_stdatamodels else 1  # primary HDU + embedded ASDF extension
+        assert len(hdul) == 3 if has_stdatamodels else 2  # primary HDU + cutout HDU + embedded ASDF extension
         assert hdul[0].name == 'PRIMARY'
+        assert hdul[1].name == 'CUTOUT'
 
         # Check ASDF extension contents (stdatamodels optional)
         if has_stdatamodels:
-            assert hdul[1].name == 'ASDF'
+            assert hdul[2].name == 'ASDF'
             from stdatamodels import asdf_in_fits
             with asdf_in_fits.open(hdul) as af:
-                check_lite_metadata(af)
+                check_lite_metadata(af, meta_only=True)
 
 
 def test_asdf_cutout_partial(test_images, center_coord, cutout_size):
@@ -404,6 +411,36 @@ def test_asdf_cutout_gwcs(test_images, center_coord):
     assert gwcs.bounding_box.intervals[1].upper == 39
 
 
+@pytest.mark.parametrize(('is_installed', 'warn_msg'), 
+                         [(True, 'not available in the correct version'), (False, 'package cannot be imported')])
+def test_asdf_cutout_stdatamodels(test_images, center_coord, cutout_size, is_installed, warn_msg):
+    """ Test that warning is emitted about ASDF-in-FITS embedding for stdatamodels issues """
+    mock_stdatamodels = None
+    if is_installed:
+        mock_stdatamodels = MagicMock()
+        mock_stdatamodels.__version__ = '1.0.0'
+        mock_stdatamodels.asdf_in_fits = MagicMock()
+    patch_dict = {'stdatamodels': mock_stdatamodels}
+
+    with patch.dict('sys.modules', patch_dict):
+        with pytest.warns(ModuleWarning, match=warn_msg):
+            cutout = ASDFCutout(test_images, center_coord, cutout_size)
+            fits_cutouts = cutout.fits_cutouts
+        assert cutout._can_embed_asdf_in_fits is False
+        assert len(fits_cutouts[0]) == 2  # primary + cutout HDU only
+
+
+def test_asdf_cutout_python_version(test_images, center_coord, cutout_size):
+    """ Test that warning is emitted about ASDF-in-FITS embedding for Python <3.11 """
+    with patch('sys.version_info', (3, 10, 0)):
+        with pytest.warns(ModuleWarning, match='requires Python 3.11 or higher'):
+            cutout = ASDFCutout(test_images, center_coord, cutout_size)
+            fits_cutouts = cutout.fits_cutouts
+        assert cutout._py311_or_higher is False
+        assert cutout._can_embed_asdf_in_fits is False
+        assert len(fits_cutouts[0]) == 2  # primary + cutout HDU only
+
+
 def test_get_center_pixel(fakedata):
     """ Test get_center_pixel function """
     # Get the fake data