Adding HAPCut support to ~mast.Cutouts

jaymedina · bsipocz · commit e77312179cf3 · 2022-12-19T21:34:29.000-08:00
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -128,6 +128,9 @@ mast
 - Add a ``flat`` option to ``Observation.download_products()`` to turn off the
   automatic creation and organizing of products into subdirectories. [#2511]
 
+- Expanding ``Cutouts`` functionality to support making Hubble Advanced Product (HAP)
+  cutouts via HAPCut. [#2613]
+
 oac
 ^^^
 
diff --git a/astroquery/mast/__init__.py b/astroquery/mast/__init__.py
@@ -30,7 +30,7 @@ class Conf(_config.ConfigNamespace):
 
 conf = Conf()
 
-from .cutouts import TesscutClass, Tesscut, ZcutClass, Zcut
+from .cutouts import TesscutClass, Tesscut, ZcutClass, Zcut, HapcutClass, Hapcut
 from .observations import Observations, ObservationsClass, MastClass, Mast
 from .collections import Catalogs, CatalogsClass
 from .missions import MastMissions, MastMissionsClass
@@ -42,5 +42,6 @@ class Conf(_config.ConfigNamespace):
            'Mast', 'MastClass',
            'Tesscut', 'TesscutClass',
            'Zcut', 'ZcutClass',
+           'Hapcut', 'HapcutClass',
            'Conf', 'conf', 'utils',
            ]
diff --git a/astroquery/mast/cutouts.py b/astroquery/mast/cutouts.py
@@ -201,7 +201,7 @@ def get_sectors(self, *, coordinates=None, radius=0*u.deg, objectname=None, movi
         return Table(sector_dict)
 
     def download_cutouts(self, *, coordinates=None, size=5, sector=None, path=".", inflate=True,
-                         objectname=None, moving_target=False, mt_type=None):
+                         objectname=None, moving_target=False, mt_type=None, verbose=False):
         """
         Download cutout target pixel file(s) around the given coordinates with indicated size.
 
@@ -301,7 +301,7 @@ def download_cutouts(self, *, coordinates=None, size=5, sector=None, path=".", i
             localpath_table['Local Path'] = [zipfile_path]
             return localpath_table
 
-        print("Inflating...")
+        if verbose: print("Inflating...")
         # unzipping the zipfile
         zip_ref = zipfile.ZipFile(zipfile_path, 'r')
         cutout_files = zip_ref.namelist()
@@ -477,7 +477,8 @@ def get_surveys(self, coordinates, *, radius="0d"):
             warnings.warn("Coordinates are not in an available deep field survey.", NoResultsWarning)
         return survey_json
 
-    def download_cutouts(self, coordinates, *, size=5, survey=None, cutout_format="fits", path=".", inflate=True, **img_params):
+    def download_cutouts(self, coordinates, *, size=5, survey=None, cutout_format="fits", path=".", inflate=True,
+                         verbose=False, **img_params):
         """
         Download cutout FITS/image file(s) around the given coordinates with indicated size.
 
@@ -560,7 +561,7 @@ def download_cutouts(self, coordinates, *, size=5, survey=None, cutout_format="f
             localpath_table['Local Path'] = [zipfile_path]
             return localpath_table
 
-        print("Inflating...")
+        if verbose: print("Inflating...")
         # unzipping the zipfile
         zip_ref = zipfile.ZipFile(zipfile_path, 'r')
         cutout_files = zip_ref.namelist()
@@ -637,3 +638,164 @@ def get_cutouts(self, coordinates, *, size=5, survey=None):
 
 
 Zcut = ZcutClass()
+
+
+class HapcutClass(MastQueryWithLogin):
+    """
+    MAST Hubble Advanced Product (HAP) cutout query class.
+    
+    Class for accessing HAP image cutouts.
+    """
+
+    def __init__(self):
+
+        super().__init__()
+
+        services = {"astrocut": {"path": "astrocut"}}
+
+        self._service_api_connection.set_service_params(services, "hapcut")
+
+    def download_cutouts(self, coordinates, *, size=5, path=".", inflate=True, verbose=False):
+        """
+        Download cutout images around the given coordinates with indicated size.
+
+        Parameters
+        ----------
+        coordinates : str or `astropy.coordinates` object
+            The target around which to search. It may be specified as a
+            string or as the appropriate `astropy.coordinates` object.
+        size : int, array-like, `~astropy.units.Quantity`
+            Optional, default 5 pixels.
+            The size of the cutout array. If ``size`` is a scalar number or
+            a scalar `~astropy.units.Quantity`, then a square cutout of ``size``
+            will be created.  If ``size`` has two elements, they should be in
+            ``(ny, nx)`` order.  Scalar numbers in ``size`` are assumed to be in
+            units of pixels. `~astropy.units.Quantity` objects must be in pixel or
+            angular units.
+        path : str
+            Optional.
+            The directory in which the cutouts will be saved.
+            Defaults to current directory.
+        inflate : bool
+            Optional, default True.
+            Cutout target pixel files are returned from the server in a zip file,
+            by default they will be inflated and the zip will be removed.
+            Set inflate to false to stop before the inflate step.
+
+        Returns
+        -------
+        response : `~astropy.table.Table`
+        """
+
+        # Get Skycoord object for coordinates/object
+        coordinates = parse_input_location(coordinates)
+
+        # Build initial astrocut request
+        astrocut_request = f"astrocut?ra={coordinates.ra.deg}&dec={coordinates.dec.deg}"
+
+        # Add size parameters to request
+        size_dict = _parse_cutout_size(size)
+        astrocut_request += f"&x={size_dict['x']}&y={size_dict['y']}&units={size_dict['units']}"
+
+        # Build the URL
+        astrocut_url = self._service_api_connection.REQUEST_URL + astrocut_request
+
+        # Set up the download path
+        path = os.path.join(path, '')
+        zipfile_path = "{}hapcut_{}.zip".format(path, time.strftime("%Y%m%d%H%M%S"))
+
+        # Download
+        self._download_file(astrocut_url, zipfile_path)
+        localpath_table = Table(names=["Local Path"], dtype=[str])
+
+        # Checking if we got a zip file or a json no results message
+        if not zipfile.is_zipfile(zipfile_path):
+            with open(zipfile_path, 'r') as FLE:
+                response = json.load(FLE)
+            warnings.warn(response['msg'], NoResultsWarning)
+            return localpath_table
+
+        if not inflate:  # not unzipping
+            localpath_table['Local Path'] = [zipfile_path]
+            return localpath_table
+
+        if verbose: print("Inflating...")
+        # unzipping the zipfile
+        zip_ref = zipfile.ZipFile(zipfile_path, 'r')
+        cutout_files = zip_ref.namelist()
+        zip_ref.extractall(path, members=cutout_files)
+        zip_ref.close()
+        os.remove(zipfile_path)
+
+        localpath_table['Local Path'] = [path+x for x in cutout_files]
+        return localpath_table
+
+
+    def get_cutouts(self, coordinates, *, size=5):
+        """
+        Get cutout image(s) around the given coordinates with indicated size,
+        and return them as a list of  `~astropy.io.fits.HDUList` objects.
+
+        Parameters
+        ----------
+        coordinates : str or `astropy.coordinates` object
+            The target around which to search. It may be specified as a
+            string or as the appropriate `astropy.coordinates` object.
+        size : int, array-like, `~astropy.units.Quantity`
+            Optional, default 5 pixels.
+            The size of the cutout array. If ``size`` is a scalar number or
+            a scalar `~astropy.units.Quantity`, then a square cutout of ``size``
+            will be created.  If ``size`` has two elements, they should be in
+            ``(ny, nx)`` order.  Scalar numbers in ``size`` are assumed to be in
+            units of pixels. `~astropy.units.Quantity` objects must be in pixel or
+            angular units.
+
+        Returns
+        -------
+        response : A list of `~astropy.io.fits.HDUList` objects.
+        """
+
+        # Get Skycoord object for coordinates/object
+        coordinates = parse_input_location(coordinates)
+
+        param_dict = _parse_cutout_size(size)
+
+        # Need to convert integers from numpy dtypes
+        # so we can convert the dictionary into a JSON object
+        # in service_request_async(...)
+        param_dict["x"] = float(param_dict["x"])
+        param_dict["y"] = float(param_dict["y"])
+
+        # Adding RA and DEC to parameters dictionary
+        param_dict["ra"] = coordinates.ra.deg
+        param_dict["dec"] = coordinates.dec.deg
+
+        response = self._service_api_connection.service_request_async("astrocut", param_dict, use_json=True)
+        response.raise_for_status()  # Raise any errors
+
+        try:
+            ZIPFILE = zipfile.ZipFile(BytesIO(response.content), 'r')
+        except zipfile.BadZipFile:
+            message = response.json()
+            if len(message['results']) == 0:
+                warnings.warn(message['msg'], NoResultsWarning)
+                return []
+            else:
+                raise
+        
+        # Open all the contained fits files:
+        # Since we cannot seek on a compressed zip file,
+        # we have to read the data, wrap it in another BytesIO object,
+        # and then open that using fits.open
+        cutout_hdus_list = []
+        for name in ZIPFILE.namelist():
+            CUTOUT = BytesIO(ZIPFILE.open(name).read())
+            cutout_hdus_list.append(fits.open(CUTOUT))
+
+            # preserve the original filename in the fits object
+            cutout_hdus_list[-1].filename = name
+
+        return cutout_hdus_list
+
+
+Hapcut = HapcutClass()
diff --git a/astroquery/mast/tests/test_mast_remote.py b/astroquery/mast/tests/test_mast_remote.py
@@ -1088,3 +1088,76 @@ def test_zcut_get_cutouts(self):
         assert isinstance(cutout_list, list)
         assert len(cutout_list) == 1
         assert isinstance(cutout_list[0], fits.HDUList)
+
+    ###################
+    # HapcutClass tests #
+    ###################
+
+    def test_hapcut_download_cutouts(self, tmpdir):
+        
+        # Test 1: Simple API call with expected results
+        coord = SkyCoord(351.347812, 28.497808, unit="deg")
+
+        cutout_table = mast.Hapcut.download_cutouts(coordinates=coord, size=5, path=str(tmpdir))
+        assert isinstance(cutout_table, Table)
+        assert len(cutout_table) >= 1
+        for row in cutout_table:
+            assert os.path.isfile(row['Local Path'])
+            if 'fits' in os.path.basename(row['Local Path']):
+                assert fits.getdata(row['Local Path']).shape == (5, 5)
+
+        # Test 2: Make input size a list
+        cutout_table = mast.Hapcut.download_cutouts(coordinates=coord, size=[2, 3], path=str(tmpdir))
+        assert isinstance(cutout_table, Table)
+        assert len(cutout_table) >= 1
+        for row in cutout_table:
+            assert os.path.isfile(row['Local Path'])
+            if 'fits' in os.path.basename(row['Local Path']):
+                assert fits.getdata(row['Local Path']).shape == (3, 2)
+
+        # Test 3: Specify unit for input size
+        cutout_table = mast.Hapcut.download_cutouts(coordinates=coord, size=5*u.arcsec, path=str(tmpdir))
+        assert isinstance(cutout_table, Table)
+        assert len(cutout_table) >= 1
+        for row in cutout_table:
+            assert os.path.isfile(row['Local Path'])
+
+        # Test 4: Intentional API call with no results
+        bad_coord = SkyCoord(102.7, 70.50, unit="deg")
+        with pytest.warns(NoResultsWarning, match='Missing HAP files for input target. Cutout not performed.'):
+            cutout_table = mast.Hapcut.download_cutouts(coordinates=bad_coord, size=5, path=str(tmpdir))
+            assert isinstance(cutout_table, Table)
+            assert len(cutout_table) == 0
+
+    def test_hapcut_get_cutouts(self):
+
+        # Test 1: Simple API call with expected results
+        coord = SkyCoord(351.347812, 28.497808, unit="deg")
+
+        cutout_list = mast.Hapcut.get_cutouts(coordinates=coord)
+        assert isinstance(cutout_list, list)
+        assert len(cutout_list) >= 1
+        assert isinstance(cutout_list[0], fits.HDUList)
+        assert cutout_list[0][1].data.shape == (5, 5)
+
+        # Test 2: Make input size a list
+        cutout_list = mast.Hapcut.get_cutouts(coordinates=coord, size=[2, 3])
+        assert isinstance(cutout_list, list)
+        assert len(cutout_list) >= 1
+        assert isinstance(cutout_list[0], fits.HDUList)
+        assert cutout_list[0][1].data.shape == (3, 2)
+
+        # Test 3: Specify unit for input size
+        cutout_list = mast.Hapcut.get_cutouts(coordinates=coord, size=5*u.arcsec)
+        assert isinstance(cutout_list, list)
+        assert len(cutout_list) >= 1
+        assert isinstance(cutout_list[0], fits.HDUList)
+        assert cutout_list[0][1].data.shape == (42, 42)
+
+        # Test 4: Intentional API call with no results
+        bad_coord = SkyCoord(102.7, 70.50, unit="deg")
+
+        with pytest.warns(NoResultsWarning, match='Missing HAP files for input target. Cutout not performed.'):
+            cutout_list = mast.Hapcut.get_cutouts(coordinates=bad_coord)
+            assert isinstance(cutout_list, list)
+            assert len(cutout_list) == 0
diff --git a/docs/mast/mast.rst b/docs/mast/mast.rst
@@ -1120,6 +1120,62 @@ To list the available deep field surveys at a particular location there is `~ast
    ['candels_gn_60mas', 'candels_gn_30mas', 'goods_north']
 
 
+HAPCut
+======
+
+
+HAPCut for MAST allows users to request cutouts from various Hubble Advance Products (HAPs). The cutouts can
+be returned as fits files (image files are not currently supported). This tool can be accessed in
+Astroquery by using the Hapcut class. Documentation for the supported HAPCut API can be found here:
+https://mast.stsci.edu/hapcut/
+
+
+Cutouts
+-------
+
+The `~astroquery.mast.HapcutClass.get_cutouts` function takes a coordinate and cutout size (in pixels or
+an angular quantity) and returns the cutout FITS file(s) as a list of `~astropy.io.fits.HDUList` objects.
+
+If the given coordinate appears in more than one product, a FITS file will be produced for each.
+
+.. doctest-remote-data::
+
+   >>> from astroquery.mast import Hapcut
+   >>> from astropy.coordinates import SkyCoord
+   ...
+   >>> cutout_coord = SkyCoord(351.347812, 28.497808, unit="deg")
+   >>> hdulist = Hapcut.get_cutouts(coordinates=cutout_coord, size=5)
+   >>> hdulist[0].info()    # doctest: +IGNORE_OUTPUT
+   Filename: <class '_io.BytesIO'>
+   No.    Name      Ver    Type      Cards   Dimensions   Format
+   0  PRIMARY       1 PrimaryHDU     754   ()      
+   1  SCI           1 ImageHDU       102   (5, 5)   float32   
+   2  WHT           1 ImageHDU        56   (5, 5)   float32
+
+
+The `~astroquery.mast.HapcutClass.download_cutouts` function takes a coordinate and cutout size (in pixels or
+an angular quantity) and downloads the cutout fits file(s) as fits files.
+
+If the given coordinate appears in more than one product, a cutout will be produced for each.
+
+.. doctest-remote-data::
+
+   >>> from astroquery.mast import Hapcut
+   >>> from astropy.coordinates import SkyCoord
+   ...
+   >>> cutout_coord = SkyCoord(351.347812, 28.497808, unit="deg")
+   >>> manifest = Hapcut.download_cutouts(coordinates=cutout_coord, size=[50, 100])    # doctest: +IGNORE_OUTPUT
+   Downloading URL https://mast.stsci.edu/hapcut/api/v0.1/astrocut?ra=351.347812&dec=28.497808&x=100&y=50&units=px to ./hapcut_20221130112710.zip ... [Done]
+   Inflating...
+   ...
+   >>> print(manifest)    # doctest: +IGNORE_OUTPUT
+                                 Local Path                                   
+   ---------------------------------------------------------------------------------
+   ./hst_cutout_skycell-p2007x09y05-ra351d3478-decn28d4978_wfc3_ir_f160w_coarse.fits
+   ./hst_cutout_skycell-p2007x09y05-ra351d3478-decn28d4978_wfc3_uvis_f606w.fits
+   ./hst_cutout_skycell-p2007x09y05-ra351d3478-decn28d4978_wfc3_uvis_f814w.fits
+
+
 Accessing Proprietary Data
 ==========================
 
