Python astropy.io.fits.BinTableHDU() Examples

def test_extend_variable_length_array(self):
        """Regression test for https://aeon.stsci.edu/ssb/trac/pyfits/ticket/54"""

        def test(format_code):
            arr = [[1] * 10] * 10
            col1 = fits.Column(name='TESTVLF', format=format_code, array=arr)
            col2 = fits.Column(name='TESTSCA', format='J', array=[1] * 10)
            tb_hdu = fits.BinTableHDU.from_columns([col1, col2], nrows=15)
            # This asserts that the normal 'scalar' column's length was extended
            assert len(tb_hdu.data['TESTSCA']) == 15
            # And this asserts that the VLF column was extended in the same manner
            assert len(tb_hdu.data['TESTVLF']) == 15
            # We can't compare the whole array since the _VLF is an array of
            # objects, but comparing just the edge case rows should suffice
            assert (tb_hdu.data['TESTVLF'][0] == arr[0]).all()
            assert (tb_hdu.data['TESTVLF'][9] == arr[9]).all()
            assert (tb_hdu.data['TESTVLF'][10] == ([0] * 10)).all()
            assert (tb_hdu.data['TESTVLF'][-1] == ([0] * 10)).all()

        for code in ('PJ()', 'QJ()'):
            test(code) 

def test_table_from_bool_fields(self):
        """
        Regression test for https://aeon.stsci.edu/ssb/trac/pyfits/ticket/113

        Tests creating a table from a recarray containing numpy.bool columns.
        """

        array = np.rec.array([(True, False), (False, True)], formats='|b1,|b1')
        thdu = fits.BinTableHDU.from_columns(array)
        assert thdu.columns.formats == ['L', 'L']
        assert comparerecords(thdu.data, array)

        # Test round trip
        thdu.writeto(self.temp('table.fits'))
        data = fits.getdata(self.temp('table.fits'), ext=1)
        assert thdu.columns.formats == ['L', 'L']
        assert comparerecords(data, array) 

def test_dump_load_array_colums(self):
        """
        Regression test for https://github.com/spacetelescope/PyFITS/issues/22

        Ensures that a table containing a multi-value array column can be
        dumped and loaded successfully.
        """

        data = np.rec.array([('a', [1, 2, 3, 4], 0.1),
                             ('b', [5, 6, 7, 8], 0.2)],
                            formats='a1,4i4,f8')
        tbhdu = fits.BinTableHDU.from_columns(data)
        datafile = self.temp('data.txt')
        cdfile = self.temp('coldefs.txt')
        hfile = self.temp('header.txt')

        tbhdu.dump(datafile, cdfile, hfile)
        new_tbhdu = fits.BinTableHDU.load(datafile, cdfile, hfile)
        assert comparerecords(tbhdu.data, new_tbhdu.data)
        assert str(tbhdu.header) == str(new_tbhdu.header) 

def test_dump_load_round_trip(self):
        """
        A simple test of the dump/load methods; dump the data, column, and
        header files and try to reload the table from them.
        """

        hdul = fits.open(self.data('table.fits'))
        tbhdu = hdul[1]
        datafile = self.temp('data.txt')
        cdfile = self.temp('coldefs.txt')
        hfile = self.temp('header.txt')

        tbhdu.dump(datafile, cdfile, hfile)

        new_tbhdu = fits.BinTableHDU.load(datafile, cdfile, hfile)

        assert comparerecords(tbhdu.data, new_tbhdu.data)

        # Double check that the headers are equivalent
        assert str(tbhdu.header) == str(new_tbhdu.header)

        hdul.close() 

