1654

pending completion

Build # 1654

Build Type

Pull #257

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/LCOGT/banzai/commit/<a class=hub.com/LCOGT/banzai/commit/033334d626ac49606cf29c28e6639936f1a21d10">033334d62<a href="https://github.com/LCOGT/banzai/commit/033334d626ac49606cf29c28e6639936f1a21d10">">Merge </a><a class="double-link" href="https://github.com/LCOGT/banzai/commit/<a class="double-link" href="https://github.com/LCOGT/banzai/commit/f2bf122052e804ec74901ef68d9f233841b45130">f2bf12205</a>">f2bf12205</a><a href="https://github.com/LCOGT/banzai/commit/033334d626ac49606cf29c28e6639936f1a21d10"> into 1bf5aec56">1bf5aec56</a>

Pull Request Pull Request #257: Seek cursor to beginning of buffer before returning.

Run Details

1 of 1 new or added line in 1 file covered. (100.0%)

1769 of 2807 relevant lines covered (63.02%)

0.63 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

42.28

/banzai/utils/fits_utils.py

import logging
from typing import Optional
import requests

from banzai import logs

import numpy as np
from astropy.io import fits
from astropy.coordinates import SkyCoord
from astropy import units
from tenacity import retry, wait_exponential, stop_after_attempt
import io
import os

logger = logging.getLogger('banzai')

FITS_MANDATORY_KEYWORDS = ['SIMPLE', 'BITPIX', 'NAXIS', 'EXTEND', 'COMMENT', 'CHECKSUM', 'DATASUM']


def sanitize_header(header):
    # Remove the mandatory keywords from a header so it can be copied to a new
    # image.
    header = header.copy()

    # Let the new data decide what these values should be
    for i in FITS_MANDATORY_KEYWORDS:
        if i in header.keys():
            header.pop(i)

    return header


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 parse_ra_dec(header):
    try:
        coord = SkyCoord(header.get('CRVAl1'), header.get('CRVAL2'), unit=(units.degree, units.degree))
        ra = coord.ra.deg
        dec = coord.dec.deg
    except (ValueError, TypeError):
        # Fallback to RA and DEC
        try:
            coord = SkyCoord(header.get('RA'), header.get('DEC'), unit=(units.hourangle, units.degree))
            ra = coord.ra.deg
            dec = coord.dec.deg
        except (ValueError, TypeError):
            # Fallback to Cat-RA and CAT-DEC
            try:
                coord = SkyCoord(header.get('CAT-RA'), header.get('CAT-DEC'), unit=(units.hourangle, units.degree))
                ra = coord.ra.deg
                dec = coord.dec.deg
            except (ValueError, TypeError) as e:
                logger.error('Could not get initial pointing guess. {0}'.format(e),
                             extra_tags={'filename': header.get('ORIGNAME')})
                ra, dec = np.nan, np.nan
    return ra, dec


def get_primary_header(filename) -> Optional[fits.Header]:
    try:
        header = fits.getheader(filename, ext=0)
        for keyword in header:
            if keyword not in FITS_MANDATORY_KEYWORDS:
                return header
        return fits.getheader(filename, ext=1)

    except Exception:
        logger.error("Unable to open fits file: {}".format(logs.format_exception()), extra_tags={'filename': filename})
        return None


# Stop after 4 attempts, and back off exponentially with a minimum wait time of 4 seconds, and a maximum of 10.
# If it fails after 4 attempts, "reraise" the original exception back up to the caller.
@retry(wait=wait_exponential(multiplier=2, min=4, max=10), stop=stop_after_attempt(4), reraise=True)
def download_from_s3(file_info, context, is_raw_frame=False):
    frame_id = file_info.get('frameid')

    logger.info(f"Downloading file {file_info.get('filename')} from archive. ID: {frame_id}.",
                extra_tags={'filename': file_info.get('filename'),
                            'attempt_number': download_from_s3.retry.statistics['attempt_number']})

    if is_raw_frame:
        url = f'{context.RAW_DATA_FRAME_URL}/{frame_id}'
        archive_auth_token = context.RAW_DATA_AUTH_TOKEN
    else:
        url = f'{context.ARCHIVE_FRAME_URL}/{frame_id}'
        archive_auth_token = context.ARCHIVE_AUTH_TOKEN
    response = requests.get(url, headers=archive_auth_token).json()
    buffer = io.BytesIO()
    buffer.write(requests.get(response['url'], stream=True).content)
    buffer.seek(0)
    return buffer


