Python astropy.io.fits.HDUList() Examples

def test_compress():
    """Test the compression functionality"""
    # return None when the factor is not a positive integer
    if not (fits_interp.compress(None, factor=-1) is None): raise AssertionError()
    if not (fits_interp.compress(None, factor=0.3) is None): raise AssertionError()
    if not (fits_interp.compress(None, factor=0) is None): raise AssertionError()
    # test with factor = 10 for speed
    fname = 'tests/test_files/1904-66_AIT.fits'
    hdulist = fits.open(fname)
    cd1 = hdulist[0].header['CDELT1']
    cd2 = hdulist[0].header['CDELT2']
    # compress using CDELT1 and CDELT2
    if not (isinstance(fits_interp.compress(hdulist, factor=10), fits.HDUList)): raise AssertionError()
    hdulist[0].header['CD1_1'] = cd1
    del hdulist[0].header['CDELT1']
    hdulist[0].header['CD2_2'] = cd2
    del hdulist[0].header['CDELT2']
    # compress using CD1_1 and CDELT2_2 instead
    if not (isinstance(fits_interp.compress(hdulist, factor=10), fits.HDUList)): raise AssertionError()
    # now strip CD2_2 and we should get error
    del hdulist[0].header['CD2_2']
    if not (fits_interp.compress(hdulist, factor=10) is None): raise AssertionError()
    # same for CD1_1
    del hdulist[0].header['CD1_1']
    if not (fits_interp.compress(hdulist, factor=10) is None): raise AssertionError() 

def _load_cube_from_file(self, data=None):
        """Initialises a cube from a file."""

        if data is not None:
            assert isinstance(data, fits.HDUList), 'data is not an HDUList object'
        else:
            try:
                with gunzip(self.filename) as gg:
                    self.data = fits.open(gg.name)
            except (IOError, OSError) as err:
                raise OSError('filename {0} cannot be found: {1}'.format(self.filename, err))

        self.header = self.data[1].header
        self.wcs = WCS(self.header)

        self._check_file(self.data[0].header, self.data, 'Cube')

        self._wavelength = self.data['WAVE'].data
        self._shape = (self.header['NAXIS2'], self.header['NAXIS1'])

        self._do_file_checks(self) 

def test_is_fits_image_file_table_img(self, dp_mock):
        with tempfile.TemporaryDirectory() as tmpdir:
            with self.settings(MEDIA_ROOT=tmpdir):
                img_file = os.path.join(tmpdir, 'img.blah')  # file name should be irrelevant
                img = fits.PrimaryHDU(np.arange(100))
                img.header['EXTNAME'] = 'SCI'
                hdul = fits.HDUList([img])
                hdul.writeto(img_file)

                tabimg_file = os.path.join(tmpdir, 'both.fits')
                table = fits.BinTableHDU.from_columns([
                    fits.Column(name='col1', format='I', array=np.array([1, 2, 3])),
                    fits.Column(name='col2', format='I', array=np.array([4, 5, 6]))
                ])
                hdul = fits.HDUList([img, table])
                hdul.writeto(tabimg_file)
                self.data_product.data = tabimg_file
                self.assertTrue(is_fits_image_file(self.data_product.data.file)) 

def __init__(self, primary_hdu=None, fits_extensions=None):

        hdu_list = []

        if primary_hdu is None:

            primary_hdu = fits.PrimaryHDU()

        else:

            assert isinstance(primary_hdu, fits.PrimaryHDU)

        hdu_list.append(primary_hdu)

        if fits_extensions is not None:

            fits_extensions = list(fits_extensions)

            hdu_list.extend([x.hdu for x in fits_extensions])

        # We embed instead of subclassing because the HDUList class has some weird interaction with the
        # __init__ and __new__ methods which makes difficult to do so (we couldn't figure it out)

        self._hdu_list = fits.HDUList(hdus=hdu_list) 

def saveScrns(self, DIR):
        """
        Saves the currently loaded phase screens to file,
        saving the r0 value in the fits header (in units of pixels). 
        Saved phase data is in radians @500nm

        Args:
            DIR (string): The directory to save the screens
        """

        for scrn in range(self.scrnNo):
            logger.info("Write Sreen {}....".format(scrn))
            hdu = fits.PrimaryHDU(self.wholeScrns[scrn])
            hdu.header["R0"] = self.wholeScrnR0 / self.pixel_scale
            hdulist = fits.HDUList([hdu])
            hdulist.writeto(DIR+"/scrn{}.fits".format(scrn))
            hdulist.close()

            logger.info("Done!") 

