Merge pull request #607 from Eixalde/save_results_fits

Save results fits

.gitignore

@@ -30,6 +30,8 @@ htmlcov/
 .settings
 .project
 .vscode/
+.vipenv/
+justfile
 
 # scripts
 _pypi.sh

vip_hci/fits/fits.py

@@ -1,26 +1,30 @@
 #! /usr/bin/env python
-
 """
 Module with various fits handling functions.
 """
 
 
-__author__ = 'Carlos Alberto Gomez Gonzalez'
-__all__ = ['open_fits',
-           'info_fits',
-           'write_fits',
-           'verify_fits',
-           'byteswap_array']
+__author__ = "Carlos Alberto Gomez Gonzalez"
+__all__ = ["open_fits", "info_fits", "write_fits", "verify_fits", "byteswap_array"]
 
 
 import os
 import numpy as np
 from astropy.io import fits as ap_fits
-
-
-def open_fits(fitsfilename, n=0, header=False, ignore_missing_end=False,
-              precision=np.float32, return_memmap=False, verbose=True,
-              **kwargs):
+from astropy.io.fits import HDUList
+from ..var.paramenum import ALL_FITS
+
+
+def open_fits(
+    fitsfilename,
+    n=0,
+    get_header=False,
+    ignore_missing_end=False,
+    precision=np.float32,
+    return_memmap=False,
+    verbose=True,
+    **kwargs,
+):
     """
     Load a fits file into memory as numpy array.
 
@@ -29,7 +33,8 @@ def open_fits(fitsfilename, n=0, header=False, ignore_missing_end=False,
     fitsfilename : string or pathlib.Path
         Name of the fits file or ``pathlib.Path`` object
     n : int, optional
-        It chooses which HDU to open. Default is the first one.
+        It chooses which HDU to open. Default is the first one. If n is equal
+        to -2, opens and returns all extensions.
     header : bool, optional
         Whether to return the header along with the data or not.
     precision : numpy dtype, optional
@@ -50,42 +55,99 @@ def open_fits(fitsfilename, n=0, header=False, ignore_missing_end=False,
 
     Returns
     -------
-    hdulist : hdulist
-        [memmap=True] FITS file ``n`` hdulist.
-    data : numpy ndarray
-        [memmap=False] Array containing the frames of the fits-cube.
-    header : dict
+    hdulist : HDU or HDUList
+        [memmap=True] FITS file ``n`` hdulist. If n equals -2, returns the whole
+        hdulist.
+    data : numpy ndarray or list of numpy ndarrays
+        [memmap=False] Array containing the frames of the fits-cube. If n
+        equals -2, returns a list of all arrays.
+    header : dict or list of dict
         [memmap=False, header=True] Dictionary containing the fits header.
+        If n equals -2, returns a list of all dictionnaries.
 
     """
     fitsfilename = str(fitsfilename)
     if not os.path.isfile(fitsfilename):
-        fitsfilename += '.fits'
+        fitsfilename += ".fits"
+
+    hdulist = ap_fits.open(
+        fitsfilename, ignore_missing_end=ignore_missing_end, memmap=True, **kwargs
+    )
+
+    # Opening all extensions in a MEF
+    if n == ALL_FITS:
+        data_list = []
+        header_list = []
+        if return_memmap:
+            return hdulist
+
+        for index, element in enumerate(hdulist):
+            data, header = _return_data_fits(
+                hdulist=hdulist,
+                index=index,
+                precision=precision,
+                verbose=verbose,
+            )
+            data_list.append(data)
+            header_list.append(header)
 
-    hdulist = ap_fits.open(fitsfilename, ignore_missing_end=ignore_missing_end,
-                           memmap=True, **kwargs)
-
-    if return_memmap:
-        return hdulist[n]
-    else:
-        data = hdulist[n].data
-        data = np.array(data, dtype=precision)
         hdulist.close()
-        if header:
-            header = hdulist[n].header
+        if get_header:
             if verbose:
-                print("Fits HDU-{} data and header successfully loaded. "
-                      "Data shape: {}".format(n, data.shape))
-            return data, header
+                print(f"All fits HDU data and headers succesfully loaded.")
+            return data_list, header_list
         else:
             if verbose:
-                print("Fits HDU-{} data successfully loaded. "
-                      "Data shape: {}".format(n, data.shape))
+                print(f"All fits HDU data succesfully loaded.")
+            return data_list
+    # Opening only a specified extension
+    else:
+        if return_memmap:
+            return hdulist[n]
+
+        data, header = _return_data_fits(
+            hdulist=hdulist,
+            index=n,
+            precision=precision,
+            verbose=verbose,
+        )
+        hdulist.close()
+        if get_header:
+            return data, header
+        else:
             return data
 
 
+def _return_data_fits(
+    hdulist: HDUList,
+    index: int,
+    precision=np.float32,
+    verbose: bool = True,
+):
+    """
+    Subfunction used to return data (and header) from a given index.
+
+    Parameters
+    ----------
+    hdulist : HDUList
+        List of FITS cubes with their headers.
+    index : int
+        The wanted index to extract.
+    """
+    data = hdulist[index].data
+    data = np.array(data, dtype=precision)
+    header = hdulist[index].header
+
+    if verbose:
+        print(
+            f"Fits HDU-{index} data successfully loaded, header available. "
+            f"Data shape: {data.shape}"
+        )
+    return data, header
+
+
 def byteswap_array(array):
-    """ FITS files are stored in big-endian byte order. All modern CPUs are
+    """FITS files are stored in big-endian byte order. All modern CPUs are
     little-endian byte order, so at some point you have to byteswap the data.
     Some FITS readers (cfitsio, the fitsio python module) do the byteswap when
     reading the data from disk to memory, so we get numpy arrays in native
@@ -153,8 +215,14 @@ def verify_fits(fitsfilename):
             f.verify()
 
 
-def write_fits(fitsfilename, array, header=None, output_verify='exception',
-               precision=np.float32, verbose=True):
+def write_fits(
+    fitsfilename,
+    array,
+    header=None,
+    output_verify="exception",
+    precision=np.float32,
+    verbose=True,
+):
     """
     Write array and header into FITS file.
 
@@ -181,20 +249,20 @@ def write_fits(fitsfilename, array, header=None, output_verify='exception',
 
     """
 
-    if not fitsfilename.endswith('.fits'):
-        fitsfilename += '.fits'
+    if not fitsfilename.endswith(".fits"):
+        fitsfilename += ".fits"
 
     res = "saved"
     if os.path.exists(fitsfilename):
         os.remove(fitsfilename)
-        res = 'overwritten'
+        res = "overwritten"
 
     if isinstance(array, tuple):
         new_hdul = ap_fits.HDUList()
         if header is None:
-            header = [None]*len(array)
+            header = [None] * len(array)
         elif not isinstance(header, tuple):
-            header = [header]*len(array)
+            header = [header] * len(array)
         elif len(header) != len(array):
             msg = "If input header is a tuple, it should have the same length "
             msg += "as tuple of arrays."