""" This test file uses the https://github.com/esheldon/fitsio package to verify our compression and decompression routines against the implementation in cfitsio. *Note*: The fitsio library is GPL licensed, therefore it could be interpreted that so is this test file. Given that this test file isn't imported anywhere else in the code this shouldn't cause us any issues. Please bear this in mind when editing this file. """ import os import numpy as np import pytest from astropy.io import fits from .conftest import _expand, fitsio_param_to_astropy_param # This is so that tox can force this file to be run, and not be silently # skipped on CI, but in all other test runs it's skipped if fitsio isn't present. if "ASTROPY_ALWAYS_TEST_FITSIO" in os.environ: import fitsio else: fitsio = pytest.importorskip("fitsio") @pytest.fixture( scope="module", params=_expand( [((10,),), ((5,), (1,), (3,))], [((12, 12),), ((1, 12), (4, 5), (6, 6), None)], [((15, 15),), ((1, 15), (5, 1), (5, 5))], [ ((15, 15, 15),), ((5, 5, 1), (5, 7, 1), (1, 5, 4), (1, 1, 15), (15, 1, 5)), ], # Test the situation where the tile shape is passed larger than the # array shape [ ( (4, 4, 5), (5, 5, 5), ), ( (5, 5, 1), None, ), ], # Test shapes which caused errors # This one we can't test here as it causes segfaults in cfitsio # It is tested in test_roundtrip_high_D though. # [ # ((3, 4, 5),), # ((1, 2, 3),), # ], # >3D Data are not currently supported by cfitsio ), ids=lambda x: f"shape: {x[0]} tile_dims: {x[1]}", ) def array_shapes_tile_dims(request, compression_type): shape, tile_dims = request.param # H_COMPRESS needs >=2D data and always 2D tiles if compression_type == "HCOMPRESS_1": if ( # We don't have at least a 2D image len(shape) < 2 or # We don't have 2D tiles np.count_nonzero(np.array(tile_dims) != 1) != 2 or # TODO: The following restrictions can be lifted with some extra work. # The tile is not the first two dimensions of the data tile_dims[0] == 1 or tile_dims[1] == 1 or # The tile dimensions not an integer multiple of the array dims np.count_nonzero(np.array(shape[:2]) % tile_dims[:2]) != 0 ): pytest.xfail( "HCOMPRESS requires 2D tiles, from the first two" "dimensions, and an integer number of tiles along the first two" "axes." ) return shape, tile_dims @pytest.fixture(scope="module") def tile_dims(array_shapes_tile_dims): return array_shapes_tile_dims[1] @pytest.fixture(scope="module") def data_shape(array_shapes_tile_dims): return array_shapes_tile_dims[0] @pytest.fixture(scope="module") def base_original_data(data_shape, dtype, numpy_rng, compression_type): random = numpy_rng.uniform(high=255, size=data_shape) # Set first value to be exactly zero as zero values require special treatment # for SUBTRACTIVE_DITHER_2 random.ravel()[0] = 0.0 # There seems to be a bug with the fitsio library where HCOMPRESS doesn't # work with int16 random data, so use a bit for structured test data. if compression_type.startswith("HCOMPRESS") and "i2" in dtype or "u1" in dtype: random = np.arange(np.prod(data_shape)).reshape(data_shape) return random.astype(dtype) @pytest.fixture(scope="module") def fitsio_compressed_file_path( tmp_path_factory, comp_param_dtype, base_original_data, data_shape, # For debugging tile_dims, ): compression_type, param, dtype = comp_param_dtype if ( base_original_data.ndim > 2 and "u1" in dtype and compression_type == "HCOMPRESS_1" ): pytest.xfail("fitsio won't write these") if compression_type == "PLIO_1" and "f" in dtype: # fitsio fails with a compression error pytest.xfail("fitsio fails to write these") if compression_type == "NOCOMPRESS": pytest.xfail("fitsio does not support NOCOMPRESS") if ( compression_type == "HCOMPRESS_1" and "f" in dtype and param.get("qmethod", None) == 2 ): # fitsio writes these files with very large/incorrect zzero values, whereas # qmethod == 1 works (and the two methods should be identical except for the # treatment of zeros) pytest.xfail("fitsio writes these files with very large/incorrect zzero values") tmp_path = tmp_path_factory.mktemp("fitsio") original_data = base_original_data.astype(dtype) filename = tmp_path / f"{compression_type}_{dtype}.fits" fits = fitsio.FITS(filename, "rw") fits.write(original_data, compress=compression_type, tile_dims=tile_dims, **param) return filename @pytest.fixture(scope="module") def astropy_compressed_file_path( comp_param_dtype, tmp_path_factory, base_original_data, data_shape, # For debugging tile_dims, ): compression_type, param, dtype = comp_param_dtype original_data = base_original_data.astype(dtype) tmp_path = tmp_path_factory.mktemp("astropy") filename = tmp_path / f"{compression_type}_{dtype}.fits" param = fitsio_param_to_astropy_param(param) hdu = fits.CompImageHDU( data=original_data, compression_type=compression_type, tile_shape=None if tile_dims is None else tile_dims[::-1], **param, ) hdu.writeto(filename) return filename def test_decompress( fitsio_compressed_file_path, comp_param_dtype, ): compression_type, param, dtype = comp_param_dtype with fits.open(fitsio_compressed_file_path) as hdul: data = hdul[1].data assert hdul[1]._header["ZCMPTYPE"].replace("ONE", "1") == compression_type assert hdul[1].data.dtype.kind == np.dtype(dtype).kind assert hdul[1].data.dtype.itemsize == np.dtype(dtype).itemsize # The data might not always match the original data exactly in the case of # lossy compression so instead of comparing the array read by astropy to the # original data, we compare it to the data read in by fitsio (as those # should match) fts = fitsio.FITS(fitsio_compressed_file_path) data2 = fts[1].read() np.testing.assert_allclose(data, data2) # The first value should be exactly equal to zero when using SUBTRACTIVE_DITHER_2 if param.get("qmethod", None) == 2: assert data.ravel()[0] == 0.0 def test_compress( astropy_compressed_file_path, compression_type, dtype, ): if compression_type == "NOCOMPRESS": pytest.xfail("fitsio does not support NOCOMPRESS") fts = fitsio.FITS(astropy_compressed_file_path, "r") header = fts[1].read_header() data = fts[1].read() assert header["ZCMPTYPE"] == compression_type assert data.dtype.kind == np.dtype(dtype).kind assert data.dtype.itemsize == np.dtype(dtype).itemsize # The data might not always match the original data exactly in the case of # lossy compression so instead of comparing the array read by fitsio to the # original data, we compare it to the data read in by astropy (as those # should match) with fits.open(astropy_compressed_file_path) as hdul: np.testing.assert_allclose(data, hdul[1].data)