File: fits_embed.py

package info (click to toggle)
python-asdf 2.7.2-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye, sid
  • size: 3,248 kB
  • sloc: python: 13,104; makefile: 125
file content (330 lines) | stat: -rw-r--r-- 12,327 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-


"""
Utilities for embedded ADSF files in FITS.
"""
import io
import re

import numpy as np

from . import asdf
from . import block
from . import util
from . import generic_io

try:
    from astropy.io import fits
except ImportError:
    raise ImportError("AsdfInFits requires astropy")


ASDF_EXTENSION_NAME = 'ASDF'
FITS_SOURCE_PREFIX = 'fits:'


__all__ = ['AsdfInFits']


class _FitsBlock:
    def __init__(self, hdu):
        self._hdu = hdu

    def __repr__(self):
        return '<FitsBlock {0},{1}>'.format(self._hdu.name, self._hdu.ver)

    def __len__(self):
        return self._hdu.data.nbytes

    @property
    def data(self):
        return self._hdu.data

    @property
    def readonly(self):
        return False

    @property
    def array_storage(self):
        return 'fits'

    def override_byteorder(self, byteorder):
        # FITS data is always stored in big-endian byte order.
        # The data array may not report big-endian, but we want
        # the value written to the tree to match the actual
        # byte order on disk.
        return 'big'

    @property
    def trust_data_dtype(self):
        # astropy.io.fits returns arrays in native byte order
        # when it has to apply scaling.  In that case, we don't
        # want to interpret the bytes as big-endian, since astropy
        # has already converted them properly.
        return True


class _EmbeddedBlockManager(block.BlockManager):
    def __init__(self, hdulist, asdffile):
        self._hdulist = hdulist

        super(_EmbeddedBlockManager, self).__init__(asdffile)

    def get_block(self, source):
        if (isinstance(source, str) and source.startswith(FITS_SOURCE_PREFIX)):
            parts = re.match(
                # All printable ASCII characters are allowed in EXTNAME
                '((?P<name>[ -~]+),)?(?P<ver>[0-9]+)',
                source[len(FITS_SOURCE_PREFIX):])
            if parts is not None:
                ver = int(parts.group('ver'))
                if parts.group('name'):
                    pair = (parts.group('name'), ver)
                else:
                    pair = ver
                return _FitsBlock(self._hdulist[pair])
            else:
                raise ValueError("Can not parse source '{0}'".format(source))

        return super(_EmbeddedBlockManager, self).get_block(source)

    def get_source(self, block):
        if isinstance(block, _FitsBlock):
            for i, hdu in enumerate(self._hdulist):
                if hdu is block._hdu:
                    if hdu.name == '':
                        return '{0}{1}'.format(
                            FITS_SOURCE_PREFIX, i)
                    else:
                        return '{0}{1},{2}'.format(
                            FITS_SOURCE_PREFIX, hdu.name, hdu.ver)
            raise ValueError("FITS block seems to have been removed")

        return super(_EmbeddedBlockManager, self).get_source(block)

    def find_or_create_block_for_array(self, arr, ctx):
        from .tags.core import ndarray

        if not isinstance(arr, ndarray.NDArrayType):
            base = util.get_array_base(arr)
            for hdu in self._hdulist:
                if hdu.data is None:
                    continue
                if base is util.get_array_base(hdu.data):
                    return _FitsBlock(hdu)

        return super(
            _EmbeddedBlockManager, self).find_or_create_block_for_array(arr, ctx)


