1681

pending completion

Build # 1681

Build Type

Pull #253

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/LCOGT/banzai/commit/<a class=hub.com/LCOGT/banzai/commit/7407598adc0bdea09ea3da5ced520ddec9bad0e0">7407598ad<a href="https://github.com/LCOGT/banzai/commit/7407598adc0bdea09ea3da5ced520ddec9bad0e0">">Merge </a><a class="double-link" href="https://github.com/LCOGT/banzai/commit/<a class="double-link" href="https://github.com/LCOGT/banzai/commit/a5afa896f5b3ce6a6dac42d78dcabad04eb655e9">a5afa896f</a>">a5afa896f</a><a href="https://github.com/LCOGT/banzai/commit/7407598adc0bdea09ea3da5ced520ddec9bad0e0"> into 4ed7aa072">4ed7aa072</a>

Pull Request Pull Request #253: Merge dev into master

Run Details

1399 of 1399 new or added lines in 31 files covered. (100.0%)

1813 of 2834 relevant lines covered (63.97%)

0.64 hits per line

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

51.28

/banzai/utils/fits_utils.py

import logging
from typing import Optional, List
import requests

from banzai import logs
from banzai.context import Context

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)


def reorder_hdus(hdu_list: List,  obstype: str, context: Context):
    for idx, extension_name in enumerate(context.REDUCED_DATA_EXTENSION_ORDERING.get(obstype)):
        if hdu_list[idx].name != extension_name:
            hdu = hdu_list[extension_name]
            hdu_list.remove(hdu)
            hdu_list.insert(idx, hdu)

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

LCOGT / banzai / 1681

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