diff --git a/.github/workflows/python-package-pr.yml b/.github/workflows/python-package-pr.yml index 117c89d..b7ea545 100644 --- a/.github/workflows/python-package-pr.yml +++ b/.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 diff --git a/.github/workflows/python-package-publish.yml b/.github/workflows/python-package-publish.yml index 1d5dc8f..bffe4b9 100644 --- a/.github/workflows/python-package-publish.yml +++ b/.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 diff --git a/.github/workflows/python-package-push.yml b/.github/workflows/python-package-push.yml index 2d08818..f16f1ad 100644 --- a/.github/workflows/python-package-push.yml +++ b/.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 diff --git a/README.md b/README.md index 34417f7..5f248e7 100644 --- a/README.md +++ b/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: diff --git a/healsparse/fits_shim.py b/healsparse/fits_shim.py index c3775c2..a2623be 100644 --- a/healsparse/fits_shim.py +++ b/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,7 +230,14 @@ 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) @@ -200,7 +245,7 @@ def append_extension(self, extension, data): # 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): diff --git a/healsparse/io_coverage_fits.py b/healsparse/io_coverage_fits.py index 4ccdc78..9bfd7ed 100644 --- a/healsparse/io_coverage_fits.py +++ b/healsparse/io_coverage_fits.py @@ -23,9 +23,13 @@ def _read_coverage_fits(coverage_class, filename_or_fits): else: fits = filename_or_fits + # HealSparse map files are required to have a COV extension. + # The exceptions listed here are those that are raised by + # the different fits backends when the extension cannot + # be found. try: cov_index_map = fits.read_ext_data('COV') - except (OSError, KeyError): + except (OSError, KeyError, ValueError): raise RuntimeError("File is not a HealSparseMap") s_hdr = fits.read_ext_header('SPARSE') diff --git a/healsparse/io_map.py b/healsparse/io_map.py index cb3f38d..2176baf 100644 --- a/healsparse/io_map.py +++ b/healsparse/io_map.py @@ -56,9 +56,10 @@ def _read_map(healsparse_class, filename, nside_coverage=None, pixels=None, head fits = HealSparseFits(filename) is_fits_file = True fits.close() - except (OSError, UnicodeDecodeError): + except (OSError, UnicodeDecodeError, ValueError): pass # UnicodeDecodeError occurs when trying to read an hdf5 file as if it's FITS + # ValueError occurs from rustfits when trying to read an hdf5 file as if it's FITS if not is_fits_file: is_parquet_file = check_parquet_dataset(filename) diff --git a/requirements.txt b/requirements.txt index cffb5f2..ebaef61 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ numpy>=1.16.0 hpgeom>=1.2.0 -astropy setuptools_scm setuptools_scm_git_archive diff --git a/setup.cfg b/setup.cfg index a96d111..d1e81b8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -26,7 +26,6 @@ python_requires = >=3.7 install_requires = numpy hpgeom>=1.2.0 - astropy zip_safe = True [options.extras_require] @@ -35,10 +34,15 @@ test = flake8 parquet = pyarrow>=5.0.0 + astropy>=5.3 test_with_healpy = healpy +astropy = + astropy>=5.3 fitsio = - fitsio + fitsio>=1.4.0 +rustfits = + rustfits>=0.1.5 hdf5 = h5py diff --git a/tests/test_cat_files.py b/tests/test_cat_files.py index 0e0e5ac..8afa207 100644 --- a/tests/test_cat_files.py +++ b/tests/test_cat_files.py @@ -96,7 +96,7 @@ def test_cat_maps(self): outfile = os.path.join(self.test_dir, 'test_%s_combined_%d.hs' % (t, int(in_mem))) - if not healsparse.fits_shim.use_fitsio and not in_mem: + if not (healsparse.fits_shim.use_rustfits or healsparse.fits_shim.use_fitsio) and not in_mem: # We cannot use out-of-memory option with astropy.io.fits self.assertRaises(RuntimeError, cat_healsparse_files, file_list, outfile, in_memory=in_mem, @@ -209,7 +209,7 @@ def test_cat_maps_overlap(self): outfile = os.path.join(self.test_dir, 'test_%s_combined_%d.hs' % (t, int(in_mem))) - if not healsparse.fits_shim.use_fitsio and not in_mem: + if not (healsparse.fits_shim.use_rustfits or healsparse.fits_shim.use_fitsio) and not in_mem: # We cannot use out-of-memory option with astropy.io.fits self.assertRaises(RuntimeError, cat_healsparse_files, file_list, outfile, in_memory=in_mem, diff --git a/tests/test_fits_shim.py b/tests/test_fits_shim.py index 94c805d..4b56a64 100644 --- a/tests/test_fits_shim.py +++ b/tests/test_fits_shim.py @@ -142,7 +142,7 @@ def test_append(self): self.test_dir = tempfile.mkdtemp(dir='./', prefix='TestHealSparse-') header = None - if healsparse.fits_shim.use_fitsio: + if healsparse.fits_shim.use_rustfits or healsparse.fits_shim.use_fitsio: # Test appending to image and recarray for t in ['array', 'recarray', 'widemask']: filename = os.path.join(self.test_dir, 'test_%s.fits' % (t)) @@ -184,24 +184,16 @@ def test_append(self): self.assertRaises(RuntimeError, HealSparseFits, filename, mode='rw') - def write_testfile_unused(self, filename, data0, data1, header): - """ - Write a testfile, using astropy.io.fits only. This is in place - until we get full compression support working in both. - """ - _header = healsparse.fits_shim._make_header(header) - _header['EXTNAME'] = 'COV' - healsparse.fits_shim.fits.writeto(filename, data0, - header=_header) - _header['EXTNAME'] = 'SPARSE' - healsparse.fits_shim.fits.append(filename, data1, - header=_header, overwrite=False) - def write_testfile(self, filename, data0, data1, header): """ Write a testfile. """ - if healsparse.fits_shim.use_fitsio: + if healsparse.fits_shim.use_rustfits: + healsparse.fits_shim.rustfits.write(filename, data0, + header=header, extname='COV') + healsparse.fits_shim.rustfits.write(filename, data1, mode='r+', + header=header, extname='SPARSE') + elif healsparse.fits_shim.use_fitsio: healsparse.fits_shim.fitsio.write(filename, data0, header=header, extname='COV') healsparse.fits_shim.fitsio.write(filename, data1, @@ -209,11 +201,18 @@ def write_testfile(self, filename, data0, data1, header): else: _header = healsparse.fits_shim._make_header(header) _header['EXTNAME'] = 'COV' - healsparse.fits_shim.fits.writeto(filename, data0, - header=_header) + healsparse.fits_shim.astropy_fits.writeto( + filename, + data0, + header=_header, + ) _header['EXTNAME'] = 'SPARSE' - healsparse.fits_shim.fits.append(filename, data1, - header=_header, overwrite=False) + healsparse.fits_shim.astropy_fits.append( + filename, + data1, + header=_header, + overwrite=False, + ) def setUp(self): self.test_dir = None