def save_single_fits(image, name, key_dict=None, image2=None, image2type=None):
        # Save an array into the first extension of a fits file
        h0 = fits.PrimaryHDU()
        h1 = fits.ImageHDU(image, name='DATA')
        if image2 is not None:
            h2 = fits.ImageHDU(image2)
            if image2type is not None:
                h2.header['EXTNAME'] = image2type

        # if a keyword dictionary is provided, put the
        # keywords into the 0th and 1st extension headers
        if key_dict is not None:
            for key in key_dict:
                h0.header[key] = key_dict[key]
                h1.header[key] = key_dict[key]

        if image2 is None:
            hdulist = fits.HDUList([h0, h1])
        else:
            hdulist = fits.HDUList([h0, h1, h2])
        hdulist.writeto(name, overwrite=True) 

def to_fits(self, context) -> Union[fits.HDUList, list]:
        hdu = fits.BinTableHDU(self.data)
        hdu.name = self.name
        # For all TTYPE header keywords, set the header comment
        # from the table column's description.
        for k in self.meta.keys():
            if 'TTYPE' in k:
                column_name = self.meta[k]
                description = self.data[column_name].description
                hdu.header[k] = (column_name.upper(), description)
                # Get the value of n in TTYPEn
                n = k[5:]
                # Also add the TCOMMn header keyword with the description of the table column
                hdu.header['TCOMM{0}'.format(n)] = description

        return [hdu] 

def pack(uncompressed_hdulist: fits.HDUList) -> fits.HDUList:
    if uncompressed_hdulist[0].data is None:
        primary_hdu = fits.PrimaryHDU(header=uncompressed_hdulist[0].header)
        hdulist = [primary_hdu]
    else:
        primary_hdu = fits.PrimaryHDU()
        compressed_hdu = fits.CompImageHDU(data=np.ascontiguousarray(uncompressed_hdulist[0].data),
                                           header=uncompressed_hdulist[0].header, quantize_level=64,
                                           dither_seed=2048, quantize_method=1)
        hdulist = [primary_hdu, compressed_hdu]

    for hdu in uncompressed_hdulist[1:]:
        if isinstance(hdu, fits.ImageHDU):
            compressed_hdu = fits.CompImageHDU(data=np.ascontiguousarray(hdu.data), header=hdu.header,
                                               quantize_level=64, quantize_method=1)
            hdulist.append(compressed_hdu)
        else:
            hdulist.append(hdu)
    return fits.HDUList(hdulist) 

def write_fits(data, header, file_name):
    """
    Combine data and a fits header to write a fits file.

    Parameters
    ----------
    data : numpy.ndarray
        The data to be written.

    header : astropy.io.fits.hduheader
        The header for the fits file.

    file_name : string
        The file to write

    Returns
    -------
    None
    """
    hdu = fits.PrimaryHDU(data)
    hdu.header = header
    hdulist = fits.HDUList([hdu])
    hdulist.writeto(file_name, overwrite=True)
    logging.info("Wrote {0}".format(file_name))
    return 

def load_file_or_hdu(filename):
    """
    Load a file from disk and return an HDUList
    If filename is already an HDUList return that instead

    Parameters
    ----------
    filename : str or HDUList
        File or HDU to be loaded

    Returns
    -------
    hdulist : HDUList
    """
    if isinstance(filename, fits.HDUList):
        hdulist = filename
    else:
        hdulist = fits.open(filename, ignore_missing_end=True)
    return hdulist 

def from_tree(cls, data, ctx):
        hdus = []
        first = True
        for hdu_entry in data:
            header = fits.Header([fits.Card(*x) for x in hdu_entry['header']])
            data = hdu_entry.get('data')
            if data is not None:
                try:
                    data = data.__array__()
                except ValueError:
                    data = None
            if first:
                hdu = fits.PrimaryHDU(data=data, header=header)
                first = False
            elif data.dtype.names is not None:
                hdu = fits.BinTableHDU(data=data, header=header)
            else:
                hdu = fits.ImageHDU(data=data, header=header)
            hdus.append(hdu)
        hdulist = fits.HDUList(hdus)
        return hdulist 

