def rnd_test2(mdl, ntrials=1000, noise=1., niter=500, gain=0.1, cutoff=0.):
"""
"""
nu = gennu(StartFreq, StepFreq, NFrequencies)
# define image space
nphi = 400
dphi = 5.0
phi = -0.5 * nphi * dphi + numpy.arange(nphi) * dphi
cpks1 = numpy.zeros(ntrials)
dpks1 = numpy.zeros(ntrials)
cpks2 = numpy.zeros(ntrials)
dpks2 = numpy.zeros(ntrials)
rmsyn = R.RMSynth(nu, StepFreq, phi)
rmcl = R.RMClean(rmsyn, niter, gain, cutoff)
progress(20, 0.)
for i in range(ntrials):
ccloc, ccval, rmc, di = do_test(mdl, noise, niter, gain, cutoff,
False, rms=rmsyn, rmc=rmcl)
cpks1[i] = abs(rmc.clean_map[170])
dpks1[i] = abs(di[170])
cpks2[i] = abs(rmc.clean_map[230])
dpks2[i] = abs(di[230])
progress(20, ((i + 1.) / ntrials) * 100.)
plot_rnd_test(dpks1, cpks1, [mdl[1]])
plot_rnd_test(dpks2, cpks2, [mdl[0]])
def do_test(mdl, noise=0., niter=500, gain=0.1, cutoff=0.,
doplot=True, rms=None, rmc=None):
"""
A routine to test the rm_tools RMClean functionality
input
-----
mdl: a list of (phi, Q + iU) values
noise: the stdv of the Gaussian (white) noise to add to each channel of
the data vector
return
------
ccloc: a list of CLEAN component locations
ccval: a list of CLEAN component values
map: the restored CLEAN image
phi: coordinates for which the map are defined
"""
# define nu range
nu = gennu(StartFreq, StepFreq, NFrequencies)
# define image space
nphi = 400
dphi = 5.0
phi = -0.5 * nphi * dphi + numpy.arange(nphi) * dphi
# compute l2
l2 = convert_nu_to_l2(nu, StepFreq)
l2 = numpy.flipud(l2)
data = dft(mdl, noise, l2)
# do rmsynth
if rms is None:
rms = R.RMSynth(nu, StepFreq, phi)
di = rms.compute_dirty_image(data)
if rmc is None:
rmc = R.RMClean(rms, niter, gain, cutoff)
else:
rmc.reset()
rmc.perform_clean(data)
rmc.restore_clean_map()
[ccloc, ccval] = sort_cc_list(rmc.cc_phi_list, rmc.cc_val_list)
if doplot:
plot_test(ccloc, ccval, rmc.clean_map, phi, mdl, di)
return ccloc, ccval, rmc, di
def main():
""" Handle all parsing here if started from the command line"""
parser = argparse.ArgumentParser(description="Rotation Measure Synthesis tool.")
parser.add_argument("-o", "--outname", type=str, help="output prefix")
parser.add_argument("-p", "--save-pol-cube", action="store_true",
dest="save_pol_cube", help="Save Pol cube", default=False)
parser.add_argument("-q", "--save-qu-cubes", action="store_true",
dest="save_qu_cubes", default=False,
help="Save derotated Q and U cubes")
parser.add_argument("-r", "--save-residual-cubes", action="store_true",
dest="save_residual_cubes", default=False,
help="Save residual cubes if cleaning")
parser.add_argument("-d", "--save-dirty-cubes", action="store_true",
dest="save_dirty_cubes", default=False,
help="Save dirty cubes if cleaning")
parser.add_argument("-a", "--auto-flag", action="store_true",
dest="auto_flag", help="auto flag data", default=False)
parser.add_argument("-m", "--auto-mask", dest="auto_mask", type=str, nargs=2,
help="Use stokes I map and cutoff value in Jy. This is in addition to a regular mask. I.E. it is OR'd with the regular mask if provided. Eg. map_i.fits 0.001")
parser.add_argument("-x", "--exclude-phi", dest="exclude_phi",
metavar='phi_range', nargs=2, type=float, default=(0,0),
help="Exclude this Phi range from 2D maps. Eg: -3 1.5")
parser.add_argument("-n", "--phi-rms", dest="phi_rms", type=float,
metavar='phi_rms_range', nargs=2, default=(0,0),
help="Make a Phi RMS image from this range. Eg: 100 115")
parser.add_argument("--single", action="store_true", default=False,
help="Use single precision floats internally. Default: False")
parser.add_argument("--double", action="store_true", default=False,
help="Output double precision floats. Default: False. If neither --single or --double is specified, use double precision internally and write out single precision.")
parser.add_argument("--rmclean-mask", type=str, help="Input mask for RMCLEAN")
parser.add_argument("--rmclean-sigma", dest="rmclean_sigma",
type=float, default=0.0,
help="Clean to RMCLEAN_SIGMA times the phi noise. This is used in combination with the cutoff value. Eg. 10")
parser.add_argument("param_file", type=str, help="Input Parameter file")
parser.add_argument("fits_cube", type=str, help="QU FITS cube")
args = parser.parse_args()
print("Parsing parameter file...")
params = parse_input_file(args.param_file)
if args.outname:
params.outputfn = args.outname
if args.single:
in_dtype = np.complex64
else:
in_dtype = np.complex128
if args.double:
out_dtype = np.complex128
else:
out_dtype = np.complex64
print("Loading data...")
hdu = fits.open(args.fits_cube)[0]
data = hdu.data
header = hdu.header
check_cube_format(header)
if params.imagemask:
mask = fits.open(params.imagemask)[0].data.squeeze().astype(bool)
else:
mask = np.ones(shape=data.shape[2:], dtype=bool)
if args.auto_mask:
stokesI_map = fits.open(args.auto_mask[0])[0].data.squeeze()
cutoff = float(args.auto_mask[1])
stokesI_mask = np.zeros_like(stokesI_map, dtype=np.uint8)
stokesI_mask[np.where(stokesI_map > cutoff)] = 1
if params.imagemask:
mask = np.logical_or(mask, stokesI_mask)
else:
mask = stokesI_mask
del stokesI_map, stokesI_mask
if args.rmclean_mask:
rmclean_mask = fits.open(args.rmclean_mask)[0].data.squeeze().astype(bool)
params.do_clean = True
else:
rmclean_mask = np.ones(shape=data.shape[2:], dtype=bool)
if args.auto_flag:
bad_chans = channel_flags(data)
flag_weights = np.logical_not(bad_chans).astype(float)
if params.weight is None:
params.weight = flag_weights
else:
params.weight = np.minimum(params.weight, flag_weights)
if args.exclude_phi[0] != args.exclude_phi[1]:
exclude_phi_range = phi_range_to_channel_range(args.exclude_phi, params)
print("Excluding phi range from 2D maps: {} => {}".format(args.exclude_phi, exclude_phi_range))
else:
exclude_phi_range = None
if args.phi_rms[0] != args.phi_rms[1]:
phi_rms_range = phi_range_to_channel_range(args.phi_rms, params)
print("RMS phi range: {} => {}".format(args.phi_rms, phi_rms_range))
else:
phi_rms_range = None
if args.rmclean_sigma != 0:
if phi_rms_range is None:
args.error("--phi-rms required for --rmclean-sigma")
# Survey cubes will be ordered: [Freqs, Stokes, Dec, RA]
data = np.moveaxis(data, 0, -1) # Move Freq to end
data = np.moveaxis(data, 0, -1) # Move Stokes to end
# Now order is [Dec, RA, Freqs, Stokes]
(DEC,RA,FREQ,STOKES) = range(4)
data_selection = np.where(mask)
non_masked_data = data[data_selection]
# Unflagged NaNs will cause RMSynth to return NaNs
non_masked_data = np.nan_to_num(non_masked_data.view(dtype=">c8").squeeze()).astype(in_dtype)
clean_regions = rmclean_mask[data_selection].squeeze()
params.nu = freqs(header, 4)
params.dnu = params.nu[1] - params.nu[0]
rms = R.RMSynth(params.nu, params.dnu, params.phi, params.weight)
if params.do_clean:
rmc = R.RMClean(rms, params.niter, params.gain, params.cutoff)
rmsfout = numpy.zeros((len(rms.rmsf), 3))
rmsfout[:, 0] = rms.rmsf_phi
rmsfout[:, 1] = rms.rmsf.real
rmsfout[:, 2] = rms.rmsf.imag
numpy.savetxt(params.outputfn + "_rmsf.txt", rmsfout)
data_out = np.empty(shape=(len(non_masked_data), params.nphi), dtype=out_dtype)
if params.do_clean:
dirty_out = np.empty(shape=(len(non_masked_data), params.nphi), dtype=out_dtype)
residual_out = np.empty(shape=(len(non_masked_data), params.nphi), dtype=out_dtype)
print("Doing RM synthesis")
n = len(non_masked_data)
width = len(str(n))
for i in range(len(non_masked_data)):
if params.do_clean and clean_regions[i]:
rmc.reset()
rmc.residual_map = rms.compute_dirty_image(non_masked_data[i])
rmc.dirty_image = rmc.residual_map
if args.rmclean_sigma != 0:
q_rms = rmc.residual_map[phi_rms_range[0]:phi_rms_range[1]].real.std()
u_rms = rmc.residual_map[phi_rms_range[0]:phi_rms_range[1]].imag.std()
cutoff = args.rmclean_sigma * 0.5 * (q_rms + u_rms)
rmc.cutoff = max(cutoff, params.cutoff)
rmc.perform_clean()
rmc.restore_clean_map()
data_out[i] = rmc.clean_map
residual_out[i] = rmc.residual_map
dirty_out[i] = rmc.dirty_image
else:
data_out[i] = rms.compute_dirty_image(non_masked_data[i])
if i % 100 == 0:
progress(20, 100.0 * i / n)
print("")
header_axes_attributes = ["NAXIS", "CTYPE", "CRVAL", "CRPIX", "CDELT", "CUNIT", "CROTA"]
for attr in header_axes_attributes:
attr += "4"
if attr in header:
del header[attr]
header["CTYPE3"] = "Phi"
header["CUNIT3"] = "rad/m/m"
header["CRPIX3"] = 1
header["CRVAL3"] = params.phi_min
header["CDELT3"] = params.dphi
if args.save_pol_cube or args.save_qu_cubes or params.do_clean and (args.save_residual_cubes or args.save_dirty_cubes):
cube_out = np.full(fill_value=np.NaN, shape=(data.shape[DEC], data.shape[RA], params.nphi), dtype=data_out.real.dtype)
else:
cube_out = None
if args.save_pol_cube:
print("Saving phi cube")
if not params.do_clean:
write_cube(cube_out, data_selection, abs(data_out), header, params.outputfn + "_di_p.fits")
else:
write_cube(cube_out, data_selection, abs(data_out), header, params.outputfn + "_clean_p.fits")
if args.save_dirty_cubes:
print("Saving phi dirty cube")
write_cube(cube_out, data_selection, abs(dirty_out), header, params.outputfn + "_di_p.fits")
if args.save_residual_cubes:
print("Saving phi residual cube")
write_cube(cube_out, data_selection, abs(residual_out), header, params.outputfn + "_residual_p.fits")
if args.save_qu_cubes:
print("Saving Q & U cubes")
if not params.do_clean:
write_cube(cube_out, data_selection, data_out.real, header, params.outputfn + "_di_q.fits")
write_cube(cube_out, data_selection, data_out.imag, header, params.outputfn + "_di_u.fits")
else:
write_cube(cube_out, data_selection, data_out.real, header, params.outputfn + "_clean_q.fits")
write_cube(cube_out, data_selection, data_out.imag, header, params.outputfn + "_clean_u.fits")
if args.save_dirty_cubes:
print("Saving Q & U dirty cubes")
write_cube(cube_out, data_selection, dirty_out.real, header, params.outputfn + "_di_q.fits")
write_cube(cube_out, data_selection, dirty_out.imag, header, params.outputfn + "_di_u.fits")
if args.save_residual_cubes:
print("Saving Q & U residual cubes")
write_cube(cube_out, data_selection, residual_out.real, header, params.outputfn + "_residual_q.fits")
write_cube(cube_out, data_selection, residual_out.imag, header, params.outputfn + "_residual_u.fits")
del cube_out
if params.do_clean:
del dirty_out, residual_out
for attr in header_axes_attributes:
attr += "3"
if attr in header:
del header[attr]
hdu.data = np.full(fill_value=np.NaN, shape=(data.shape[DEC], data.shape[RA]), dtype=data_out.real.dtype)
# RMS Image
if phi_rms_range != None:
q_rms = data_out[:,phi_rms_range[0]:phi_rms_range[1]].real.std(axis=1)
u_rms = data_out[:,phi_rms_range[0]:phi_rms_range[1]].imag.std(axis=1)
p_rms = 0.5 * (q_rms + u_rms)
hdu.data[data_selection] = p_rms
print("Saving rms map")
hdu.writeto(params.outputfn + "_rms.fits", **overwrite)
if exclude_phi_range != None:
data_out[:,exclude_phi_range[0]:exclude_phi_range[1]] = 0
# 2D maps
hdu.data[data_selection] = np.amax(abs(data_out), axis=-1)
print("Saving int pol map")
hdu.writeto(params.outputfn + "_polint.fits", **overwrite)
indx_max = np.argmax(abs(data_out), axis=-1)
hdu.data[data_selection] = data_out.real[range(len(data_out)), indx_max]
print("Saving q map")
hdu.writeto(params.outputfn + "_qmap.fits", **overwrite)
hdu.data[data_selection] = data_out.imag[range(len(data_out)), indx_max]
print("Saving u map")
hdu.writeto(params.outputfn + "_umap.fits", **overwrite)
header["BUNIT"] = "rad/m/m"
hdu.data[data_selection] = params.phi[indx_max]
print("Saving phi map")
hdu.writeto(params.outputfn + "_phi.fits", **overwrite)
#------------------------------------------------------------------------ # specify the range of fd space: phi = np.linspace(-1000, 1000, 4000) # get the input data: #inputdir = "./data" nu, los = get_1d_data(input_dir) dnu = nu[1] - nu[0] # initialise the pyrmsynth class: rms = R.RMSynth(nu, dnu, phi) # calculate the resolution in fd space and the maximum recoverable scale: res = calc_phi_res(nu) maxscale = calc_maxscale(nu, dnu) print "\n" print "Max f.d. resolution: " + str(round(res)) + " rad/m^2" print "Max f.d. scale " + str(round(maxscale)) + " rad/m^2" print "\n"
def rmsynthesis(params, options, manual=False):
"""
rmsynthesis(params, manual=False)
Conducts RM synthesis based on the parameters provided in the params class
for the data contained w/in the current directory. This may be either a
fits file containing a cube of data (indexed z,x,y when read in using
pyfits) or a text file (freq, q,u)
if manual=true returns the rmsynth cube class and the pol cube.
"""
qsubdir = '/stokes_q/'
usubdir = '/stokes_u/'
vcube_mmfn = 'stokesv.dat'
incube_mmfn = 'incube.dat'
outcube_mmfn = 'outcube.dat'
rescube_mmfn = 'outcube_res.dat'
cleancube_mmfn = 'outcube_clean.dat'
cccube_mmfn = 'outcube_cc.dat'
pmap_mmfn = 'outmap_p.dat'
phimap_mmfn = 'outmap_phi.dat'
qmap_mmfn = 'outmap_q.dat'
umap_mmfn = 'outmap_u.dat'
if options.freq_last:
freq_axnum = '4'
stokes_axnum = '3'
else:
freq_axnum = '3'
stokes_axnum = '4' # lint:ok
if params.phi is None:
msg = 'Parameters not specified correctly. No phi axis defined! ' + \
'Please double check and try again.'
raise ValueError(msg)
if options.stokes_v and options.separate_stokes:
print "Stokes V output is not available when the Stokes images are " +\
"stored in separate files. Turning Stokes V output off."
options.stokes_v = False
# Gather the input file names from the directory in the parset
if not options.separate_stokes:
# In this case, each file contains all Stokes parameters.
fns = os.listdir(params.input_dir)
if len(fns) == 0:
raise IOError('No valid files found in this directory!')
nfns = len(fns)
fns_copy = list(fns)
for indx in range(nfns): # remove any non fits files
fn = fns_copy[indx]
if fn[len(fn) - 4:len(fn)].lower() != 'fits':
fns.remove(fn)
# If a mask file is specified, above lines will add the mask fits
# to fns. Removing it from the list of valid fits files
if params.imagemask != '':
fns.remove(params.imagemask.split('/')[-1])
if len(fns) == 0:
raise IOError('No valid files found in this directory!')
fns.sort()
tdata = pyfits.getdata(params.input_dir + fns[0])
thead = pyfits.getheader(params.input_dir + fns[0])
else:
# In this case, each file contains only a single Stokes parameter.
# This is mainly here for compatability with Michiel's code.
try:
fns_q = os.listdir(params.input_dir + qsubdir)
fns_u = os.listdir(params.input_dir + usubdir)
except OSError:
print "No Stokes sub-directories found! Please put the Stokes " +\
"Q and Stokes U images into the stokes_q and stokes_u " + \
"subfolders of the input directory listed in the parset file."
raise
if len(fns_q) == 0 or len(fns_u) == 0:
raise IOError('No valid files found in this directory')
nfnsq = len(fns_q)
nfnsu = len(fns_u)
fns_q_copy = list(fns_q)
fns_u_copy = list(fns_u)
for indx in range(nfnsq): # remove any non fits files
fn = fns_q_copy[indx]
if fn[len(fn) - 4:len(fn)].lower() != 'fits':
fns_q.remove(fn)
for indx in range(nfnsu): # remove any non fits files
fn = fns_u_copy[indx]
if fn[len(fn) - 4:len(fn)].lower() != 'fits':
fns_u.remove(fn)
if len(fns_q) == 0 or len(fns_u) == 0:
raise IOError('No valid files found in this directory.')
if len(fns_q) != len(fns_u):
raise IOError('The number of Stokes Q and U images are different.')
fns_q.sort()
fns_u.sort()
tdata = pyfits.getdata(params.input_dir + qsubdir + fns_q[0])
thead = pyfits.getheader(params.input_dir + qsubdir + fns_q[0])
# indices 1=freq, 2=dec, 3=ra - stored values are complex
# Validate the image bounds
decsz = [0, len(tdata[0, 0])]
if params.dec_lim[0] != -1:
decsz[0] = params.dec_lim[0]
else:
params.dec_lim[0] = 0
if params.dec_lim[1] != -1:
decsz[1] = params.dec_lim[1]
else:
params.dec_lim[1] = len(tdata[0, 0])
if decsz[0] >= decsz[1]:
raise Exception('Invalid image bounds')
rasz = [0, len(tdata[0, 0, 0])]
if params.ra_lim[0] != -1:
rasz[0] = params.ra_lim[0]
else:
params.ra_lim[0] = 0
if params.ra_lim[1] != -1:
rasz[1] = params.ra_lim[1]
else:
params.ra_lim[1] = len(tdata[0, 0, 0])
if rasz[0] >= rasz[1]:
raise Exception('Invalid image bounds')
# Get the appropriate mask
if params.imagemask == '':
# If no mask is specified, assign 1's to all pixels in the mask image
ra_size = params.ra_lim[1] - params.ra_lim[0]
dec_size = params.dec_lim[1] - params.dec_lim[0]
maskimage = numpy.ones((dec_size, ra_size), dtype=numpy.int_)
else:
# Read in the mask file
try:
# mask should be in the input directory
maskdata = pyfits.getdata(params.input_dir + params.imagemask)
maskhead = pyfits.getheader(params.input_dir + params.imagemask)
except:
raise Exception('Unable to read the mask image.')
# Check if the mask has the same shape as the input image
if thead['NAXIS1'] != maskhead['NAXIS1'] or \
thead['NAXIS2'] != maskhead['NAXIS2']:
raise Exception(
'Mask and Stokes images have different image sizes.')
else:
# Check the FITS format of mask file;take care of possible extra
# axes like correlation or frequency
if (len(numpy.shape(maskdata)) > 2):
try:
maskdata = maskdata.reshape(params.dec_lim[1] - \
params.dec_lim[0],params.ra_lim[1] - params.ra_lim[0])
except:
raise Exception(
'Incompatible number of axes in Mask file.')
# Extract the submask as specified by ra, dec limits
maskimage = maskdata[decsz[0]:decsz[1], rasz[0]:rasz[1]]
nchan = thead.get('NAXIS' + freq_axnum)
if options.rest_freq and nchan != 1:
raise Exception('rest_freq option is only valid for files with one \
frequency.')
if options.stokes_v:
vcube = create_memmap_file_and_array(
vcube_mmfn,
(len(fns) * nchan, decsz[1] - decsz[0], rasz[1] - rasz[0]),
numpy.dtype('float64'))
if options.separate_stokes:
nsb = len(fns_q)
else:
nsb = len(fns)
cube = create_memmap_file_and_array(
incube_mmfn, (nsb * nchan, decsz[1] - decsz[0], rasz[1] - rasz[0]),
numpy.dtype('complex128'))
params.nu = numpy.zeros(nsb * nchan)
# This gets overwritten later for rest_freq mode!
# store the channel width, in Hz
params.dnu = thead.get('CDELT' + freq_axnum)
if (params.weight is not None and len(params.weight) != len(params.nu)):
raise Exception('number of frequency channels in weight list is not ' +
'compatible with input visibilities.')
if params.weight is None:
params.weight = numpy.ones(len(params.nu))
params()
print ' '
print '-----------------------'
print ' '
print 'Reading fits files into one big cube, ' +\
'converting Q & U to complex...'
progress(20, 0)
if not options.separate_stokes:
for indx in range(len(fns)):
fn = fns[indx]
tdata = pyfits.getdata(params.input_dir + fn)
thead = pyfits.getheader(params.input_dir + fn)
base_indx = nchan * indx
if not options.freq_last:
if tdata.shape[0] == 4:
# Read Stokes Q and U into the complex cube
cube.real[base_indx:base_indx + nchan] = \
tdata[1, :, decsz[0]:decsz[1], rasz[0]:rasz[1]]
cube.imag[base_indx:base_indx + nchan] = \
tdata[2, :, decsz[0]:decsz[1], rasz[0]:rasz[1]]
# XXX: The commented code above isn't the appropriate way
# to deal with weights. Hand params.weight to the
# RMSynth class init instead. This has been fixed
else:
# There are not 4 stokes parameters, and that is not OK
raise Exception('Invalid data shape!')
if options.stokes_v:
vcube[base_indx:base_indx + nchan] = \
tdata[3, :, decsz[0]:decsz[1], rasz[0]:rasz[1]]
else:
if tdata.shape[1] == 4:
cube.real[base_indx:base_indx + nchan] = \
tdata[:, 1, decsz[0]:decsz[1], rasz[0]:rasz[1]]
cube.imag[base_indx:base_indx + nchan] = \
tdata[:, 2, decsz[0]:decsz[1], rasz[0]:rasz[1]]
else:
raise Exception('Invalid data shape!')
if options.stokes_v:
vcube[base_indx:base_indx + nchan] = \
tdata[:, 3, decsz[0]:decsz[1], rasz[0]:rasz[1]]
if options.rest_freq:
if nchan != 1:
msg = 'When the rest frequency option is selected, ' + \
'only one channel is allowed per file!'
raise Exception(msg)
params.nu[base_indx] = thead.get('RESTFREQ')
else:
flist = numpy.arange(nchan) * thead.get('CDELT' + freq_axnum) \
+ thead.get('CRVAL' + freq_axnum)
params.nu[base_indx:base_indx + nchan] = flist
pcent = 100. * (indx + 1.) / len(fns)
progress(20, pcent)
del tdata
else:
# Separate Stokes has been requested...
for indx in range(len(fns_q)):
fnq = fns_q[indx]
fnu = fns_u[indx]
tdata_q = pyfits.getdata(params.input_dir + qsubdir + fnq)
thead_q = pyfits.getheader(params.input_dir + qsubdir + fnq)
tdata_u = pyfits.getdata(params.input_dir + usubdir + fnu)
thead_u = pyfits.getheader(params.input_dir + usubdir + fnu)
if options.rest_freq:
try:
q_restfreq = thead_q['RESTFREQ']
u_restfreq = thead_u['RESTFREQ']
except KeyError:
q_restfreq = thead_q['RESTFRQ']
u_restfreq = thead_u['RESTFRQ']
if q_restfreq != u_restfreq:
raise Exception(
'Something went wrong! I am trying ' +
'to combine Q & U images at different frequencies!')
else:
if options.freq_last:
if thead_q['CRVAL4'] != thead_u['CRVAL4']:
raise Exception(
'Something went wrong! I am ' +
'trying to combine Q & U images at different ' +
'frequencies!')
else:
if thead_q['CRVAL3'] != thead_u['CRVAL3']:
raise Exception(
'Something went wrong! I am trying ' +
'to combine Q & U images at different ' +
'frequencies!')
base_indx = nchan * indx
if not options.freq_last:
if tdata_q.shape[0] == 1:
cube.real[base_indx:base_indx + nchan] = \
tdata_q[0, :, decsz[0]:decsz[1], rasz[0]:rasz[1]]
cube.imag[base_indx:base_indx + nchan] = \
tdata_u[0, :, decsz[0]:decsz[1], rasz[0]:rasz[1]]
else:
raise Exception('Invalid data shape!')
else:
if tdata_q.shape[1] == 1:
cube.real[base_indx:base_indx + nchan] = \
tdata_q[:, 0, decsz[0]:decsz[1], rasz[0]:rasz[1]]
cube.imag[base_indx:base_indx + nchan] = \
tdata_u[:, 0, decsz[0]:decsz[1], rasz[0]:rasz[1]]
else:
raise Exception('Invalid data shape!')
if options.rest_freq:
if nchan != 1:
msg = 'When the rest frequency option is selected, ' + \
'only one channel is allowed per file!'
raise Exception(msg)
try:
params.nu[base_indx] = thead_q['RESTFREQ']
except KeyError:
params.nu[base_indx] = thead_q['RESTFRQ']
else:
if nchan != 1:
flist = numpy.arange(nchan) *\
thead_q.get('CDELT' + freq_axnum) +\
thead_q.get('CRVAL' + freq_axnum)
params.nu[base_indx:base_indx + nchan] = flist
else:
# When each fits file has a separate Q or U image, CDELT
# cannot be used to estimate the frequency values.
params.nu[base_indx] = thead_q.get('CRVAL' + freq_axnum)
pcent = 100. * (indx + 1.) / len(fns_q)
progress(20, pcent)
del tdata_q
del tdata_u
if options.separate_stokes:
thead = thead_q
# SARRVESH'S EDIT
if nchan == 1:
params.dnu = thead_q['CDELT' + freq_axnum]
if options.rest_freq:
# SARRVESH'S EDIT
# dnu can be estimated using optical velocity information in
# the FITS header. But this works only with images produced
# by AWImager and WSClean. Have to be tested with other imagers
if thead_q['CTYPE' +
freq_axnum] == 'VOPT' and thead_q['CUNIT' +
freq_axnum] == 'm/s':
velocity = thead_q['CDELT' + freq_axnum]
C = 299792458. # light speed in m/s
params.dnu = (params.nu[0]) / (1 - (velocity / C)) - params.nu[0]
else:
# FIXME: This isn't very general and could lead to problems.
params.dnu = params.nu[1] - params.nu[0]
# Print out basic parameters
C2 = 8.98755179e16
nus = numpy.sort(params.nu)
dnu = params.dnu
delta_l2 = C2 * (nus[0]**(-2) - nus[len(nus) - 1]**(-2))
l2min = 0.5 * C2 * ((nus[len(nus) - 1] + dnu)**(-2) +
(nus[len(nus) - 1] - dnu)**(-2))
res = 2. * math.sqrt(3) / delta_l2
maxscale = numpy.pi / l2min
print "\n"
print "The maximum theroretical resolution for the given" +\
" set of parameters is " +str(round(res)) + " rad/m^2"
print "The maximum observable scale for the given set of parameters" +\
" is " +str(round(maxscale)) + " rad/m^2"
print "\n"
# If requested, produce an array containing the spectral index information
# and apply it
if params.alpha:
if params.alphafromfile:
alphadata = pyfits.getdata(params.alpha)
if numpy.shape(alphadata) != (decsz[1] - decsz[0],
rasz[1] - rasz[0]):
try:
alpha = alphadata.reshape(
(decsz[1] - decsz[0], rasz[1] - rasz[0]))
except ValueError:
raise Exception(
'Size of the spectral index image is inconsestent with parameter inputs.'
)
else:
alpha = alphadata
else:
alpha = params.alpha
# set rest frequency to mean frequency across band
if params.ref_freq == None:
#params.ref_freq = numpy.mean(nus)
params.ref_freq = nus[0]
# Rescale cube with spectral dependence function s, assuming
# separability of the spectrally dependent Faraday spectrum,
# following de Bruyn & Brentjens 2005. Division already by
# lambda2 dependence despite the array still being ordered
# in frequency.
#PROBLEM: HOW TO TRANSLATE THE SPECTRUM? CORRECT THIS WAY?
for k in range(len(nus)):
cube[k] /= (nus[k] / params.ref_freq)**(-alpha)
# initialize the RMSynth class that does all the work
rms = R.RMSynth(params.nu, params.dnu, params.phi, params.weight)
print "Done!"
# Write out the RMSF to a text file
rmsfout = numpy.zeros((len(rms.rmsf), 3))
rmsfout[:, 0] = rms.rmsf_phi
rmsfout[:, 1] = rms.rmsf.real
rmsfout[:, 2] = rms.rmsf.imag
numpy.savetxt(params.outputfn + '_rmsf.txt', rmsfout)
if options.stokes_v:
print 'Writing Stokes V cube...'
hdu_v = pyfits.PrimaryHDU(vcube)
try:
generate_v_header(hdu_v, thead, params)
except:
print "Warning: There was a problem generating the header, " + \
"no header information stored!"
print "Unexpected error:", sys.exc_info()[0]
hdu_v_list = pyfits.HDUList([hdu_v])
hdu_v_list.writeto(params.outputfn + '_v.fits', clobber=True)
f = open(params.outputfn + '_freqlist.txt', 'wb')
Writer = csv.writer(f, delimiter=' ')
for i in range(len(fns) * nchan):
Writer.writerow([str(params.nu[i]), fns[i / nchan]])
f.close()
print 'Done!'
print 'Cleaning up Stokes V temp files...'
del vcube
os.remove(vcube_mmfn)
if options.plot_rmsf:
plot_rmsf(rms)
# in case i just want the RMS class and raw data back to e.g. analyze
# single lines of sight
if manual:
return rms, cube
# dirty image
dicube = create_memmap_file_and_array(
outcube_mmfn, (len(params.phi), len(cube[0]), len(cube[0][0])),
numpy.dtype('complex128'))
if params.do_clean:
# To store a master list of clean components for the entire cube
cclist = list()
rescube = create_memmap_file_and_array(
rescube_mmfn, (len(params.phi), len(cube[0]), len(cube[0][0])),
numpy.dtype('complex128'))
cleancube = create_memmap_file_and_array(
cleancube_mmfn, (len(params.phi), len(cube[0]), len(cube[0][0])),
numpy.dtype('complex128'))
cccube = create_memmap_file_and_array(
cccube_mmfn, (len(params.phi), len(cube[0]), len(cube[0][0])),
numpy.dtype('complex128'))
print 'Performing synthesis...'
progress(20, 0)
if params.do_clean:
# initialize the CLEAN class once, reuse for each LOS
# In doing so, the CLEAN beam convolution kernel is only computed once
rmc = R.RMClean(rms, params.niter, params.gain, params.cutoff)
for indx in range(decsz[1] - decsz[0]):
for jndx in range(rasz[1] - rasz[0]):
if maskimage[indx, jndx] != 0:
los = cube[:, indx, jndx]
if params.do_clean:
rmc.reset()
rmc.perform_clean(los)
rmc.restore_clean_map()
cleancube[:, indx, jndx] = rmc.clean_map.copy()
rescube[:, indx, jndx] = rmc.residual_map.copy()
dicube[:, indx, jndx] = rmc.dirty_image.copy()
for kndx in range(len(rmc.cc_phi_list)):
cclist.append([
rmc.cc_phi_list[kndx][0], indx, jndx,
rmc.cc_val_list[kndx].real,
rmc.cc_val_list[kndx].imag
])
# SARRVESH'S EDIT
cccube[:, indx, jndx] = rmc.cc_add_list.copy()
else:
dicube[:, indx, jndx] = rms.compute_dirty_image(los)
else:
dicube[:, indx, jndx] = numpy.zeros((params.nphi, ))
if params.do_clean:
cleancube[:, indx, jndx] = numpy.zeros((params.nphi, ))
rescube[:, indx, jndx] = numpy.zeros((params.nphi, ))
pcent = 100. * (indx + 1.) * (jndx + 1.) / (rasz[1] - rasz[0]) /\
(decsz[1] - decsz[0])
progress(20, pcent)
if params.do_clean:
print '\n'
print "The fitted FWHM of the clean beam is " + str(round(
rmc.sdev, 2)) + " rad/m^2"
print '\n'
print 'RM synthesis done! Writing out FITS files...'
write_output_files(dicube, params, thead, 'di')
if params.do_clean:
write_output_files(rescube, params, thead, 'residual')
write_output_files(cleancube, params, thead, 'clean')
write_output_files(cccube, params, thead, 'cc')
print 'Writing out CC list...'
write_output_files(cccube, params, thead, 'cc')
# TODO: need to make this usable!
# it doesn't work right now because there are just way too many CCs
#write_cc_list(cclist, params.outputfn+"_cc.txt")
print "polint"
polintmap = create_memmap_file_and_array(pmap_mmfn, \
(decsz[1]-decsz[0], rasz[1]-rasz[0]), numpy.dtype('float64'))
#print "phimap"
#phimap = create_memmap_file_and_array(phimap_mmfn, \
# (decsz[1]-decsz[0], rasz[1]-rasz[0]), numpy.dtype('float64'))
#print "qmap"
#qmap = create_memmap_file_and_array(qmap_mmfn, \
# (decsz[1]-decsz[0], rasz[1]-rasz[0]), numpy.dtype('float64'))
#print "umap"
#umap = create_memmap_file_and_array(umap_mmfn, \
# (decsz[1]-decsz[0], rasz[1]-rasz[0]), numpy.dtype('float64'))
print 'Computing polarized intensity and Faraday depth maps'
for indx in range(decsz[1] - decsz[0]):
for jndx in range(rasz[1] - rasz[0]):
if maskimage[indx, jndx] != 0:
if params.do_clean:
q_los = cleancube[:, indx, jndx].real
u_los = cleancube[:, indx, jndx].imag
else:
q_los = dicube[:, indx, jndx].real
u_los = dicube[:, indx, jndx].imag
p_los = numpy.sqrt(
numpy.add(numpy.square(q_los), numpy.square(u_los)))
if numpy.amax(p_los) > params.threshold:
polintmap[indx, jndx] = numpy.amax(p_los)
indx_max_polint = numpy.argmax(p_los)
#phimap[indx, jndx] = params.phi[indx_max_polint]
#qmap[indx, jndx] = q_los[indx_max_polint]
#umap[indx, jndx] = u_los[indx_max_polint]
else:
polintmap[indx, jndx] = numpy.nan
#phimap[indx, jndx] = numpy.nan
#qmap[indx, jndx] = numpy.nan
#umap[indx, jndx] = numpy.nan
else:
polintmap[indx, jndx] = numpy.nan
#phimap[indx, jndx] = numpy.nan
#qmap[indx, jndx] = numpy.nan
#umap[indx, jndx] = numpy.nan
print 'Writing polarized intensity and Faraday depth maps'
write_output_maps(polintmap, params, thead, 'polint')
#write_output_maps(phimap, params, thead, 'phi', 'rad/m/m')
#write_output_maps(qmap, params, thead, 'qmap')
#write_output_maps(umap, params, thead, 'umap')
print 'Cleaning up temp files...'
del dicube
del cube
if params.do_clean:
del cleancube
del rescube
os.remove(incube_mmfn)
os.remove(outcube_mmfn)
os.remove(pmap_mmfn)
#os.remove(phimap_mmfn)
#os.remove(qmap_mmfn)
#os.remove(umap_mmfn)
if params.do_clean:
os.remove(cleancube_mmfn)
os.remove(rescube_mmfn)
os.remove(cccube_mmfn)
print 'Done!'
