def restoreSources(fits_hdu,
sources,
gmaj,
gmin=None,
grot=0,
freq=None,
primary_beam=None,
apply_beamgain=False,
ignore_nobeam=False):
"""Restores sources (into the given FITSHDU) using a Gaussian PSF given by gmaj/gmin/grot, in radians.
gmaj/gmin is major/minor sigma parameter; grot is PA in the North thru East convention (PA=0 is N).
If gmaj=0, uses delta functions instead.
If freq is specified, converts flux to the specified frequency.
If primary_beam is specified, uses it to apply a PB gain to each source. This must be a function of two arguments:
r and freq, returning the power beam gain.
If apply_beamgain is true, applies beamgain atribute instead, if this exists.
Source tagged 'nobeam' will not have the PB gain applied, unless ignore_nobeam=True
"""
hdr = fits_hdu.header
data, stokes, extra_data_axes, dum = getImageCube(fits_hdu)
# create projection object, using pixel coordinates
proj = Projection.FITSWCSpix(hdr)
naxis = len(data.shape)
nx = data.shape[0]
ny = data.shape[1]
dprintf(1, "Read image of shape %s\n", data.shape)
# Now we make "indexer" tuples. These use the numpy.newarray index to turn elementary vectors into
# full arrays of the same number of dimensions as 'data' (data can be 2-, 3- or 4-dimensional, so we need
# a general solution.)
# For e.g. a nfreq x nstokes x ny x nx array, the following objects are created:
# x_indexer turns n-vector vx into a _,_,_,n array
# y_indexer turns m-vector vy into a _,_,m,_ array
# stokes_indexer turns the stokes vector into a _,nst,_,_ array
# ...where "_" is numpy.newaxis.
# The happy result of all this is that we can add a Gaussian into the data array at i1:i2,j1:j2 as follows:
# 1. form up vectors of world coordinates (vx,vy) corresponding to pixel coordinates i1:i2 and j1:j2
# 2. form up vector of Stokes parameters
# 3. g = Gauss(vx[x_indexer],vy[y_indexer])*stokes[stokes_indexer]
# 4. Just say data[j1:j2,i1:2,...] += g
# This automatically expands all array dimensions as needed.
# This is a helper function, returns an naxis-sized tuple, with slice(None) in the Nth
# position, and elem_index elsewhere.
def make_axis_indexer(n, elem_index=numpy.newaxis):
indexer = [elem_index] * naxis
indexer[n] = slice(None)
return tuple(indexer)
x_indexer = make_axis_indexer(0)
y_indexer = make_axis_indexer(1)
# figure out stokes
nstokes = len(stokes)
stokes_vec = numpy.zeros((nstokes, ))
stokes_indexer = make_axis_indexer(2)
dprint(2, "Stokes are", stokes)
dprint(2, "Stokes indexing vector is", stokes_indexer)
# get pixel sizes, in radians
# gmaj != 0: use gaussian. Estimate PSF box size. We want a +/-5 sigma box
if gmaj > 0:
# convert grot from N-E to W-N (which is the more conventional mathematical definition of these things), so X is major axis
grot += math.pi / 2
if gmin == 0:
gmin = gmaj
cos_rot = math.cos(grot)
sin_rot = math.sin(-grot)
# rotation is N->E, so swap the sign
else:
gmaj = gmin = grot = 0
conv_kernels = {}
# loop over sources in model
for src in sources:
# get normalized intensity, if spectral info is available
if freq is not None and getattr(src, 'spectrum', None):
ni = src.spectrum.normalized_intensity(freq)
dprintf(3, "Source %s: normalized spectral intensity is %f\n",
src.name, ni)
else:
ni = 1
# multiply that by PB gain, if given
if ignore_nobeam or not getattr(src, 'nobeam', False):
if apply_beamgain and hasattr(src, 'beamgain'):
ni *= getattr(src, 'beamgain')
elif primary_beam:
r = getattr(src, 'r', None)
if r is not None:
pb = primary_beam(r, freq)
ni *= pb
dprintf(3,
"Source %s: r=%g pb=%f, normalized intensity is %f\n",
src.name, r, pb, ni)
# process point sources
if src.typecode in ('pnt', 'Gau'):
# pixel coordinates of source
xsrc, ysrc = proj.lm(src.pos.ra, src.pos.dec)
# form up stokes vector
for i, st in enumerate(stokes):
stokes_vec[i] = getattr(src.flux, st, 0) * ni
dprintf(3, "Source %s, %s Jy, at pixel %f,%f\n", src.name,
stokes_vec, xsrc, ysrc)
# for gaussian sources, convolve with beam
if src.typecode == 'Gau':
pa0 = src.shape.pa + math.pi / 2
# convert PA from N->E to conventional W->N
ex0, ey0 = src.shape.ex / FWHM, src.shape.ey / FWHM
# convert extents from FWHM to sigmas, since gmaj/gmin is in same scale
if gmaj > 0:
ex, ey, pa = convolveGaussian(ex0, ey0, pa0, gmaj, gmin,
grot)
# normalize flux by beam/extent ratio
stokes_vec *= (gmaj * gmin) / (ex * ey)
# print "%3dx%-3d@%3d * %3dx%-3d@%3d -> %3dx%-3d@%3d"%(
# ex0 *FWHM*ARCSEC,ey0 *FWHM*ARCSEC,(pa0-math.pi/2)*DEG,
# gmaj*FWHM*ARCSEC,gmin*FWHM*ARCSEC,(grot-math.pi/2)*DEG,
# ex *FWHM*ARCSEC,ey *FWHM*ARCSEC,(pa-math.pi/2)*DEG)
else:
# normalize flux by pixel/extent ratio
ex, ey, pa = ex0, ey0, pa0
stokes_vec *= (abs(proj.xscale * proj.yscale)) / (ex * ey)
else:
ex, ey, pa = gmaj, gmin, grot
# gmaj != 0: use gaussian.
if ex > 0 or ey > 0:
# work out restoring box
box_radius = 5 * (max(ex, ey)) / min(abs(proj.xscale),
abs(proj.yscale))
dprintf(
2, "Will use a box of radius %f pixels for restoration\n",
box_radius)
cos_pa = math.cos(pa)
sin_pa = math.sin(-pa)
# rotation is N->E, so swap the sign
# pixel coordinates of box around source in which we evaluate the gaussian
i1 = max(0, int(math.floor(xsrc - box_radius)))
i2 = min(nx, int(math.ceil(xsrc + box_radius)))
j1 = max(0, int(math.floor(ysrc - box_radius)))
j2 = min(ny, int(math.ceil(ysrc + box_radius)))
# skip sources if box doesn't overlap image
if i1 >= i2 or j1 >= j2:
continue
# now we convert pixel indices within the box into world coordinates, relative to source position
xi = (numpy.arange(i1, i2) - xsrc) * proj.xscale
yj = (numpy.arange(j1, j2) - ysrc) * proj.yscale
# work out rotated coordinates
xi1 = (xi * cos_pa)[x_indexer] - (yj * sin_pa)[y_indexer]
yj1 = (xi * sin_pa)[x_indexer] + (yj * cos_pa)[y_indexer]
# evaluate gaussian at these, scale up by stokes vector
gg = stokes_vec[stokes_indexer] * numpy.exp(-(
(xi1 / ex)**2 + (yj1 / ey)**2) / 2.)
# add into data
data[i1:i2, j1:j2, ...] += gg
# else gmaj=0: use delta functions
else:
xsrc = int(round(xsrc))
ysrc = int(round(ysrc))
# skip sources outside image
if xsrc < 0 or xsrc >= nx or ysrc < 0 or ysrc >= ny:
continue
xdum = numpy.array([1])
ydum = numpy.array([1])
data[xsrc:xsrc + 1, ysrc:ysrc + 1, ...] += stokes_vec[
stokes_indexer] * xdum[x_indexer] * ydum[y_indexer]
# process model images -- convolve with PSF and add to data
elif src.typecode == "FITS":
modelff = pyfits.open(src.shape.filename)
model, model_stokes, extra_model_axes, removed_model_axes = \
getImageCube(modelff[0], src.shape.filename, extra_axes=extra_data_axes)
modelproj = Projection.FITSWCSpix(modelff[0].header)
# map Stokes planes: at least the first one ("I", presumably) must be present
# The rest are represented by indices in model_stp. Thus e.g. for an IQUV data image and an IV model,
# model_stp will be [0,-1,-1,1]
model_stp = [(model_stokes.index(st) if st in model_stokes else -1)
for st in stokes]
if model_stp[0] < 0:
print("Warning: model image %s lacks Stokes %s, skipping." %
(src.shape.filename, model_stokes[0]))
continue
# figure out whether the images overlap at all
# in the trivial case, both images have the same WCS, so no resampling is needed
if model.shape[:2] == data.shape[:2] and modelproj == proj:
model_resampler = lambda x: x
data_x_slice = data_y_slice = slice(None)
dprintf(
3,
"Source %s: same resolution as output, no interpolation needed\n",
src.shape.filename)
# else make a resampler engine
else:
model_resampler = ImageResampler(
modelproj, proj, numpy.arange(model.shape[0], dtype=float),
numpy.arange(model.shape[1], dtype=float),
numpy.arange(data.shape[0], dtype=float),
numpy.arange(data.shape[1], dtype=float))
data_x_slice, data_y_slice = model_resampler.targetSlice()
dprintf(3, "Source %s: resampling into image at %s, %s\n",
src.shape.filename, data_x_slice, data_y_slice)
# skip this source if no overlap
if data_x_slice is None or data_y_slice is None:
continue
# warn about ignored model axes (e.g. when model has frequency and our output doesn't)
if removed_model_axes:
print(
"Warning: model image %s has one or more axes that are not present in the output image:"
% src.shape.filename)
print(" taking the first plane along (%s)." %
(",".join(removed_model_axes)))
# evaluate convolution kernel for this model scale, if not already cached
conv_kernel = conv_kernels.get(
(modelproj.xscale, modelproj.yscale), None)
if conv_kernel is None:
box_radius = 5 * (max(gmaj, gmin)) / min(
abs(modelproj.xscale), abs(modelproj.yscale))
radius = int(round(box_radius))
# convert pixel coordinates into world coordinates relative to 0
xi = numpy.arange(-radius, radius + 1) * modelproj.xscale
yj = numpy.arange(-radius, radius + 1) * modelproj.yscale
# work out rotated coordinates
xi1 = (xi * cos_rot)[:, numpy.newaxis] - (
yj * sin_rot)[numpy.newaxis, :]
yj1 = (xi * sin_rot)[:, numpy.newaxis] + (
yj * cos_rot)[numpy.newaxis, :]
# evaluate convolution kernel
conv_kernel = numpy.exp(-((xi1 / gmaj)**2 + (yj1 / gmin)**2) /
2.)
conv_kernels[modelproj.xscale, modelproj.yscale] = conv_kernel
# Work out data slices that we need to loop over.
# For every 2D slice in the data image cube (assuming other axes besides x/y), we need to apply a
# convolution to the corresponding model slice, and add it in to the data slice. The complication
# is that any extra axis may be of length 1 in the model and of length N in the data (e.g. frequency axis),
# in which case we need to add the same model slice to all N data slices. The loop below puts together a series
# of index tuples representing each per-slice operation.
# These two initial slices correspond to the x/y axes. Additional indices will be appended to these in a loop
slices0 = [([data_x_slice,
data_y_slice], [slice(None), slice(None)])]
# work out Stokes axis
sd0 = [data_x_slice, data_y_slice]
sm0 = [slice(None), slice(None)]
slices = []
slices = [(sd0 + [dst], sm0 + [mst])
for dst, mst in enumerate(model_stp) if mst >= 0]
# for dst,mst in enumerate(model_stp):
# if mst >= 0:
# slices = [ (sd0+[dst],sm0+[mst]) for sd0,sm0 in slices ]
# now loop over extra axes
for axis in range(3, len(extra_data_axes) + 3):
# list of data image indices to iterate over for this axis, 0...N-1
indices = [[x] for x in range(data.shape[axis])]
# list of model image indices to iterate over
if model.shape[axis] == 1:
model_indices = [[0]] * len(indices)
# shape-n: must be same as data, in which case 0..N-1 is assigned to 0..N-1
elif model.shape[axis] == data.shape[axis]:
model_indices = indices
# else error
else:
raise RuntimeError("axis %s of model image %s doesn't match that of output image" % \
(extra_data_axes[axis - 3], src.shape.filename))
# update list of slices
slices = [(sd0 + sd, si0 + si) for sd0, si0 in slices
for sd, si in zip(indices, model_indices)]
# now loop over slices and assign
for sd, si in slices:
conv = convolve(model[tuple(si)], conv_kernel)
data[tuple(sd)] += model_resampler(conv)
def fitPsf(filename, cropsize=None):
"""Fits a Gaussian PSF to the FITS file given by 'filename'.
If cropsize is specified, crops the central cropsize X cropsize pixels before fitting.
Else determines cropsize by looking for the first negative sidelobe from the centre outwards.
Returns maj_sigma,min_sigma,pa_NE (in radians)
"""
# read PSF from file
psf = pyfits.open(filename)[0]
hdr = psf.header
psf = psf.data
dprintf(2, "Read PSF of shape %s from file %s\n", psf.shape, filename)
# remove stokes and freq axes
if len(psf.shape) == 4:
psf = psf[0, 0, :, :]
elif len(psf.shape) == 3:
psf = psf[0, :, :]
else:
raise RuntimeError("illegal PSF shape %s" + psf.shape)
nx, ny = psf.shape
# crop the central region
if cropsize:
size = cropsize
psf = psf[(nx - size) // 2:(nx + size) // 2,
(ny - size) // 2:(ny + size) // 2]
# if size not specified, then auto-crop by looking for the first negative value starting from the center
# this will break on very extended diagonal PSFs, but that's a pathological case
else:
ix = numpy.where(psf[:, ny // 2] < 0)[0]
ix0 = max(ix[ix < nx // 2])
ix1 = min(ix[ix > nx // 2])
iy = numpy.where(psf[nx // 2, :] < 0)[0]
iy0 = max(iy[iy < ny // 2])
iy1 = min(iy[iy > ny // 2])
print(ix0, ix1, iy0, iy1)
psf = psf[ix0:ix1, iy0:iy1]
psf[psf < 0] = 0
# estimate gaussian parameters, then fit
from . import gaussfitter2
parms0 = gaussfitter2.moments(psf, circle=0, rotate=1, vheight=0)
print(parms0)
dprint(2, "Estimated parameters are", parms0)
parms = gaussfitter2.gaussfit(psf,
None,
parms0,
autoderiv=1,
return_all=0,
circle=0,
rotate=1,
vheight=0)
dprint(0, "Fitted parameters are", parms)
# now swap x and y around, since our axes are in reverse order
ampl, y0, x0, sy, sx, rot = parms
# get pixel sizes in radians (by constructing a projection object)
proj = Projection.FITSWCSpix(hdr)
sx_rad = abs(sx * proj.xscale)
sy_rad = abs(sy * proj.yscale)
rot -= 90 # convert West through North PA into the conventional North through East
if sx_rad < sy_rad:
sx_rad, sy_rad = sy_rad, sx_rad
rot -= 90
rot %= 180
dprintf(
1,
"Fitted gaussian PSF FWHM of %f x %f pixels (%f x %f arcsec), PA %f deg\n",
sx * FWHM, sy * FWHM, sx_rad * FWHM * ARCSEC, sy_rad * FWHM * ARCSEC,
rot)
return sx_rad, sy_rad, rot / DEG