def savefits(self, array, ofile):
        """Save the given array in a fits file

        Parameters
        ----------
        array : numpy.ndarray
            2D array containing the data to be saved

        ofile : str
            Output filename to use
        """
        h0 = fits.PrimaryHDU()
        h1 = fits.ImageHDU()
        h1.data = array

        h1.header['INSTRUME'] = self.instrument
        h1.header['CDELT1'] = self.mosaic_scale_x / 3600.
        h1.header['CDELT2'] = self.mosaic_scale_y / 3600.
        h1.header['CRPIX1'] = self.crpix1
        h1.header['CRPIX2'] = self.crpix2
        h1.header['CRVAL1'] = self.center_ra
        h1.header['CRVAL2'] = self.center_dec
        h1.header['CTYPE1'] = self.mosaic_x_type
        h1.header['CTYPE2'] = self.mosaic_y_type
        h1.header['CUNIT1'] = 'deg'
        h1.header['CUNIT2'] = 'deg'
        h1.header['RA_REF'] = self.center_ra
        h1.header['DEC_REF'] = self.center_dec
        h1.header['ROLL_REF'] = self.mosaic_roll * 180./np.pi
        h1.header['WCSAXES'] = 2
        h1.header['EXTNAME'] = 'SCI    '
        h1.header['EXTVER'] = 1
        h1.header['PC1_1'] = self.cd11 / (self.mosaic_scale_x / 3600.)
        h1.header['PC1_2'] = self.cd12 / (self.mosaic_scale_x / 3600.)
        h1.header['PC2_1'] = self.cd21 / (self.mosaic_scale_y / 3600.)
        h1.header['PC2_2'] = self.cd22 / (self.mosaic_scale_y / 3600.)

        hlist = fits.HDUList([h0, h1])
        hlist.writeto(ofile, overwrite=True) 

def test_exten_num(self, tmpdir):
        filename = tmpdir.join(self.filename).strpath
        hdu1 = fits.PrimaryHDU()
        hdu2 = fits.ImageHDU(self.array)
        hdulist = fits.HDUList([hdu1, hdu2])
        hdulist.writeto(filename)
        main([filename, '-e', '1']) 

def test_initialize_from_fits_with_extension(tmpdir):
    fake_img1 = np.zeros([2, 2])
    fake_img2 = np.arange(4).reshape(2, 2)
    hdu0 = fits.PrimaryHDU()
    hdu1 = fits.ImageHDU(fake_img1)
    hdu2 = fits.ImageHDU(fake_img2)
    hdus = fits.HDUList([hdu0, hdu1, hdu2])
    filename = tmpdir.join('afile.fits').strpath
    hdus.writeto(filename)
    ccd = CCDData.read(filename, hdu=2, unit='adu')
    # ccd should pick up the unit adu from the fits header...did it?
    np.testing.assert_array_equal(ccd.data, fake_img2) 

def test_fits_table(tmpdir):
    a = np.array(
        [(0, 1), (2, 3)],
        dtype=[('A', int), ('B', int)])

    h = fits.HDUList()
    h.append(fits.BinTableHDU.from_columns(a))
    tree = {'fits': h}

    def check_yaml(content):
        assert b'!<tag:astropy.org:astropy/table/table-1.0.0>' in content

    helpers.assert_roundtrip_tree(tree, tmpdir, raw_yaml_check_func=check_yaml) 

def test_exten_name(self, tmpdir):
        filename = tmpdir.join(self.filename).strpath
        hdu1 = fits.PrimaryHDU()
        extname = 'SCI'
        hdu2 = fits.ImageHDU(self.array)
        hdu2.header['EXTNAME'] = extname
        hdulist = fits.HDUList([hdu1, hdu2])
        hdulist.writeto(filename)
        main([filename, '-e', extname]) 

def test_read_returns_image(tmpdir):
    # Test if CCData.read returns a image when reading a fits file containing
    # a table and image, in that order.
    tbl = Table(np.ones(10).reshape(5, 2))
    img = np.ones((5, 5))
    hdul = fits.HDUList(hdus=[fits.PrimaryHDU(), fits.TableHDU(tbl.as_array()),
                              fits.ImageHDU(img)])
    filename = tmpdir.join('table_image.fits').strpath
    hdul.writeto(filename)
    ccd = CCDData.read(filename, unit='adu')
    # Expecting to get (5, 5), the size of the image
    assert ccd.data.shape == (5, 5)