def test_bintable_to_asciitable(self):
        """Tests initializing a TableHDU with the data from a BinTableHDU."""

        with fits.open(self.data('tb.fits')) as hdul:
            tbdata = hdul[1].data
            tbhdu = fits.TableHDU(data=tbdata)
            with ignore_warnings():
                tbhdu.writeto(self.temp('test.fits'), overwrite=True)
            with fits.open(self.temp('test.fits')) as hdul2:
                tbdata2 = hdul2[1].data
                assert np.all(tbdata['c1'] == tbdata2['c1'])
                assert np.all(tbdata['c2'] == tbdata2['c2'])
                # c3 gets converted from float32 to float64 when writing
                # test.fits, so cast to float32 before testing that the correct
                # value is retrieved
                assert np.all(tbdata['c3'].astype(np.float32) ==
                              tbdata2['c3'].astype(np.float32))
                # c4 is a boolean column in the original table; we want ASCII
                # columns to convert these to columns of 'T'/'F' strings
                assert np.all(np.where(tbdata['c4'], 'T', 'F') ==
                              tbdata2['c4']) 

def test_dim_column_byte_order_mismatch(self):
        """
        When creating a table column with non-trivial TDIMn, and
        big-endian array data read from an existing FITS file, the data
        should not be unnecessarily byteswapped.

        Regression test for https://github.com/astropy/astropy/issues/3561
        """

        data = fits.getdata(self.data('random_groups.fits'))['DATA']
        col = fits.Column(name='TEST', array=data, dim='(3,1,128,1,1)',
                          format='1152E')
        thdu = fits.BinTableHDU.from_columns([col])
        thdu.writeto(self.temp('test.fits'))

        with fits.open(self.temp('test.fits')) as hdul:
            assert np.all(hdul[1].data['TEST'] == data) 

def test_bin_table_with_logical_array(self):
        c1 = fits.Column(name='flag', format='2L',
                         array=[[True, False], [False, True]])
        coldefs = fits.ColDefs([c1])

        tbhdu1 = fits.BinTableHDU.from_columns(coldefs)

        assert (tbhdu1.data.field('flag')[0] ==
                np.array([True, False], dtype=bool)).all()
        assert (tbhdu1.data.field('flag')[1] ==
                np.array([False, True], dtype=bool)).all()

        tbhdu = fits.BinTableHDU.from_columns(tbhdu1.data)

        assert (tbhdu.data.field('flag')[0] ==
                np.array([True, False], dtype=bool)).all()
        assert (tbhdu.data.field('flag')[1] ==
                np.array([False, True], dtype=bool)).all() 

def test_column_with_scaling(self):
        """Check that a scaled column if correctly saved once it is modified.
        Regression test for https://github.com/astropy/astropy/issues/6887
        """
        c1 = fits.Column(name='c1', array=np.array([1], dtype='>i2'),
                         format='1I', bscale=1, bzero=32768)
        S = fits.HDUList([fits.PrimaryHDU(),
                          fits.BinTableHDU.from_columns([c1])])

        # Change value in memory
        S[1].data['c1'][0] = 2
        S.writeto(self.temp("a.fits"))
        assert S[1].data['c1'] == 2

        # Read and change value in memory
        with fits.open(self.temp("a.fits")) as X:
            X[1].data['c1'][0] = 10
            assert X[1].data['c1'][0] == 10

            # Write back to file
            X.writeto(self.temp("b.fits"))

        # Now check the file
        with fits.open(self.temp("b.fits")) as hdul:
            assert hdul[1].data['c1'][0] == 10 

def test_fits_record_len(self):
        counts = np.array([312, 334, 308, 317])
        names = np.array(['NGC1', 'NGC2', 'NGC3', 'NCG4'])
        c1 = fits.Column(name='target', format='10A', array=names)
        c2 = fits.Column(name='counts', format='J', unit='DN', array=counts)
        c3 = fits.Column(name='notes', format='A10')
        c4 = fits.Column(name='spectrum', format='5E')
        c5 = fits.Column(name='flag', format='L', array=[1, 0, 1, 1])
        coldefs = fits.ColDefs([c1, c2, c3, c4, c5])
        tbhdu = fits.BinTableHDU.from_columns(coldefs)
        tbhdu.writeto(self.temp('table1.fits'))

        t1 = fits.open(self.temp('table1.fits'))

        assert len(t1[1].data[0]) == 5
        assert len(t1[1].data[0][0:4]) == 4
        assert len(t1[1].data[0][0:5]) == 5
        assert len(t1[1].data[0][0:6]) == 5
        assert len(t1[1].data[0][0:7]) == 5
        assert len(t1[1].data[0][1:4]) == 3
        assert len(t1[1].data[0][1:5]) == 4
        assert len(t1[1].data[0][1:6]) == 4
        assert len(t1[1].data[0][1:7]) == 4

        t1.close() 

