def test_beam_grids(fits_header, header_l, header_m, l_axis, m_axis):
from africanus.util.beams import beam_grids, axis_and_sign
hdr = fits_header
hdr['CTYPE1'] = header_l
hdr['CTYPE2'] = header_m
l_ax, l_sgn = axis_and_sign(l_axis, "L")
m_ax, m_sgn = axis_and_sign(m_axis, "M")
# Extract l, m and frequency axes and grids
(l, l_grid), (m, m_grid), (freq,
freq_grid) = beam_grids(fits_header, l_axis,
m_axis)
# Check expected L
assert hdr['CTYPE%d' % l] == header_l
crval = hdr['CRVAL%d' % l]
cdelt = hdr['CDELT%d' % l]
crpix = hdr['CRPIX%d' % l] - 1 # C-indexing
R = np.arange(0.0, float(hdr['NAXIS%d' % l]))
exp_l = (R - crpix) * cdelt + crval
exp_l = np.deg2rad(exp_l) * l_sgn
assert_array_almost_equal(exp_l, l_grid)
assert hdr['CTYPE%d' % m] == header_m
crval = hdr['CRVAL%d' % m]
cdelt = hdr['CDELT%d' % m]
crpix = hdr['CRPIX%d' % m] - 1 # C-indexing
R = np.arange(0.0, float(hdr['NAXIS%d' % m]))
exp_m = (R - crpix) * cdelt + crval
exp_m = np.deg2rad(exp_m) * m_sgn
assert_array_almost_equal(exp_m, m_grid)
# GFREQS used for the frequency grid
gfreqs = [
fits_header.get('GFREQ%d' % (i + 1))
for i in range(fits_header['NAXIS3'])
]
assert_array_almost_equal(freq_grid, gfreqs)
def test_beam_grids(fits_header):
from africanus.util.beams import beam_grids
hdr = fits_header
# Extract l, m and frequency axes and grids
(l, l_grid), (m, m_grid), (freq, freq_grid) = beam_grids(fits_header)
# Check expected L
crval = hdr['CRVAL%d' % l]
cdelt = hdr['CDELT%d' % l]
crpix = hdr['CRPIX%d' % l] - 1 # C-indexing
R = np.arange(0.0, float(hdr['NAXIS%d' % l]))
exp_l = (R - crpix) * cdelt + crval
exp_l = np.deg2rad(exp_l)
assert np.allclose(exp_l, l_grid)
crval = hdr['CRVAL%d' % m]
cdelt = hdr['CDELT%d' % m]
crpix = hdr['CRPIX%d' % m] - 1 # C-indexing
R = np.arange(0.0, float(hdr['NAXIS%d' % m]))
# Check expected M. It's -M in the FITS header
# so there's a flip in direction here
exp_m = (R - crpix) * cdelt + crval
exp_m = np.deg2rad(exp_m)
exp_m = np.flipud(exp_m)
assert np.allclose(exp_m, m_grid)
# GFREQS used for the frequency grid
gfreqs = [
fits_header.get('GFREQ%d' % (i + 1))
for i in range(fits_header['NAXIS3'])
]
assert np.allclose(freq_grid, gfreqs)
def load_beams(beam_file_schema, corr_types):
class FITSFile(object):
""" Exists so that fits file is closed when last ref is gc'd """
def __init__(self, filename):
self.hdul = hdul = fits.open(filename)
assert len(hdul) == 1
self.__del_ref = weakref.ref(self, lambda r: hdul.close())
# Open files and get headers
beam_files = []
headers = []
for corr, (re, im) in beam_filenames(beam_file_schema, corr_types).items():
re_f = FITSFile(re)
im_f = FITSFile(im)
beam_files.append((corr, (re_f, im_f)))
headers.append((corr, (re_f.hdul[0].header, im_f.hdul[0].header)))
# All FITS headers should agree
flat_headers = [d for k, v in headers for d in v]
if not all(flat_headers[0] == h for h in flat_headers[1:]):
raise ValueError("BEAM FITS Header Files differ")
# Map FITS header type to NumPy type
BITPIX_MAP = {8: np.dtype('uint8').type, 16: np.dtype('int16').type,
32: np.dtype('int32').type, -32: np.dtype('float32').type,
-64: np.dtype('float64').type}
header = flat_headers[0]
bitpix = header['BITPIX']
try:
dtype = BITPIX_MAP[bitpix]
except KeyError:
raise ValueError("No mapping from BITPIX %s to a numpy type" % bitpix)
else:
dtype = np.result_type(dtype, np.complex64)
if not header['NAXIS'] == 3:
raise ValueError("FITS must have exactly three axes. "
"L or X, M or Y and FREQ. NAXIS != 3")
(l_ax, l_grid), (m_ax, m_grid), (nu_ax, nu_grid) = beam_grids(header)
# Shape of each correlation
shape = (l_grid.shape[0], m_grid.shape[0], nu_grid.shape[0])
# Axis tranpose, FITS is FORTRAN ordered
ax = (nu_ax - 1, m_ax - 1, l_ax - 1)
def _load_correlation(re, im, ax):
# Read real and imaginary for each correlation
return (re.hdul[0].data.transpose(ax) +
im.hdul[0].data.transpose(ax)*1j)
# Create delayed loads of the beam
beam_loader = dask.delayed(_load_correlation)
beam_corrs = [beam_loader(re, im, ax)
for c, (corr, (re, im)) in enumerate(beam_files)]
beam_corrs = [da.from_delayed(bc, shape=shape, dtype=dtype)
for bc in beam_corrs]
# Stack correlations and rechunk to one great big block
beam = da.stack(beam_corrs, axis=3)
beam = beam.rechunk(shape + (len(corr_types),))
# Dask arrays for the beam extents and beam frequency grid
beam_lm_ext = np.array([[l_grid[0], l_grid[-1]], [m_grid[0], m_grid[-1]]])
beam_lm_ext = da.from_array(beam_lm_ext, chunks=beam_lm_ext.shape)
beam_freq_grid = da.from_array(nu_grid, chunks=nu_grid.shape)
return beam, beam_lm_ext, beam_freq_grid