# https://github.com/astropy/astropy/issues/9664 

def test_to_hdu():
    ccd_data = create_ccd_data()
    ccd_data.meta = {'observer': 'Edwin Hubble'}
    fits_hdulist = ccd_data.to_hdu()
    assert isinstance(fits_hdulist, fits.HDUList)
    for k, v in ccd_data.meta.items():
        assert fits_hdulist[0].header[k] == v
    np.testing.assert_array_equal(fits_hdulist[0].data, ccd_data.data) 

def test_fromMEF(tmpdir):
    ccd_data = create_ccd_data()
    hdu = fits.PrimaryHDU(ccd_data)
    hdu2 = fits.PrimaryHDU(2 * ccd_data.data)
    hdulist = fits.HDUList(hdu)
    hdulist.append(hdu2)
    filename = tmpdir.join('afile.fits').strpath
    hdulist.writeto(filename)
    # by default, we reading from the first extension
    cd = CCDData.read(filename, unit=u.electron)
    np.testing.assert_array_equal(cd.data, ccd_data.data)
    # but reading from the second should work too
    cd = CCDData.read(filename, hdu=1, unit=u.electron)
    np.testing.assert_array_equal(cd.data, 2 * ccd_data.data) 

def test_sliced_ccdata_to_hdu():
    wcs = WCS(naxis=2)
    wcs.wcs.crpix = 10, 10
    ccd = CCDData(np.ones((10, 10)), wcs=wcs, unit='pixel')
    trimmed = ccd[2:-2, 2:-2]
    hdul = trimmed.to_hdu()
    assert isinstance(hdul, fits.HDUList)
    assert hdul[0].header['CRPIX1'] == 8
    assert hdul[0].header['CRPIX2'] == 8 

def _get_extension_data(self, extension, data, data_origin='file'):
        """Returns the value of an extension for this fibre, either from file or API.

        Parameters
        ----------
        extension : datamodel object
            The datamodel object containing the information for the extension
            we want to retrieve.
        data : ~astropy.io.fits.HDUList
            An `~astropy.io.fits.HDUList` object containing the RSS
            information.

        """

        # Determine if this is an RSS datamodel object or an spectrum.
        # If the origin is the API, the extension data contains a single spectrum,
        # not a row-stacked array, so we consider it a 1D array.
        is_extension_data_1D = isinstance(extension, SpectrumDataModel) or data_origin == 'api'

        value = data[extension.fits_extension()].data

        if extension.has_mask():
            mask = data[extension.fits_extension('mask')].data
        else:
            mask = None

        if hasattr(extension, 'has_ivar') and extension.has_ivar():
            ivar = data[extension.fits_extension('ivar')].data
        elif hasattr(extension, 'has_std') and extension.has_std():
            std = data[extension.fits_extension('std')].data
            ivar = 1. / (std**2)
        else:
            ivar = None

        # If this is an RSS, gets the right row in the stacked spectra.
        if not is_extension_data_1D:
            value = value[self.fiberid, :]
            mask = mask[self.fiberid, :] if mask is not None else None
            ivar = ivar[self.fiberid, :] if ivar is not None else None

        return value, ivar, mask 

def _load_itm_library(library_file):
    """Load ITM FITS file

    Parameters
    ----------
    library_path : str
        Path pointing to the location of the PSF library

    Returns
    -------
    library : photutils.griddedPSFModel
        Object containing PSF library
    """
    data = fits.getdata(library_file)
    hdr = fits.getheader(library_file)
    if data.shape == (2048, 2048):
        # Normalize the data
        data /= np.sum(data)

        # Add PSF location and oversampling keywords
        hdr['DET_YX0'] = ('(1023, 1023)', "The #0 PSF's (y,x) detector pixel position")
        hdr['OVERSAMP'] = (1, 'Oversampling factor for FFTs in computation')

        # Convert to HDUList and create library
        phdu = fits.PrimaryHDU(data, hdr)
        hdulist = fits.HDUList(phdu)
        library = to_griddedpsfmodel(hdulist)

        return library
    else:
        raise ValueError('Expecting ITM data of size (2048, 2048), not {}'.format(data.shape)) 

