Adding cube_cut function

Author: snbianco

astrocut/__init__.py

@@ -29,7 +29,7 @@
 # Import key submodules and functions if not in setup mode
 if not _ASTROPY_SETUP_:  # noqa
     from .make_cube import CubeFactory, TicaCubeFactory  # noqa
-    from .cube_cut import CutoutFactory  # noqa
+    from .cube_cut import CutoutFactory, cube_cut  # noqa
     from .cutouts import fits_cut, img_cut, normalize_img  # noqa
     from .cutout_processing import (  # noqa
         path_to_footprints, center_on_path, CutoutsCombiner, build_default_combine_function  # noqa

astrocut/cube_cut.py

@@ -2,7 +2,13 @@
 
 """This module implements the cutout functionality."""
 
-from typing import Literal, Union
+from pathlib import Path
+from typing import Literal, Optional, Union, List, Tuple
+
+import astropy.units as u
+import numpy as np
+from astropy.coordinates import SkyCoord
+from s3path import S3Path
 
 from .TessCubeCutout import TessCubeCutout
 
@@ -20,18 +26,11 @@ class CutoutFactory():
     `~astrocut.TessCubeCutout` class.
     """
 
-    def cube_cut(
-        self,
-        cube_file,
-        coordinates,
-        cutout_size,
-        product="SPOC",
-        target_pixel_file=None,
-        output_path=".",
-        memory_only=False,
-        threads: Union[int, Literal["auto"]] = 1,
-        verbose=False,
-    ):
+    def cube_cut(self, cube_file: Union[str, Path, S3Path], coordinates: Union[SkyCoord, str],
+                 cutout_size: Union[int, np.ndarray, u.Quantity, List[int], Tuple[int]], 
+                 product: str = 'SPOC', target_pixel_file: Optional[str] = None, 
+                 output_path: Union[str, Path] = '.', memory_only: bool = False, 
+                 threads: Union[int, Literal["auto"]] = 1, verbose: bool = False):
         """
         Takes a cube file (as created by `~astrocut.CubeFactory`), and makes a cutout target pixel
         file of the given size around the given coordinates. The target pixel file is formatted like
@@ -79,9 +78,10 @@ def cube_cut(
 
         Returns
         -------
-        response: string or None
-            If successful, returns the path to the target pixel file,
-            if unsuccessful returns None.
+        response: string or `~astropy.io.fits.HDUList` or None
+            If successful, returns the target pixel file as an `~astropy.io.fits.HDUList` object,
+            or the path to the target pixel file if saved to disk.
+            If unsuccessful returns None.
         """
         cube_cutout = TessCubeCutout(input_files=cube_file,
                                      coordinates=coordinates,
@@ -102,3 +102,75 @@ def cube_cut(
 
         return cube_cutout.write_as_tpf(output_dir=output_path, 
                                         output_file=target_pixel_file)[0]
+    
+
+
+def cube_cut(cube_file: Union[str, Path, S3Path], coordinates: Union[SkyCoord, str],
+             cutout_size: Union[int, np.ndarray, u.Quantity, List[int], Tuple[int]], 
+             product: str = 'SPOC', target_pixel_file: Optional[str] = None, 
+             output_path: Union[str, Path] = '.', memory_only: bool = False, 
+             threads: Union[int, Literal["auto"]] = 1, verbose: bool = False):
+    """
+    Takes a cube file (as created by `~astrocut.CubeFactory`), and makes a cutout target pixel
+    file of the given size around the given coordinates. The target pixel file is formatted like
+    a TESS pipeline target pixel file.
+
+    This function is maintained for backwards compatibility. For maximum flexibility, we recommend using the
+    `~astrocut.TessCubeCutout` class.
+
+    Parameters
+    ----------
+    cube_file : str
+        The cube file containing all the images to be cutout.
+        Must be in the format returned by ~astrocut.make_cube.
+    coordinates : str or `astropy.coordinates.SkyCoord` object
+        The position around which to cutout.
+        It may be specified as a string ("ra dec" in degrees)
+        or as the appropriate `~astropy.coordinates.SkyCoord` object.
+    cutout_size : int, array-like, `~astropy.units.Quantity`
+        The size of the cutout array. If ``cutout_size``
+        is a scalar number or a scalar `~astropy.units.Quantity`,
+        then a square cutout of ``cutout_size`` will be created.  If
+        ``cutout_size`` has two elements, they should be in ``(ny, nx)``
+        order.  Scalar numbers in ``cutout_size`` are assumed to be in
+        units of pixels. `~astropy.units.Quantity` objects must be in pixel or
+        angular units.
+    product : str
+        The product type to make the cutouts from.
+        Can either be 'SPOC' or 'TICA' (default is 'SPOC').
+    target_pixel_file : str
+        Optional. The name for the output target pixel file.
+        If no name is supplied, the file will be named:
+        ``<cube_file_base>_<ra>_<dec>_<cutout_size>_astrocut.fits``
+    output_path : str
+        Optional. The path where the output file is saved.
+        The current directory is default.
+    memory_only : bool
+        Optional. If true, the cutout is made in memory only and not saved to disk.
+        Default is False.
+    threads : int, "auto", default=1
+        Number of threads to use when making remote (e.g. s3) cutouts, will not use threads for local access
+        <=1 disables the threadpool, >1 sets threadpool to the specified number of threads,
+        "auto" uses `concurrent.futures.ThreadPoolExecutor`'s default: cpu_count + 4, limit to max of 32
+    verbose : bool
+        Optional. If true intermediate information is printed.
+
+    Returns
+    -------
+    response: string or `~astropy.io.fits.HDUList` or None
+        If successful, returns the target pixel file as an `~astropy.io.fits.HDUList` object,
+        or the path to the target pixel file if saved to disk.
+        If unsuccessful, returns None.
+    """
+    cube_cutout = TessCubeCutout(input_files=cube_file,
+                                 coordinates=coordinates,
+                                 cutout_size=cutout_size,
+                                 product=product,
+                                 threads=threads,
+                                 verbose=verbose)
+            
+    if memory_only:
+        return cube_cutout.tpf_cutouts[0]
+    
+    return cube_cutout.write_as_tpf(output_dir=output_path, 
+                                    output_file=target_pixel_file)[0]

astrocut/tests/test_cube_cut.py

@@ -12,7 +12,7 @@
 from astropy.time import Time
 from astropy.wcs import FITSFixedWarning
 
-from ..cube_cut import CutoutFactory
+from ..cube_cut import CutoutFactory, cube_cut
 from ..exceptions import InputWarning
 from ..make_cube import CubeFactory, TicaCubeFactory
 from .utils_for_test import create_test_ffis
@@ -123,17 +123,15 @@ def check_flux(flux, x1, x2, y1, y2, ecube, label, hdulist):
 
 
 @pytest.mark.parametrize("ffi_type", ["SPOC", "TICA"])
-def test_cube_cutout(cube_file, ffi_files, ffi_type, tmp_path):
+@pytest.mark.parametrize("use_factory", [True, False])
+def test_cube_cutout(cube_file, ffi_files, ffi_type, use_factory, tmp_path):
     """
     Testing the cube cutout functionality.
     """
     tmpdir = str(tmp_path)
-
     img_sz = 10
     num_im = 100
 
-    cutout_maker = CutoutFactory()
-
     # Read one of the input images to get the WCS
     n = 1 if ffi_type == 'SPOC' else 0
     img_header = fits.getheader(ffi_files[0], n)
@@ -159,8 +157,14 @@ def test_cube_cutout(cube_file, ffi_files, ffi_type, tmp_path):
     csize[-1] = img_sz+5
     for i, v in enumerate(world_coords):
         coord = SkyCoord(v[0], v[1], frame='icrs', unit='deg')
-        cutout_maker.cube_cut(cube_file, coord, csize[i], ffi_type, target_pixel_file=cutlist[i],
-                              output_path=tmpdir, verbose=False)
+        if use_factory:
+            cutout_maker = CutoutFactory()
+            cutout_maker.cube_cut(cube_file, coord, csize[i], ffi_type, target_pixel_file=cutlist[i],
+                                  output_path=tmpdir, verbose=False)
+        else:
+            cube_cut(cube_file, coord, csize[i], ffi_type, target_pixel_file=cutlist[i],
+                     output_path=tmpdir, verbose=False)
+
 
     # expected values for cube
     ecube = np.zeros((img_sz, img_sz, num_im, 2))
@@ -470,7 +474,7 @@ def test_tica_cutout_error(tmp_path):
 
     ffi_type = 'TICA'
     # Test cutouts from TICA cube with error array
-    test_cube_cutout(tica_cube_error_file, tica_ffi_files, ffi_type, tmp_path)
+    test_cube_cutout(tica_cube_error_file, tica_ffi_files, ffi_type, True, tmp_path)
     test_header_keywords_quality(tica_cube_error_file, ffi_type, tmp_path)
     test_header_keywords_diffs(tica_cube_error_file, ffi_type, tmp_path)
     test_target_pixel_file(tica_cube_error_file, ffi_type, tmp_path)