def test_mask_array(self):
        t = fits.open(self.data('table.fits'))
        tbdata = t[1].data
        mask = tbdata.field('V_mag') > 12
        newtbdata = tbdata[mask]
        hdu = fits.BinTableHDU(newtbdata)
        hdu.writeto(self.temp('newtable.fits'))

        hdul = fits.open(self.temp('newtable.fits'))

        # match to a regex rather than a specific string.
        expect = r"\[\('NGC1002',\s+12.3[0-9]*\) \(\'NGC1003\',\s+15.[0-9]+\)\]"
        assert re.match(expect, str(hdu.data))
        assert re.match(expect, str(hdul[1].data))

        t.close()
        hdul.close() 

def test_modify_column_attributes(self):
        """Regression test for https://github.com/astropy/astropy/issues/996

        This just tests one particular use case, but it should apply pretty
        well to other similar cases.
        """

        NULLS = {'a': 2, 'b': 'b', 'c': 2.3}

        data = np.array(list(zip([1, 2, 3, 4],
                                 ['a', 'b', 'c', 'd'],
                                 [2.3, 4.5, 6.7, 8.9])),
                        dtype=[('a', int), ('b', 'S1'), ('c', float)])

        b = fits.BinTableHDU(data=data)
        for col in b.columns:
            col.null = NULLS[col.name]

        b.writeto(self.temp('test.fits'), overwrite=True)

        with fits.open(self.temp('test.fits')) as hdul:
            header = hdul[1].header
            assert header['TNULL1'] == 2
            assert header['TNULL2'] == 'b'
            assert header['TNULL3'] == 2.3 

def test_variable_length_columns(self):
        def test(format_code):
            col = fits.Column(name='QUAL_SPE', format=format_code,
                              array=[[0] * 1571] * 225)
            tb_hdu = fits.BinTableHDU.from_columns([col])
            pri_hdu = fits.PrimaryHDU()
            hdu_list = fits.HDUList([pri_hdu, tb_hdu])
            with ignore_warnings():
                hdu_list.writeto(self.temp('toto.fits'), overwrite=True)

            with fits.open(self.temp('toto.fits')) as toto:
                q = toto[1].data.field('QUAL_SPE')
                assert (q[0][4:8] ==
                        np.array([0, 0, 0, 0], dtype=np.uint8)).all()
                assert toto[1].columns[0].format.endswith('J(1571)')

        for code in ('PJ()', 'QJ()'):
            test(code) 

def test_variable_length_table_format_pd_from_list(self):
        def test(format_code):
            a = [np.array([7.2e-20, 7.3e-20]), np.array([0.0]),
                 np.array([0.0])]
            acol = fits.Column(name='testa', format=format_code, array=a)
            tbhdu = fits.BinTableHDU.from_columns([acol])
            with ignore_warnings():
                tbhdu.writeto(self.temp('newtable.fits'), overwrite=True)

            with fits.open(self.temp('newtable.fits')) as tbhdu1:
                assert tbhdu1[1].columns[0].format.endswith('D(2)')
                for j in range(3):
                    for i in range(len(a[j])):
                        assert tbhdu1[1].data.field(0)[j][i] == a[j][i]

        for code in ('PD()', 'QD()'):
            test(code) 

