Add support for rustfits fits backend.

.github/workflows/python-package-pr.yml

@@ -13,7 +13,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.11", "3.12", "3.13", "3.14"]
 
     steps:
     - uses: actions/checkout@v3
@@ -24,7 +24,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install .[parquet,test,test_with_healpy,hdf5]
+        pip install .[parquet,test,test_with_healpy,hdf5,astropy]
     - name: Lint with flake8
       run: |
         flake8
@@ -34,3 +34,5 @@ jobs:
         pytest
         pip install fitsio>=1.0.5
         pytest
+        pip install rustfits
+        pytest

.github/workflows/python-package-publish.yml

@@ -25,7 +25,7 @@ jobs:
       - name: Build and install
         run: |
           python -m pip install --upgrade pip setuptools
-          python -m pip install .[parquet,test,test_with_healpy]
+          python -m pip install .[parquet,test,test_with_healpy,rustfits]
       - name: Run tests
         run: |
           cd tests

.github/workflows/python-package-push.yml

@@ -24,7 +24,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install .[parquet,test,test_with_healpy]
+        python -m pip install .[parquet,test,test_with_healpy,astropy]
     - name: Lint with flake8
       run: |
         flake8
@@ -34,3 +34,5 @@ jobs:
         pytest
         pip install fitsio>=1.0.5
         pytest
+        pip install rustfits
+        pytest

README.md

