-
Notifications
You must be signed in to change notification settings - Fork 2
/
fits.py
93 lines (76 loc) · 2.36 KB
/
fits.py
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
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
FITS file related functions.
'''
# System modules
from astropy import log
from astropy.io import fits
try:
import fitsio
FAST = True
except ImportError:
log.warning("fitsio module not found: cannot use fast mode.")
FAST = False
# Local modules
def choose_hdu(filename, fast=False):
'''
Detect whether the fits file is compressed with
fpack, and choose the right HDU.
fast: Alternative mode based on fitsio
'''
if fast:
finfo = fitsio.FITS(filename) # Object
finfo_list = [f.get_extnum() for f in finfo if f.is_compressed()]
else:
finfo = fits.info(filename, output=False) # List of tuples.
finfo_list = [f[0] for f in finfo if 'COMPRESSED_IMAGE' in f]
return 0 if not finfo_list else 1 # finfo=0 # finfo_list[0]=1
def get_fits_header(filename, fast=False):
'''
Return the header of the fits file.
fast uses fitsio.
'''
which_hdu = choose_hdu(filename, fast=fast)
if fast:
#header = fitsio.read_header(filename, which_hdu)
with fitsio.FITS(filename) as file_:
header = file_[which_hdu].read_header()
else:
header = fits.getheader(filename, which_hdu)
log.debug("Getting header from {filename}", filename=filename)
return header
def get_fits_data(filename, fast=FAST):
'''
Return the data of the fits file.
If fitsio=True, use fitsio.
'''
which_hdu = choose_hdu(filename, fast=fast)
if fast:
#data = fitsio.read(filename, which_hdu)
with fitsio.FITS(filename) as file_:
data = file_[which_hdu].read()
else:
data = fits.getdata(filename, which_hdu)
log.debug("Getting data from {filename}", filename=filename)
return data
def write_fits(data, output_file, header=None, fast=False):
'''
Write a fits file.
It adds a checksum keyword.
'''
if fast:
hdu = fitsio.FITS(output_file, 'rw')
if header:
hdu.write(data=data, header=header)
else:
hdu.write(data=data)
hdu.close()
else:
if header:
hdu = fits.PrimaryHDU(data, header=header)
else:
hdu = fits.PrimaryHDU(data)
hdu.writeto(output_file, overwrite=True, checksum=True)
log.info(f"Writing fits file to {output_file}")
return hdu