class AsdfInFits(asdf.AsdfFile):
    """
    Embed ASDF tree content in a FITS file.

    The YAML rendering of the tree is stored in a special FITS
    extension with the EXTNAME of ``ASDF``.  Arrays in the ASDF tree
    may refer to binary data in other FITS extensions by setting
    source to a string with the prefix ``fits:`` followed by an
    ``EXTNAME``, ``EXTVER`` pair, e.g. ``fits:SCI,0``.

    Examples
    --------
    Create a FITS file with ASDF structure, based on an existing FITS
    file::

        from astropy.io import fits

        hdulist = fits.HDUList()
        hdulist.append(fits.ImageHDU(np.arange(512, dtype=float), name='SCI'))
        hdulist.append(fits.ImageHDU(np.arange(512, dtype=float), name='DQ'))

        tree = {
            'model': {
                'sci': {
                    'data': hdulist['SCI'].data,
                    'wcs': 'WCS info'
                },
                'dq': {
                    'data': hdulist['DQ'].data,
                    'wcs': 'WCS info'
                }
            }
        }

        ff = fits_embed.AsdfInFits(hdulist, tree)
        ff.write_to('test.fits')  # doctest: +SKIP
    """

    def __init__(self, hdulist=None, tree=None, **kwargs):
        if hdulist is None:
            hdulist = fits.HDUList()
        super(AsdfInFits, self).__init__(tree=tree, **kwargs)
        self._blocks = _EmbeddedBlockManager(hdulist, self)
        self._hdulist = hdulist
        self._close_hdulist = False

    def __exit__(self, type, value, traceback):
        super(AsdfInFits, self).__exit__(type, value, traceback)
        if self._close_hdulist:
            self._hdulist.close()
        self._tree = {}

    def close(self):
        super(AsdfInFits, self).close()
        if self._close_hdulist:
            self._hdulist.close()
        self._tree = {}

    @classmethod
    def open(cls, fd, uri=None, validate_checksums=False, extensions=None,
             ignore_version_mismatch=True, ignore_unrecognized_tag=False,
             strict_extension_check=False, ignore_missing_extensions=False,
             validate_on_read=True):
        """Creates a new AsdfInFits object based on given input data

        Parameters
        ----------
        fd : FITS HDUList instance, URI string, or file-like object
            May be an already opened instance of a FITS HDUList instance,
            string ``file`` or ``http`` URI, or a Python file-like object.

        uri : str, optional
            The URI for this ASDF file.  Used to resolve relative
            references against.  If not provided, will be
            automatically determined from the associated file object,
            if possible and if created from `asdf.open`.

        validate_checksums : bool, optional
            If `True`, validate the blocks against their checksums.
            Requires reading the entire file, so disabled by default.

        extensions : list of AsdfExtension, optional
            A list of extensions to the ASDF to support when reading
            and writing ASDF files.  See `asdf.types.AsdfExtension` for
            more information.

        ignore_version_mismatch : bool, optional
            When `True`, do not raise warnings for mismatched schema versions.

        strict_extension_check : bool, optional
            When `True`, if the given ASDF file contains metadata about the
            extensions used to create it, and if those extensions are not
            installed, opening the file will fail. When `False`, opening a file
            under such conditions will cause only a warning. Defaults to
            `False`.

        ignore_missing_extensions : bool, optional
            When `True`, do not raise warnings when a file is read that
            contains metadata about extensions that are not available. Defaults
            to `False`.

        validate_on_read : bool, optional
            When `True`, validate the newly opened file against tag and custom
            schemas.  Recommended unless the file is already known to be valid.
        """
        return cls._open_impl(fd, uri=uri,
                       validate_checksums=validate_checksums,
                       extensions=extensions,
                       ignore_version_mismatch=ignore_version_mismatch,
                       ignore_unrecognized_tag=ignore_unrecognized_tag,
                       strict_extension_check=strict_extension_check,
                       ignore_missing_extensions=ignore_missing_extensions,
                       validate_on_read=validate_on_read)

    @classmethod
    def _open_impl(cls, fd, uri=None, validate_checksums=False, extensions=None,
             ignore_version_mismatch=True, ignore_unrecognized_tag=False,
             strict_extension_check=False, _extension_metadata=None,
             ignore_missing_extensions=False, validate_on_read=True):

        close_hdulist = False
        if isinstance(fd, fits.hdu.hdulist.HDUList):
            hdulist = fd
        else:
            uri = generic_io.get_uri(fd)
            try:
                hdulist = fits.open(fd)
                # Since we created this HDUList object, we need to be
                # responsible for cleaning up upon close() or __exit__
                close_hdulist = True
            except IOError:
                msg = "Failed to parse given file '{}'. Is it FITS?"
                raise ValueError(msg.format(uri))

        self = cls(hdulist, uri=uri, extensions=extensions,
                   ignore_version_mismatch=ignore_version_mismatch,
                   ignore_unrecognized_tag=ignore_unrecognized_tag)
        if _extension_metadata is not None:
            self._extension_metadata = _extension_metadata

        self._close_hdulist = close_hdulist

        try:
            asdf_extension = hdulist[ASDF_EXTENSION_NAME]
        except (KeyError, IndexError, AttributeError):
            # This means there is no ASDF extension
            return self

        buff = io.BytesIO(asdf_extension.data)

        try:
            return cls._open_asdf(self, buff, uri=uri, mode='r',
                              validate_checksums=validate_checksums,
                              strict_extension_check=strict_extension_check,
                              ignore_missing_extensions=ignore_missing_extensions,
                              validate_on_read=validate_on_read)
        except RuntimeError:
            self.close()
            raise

    def _create_hdu(self, buff, use_image_hdu):
        # Allow writing to old-style ImageHDU for backwards compatibility
        if use_image_hdu:
            array = np.frombuffer(buff.getvalue(), np.uint8)
            return fits.ImageHDU(array, name=ASDF_EXTENSION_NAME)
        else:
            data = np.array(buff.getbuffer(), dtype=np.uint8)[None, :]
            fmt = '{}B'.format(len(data[0]))
            column = fits.Column(array=data, format=fmt, name='ASDF_METADATA')
            return fits.BinTableHDU.from_columns([column], name=ASDF_EXTENSION_NAME)

    def _update_asdf_extension(self, all_array_storage=None,
                               all_array_compression=None, auto_inline=None,
                               pad_blocks=False, use_image_hdu=False):
        if self.blocks.streamed_block is not None:
            raise ValueError(
                "Can not save streamed data to ASDF-in-FITS file.")

        buff = io.BytesIO()
        super(AsdfInFits, self).write_to(
            buff, all_array_storage=all_array_storage,
            all_array_compression=all_array_compression,
            auto_inline=auto_inline, pad_blocks=pad_blocks,
            include_block_index=False)

        if ASDF_EXTENSION_NAME in self._hdulist:
            del self._hdulist[ASDF_EXTENSION_NAME]
        self._hdulist.append(self._create_hdu(buff, use_image_hdu))

    def write_to(self, filename, all_array_storage=None,
                 all_array_compression=None, auto_inline=None,
                 pad_blocks=False, use_image_hdu=False, *args, **kwargs):
        self._update_asdf_extension(
            all_array_storage=all_array_storage,
            all_array_compression=all_array_compression,
            auto_inline=auto_inline, pad_blocks=pad_blocks,
            use_image_hdu=use_image_hdu)

        self._hdulist.writeto(filename, *args, **kwargs)

    def update(self, all_array_storage=None, all_array_compression=None,
               auto_inline=None, pad_blocks=False):
        raise NotImplementedError(
            "In-place update is not currently implemented for ASDF-in-FITS")

        self._update_asdf_extension(
            all_array_storage=all_array_storage,
            all_array_compression=all_array_compression,
            auto_inline=auto_inline, pad_blocks=pad_blocks)