def test_column_endianness(self):
        """
        Regression test for https://aeon.stsci.edu/ssb/trac/pyfits/ticket/77
        (Astropy doesn't preserve byte order of non-native order column arrays)
        """

        a = [1., 2., 3., 4.]
        a1 = np.array(a, dtype='<f8')
        a2 = np.array(a, dtype='>f8')

        col1 = fits.Column(name='a', format='D', array=a1)
        col2 = fits.Column(name='b', format='D', array=a2)
        cols = fits.ColDefs([col1, col2])
        tbhdu = fits.BinTableHDU.from_columns(cols)

        assert (tbhdu.data['a'] == a1).all()
        assert (tbhdu.data['b'] == a2).all()

        # Double check that the array is converted to the correct byte-order
        # for FITS (big-endian).
        tbhdu.writeto(self.temp('testendian.fits'), overwrite=True)
        with fits.open(self.temp('testendian.fits')) as hdul:
            assert (hdul[1].data['a'] == a2).all()
            assert (hdul[1].data['b'] == a2).all() 

def test_table_to_hdu():
    from astropy.table import Table
    table = Table([[1, 2, 3], ['a', 'b', 'c'], [2.3, 4.5, 6.7]],
                    names=['a', 'b', 'c'], dtype=['i', 'U1', 'f'])
    table['a'].unit = 'm/s'
    table['b'].unit = 'not-a-unit'
    table.meta['foo'] = 'bar'

    with catch_warnings() as w:
        hdu = fits.BinTableHDU(table, header=fits.Header({'TEST': 1}))
        assert len(w) == 1
        assert str(w[0].message).startswith("'not-a-unit' did not parse as"
                                            " fits unit")

    for name in 'abc':
        assert np.array_equal(table[name], hdu.data[name])

    # Check that TUNITn cards appear in the correct order
    # (https://github.com/astropy/astropy/pull/5720)
    assert hdu.header.index('TUNIT1') < hdu.header.index('TTYPE2')

    assert hdu.header['FOO'] == 'bar'
    assert hdu.header['TEST'] == 1 

def parse_input_healpix_data(input_data, field=0, hdu_in=None, nested=None):
    """
    Parse input HEALPIX data to return a Numpy array and coordinate frame object.
    """

    if isinstance(input_data, (TableHDU, BinTableHDU)):
        data = input_data.data
        header = input_data.header
        coordinate_system_in = parse_coord_system(header['COORDSYS'])
        array_in = data[data.columns[field].name].ravel()
        if 'ORDERING' in header:
            nested = header['ORDERING'].lower() == 'nested'
    elif isinstance(input_data, str):
        hdu = fits.open(input_data)[hdu_in or 1]
        return parse_input_healpix_data(hdu, field=field)
    elif isinstance(input_data, tuple) and isinstance(input_data[0], np.ndarray):
        array_in = input_data[0]
        coordinate_system_in = parse_coord_system(input_data[1])
    else:
        raise TypeError("input_data should either be an HDU object or a tuple of (array, frame)")

    return array_in, coordinate_system_in, nested 

def table_to_fits(table):
    """
    Convert an astropy table to a fits binary table HDU
    :param table: astropy table
    :return: fits BinTableHDU
    """
    hdu = fits.BinTableHDU(table)
    # Put in the description keywords
    for k in hdu.header.keys():
        if 'TTYPE' in k:
            column_name = hdu.header[k].lower()
            description = table[column_name].description
            hdu.header[k] = (column_name.upper(), description)
            # Get the value of n in TTYPE
            n = k[5:]
            hdu.header['TCOMM{0}'.format(n)] = description
    return hdu 

