"""fits related utilities.
"""
from astropy.io import fits
import numpy as np
from typing import Tuple, List, Optional
[docs]def read_image_fits(fn: str) -> Tuple[np.ndarray, fits.Header]:
with fits.open(fn) as hdulist:
image = hdulist[0].data[0, 0]
header = hdulist[0].header
return image, header
[docs]def write_image_fits(fn, header, data, overwrite=False):
fits.PrimaryHDU(np.reshape(data, newshape=(1, 1, *data.shape)), header=header).writeto(
fn, overwrite=overwrite)
[docs]def write_fits_mask_with_box_xy_coordindates(output_fits_path: str, imsize: int,
center_list: List[Tuple[int, int]], width_list: List[int]) -> str:
"""Writes a fits mask with a list of boxes to file.
Writes a fits mask with a list of boxes to file. Both im+list and center_list are in
WCS X Y coordindates (i.e. transpose of the numpy array
from astropy.fits. The fits will NOT have a header that is accurate since only the pixels will be used.
If you only need one box you can use one-element lists for the arguments.
Args:
output_fits_path: The fits file path.
imsize: size of the image.
center_list: A list of enters of the boxes. Must be in XY index (i.e. what ds9 shows).
width_list: A list of widths of the boxes in pixels, with each element being the width of the corresponding
center_list element.
Returns: The fits mask path.
"""
image = np.zeros(shape=(imsize, imsize), dtype='>f4')
for i, center in enumerate(center_list):
width = width_list[i]
assert width % 2 == 1, 'width must be an odd number'
image[(center[0] - width//2):(center[0] + width//2 + 1), (center[1] - width//2):(center[1] + width//2 + 1)] = \
np.ones(shape=(width, width))
write_image_fits(output_fits_path, get_sample_header(), image.T, overwrite=True)
return output_fits_path
[docs]def co_add(fits_list: List[str], output_fits_path: str, header_index: Optional[int] = None) -> str:
im, header = read_image_fits(fits_list[header_index if header_index else len(fits_list) // 2])
averaged_im = co_add_arr(fits_list, im.shape)
write_image_fits(output_fits_path, header, averaged_im)
return output_fits_path
[docs]def co_add_arr(fits_list, dims):
averaged_im = np.zeros(shape=dims)
n = len(fits_list)
for fn in fits_list:
im, _ = read_image_fits(fn)
averaged_im += (im / n)
return averaged_im