def save_dispersed_seed_image(self, seed_image):
        """Save the dispersed seed image"""
        primary_hdu = fits.PrimaryHDU()
        image_hdu = fits.ImageHDU(seed_image)
        hdu_list = fits.HDUList([primary_hdu, image_hdu])
        hdu_list[0].header['units'] = 'e/sec'
        hdu_list[1].header['units'] = 'e/sec'
        if self.disp_seed_filename is None:
            self.disp_seed_filename = self.default_dispersed_filename
        hdu_list.writeto(self.disp_seed_filename, overwrite=True)
        print(("Dispersed seed image saved to {}".format(self.disp_seed_filename))) 

def reorder_hdus(hdu_list: fits.HDUList, extensions: list):
    """
    Re-order HDUs in an HDUList before writing to disk
    :param hdu_list: Astropy fits.HDUList
    :param extensions: Ordered list of extensions by EXTNAME
    """
    if extensions is None:
        extensions = []
    extensions += [hdu.name for hdu in hdu_list]
    # Use an ordered dict to get unique elements
    extensions = list(OrderedDict.fromkeys(extensions))
    hdu_list.sort(key=lambda x: extensions.index(x.name)) 

def MakeFITS(model, fitsfile=None):
    '''
    Generate a FITS file for a given :py:mod:`everest` run.

    :param model: An :py:mod:`everest` model instance

    '''

    # Get the fits file name
    if fitsfile is None:
        outfile = os.path.join(model.dir, model._mission.FITSFile(
            model.ID, model.season, model.cadence))
    else:
        outfile = os.path.join(model.dir, fitsfile)
    if os.path.exists(outfile) and not model.clobber:
        return
    elif os.path.exists(outfile):
        os.remove(outfile)

    log.info('Generating FITS file...')

    # Create the HDUs
    primary = PrimaryHDU(model)
    lightcurve = LightcurveHDU(model)
    pixels = PixelsHDU(model)
    aperture = ApertureHDU(model)
    images = ImagesHDU(model)
    hires = HiResHDU(model)

    # Combine to get the HDUList
    hdulist = pyfits.HDUList(
        [primary, lightcurve, pixels, aperture, images, hires])

    # Output to the FITS file
    hdulist.writeto(outfile)

    return 

def to_fits(self, context):
        hdu_list_to_write = fits.HDUList([])
        for hdu in self._hdus:
            hdu_list_to_write += hdu.to_fits(context)
        fits_utils.reorder_hdus(hdu_list_to_write, self.hdu_order)
        if not isinstance(hdu_list_to_write[0], fits.PrimaryHDU):
            hdu_list_to_write[0] = fits.PrimaryHDU(data=hdu_list_to_write[0].data, header=hdu_list_to_write[0].header)
        fits_utils.convert_extension_datatypes(hdu_list_to_write, context.REDUCED_DATA_EXTENSION_TYPES)
        if context.fpack:
            hdu_list_to_write = fits_utils.pack(hdu_list_to_write)
        return hdu_list_to_write 

def to_fits(self, context):
        data_hdu = fits.ImageHDU(data=self.data, header=fits.Header(self.meta))
        bpm_hdu = fits_utils.to_fits_image_extension(self.mask, self.extension_name, 'BPM', context,
                                                     extension_version=self.meta.get('EXTVER'))
        uncertainty_hdu = fits_utils.to_fits_image_extension(self.uncertainty, self.extension_name, 'ERR', context,
                                                             extension_version=self.meta.get('EXTVER'))
        hdulist = fits.HDUList([data_hdu, bpm_hdu, uncertainty_hdu])
        return hdulist 

def to_fits(self, context) -> Union[fits.HDUList, list]:
        return [fits.ImageHDU(data=self.data, header=self.meta)] 

def to_fits(self, context) -> Union[fits.HDUList, list]:
        pass 

def convert_extension_datatypes(hdu_list: fits.HDUList, extension_datatypes: dict):
    """
    Convert extensions' data types into desired form.
    :param hdu_list: FITS HDUList
    :param extension_datatypes: Dictionary of desired data types, keyed by extension name
    """
    for hdu in hdu_list:
        if hdu.name in extension_datatypes:
            hdu.data = hdu.data.astype(extension_datatypes[hdu.name])