def test_a3dtable(tmpdir):
    testfile = str(tmpdir.join('test.fits'))
    hdu = fits.BinTableHDU.from_columns([
        fits.Column(name='FOO', format='J', array=np.arange(10))
    ])
    hdu.header['XTENSION'] = 'A3DTABLE'
    hdu.writeto(testfile, output_verify='ignore')

    with fits.open(testfile) as hdul:
        assert hdul[1].header['XTENSION'] == 'A3DTABLE'

        with catch_warnings() as w:
            hdul.verify('fix')

        assert str(w[0].message) == 'Verification reported errors:'
        assert str(w[2].message).endswith(
            'Converted the XTENSION keyword to BINTABLE.')

        assert hdul[1].header['XTENSION'] == 'BINTABLE' 

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 test_variable_length_table_format_pa_from_object_array(self):
        def test(format_code):
            a = np.array([np.array(['a', 'b', 'c']), np.array(['d', 'e']),
                          np.array(['f'])], 'O')
            acol = fits.Column(name='testa', format=format_code, array=a)
            tbhdu = fits.BinTableHDU.from_columns([acol])
            with ignore_warnings():
                tbhdu.writeto(self.temp('newtable.fits'), overwrite=True)

            with fits.open(self.temp('newtable.fits')) as hdul:
                assert hdul[1].columns[0].format.endswith('A(3)')
                for j in range(3):
                    for i in range(len(a[j])):
                        assert hdul[1].data.field(0)[j][i] == a[j][i]

        for code in ('PA()', 'QA()'):
            test(code) 

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 test_table_to_hdu(self):
        table = Table([[1, 2, 3], ['a', 'b', 'c'], [2.3, 4.5, 6.7]],
                      names=['a', 'b', 'c'], dtype=['i', 'U1', 'f'])
        table['a'].unit = 'm/s'
        table['b'].unit = 'not-a-unit'

        with catch_warnings() as w:
            hdu = fits.table_to_hdu(table)
            assert len(w) == 1
            assert str(w[0].message).startswith("'not-a-unit' did not parse as"
                                                " fits unit")

        # Check that TUNITn cards appear in the correct order
        # (https://github.com/astropy/astropy/pull/5720)
        assert hdu.header.index('TUNIT1') < hdu.header.index('TTYPE2')

        assert isinstance(hdu, fits.BinTableHDU)
        filename = self.temp('test_table_to_hdu.fits')
        hdu.writeto(filename, overwrite=True) 

def test_getdata_vla(self):
        """Regression test for https://aeon.stsci.edu/ssb/trac/pyfits/ticket/200"""

        def test(format_code):
            col = fits.Column(name='QUAL_SPE', format=format_code,
                              array=[np.arange(1572)] * 225)
            tb_hdu = fits.BinTableHDU.from_columns([col])
            pri_hdu = fits.PrimaryHDU()
            hdu_list = fits.HDUList([pri_hdu, tb_hdu])
            with ignore_warnings():
                hdu_list.writeto(self.temp('toto.fits'), overwrite=True)

            data = fits.getdata(self.temp('toto.fits'))

            # Need to compare to the original data row by row since the FITS_rec
            # returns an array of _VLA objects
            for row_a, row_b in zip(data['QUAL_SPE'], col.array):
                assert (row_a == row_b).all()

        for code in ('PJ()', 'QJ()'):
            test(code) 

def test_variable_length_table_format_pa_from_list(self):
        def test(format_code):
            a = ['a', 'ab', 'abc']
            acol = fits.Column(name='testa', format=format_code, array=a)
            tbhdu = fits.BinTableHDU.from_columns([acol])
            with ignore_warnings():
                tbhdu.writeto(self.temp('newtable.fits'), overwrite=True)

            with fits.open(self.temp('newtable.fits')) as hdul:
                assert hdul[1].columns[0].format.endswith('A(3)')
                for j in range(3):
                    for i in range(len(a[j])):
                        assert hdul[1].data.field(0)[j][i] == a[j][i]

        for code in ('PA()', 'QA()'):
            test(code) 

def test_constructor_copies_header(self):
        """
        Regression test for https://aeon.stsci.edu/ssb/trac/pyfits/ticket/153

        Ensure that a header from one HDU is copied when used to initialize new
        HDU.

        This is like the test of the same name in test_image, but tests this
        for tables as well.
        """

        ifd = fits.HDUList([fits.PrimaryHDU(), fits.BinTableHDU()])
        thdr = ifd[1].header
        thdr['FILENAME'] = 'labq01i3q_rawtag.fits'

        thdu = fits.BinTableHDU(header=thdr)
        ofd = fits.HDUList(thdu)
        ofd[0].header['FILENAME'] = 'labq01i3q_flt.fits'

        # Original header should be unchanged
        assert thdr['FILENAME'] == 'labq01i3q_rawtag.fits' 

def test_regression_5383():

    # Regression test for an undefined variable

    x = np.array([1, 2, 3])
    col = fits.Column(name='a', array=x, format='E')
    hdu = fits.BinTableHDU.from_columns([col])
    del hdu._header['TTYPE1']
    hdu.columns[0].name = 'b' 

def test_pseudo_unsigned_ints(self):
        """
        Tests updating a table column containing pseudo-unsigned ints.
        """

        data = np.array([1, 2, 3], dtype=np.uint32)
        col = fits.Column(name='A', format='1J', bzero=2**31, array=data)
        thdu = fits.BinTableHDU.from_columns([col])
        thdu.writeto(self.temp('test.fits'))

        # Test that the file wrote out correctly
        with fits.open(self.temp('test.fits'), uint=True) as hdul:
            hdu = hdul[1]
            assert 'TZERO1' in hdu.header
            assert hdu.header['TZERO1'] == 2**31
            assert hdu.data['A'].dtype == np.dtype('uint32')
            assert np.all(hdu.data['A'] == data)

            # Test updating the unsigned int data
            hdu.data['A'][0] = 99
            hdu.writeto(self.temp('test2.fits'))

        with fits.open(self.temp('test2.fits'), uint=True) as hdul:
            hdu = hdul[1]
            assert 'TZERO1' in hdu.header
            assert hdu.header['TZERO1'] == 2**31
            assert hdu.data['A'].dtype == np.dtype('uint32')
            assert np.all(hdu.data['A'] == [99, 2, 3]) 

def test_empty_table(tmpdir):
    ofile = str(tmpdir.join('emptytable.fits'))
    hdu = fits.BinTableHDU(header=None, data=None, name='TEST')
    hdu.writeto(ofile)

    with fits.open(ofile) as hdul:
        assert hdul['TEST'].data.size == 0

    ofile = str(tmpdir.join('emptytable.fits.gz'))
    hdu = fits.BinTableHDU(header=None, data=None, name='TEST')
    hdu.writeto(ofile, overwrite=True)

    with fits.open(ofile) as hdul:
        assert hdul['TEST'].data.size == 0 

def test_table_from_columns_of_other_table(self):
        """Tests a rare corner case where the columns of an existing table
        are used to create a new table with the new_table function.  In this
        specific case, however, the existing table's data has not been read
        yet, so new_table has to get at it through the Delayed proxy.

        Note: Although this previously tested new_table it now uses
        BinTableHDU.from_columns directly, around which new_table is a mere
        wrapper.
        """

        hdul = fits.open(self.data('table.fits'))

        # Make sure the column array is in fact delayed...
        assert isinstance(hdul[1].columns._arrays[0], Delayed)

        # Create a new table...
        t = fits.BinTableHDU.from_columns(hdul[1].columns)

        # The original columns should no longer be delayed...
        assert not isinstance(hdul[1].columns._arrays[0], Delayed)

        t.writeto(self.temp('test.fits'))

        with fits.open(self.temp('test.fits')) as hdul2:
            assert comparerecords(hdul[1].data, hdul2[1].data)

        hdul.close() 

def test_bool_column_update(self):
        """Regression test for https://aeon.stsci.edu/ssb/trac/pyfits/ticket/139"""

        c1 = fits.Column('F1', 'L', array=[True, False])
        c2 = fits.Column('F2', 'L', array=[False, True])
        thdu = fits.BinTableHDU.from_columns(fits.ColDefs([c1, c2]))
        thdu.writeto(self.temp('table.fits'))

        with fits.open(self.temp('table.fits'), mode='update') as hdul:
            hdul[1].data['F1'][1] = True
            hdul[1].data['F2'][0] = True

        with fits.open(self.temp('table.fits')) as hdul:
            assert (hdul[1].data['F1'] == [True, True]).all()
            assert (hdul[1].data['F2'] == [True, True]).all()