Merge pull request #73 from iprafols/fitisio_migration

Fitisio migration

.github/workflows/python-package.yml

@@ -28,6 +28,12 @@ jobs:
         python -m pip install --upgrade pip
         pip install wheel pytest pytest-cov coveralls
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    # TODO: the section below can be removed if the compile step is fixed otherwise
+    - name: Reinstall fitsio to ensure the numpy version during compilation matches the one at runtime
+      run: |
+        pip uninstall fitsio -y
+        pip cache purge
+        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
     - name: Install SQUEzE
       run: pip install -e .
     - name: Test with pytest

py/squeze/boss_spectrum.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 from numpy.random import randn
-import astropy.io.fits as fits
+import fitsio
 
 from squeze.error import Error
 from squeze.spectrum import Spectrum
@@ -71,10 +71,10 @@ def __init__(self, spectrum_file, metadata, sky_mask, mask_jpas=False,
             raise Error("""The property "SPECID" must be present in metadata""")
 
         # open fits file
-        spectrum_hdu = fits.open(spectrum_file)
+        spectrum_hdul = fitsio.FITS(spectrum_file)
 
         # compute sky mask
-        self._wave = 10**spectrum_hdu[1].data["LOGLAM"].copy()
+        self._wave = 10**spectrum_hdul[1]["LOGLAM"][:]
         masklambda = sky_mask[0]
         margin = sky_mask[1]
         self.__skymask = None
@@ -87,8 +87,8 @@ def __init__(self, spectrum_file, metadata, sky_mask, mask_jpas=False,
         # store the wavelength, flux and inverse variance as arrays
         # The 1.0 mulitplying is added to change type from >4f to np.float
         # this is required by numba later on
-        self._flux = 1.0*spectrum_hdu[1].data["FLUX"].copy()
-        self._ivar = 1.0*spectrum_hdu[1].data["IVAR"].copy()
+        self._flux = 1.0*spectrum_hdul[1]["FLUX"][:]
+        self._ivar = 1.0*spectrum_hdul[1]["IVAR"][:]
         # mask pixels
         self._ivar[self.__skymask] = 0.0
         self._metadata = metadata
@@ -112,8 +112,7 @@ def __init__(self, spectrum_file, metadata, sky_mask, mask_jpas=False,
             self._ivar = self._ivar[pos].copy()
             self._flux = self._flux[pos].copy()
 
-        del spectrum_hdu[1].data
-        spectrum_hdu.close()
+        spectrum_hdul.close()
 
     def __add_noise(self, noise_amount):
         """ Adds noise to the spectrum by adding a gaussian random number of width

py/squeze/candidates.py

@@ -15,7 +15,7 @@
 import time
 
 import pandas as pd
-import astropy.io.fits as fits
+import fitsio
 from astropy.table import Table
 
 from squeze.common_functions import verboseprint
@@ -476,18 +476,14 @@ def __load_model_settings(self):
 
     def save_candidates(self):
         """ Save the candidates DataFrame. """
-        def convert_dtype(dtype):
-             if dtype == "O":
-                 return "15A"
-             else:
-                 return dtype
-
-        hdu = fits.BinTableHDU.from_columns([fits.Column(name=col,
-                                                         format=convert_dtype(dtype),
-                                                         array=self.__candidates[col])
-                                             for col, dtype in zip(self.__candidates.columns,
-                                                                   self.__candidates.dtypes)])
-        hdu.writeto(self.__name, overwrite=True)
+        results = fitsio.FITS(self.__name, 'rw', clobber=True)
+        names = [col for col in self.__candidates.columns]
+        cols = [np.array(self.__candidates[col].values, dtype=str)
+                if self.__candidates[col].dtype == "object"
+                else self.__candidates[col].values
+                for col in self.__candidates.columns]
+        results.write(cols, names=names, extname="CANDIDATES")
+        results.close()
 
     def candidates(self):
         """ Access the candidates DataFrame. """
@@ -981,22 +977,18 @@ def save_catalogue(self, filename, prob_cut):
         if filename is None:
             filename = self.__name.replace(".fits", "_catalogue.fits")
 
-        def convert_dtype(dtype):
-             if dtype == "O":
-                 return "15A"
-             else:
-                 return dtype
-
         # filter data DataFrame
         data_frame = self.__candidates[(~self.__candidates["DUPLICATED"]) &
                                        (self.__candidates["PROB"] >= prob_cut)]
 
-        hdu = fits.BinTableHDU.from_columns([fits.Column(name=col,
-                                                         format=convert_dtype(dtype),
-                                                         array=data_frame[col])
-                                             for col, dtype in zip(data_frame.columns,
-                                                                   data_frame.dtypes)])
-        hdu.writeto(filename, overwrite=True)
+        results = fitsio.FITS(filename, 'rw', clobber=True)
+        names = [col for col in data_frame.columns]
+        cols = [np.array(data_frame[col].values, dtype=str)
+                if data_frame[col].dtype == "object"
+                else data_frame[col].values
+                for col in data_frame.columns]
+        results.write(cols, names=names, extname="CANDIDATES")
+        results.close()
 
 if __name__ == '__main__':
     pass