def get_configuration_mode(header):
    configuration_mode = header.get('CONFMODE', 'default')
    # If the configuration mode is not in the header, fallback to default to support legacy data
    if configuration_mode == 'N/A' or configuration_mode == 0 or configuration_mode.lower() == 'normal':
        configuration_mode = 'default'

    header['CONFMODE'] = configuration_mode
    return configuration_mode


def open_fits_file(file_info, context, is_raw_frame=False):
    if file_info.get('path') is not None and os.path.exists(file_info.get('path')):
        buffer = open(file_info.get('path'), 'rb')
        filename = os.path.basename(file_info.get('path'))
        frame_id = None
    elif file_info.get('frameid') is not None:
        buffer = download_from_s3(file_info, context, is_raw_frame=is_raw_frame)
        filename = file_info.get('filename')
        frame_id = file_info.get('frameid')
    else:
        raise ValueError('This file does not exist and there is no frame id to get it from S3.')

    hdu_list = fits.open(buffer, memmap=False)
    uncompressed_hdu_list = unpack(hdu_list)
    hdu_list.close()
    buffer.close()
    del hdu_list
    del buffer

    return uncompressed_hdu_list, filename, frame_id


def unpack(compressed_hdulist: fits.HDUList) -> fits.HDUList:
    # If the primary fits header only has the mandatory keywords, then we throw away that extension
    # and extension 1 gets moved to 0
    # Otherwise the primary HDU is kept
    move_1_to_0 = True
    for keyword in compressed_hdulist[0].header:
        if keyword not in FITS_MANDATORY_KEYWORDS:
            move_1_to_0 = False
            break
    if not move_1_to_0 or not isinstance(compressed_hdulist[1], fits.CompImageHDU):
        primary_hdu = fits.PrimaryHDU(data=None, header=compressed_hdulist[0].header)
    else:
        data_type = str(compressed_hdulist[1].data.dtype)
        if 'int' == data_type[:3]:
            data_type = 'u' + data_type
            data = np.array(compressed_hdulist[1].data, getattr(np, data_type))
        else:
            data = compressed_hdulist[1].data
        primary_hdu = fits.PrimaryHDU(data=data, header=compressed_hdulist[1].header)
        if 'ZDITHER0' in primary_hdu.header:
            primary_hdu.header.pop('ZDITHER0')
    hdulist = [primary_hdu]
    if move_1_to_0:
        starting_extension = 2
    else:
        starting_extension = 1
    for hdu in compressed_hdulist[starting_extension:]:
        if isinstance(hdu, fits.CompImageHDU):
            data_type = str(hdu.data.dtype)
            if 'int' == data_type[:3]:
                data_type = getattr(np, 'u' + data_type)
                data = np.array(hdu.data, data_type)
            else:
                data = np.array(hdu.data, hdu.data.dtype)
            header = hdu.header
            hdulist.append(fits.ImageHDU(data=data, header=header))
        elif isinstance(hdu, fits.BinTableHDU):
            hdulist.append(fits.BinTableHDU(data=hdu.data, header=hdu.header))
        else:
            hdulist.append(fits.ImageHDU(data=hdu.data, header=hdu.header))
    return fits.HDUList(hdulist)


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 to_fits_image_extension(data, master_extension_name, extension_name, context, extension_version=None):
    extension_name = master_extension_name + extension_name
    for extname_to_condense in context.EXTENSION_NAMES_TO_CONDENSE:
        extension_name = extension_name.replace(extname_to_condense, '')
    header = fits.Header({'EXTNAME': extension_name})
    if extension_version is not None:
        header['EXTVER'] = extension_version
    return fits.ImageHDU(data=data, header=header)

