Skip to main content
`Ctrl`+`K`

# Source code for astropy.io.fits.hdu.compressed.section

    # Licensed under a 3-clause BSD style license - see LICENSE.rst

    import numpy as np

    from astropy.io.fits.hdu.base import BITPIX2DTYPE
    from astropy.io.fits.hdu.compressed._tiled_compression import (
        decompress_image_data_section,
    )
    from astropy.io.fits.hdu.compressed.utils import _n_tiles
    from astropy.utils.shapes import simplify_basic_index

    __all__ = ["CompImageSection"]

    class CompImageSection:
        """
        Class enabling subsets of CompImageHDU data to be loaded lazily via slicing.

        Slices of this object load the corresponding section of an image array from
        the underlying FITS file, and applies any BSCALE/BZERO factors.

        Section slices cannot be assigned to, and modifications to a section are
        not saved back to the underlying file.

        See the :ref:`astropy:data-sections` section of the Astropy documentation
        for more details.
        """

        def __init__(self, hdu):
            self.hdu = hdu

            self._data_shape = self.hdu.shape
            self._tile_shape = self.hdu.tile_shape

            self._n_dim = len(self._data_shape)
            self._n_tiles = np.array(
                _n_tiles(self._data_shape, self._tile_shape), dtype=int
            )

        @property
        def shape(self):
            return tuple(self._data_shape)

        @property
        def ndim(self):
            return len(self.hdu.shape)

        @property
        def dtype(self):
            return np.dtype(BITPIX2DTYPE[self.hdu._bitpix])

        def __getitem__(self, index):
            if self.hdu._bintable is None:
                return self.hdu.data[index]

            # Shortcut if the whole data is requested (this is used by the
            # data property, so we optimize it as it is frequently used)
            if index is Ellipsis:
                first_tile_index = np.zeros(self._n_dim, dtype=int)
                last_tile_index = self._n_tiles - 1
                data = decompress_image_data_section(
                    self.hdu._bintable.data,
                    self.hdu.compression_type,
                    self.hdu._bintable.header,
                    self.hdu._bintable,
                    self.hdu.header,
                    first_tile_index,
                    last_tile_index,
                )
                if self.hdu._do_not_scale_image_data:
                    return data
                else:
                    return self.hdu._scale_data(data)

            index = simplify_basic_index(index, shape=self._data_shape)

            # Determine for each dimension the first and last tile to extract

            first_tile_index = np.zeros(self._n_dim, dtype=int)
            last_tile_index = np.zeros(self._n_dim, dtype=int)

            final_array_index = []

            for dim, idx in enumerate(index):
                if isinstance(idx, slice):
                    if idx.step > 0:
                        first_tile_index[dim] = idx.start // self._tile_shape[dim]
                        last_tile_index[dim] = (idx.stop - 1) // self._tile_shape[dim]
                    else:
                        stop = 0 if idx.stop is None else max(idx.stop - 1, 0)
                        first_tile_index[dim] = stop // self._tile_shape[dim]
                        last_tile_index[dim] = idx.start // self._tile_shape[dim]

                    # Because slices such as slice(5, 0, 1) can exist (which
                    # would be empty) we need to make sure last_tile_index is
                    # always larger than first_tile_index
                    last_tile_index = np.maximum(last_tile_index, first_tile_index)

                    if idx.step < 0 and idx.stop is None:
                        final_array_index.append(idx)
                    else:
                        final_array_index.append(
                            slice(
                                idx.start - self._tile_shape[dim] * first_tile_index[dim],
                                idx.stop - self._tile_shape[dim] * first_tile_index[dim],
                                idx.step,
                            )
                        )
                else:
                    first_tile_index[dim] = idx // self._tile_shape[dim]
                    last_tile_index[dim] = first_tile_index[dim]
                    final_array_index.append(
                        idx - self._tile_shape[dim] * first_tile_index[dim]
                    )

            data = decompress_image_data_section(
                self.hdu._bintable.data,
                self.hdu.compression_type,
                self.hdu._bintable.header,
                self.hdu._bintable,
                self.hdu.header,
                first_tile_index,
                last_tile_index,
            )

            data = data[tuple(final_array_index)]

            if self.hdu._do_not_scale_image_data:
                return data
            else:
                return self.hdu._scale_data(data)