@@ -19,8 +19,8 @@ In addition, there is general support for any [numpy](https://github.com/numpy/n
 - [hpgeom](https://github.com/LSSTDESC/hpgeom)
 - [astropy](https://astropy.org)
 
-The following package is optional but recommended for all features including reading full maps in HEALPix format:
-- [fitsio](https://github.com/esheldon/fitsio)
+The following packages are optional but recommended for all features including reading full maps in HEALPix format:
+- [rustfits](https://github.com/esheldon/rustfits)
 - [healpy](https://github.com/healpy/healpy/)
 
 ## Install:

healsparse/fits_shim.py

@@ -1,16 +1,33 @@
+import copy
 import numpy as np
-import mmap
 from .utils import is_integer_value
-# We need this for compression before a newer version of fitsio arrives
-import astropy.io.fits as fits
 
+use_rustfits = False
 use_fitsio = False
+has_astropy = False
+
+try:
+    import rustfits
+    use_rustfits = True
+except ImportError:
+    pass
+
+if not use_rustfits:
+    try:
+        import fitsio
+        use_fitsio = True
+    except ImportError:
+        pass
+
 try:
-    import fitsio
-    use_fitsio = True
+    import astropy.io.fits as astropy_fits
+    has_astropy = True
 except ImportError:
     pass
 
+if not use_rustfits and not use_fitsio and not has_astropy:
+    raise ImportError("HealSparse requires rustfits or fitsio or astropy to be installed.")
+
 
 _image_bitpix2npy = {
     8: 'u1',
@@ -29,12 +46,12 @@
                  'ZPCOUNT', 'ZGCOUNT', 'ZTILE1', 'ZCMPTYPE',
                  'ZNAME1', 'ZVAL1', 'ZQUANTIZ',
                  'SIMPLE', 'BITPIX', 'NAXIS', 'NAXIS1', 'NAXIS2',
-                 'PCOUNT', 'GCOUNT']
+                 'PCOUNT', 'GCOUNT', 'XTENSION']
 
 
 class HealSparseFits(object):
     """
-    Wrapper class to handle fitsio or astropy.io.fits
+    Wrapper class to handle rustfits, fitsio, or astropy.io.fits
     """
     def __init__(self, filename, mode='r'):
         """
@@ -54,15 +71,27 @@ def __init__(self, filename, mode='r'):
         self._filename = filename
         self._mode = mode
 
-        if use_fitsio:
+        if use_rustfits:
+            if mode == "r":
+                rustfits_mode = "r"
+            elif mode == "rw":
+                rustfits_mode = "r+"
+
+            self.fits_object = rustfits.FITS(filename, mode=rustfits_mode)
+
+            try:
+                _ = self.fits_object[0]
+            except ValueError:
+                raise IOError("File %s does not appear to be a fits file." % (filename))
+        elif use_fitsio:
             self.fits_object = fitsio.FITS(filename, mode=mode)
-        else:
+        elif has_astropy:
             if mode == 'r':
                 fits_mode = 'readonly'
             else:
                 raise RuntimeError('Readonly is only useful mode supported for astropy.io.fits')
-            self.fits_object = fits.open(filename, memmap=True, lazy_load_hdus=True,
-                                         mode=fits_mode)
+            self.fits_object = astropy_fits.open(filename, memmap=True, lazy_load_hdus=True,
+                                                 mode=fits_mode)
 
     def read_ext_header(self, extension):
         """
@@ -80,6 +109,7 @@ def read_ext_header(self, extension):
         if use_fitsio:
             return self.fits_object[extension].read_header()
         else:
+            # This works for rustfits and astropy fits.
             return self.fits_object[extension].header
 
     def get_ext_dtype(self, extension):
@@ -95,13 +125,19 @@ def get_ext_dtype(self, extension):
         -------
         dtype : `np.dtype`
         """
-        if use_fitsio:
+        if use_rustfits:
             hdu = self.fits_object[extension]
-            if hdu.get_exttype() == 'IMAGE_HDU':
+            if self.ext_is_image(extension):
+                return hdu.dtype.str
+            else:
+                return hdu.dtype
+        elif use_fitsio:
+            hdu = self.fits_object[extension]
+            if self.ext_is_image(extension):
                 return _image_bitpix2npy[hdu.get_info()['img_equiv_type']]
             else:
                 return self.fits_object[extension].get_rec_dtype()[0]
-        else:
+        elif has_astropy:
             hdu = self.fits_object[extension]
             if hdu.is_image:
                 return _image_bitpix2npy[hdu._bitpix]
@@ -126,7 +162,7 @@ def read_ext_data(self, extension, row_range=None, col_range=None):
         -------
         data : `np.ndarray`
         """
-        if use_fitsio:
+        if use_rustfits or use_fitsio:
             hdu = self.fits_object[extension]
             if row_range is None:
                 return hdu.read()
@@ -135,7 +171,7 @@ def read_ext_data(self, extension, row_range=None, col_range=None):
             else:
                 return hdu[slice(col_range[0], col_range[1]),
                            slice(row_range[0], row_range[1])]
-        else:
+        elif has_astropy:
             # Note that for astropy this does not actually seem to work
             # read a subregion from a tile-compressed image; it reads
             # the full thing.
@@ -169,12 +205,14 @@ def ext_is_image(self, extension):
         is_image : `bool`
         """
         hdu = self.fits_object[extension]
-        if use_fitsio:
+        if use_rustfits:
+            return isinstance(hdu, rustfits.ImageHDU)
+        elif use_fitsio:
             if hdu.get_exttype() == 'IMAGE_HDU':
                 return True
             else:
                 return False
-        else:
+        elif has_astropy:
             return hdu.is_image
 
     def append_extension(self, extension, data):
@@ -192,15 +230,22 @@ def append_extension(self, extension, data):
             raise RuntimeError("Appending only allowed in rw mode")
 
         hdu = self.fits_object[extension]
-        if use_fitsio:
+        if use_rustfits:
+            if hasattr(hdu, 'extend'):
+                # We can append
+                hdu.extend(data)
+            else:
+                firstrow = (hdu.get_dims()[0], )
+                hdu.write(data, start=firstrow)
+        elif use_fitsio:
             if hasattr(hdu, 'append'):
                 # A recarray that we can append to
                 hdu.append(data)
             else:
                 # An image that we cannot append to
                 firstrow = (hdu.get_dims()[0], )
                 hdu.write(data, start=firstrow)
-        else:
+        elif has_astropy:
             raise RuntimeError("Appending is not supported by astropy.io.fits")
 
     def close(self):
@@ -260,7 +305,26 @@ def _write_filename(filename, c_hdr, s_hdr, cov_index_map, sparse_map,
             _tile_shape = (compress_tilesize, )
             s_hdr['RESHAPED'] = False
 
-    if use_fitsio and integer_map:
+    if use_rustfits:
+        with rustfits.FITS(filename, mode="w+") as f:
+            f.write_image(cov_index_map, extname=c_hdr["EXTNAME"], header=c_hdr)
+
+            if compress:
+                if compression == "GZIP_2":
+                    comp = rustfits.Gzip2(tile_shape=_tile_shape)
+                elif compression == "RICE_1":
+                    comp = rustfits.Rice1(tile_shape=_tile_shape)
+
+                f.write_image(
+                    _sparse_map,
+                    extname=s_hdr["EXTNAME"],
+                    header=s_hdr,
+                    compress=comp,
+                )
+            else:
+                f.write(sparse_map, extname=s_hdr["EXTNAME"], header=s_hdr)
+
+    elif use_fitsio:
         # Preferred because it is faster for integer writes.
         # Floating point writing with compression has only just
         # been fixed and I don't want to put a lower limit on
@@ -283,62 +347,32 @@ def _write_filename(filename, c_hdr, s_hdr, cov_index_map, sparse_map,
                 f.write(sparse_map, extname=s_hdr["EXTNAME"], header=s_hdr)
 
     else:
-        hdu_list = fits.HDUList()
+        hdu_list = astropy_fits.HDUList()
 
-        hdu = fits.PrimaryHDU(data=cov_index_map, header=fits.Header())
+        hdu = astropy_fits.PrimaryHDU(data=cov_index_map, header=astropy_fits.Header())
         _make_hierarch_header(c_hdr, hdu.header)
         hdu_list.append(hdu)
 
         if compress:
-            try:
-                # Try new tile_shape API (astropy>=5.3).
-                hdu = fits.CompImageHDU(data=_sparse_map, header=fits.Header(),
-                                        compression_type=compression,
-                                        tile_shape=_tile_shape,
-                                        quantize_level=0.0)
-            except TypeError:
-                # Fall back to old tile_size API.
-                hdu = fits.CompImageHDU(data=sparse_map, header=fits.Header(),
-                                        compression_type=compression,
-                                        tile_size=_tile_shape,
-                                        quantize_level=0.0)
+            # Try new tile_shape API (astropy>=5.3).
+            hdu = astropy_fits.CompImageHDU(
+                data=_sparse_map,
+                header=astropy_fits.Header(),
+                compression_type=compression,
+                tile_shape=_tile_shape,
+                quantize_level=0.0,
+            )
         else:
             if sparse_map.dtype.fields is not None:
-                hdu = fits.BinTableHDU(data=sparse_map, header=fits.Header())
+                hdu = astropy_fits.BinTableHDU(data=sparse_map, header=astropy_fits.Header())
             else:
-                hdu = fits.ImageHDU(data=sparse_map, header=fits.Header())
+                hdu = astropy_fits.ImageHDU(data=sparse_map, header=astropy_fits.Header())
 
         _make_hierarch_header(s_hdr, hdu.header)
         hdu_list.append(hdu)
 
         hdu_list.writeto(filename, overwrite=True)
 
-        # When writing a gzip unquantized (lossless) floating point image,
-        # current versions of astropy (4.0.1 and earlier, at least) write
-        # the ZQUANTIZ header value as NO_DITHER, while cfitsio expects
-        # this to be NONE for unquantized data.  The only way to overwrite
-        # this reserved header keyword is to manually overwrite the bytes
-        # in the file.  The following code uses mmap to overwrite the
-        # necessary header keyword without loading the full image into
-        # memory.  Note that healsparse files only have one compressed
-        # extension, so there will only be one use of ZQUANTIZ in the file.
-        if compress and not is_integer_value(sparse_map[0]):
-            with open(filename, "r+b") as f:
-                try:
-                    mm = mmap.mmap(f.fileno(), 0)
-                    loc = mm.find(b"ZQUANTIZ= 'NO_DITHER'")
-                    if loc >= 0:
-                        mm.seek(loc)
-                        mm.write(b"ZQUANTIZ= 'NONE     '")
-                except OSError:
-                    # Some systems do not have the mmap available,
-                    # we need to read in the full file.
-                    data = f.read()
-                    loc = data.find(b"ZQUANTIZ= 'NO_DITHER'")
-                    if loc >= 0:
-                        f.seek(loc)
-                        f.write(b"ZQUANTIZ= 'NONE     '")
-
 
 def _make_header(metadata, force_astropy=False):
     """
@@ -355,10 +389,18 @@ def _make_header(metadata, force_astropy=False):
     -------
     header : `fitsio.FITSHDR` or `astropy.io.fits.Header`
     """
+    if use_rustfits and not force_astropy:
+        if metadata is None:
+            return {}
+        else:
+            hdr = copy.copy(metadata)
+            for reserved in FITS_RESERVED:
+                hdr.pop(reserved, None)
+            return hdr
     if use_fitsio and not force_astropy:
         hdr = fitsio.FITSHDR(metadata)
     else:
-        hdr = fits.Header()
+        hdr = astropy_fits.Header()
 
     if metadata is not None:
         _make_hierarch_header(metadata, hdr)
@@ -368,7 +410,7 @@ def _make_header(metadata, force_astropy=False):
 
 def _write_healpix_filename(filename, hdr, output_struct):
     """
-    Write to a filename, HEALPix EXPLICIT format, using astropy.io.fits.
+    Write to a filename, HEALPix EXPLICIT format.
 
     This assumes that you want to overwrite any existing file (as should be
     checked in the calling function.)
@@ -382,14 +424,19 @@ def _write_healpix_filename(filename, hdr, output_struct):
     output_struct : `numpy.recarray`
         Correctly formatted output struct.
     """
-    hdu_list = fits.HDUList()
+    if use_rustfits:
+        rustfits.write(filename, output_struct, header=hdr, mode="w+")
+    elif use_fitsio:
+        fitsio.write(filename, output_struct, header=hdr, clobber=True)
+    elif has_astropy:
+        hdu_list = astropy_fits.HDUList()
 
-    hdu = fits.BinTableHDU(data=output_struct, header=fits.Header())
+        hdu = astropy_fits.BinTableHDU(data=output_struct, header=astropy_fits.Header())
 
-    _make_hierarch_header(hdr, hdu.header, skip_reserved=False)
-    hdu_list.append(hdu)
+        _make_hierarch_header(hdr, hdu.header, skip_reserved=False)
+        hdu_list.append(hdu)
 
-    hdu_list.writeto(filename, overwrite=True)
+        hdu_list.writeto(filename, overwrite=True)
 
 
 def _make_hierarch_header(hdr_in, hdr_out, skip_reserved=True):