1	import logging	1✔
2	from typing import Optional	1✔
3	import requests	1✔
4
5	from banzai import logs	1✔
6
7	import numpy as np	1✔
8	from astropy.io import fits	1✔
9	from astropy.coordinates import SkyCoord	1✔
10	from astropy import units	1✔
11	from tenacity import retry, wait_exponential, stop_after_attempt	1✔
12	import io	1✔
13	import os	1✔
14
15	logger = logging.getLogger('banzai')	1✔
16
17	FITS_MANDATORY_KEYWORDS = ['SIMPLE', 'BITPIX', 'NAXIS', 'EXTEND', 'COMMENT', 'CHECKSUM', 'DATASUM']	1✔
18
19
20	def sanitize_header(header):	1✔
21	# Remove the mandatory keywords from a header so it can be copied to a new
22	# image.
23	header = header.copy()	1✔
24
25	# Let the new data decide what these values should be
26	for i in FITS_MANDATORY_KEYWORDS:	1✔
27	if i in header.keys():	1✔
28	header.pop(i)	1✔
29
30	return header	1✔
31
32
33	def table_to_fits(table):	1✔
34	"""
35	Convert an astropy table to a fits binary table HDU
36	:param table: astropy table
37	:return: fits BinTableHDU
38	"""
39	hdu = fits.BinTableHDU(table)	1✔
40	# Put in the description keywords
41	for k in hdu.header.keys():	1✔
42	if 'TTYPE' in k:	1✔
43	column_name = hdu.header[k].lower()	1✔
44	description = table[column_name].description	1✔
45	hdu.header[k] = (column_name.upper(), description)	1✔
46	# Get the value of n in TTYPE
47	n = k[5:]	1✔
48	hdu.header['TCOMM{0}'.format(n)] = description	1✔
49	return hdu	1✔
50
51
52	def parse_ra_dec(header):	1✔
53	try:	1✔
54	coord = SkyCoord(header.get('CRVAl1'), header.get('CRVAL2'), unit=(units.degree, units.degree))	1✔
55	ra = coord.ra.deg	×
56	dec = coord.dec.deg	×
57	except (ValueError, TypeError):	1✔
58	# Fallback to RA and DEC
59	try:	1✔
60	coord = SkyCoord(header.get('RA'), header.get('DEC'), unit=(units.hourangle, units.degree))	1✔
61	ra = coord.ra.deg	1✔
62	dec = coord.dec.deg	1✔
63	except (ValueError, TypeError):	1✔
64	# Fallback to Cat-RA and CAT-DEC
65	try:	1✔
66	coord = SkyCoord(header.get('CAT-RA'), header.get('CAT-DEC'), unit=(units.hourangle, units.degree))	1✔
67	ra = coord.ra.deg	×
68	dec = coord.dec.deg	×
69	except (ValueError, TypeError) as e:	1✔
70	logger.error('Could not get initial pointing guess. {0}'.format(e),	1✔
71	extra_tags={'filename': header.get('ORIGNAME')})
72	ra, dec = np.nan, np.nan	1✔
73	return ra, dec	1✔
74
75
76	def get_primary_header(filename) -> Optional[fits.Header]:	1✔
77	try:	×
78	header = fits.getheader(filename, ext=0)	×
79	for keyword in header:	×
80	if keyword not in FITS_MANDATORY_KEYWORDS:	×
81	return header	×
82	return fits.getheader(filename, ext=1)	×
83
84	except Exception:	×
85	logger.error("Unable to open fits file: {}".format(logs.format_exception()), extra_tags={'filename': filename})	×
86	return None	×
87
88
89	# Stop after 4 attempts, and back off exponentially with a minimum wait time of 4 seconds, and a maximum of 10.
90	# If it fails after 4 attempts, "reraise" the original exception back up to the caller.
91	@retry(wait=wait_exponential(multiplier=2, min=4, max=10), stop=stop_after_attempt(4), reraise=True)	1✔
92	def download_from_s3(file_info, context, is_raw_frame=False):	1✔
93	frame_id = file_info.get('frameid')	×
94
95	logger.info(f"Downloading file {file_info.get('filename')} from archive. ID: {frame_id}.",	×
96	extra_tags={'filename': file_info.get('filename'),
97	'attempt_number': download_from_s3.retry.statistics['attempt_number']})
98
99	if is_raw_frame:	×
100	url = f'{context.RAW_DATA_FRAME_URL}/{frame_id}'	×
101	archive_auth_token = context.RAW_DATA_AUTH_TOKEN	×
102	else:
103	url = f'{context.ARCHIVE_FRAME_URL}/{frame_id}'	×
104	archive_auth_token = context.ARCHIVE_AUTH_TOKEN	×
105	response = requests.get(url, headers=archive_auth_token).json()	×
106	buffer = io.BytesIO()	×
107	buffer.write(requests.get(response['url'], stream=True).content)	×
108	buffer.seek(0)	×
109	return buffer	×
110
111
112	def get_configuration_mode(header):	1✔
113	configuration_mode = header.get('CONFMODE', 'default')	1✔
114	# If the configuration mode is not in the header, fallback to default to support legacy data
115	if configuration_mode == 'N/A' or configuration_mode == 0 or configuration_mode.lower() == 'normal':	1✔
116	configuration_mode = 'default'	1✔
117
118	header['CONFMODE'] = configuration_mode	1✔
119	return configuration_mode	1✔
120
121
122	def open_fits_file(file_info, context, is_raw_frame=False):	1✔
123	if file_info.get('path') is not None and os.path.exists(file_info.get('path')):	×
124	buffer = open(file_info.get('path'), 'rb')	×
125	filename = os.path.basename(file_info.get('path'))	×
126	frame_id = None	×
127	elif file_info.get('frameid') is not None:	×
128	buffer = download_from_s3(file_info, context, is_raw_frame=is_raw_frame)	×
129	filename = file_info.get('filename')	×
130	frame_id = file_info.get('frameid')	×
131	else:
132	raise ValueError('This file does not exist and there is no frame id to get it from S3.')	×
133
134	hdu_list = fits.open(buffer, memmap=False)	×
135	uncompressed_hdu_list = unpack(hdu_list)	×
136	hdu_list.close()	×
137	buffer.close()	×
138	del hdu_list	×
139	del buffer	×
140
141	return uncompressed_hdu_list, filename, frame_id	×
142
143
144	def unpack(compressed_hdulist: fits.HDUList) -> fits.HDUList:	1✔
145	# If the primary fits header only has the mandatory keywords, then we throw away that extension
146	# and extension 1 gets moved to 0
147	# Otherwise the primary HDU is kept
148	move_1_to_0 = True	×
149	for keyword in compressed_hdulist[0].header:	×
150	if keyword not in FITS_MANDATORY_KEYWORDS:	×
151	move_1_to_0 = False	×
152	break	×
153	if not move_1_to_0 or not isinstance(compressed_hdulist[1], fits.CompImageHDU):	×
154	primary_hdu = fits.PrimaryHDU(data=None, header=compressed_hdulist[0].header)	×
155	else:
156	data_type = str(compressed_hdulist[1].data.dtype)	×
157	if 'int' == data_type[:3]:	×
158	data_type = 'u' + data_type	×
159	data = np.array(compressed_hdulist[1].data, getattr(np, data_type))	×
160	else:
161	data = compressed_hdulist[1].data	×
162	primary_hdu = fits.PrimaryHDU(data=data, header=compressed_hdulist[1].header)	×
163	if 'ZDITHER0' in primary_hdu.header:	×
164	primary_hdu.header.pop('ZDITHER0')	×
165	hdulist = [primary_hdu]	×
166	if move_1_to_0:	×
167	starting_extension = 2	×
168	else:
169	starting_extension = 1	×
170	for hdu in compressed_hdulist[starting_extension:]:	×
171	if isinstance(hdu, fits.CompImageHDU):	×
172	data_type = str(hdu.data.dtype)	×
173	if 'int' == data_type[:3]:	×
174	data_type = getattr(np, 'u' + data_type)	×
175	data = np.array(hdu.data, data_type)	×
176	else:
177	data = np.array(hdu.data, hdu.data.dtype)	×
178	header = hdu.header	×
179	hdulist.append(fits.ImageHDU(data=data, header=header))	×
180	elif isinstance(hdu, fits.BinTableHDU):	×
181	hdulist.append(fits.BinTableHDU(data=hdu.data, header=hdu.header))	×
182	else:
183	hdulist.append(fits.ImageHDU(data=hdu.data, header=hdu.header))	×
184	return fits.HDUList(hdulist)	×
185
186
187	def pack(uncompressed_hdulist: fits.HDUList) -> fits.HDUList:	1✔
188	if uncompressed_hdulist[0].data is None:	×
189	primary_hdu = fits.PrimaryHDU(header=uncompressed_hdulist[0].header)	×
190	hdulist = [primary_hdu]	×
191	else:
192	primary_hdu = fits.PrimaryHDU()	×
193	compressed_hdu = fits.CompImageHDU(data=np.ascontiguousarray(uncompressed_hdulist[0].data),	×
194	header=uncompressed_hdulist[0].header, quantize_level=64,
195	dither_seed=2048, quantize_method=1)
196	hdulist = [primary_hdu, compressed_hdu]	×
197
198	for hdu in uncompressed_hdulist[1:]:	×
199	if isinstance(hdu, fits.ImageHDU):	×
200	compressed_hdu = fits.CompImageHDU(data=np.ascontiguousarray(hdu.data), header=hdu.header,	×
201	quantize_level=64, quantize_method=1)
202	hdulist.append(compressed_hdu)	×
203	else:
204	hdulist.append(hdu)	×
205	return fits.HDUList(hdulist)	×
206
207
208	def to_fits_image_extension(data, master_extension_name, extension_name, context, extension_version=None):	1✔
209	extension_name = master_extension_name + extension_name	1✔
210	for extname_to_condense in context.EXTENSION_NAMES_TO_CONDENSE:	1✔
211	extension_name = extension_name.replace(extname_to_condense, '')	1✔
212	header = fits.Header({'EXTNAME': extension_name})	1✔
213	if extension_version is not None:	1✔
214	header['EXTVER'] = extension_version	×
215	return fits.ImageHDU(data=data, header=header)	1✔

LCOGT / banzai / 1654

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous