def cl_from_im(image, clname=None, threshold=None):
from taskinit import qa,ia, cl
ia.open(image)
data = ia.getregion()
cs = ia.coordsys()
ia.done()
data = np.where(np.isnan(data), 0, data)
if threshold is None:
datanz = np.nonzero(data)
else:
datanz = np.nonzero(np.abs(data) > threshold)
modelfluxes = data[datanz]
modelpixels = np.array(datanz)
modellist = cs.convertmany(modelpixels,
unitsin=['pix', 'pix', 'pix', 'pix'],
unitsout=['rad', 'rad', '', 'Hz'])
cl.done()
for i in range(modellist.shape[1]):
x = qa.formxxx(str(modellist[0, i])+'rad', format='hms', prec=6)
y = qa.formxxx(str(modellist[1, i])+'rad', format='dms', prec=6)
pos = ' '.join(['J2000', x, y])
freq = str(modellist[3, i])+'Hz'
flux = modelfluxes[i]
cl.addcomponent(flux=flux, fluxunit='Jy', dir=pos, freq=freq)
if clname is not None:
cl.rename(clname)
cl.done()
else:
return cl
def noise(coords, nrandom = 50, imagenames=[], stampsize=32,
method = 'mean', weighting = 'simga2', maskradius=None,
psfmode = 'point'):
import stacker
import numpy as np
from taskinit import ia, qa
ia.open(imagenames[0])
beam = qa.convert(ia.restoringbeam()['major'], 'rad')['value']
ia.done()
# if coords.coord_type == 'physical':
# coords = stacker.getPixelCoords(coords, imagenames)
_allocate_buffers(imagenames, stampsize, len(coords)*len(imagenames))
dist = []
for i in range(nrandom):
random_coords = stacker.randomizeCoords(coords, beam=beam)
random_coords = stacker.getPixelCoords(random_coords, imagenames)
_load_stack(random_coords, psfmode)
if method == 'mean' and weighting == 'sigma2':
random_coords = _calculate_sigma2_weights(random_coords, maskradius)
elif method == 'mean' and weighting == 'sigma':
random_coords = _calculate_sigma_weights(random_coords, maskradius)
stacked_im = _stack_stack(method, random_coords)
dist.append(stacked_im[int(stampsize/2+0.5), int(stampsize/2+0.5),0,0])
return np.std(dist)
def noise(coords,
vis,
weighting='sigma2',
imagenames=[],
beam=None,
nrand=50,
stampsize=32,
maskradius=None):
""" Calculate noise using a Monte Carlo method, can be time consuming. """
import stacker
import stacker.image
from math import pi
if beam is None:
try:
from taskinit import ia, qa
ia.open(imagenames[0])
beam = qa.convert(ia.restoringbeam()['major'], 'rad')['value']
ia.done()
except ImportError:
beam = 1 / 3600. / 180. * pi
dist = []
for i in range(nrand):
random_coords = stacker.randomizeCoords(coords, beam=beam)
if weighting == 'sigma2':
random_coords = stacker.image.calculate_sigma2_weights(
random_coords, imagenames, stampsize, maskradius)
dist.append(stack(random_coords, vis))
return np.std(np.real(np.array(dist)))
def randomCoords(imagenames, ncoords=10):
import random
from taskinit import ia, qa
xmin, xmax = [], []
ymin, ymax = [], []
for image in imagenames:
ia.open(image)
print image, ia.boundingbox()
trc = ia.boundingbox()['trcf'].split(', ')
blc = ia.boundingbox()['blcf'].split(', ')
xmin.append(qa.convert(qa.quantity(trc[0]), 'rad')['value'])
xmax.append(qa.convert(qa.quantity(blc[0]), 'rad')['value'])
ymin.append(qa.convert(qa.quantity(blc[1]), 'rad')['value'])
ymax.append(qa.convert(qa.quantity(trc[1]), 'rad')['value'])
ia.done()
randomcoords = CoordList(imagenames)
for i in range(ncoords):
imageid = random.randint(0, len(imagenames) - 1)
x = random.uniform(xmin[imageid], xmax[imageid])
y = random.uniform(ymin[imageid], ymax[imageid])
c = Coord(x, y, 1.0)
randomcoords.append(c)
return randomcoords
def cl_from_im(image, clname=None, threshold=None):
from taskinit import qa, ia, cl
ia.open(image)
data = ia.getregion()
cs = ia.coordsys()
ia.done()
data = np.where(np.isnan(data), 0, data)
if threshold is None:
datanz = np.nonzero(data)
else:
datanz = np.nonzero(np.abs(data) > threshold)
modelfluxes = data[datanz]
modelpixels = np.array(datanz)
modellist = cs.convertmany(modelpixels,
unitsin=['pix', 'pix', 'pix', 'pix'],
unitsout=['rad', 'rad', '', 'Hz'])
cl.done()
for i in range(modellist.shape[1]):
x = qa.formxxx(str(modellist[0, i]) + 'rad', format='hms', prec=6)
y = qa.formxxx(str(modellist[1, i]) + 'rad', format='dms', prec=6)
pos = ' '.join(['J2000', x, y])
freq = str(modellist[3, i]) + 'Hz'
flux = modelfluxes[i]
cl.addcomponent(flux=flux, fluxunit='Jy', dir=pos, freq=freq)
if clname is not None:
cl.rename(clname)
cl.done()
else:
return cl
def _getPixelCoords1ImSimpleProj(coords, imagename):
from taskinit import ia
from interval import interval
ia.open(imagename)
cs = ia.coordsys()
imshape = ia.shape()
ia.done()
pixcoords = []
for coord in coords:
p = cs.convert(coordin=[coord.x, coord.y, 0, 0],
absin=[True] * 4,
unitsin=[coords.unit, coords.unit, 'pix', 'pix'],
absout=[True] * 4,
unitsout=['pix'] * 4)
x = p[0]
y = p[1]
if x in interval[0, imshape[0] - 1] and y in interval[0.,
imshape[1] - 1]:
c = Coord(x, y)
try:
c.index = coord.index
except AttributeError:
pass
pixcoords.append(c)
return pixcoords
def noise(coords, vis, weighting='sigma2', imagenames=[], beam=None, nrand=50,
stampsize=32, maskradius=None):
""" Calculate noise using a Monte Carlo method, can be time consuming. """
import stacker
import stacker.image
from math import pi
if beam is None:
try:
from taskinit import ia, qa
ia.open(imagenames[0])
beam = qa.convert(ia.restoringbeam()['major'], 'rad')['value']
ia.done()
except ImportError:
beam = 1/3600./180.*pi
dist = []
for i in range(nrand):
random_coords = stacker.randomizeCoords(coords, beam=beam)
if weighting == 'sigma2':
random_coords = stacker.image.calculate_sigma2_weights(
random_coords, imagenames, stampsize, maskradius)
dist.append(stack(random_coords, vis))
return np.std(np.real(np.array(dist)))
def _calculate_sigma_weights(coords, maxmaskradius=None):
import numpy as np
from taskinit import ia
global stampsize
if coords.physical and coords.imagenames[0]:
ia.open(coords.imagenames[0])
# beam = qa.convert(ia.restoringbeam()['major'], 'rad')['value']
if ia.restoringbeam() == {}:
masksize = 10
else:
masksize = 2*np.abs(qa.convert(ia.restoringbeam()['major'],'rad')['value']/ ia.coordsys().increment()['numeric'][0])
ia.done()
X = np.arange(0, stampsize)-stampsize/2
Y = np.arange(0, stampsize)-stampsize/2
X,Y = np.meshgrid(X,Y)
for i,coord in enumerate(coords):
tmpdata = data[i,:,:,:,:]
for j in range(tmpdata.shape[2]):
for k in range(tmpdata.shape[3]):
tmpdata[:,:,j,k] = (tmpdata[:,:,j,k]*np.double( np.sqrt(X**2+Y**2)>masksize))
sigma = np.std(tmpdata)
if sigma == 0:
coord.weight = 0.
else:
coord.weight = 1/sigma
if maxmaskradius and maxmaskradius < masksize:
masksize = maxmaskradius
return coords
def make_casa_testimage(infile, outname):
infile = str(infile)
outname = str(outname)
if not casaOK:
raise Exception("Attempted to make a CASA test image in a non-CASA "
"environment")
ia = casatools.image()
ia.fromfits(infile=infile, outfile=outname, overwrite=True)
ia.unlock()
ia.close()
ia.done()
cube = SpectralCube.read(infile)
if isinstance(cube, VaryingResolutionSpectralCube):
ia.open(outname)
# populate restoring beam emptily
ia.setrestoringbeam(
major={
'value': 1.0,
'unit': 'arcsec'
},
minor={
'value': 1.0,
'unit': 'arcsec'
},
pa={
'value': 90.0,
'unit': 'deg'
},
channel=len(cube.beams) - 1,
polarization=-1,
)
# populate each beam (hard assumption of 1 poln)
for channum, beam in enumerate(cube.beams):
casabdict = {
'major': {
'value': beam.major.to(u.deg).value,
'unit': 'deg'
},
'minor': {
'value': beam.minor.to(u.deg).value,
'unit': 'deg'
},
'positionangle': {
'value': beam.pa.to(u.deg).value,
'unit': 'deg'
}
}
ia.setrestoringbeam(beam=casabdict,
channel=channum,
polarization=0)
ia.unlock()
ia.close()
ia.done()
def getFlux(imagename):
from taskinit import ia,rg
ia.open(imagename)
cs = ia.coordsys()
x = int(cs.referencepixel()['numeric'][0])
y = int(cs.referencepixel()['numeric'][1])
ia.done()
return float(ia.getregion(region=rg.box([x,y], [x,y])))
def __init__(self, imagename, *args, **kwargs):
"""
Constructor
Keyword arguments:
imagename -- Str to casa image of primary beam
"""
super(MSPrimaryBeamModel, self).__init__(*args, **kwargs)
self.imagename = imagename
try:
from taskinit import ia
ia.open(imagename)
self.cs = ia.coordsys()
self.nx = ia.shape()[0]
self.ny = ia.shape()[1]
self.refpix_x = self.cs.referencepixel()['numeric'][0]
self.refpix_y = self.cs.referencepixel()['numeric'][1]
self.increment_x = self.cs.increment()['numeric'][0]
self.increment_y = self.cs.increment()['numeric'][1]
try:
self.frequencyaxis = self.cs.findaxisbyname('frequency')
self.nu0 = self.cs.referencevalue()['numeric'][
self.frequencyaxis]
except Exception:
self.nu0 = None
print('Some stuff!')
self.data = ia.getregion()[:, :, 0, 0]
ia.done()
except ImportError:
from pyrap.images import image
im = image(imagename)
self.nx = im.shape()[-1]
self.nx = im.shape()[-2]
self.cs = im.coordinates()
self.cs_dir = self.cs.get_coordinate('direction')
self.refpix_x = self.cs_dir.get_referencepixel()[1]
self.refpix_y = self.cs_dir.get_referencepixel()[0]
self.increment_x = self.cs_dir.get_increment()[1]
self.increment_y = self.cs_dir.get_increment()[0]
try:
self.nu0 = self.cs.get_coordinate(
'spectral').get_referencevalue()
except Exception:
self.nu0 = None
print(
'Warning! No frequency information in primary beam model.')
self.data = im.getdata()[0, 0]
def make_pbfile(vis, pbfile):
from taskinit import im, ms, ia, qa, tb
import numpy as np
from scipy.constants import c
ms.open(vis)
fields = ms.range('field_id')['field_id']
ms.done()
im.open(vis)
im.selectvis(field=fields[0])
ms.open(vis)
freq = np.mean(ms.range('chan_freq')['chan_freq'])
phase_dir = ms.range('phase_dir')['phase_dir']['direction']
ms.done()
phase_dir = phase_dir[0][0], phase_dir[1][0]
phase_dir = [
qa.formxxx(str(phase_dir[0]) + 'rad', format='hms'),
qa.formxxx(str(phase_dir[1]) + 'rad', format='dms')
]
phase_dir = 'J2000 ' + ' '.join(phase_dir)
tb.open(vis + '/ANTENNA/')
dishdia = np.min(tb.getcol('DISH_DIAMETER'))
tb.done()
# pb of 512 pix cover pb down to 0.001
# ensure largest pixel to pixel var to .01
minpb = 0.001
nx = 512
cellconv = (nx * np.sqrt(np.log(2) / np.log(1 / minpb)))**-1
beam = c / freq / dishdia
cell = {}
cell['value'] = beam * cellconv
cell['unit'] = 'rad'
# nx = int(3*3e8/freq/dishdia*1.22*180/
# math.pi*3600/qa.convert(advise['cell'],
# 'arcsec')['value'])
# Chosen as to be 3 times fwhm of primary beam,
# should include up to approximately .01 of peak flux
im.defineimage(nx=nx, ny=nx, cellx=cell, celly=cell, phasecenter=phase_dir)
im.setvp(dovp=True)
im.makeimage(type='pb', image=pbfile)
im.done()
ia.open(pbfile)
cs = ia.coordsys()
cs.setreferencevalue(type='direction', value=[0., 0.])
ia.setcoordsys(cs.torecord())
ia.maskhandler('delete', 'mask0')
ia.done()
def test_casa_mask(data_adv, tmp_path):
cube = SpectralCube.read(data_adv)
mask_array = np.array([[True, False], [False, False], [True, True]])
bool_mask = BooleanArrayMask(mask=mask_array,
wcs=cube._wcs,
shape=cube.shape)
cube = cube.with_mask(bool_mask)
make_casa_mask(cube,
str(tmp_path / 'casa.mask'),
add_stokes=False,
append_to_image=False,
overwrite=True)
ia = casatools.image()
ia.open(str(tmp_path / 'casa.mask'))
casa_mask = ia.getchunk()
coords = ia.coordsys()
ia.unlock()
ia.close()
ia.done()
# Test masks
# Mask array is broadcasted to the cube shape. Mimic this, switch to ints,
# and transpose to match CASA image.
compare_mask = np.tile(mask_array, (4, 1, 1)).astype('int16').T
assert np.all(compare_mask == casa_mask)
# Test WCS info
# Convert back to an astropy wcs object so transforms are dealt with.
casa_wcs = wcs_casa2astropy(ia, coords)
header = casa_wcs.to_header() # Invokes transform
# Compare some basic properties EXCLUDING the spectral axis
assert np.allclose(cube.wcs.wcs.crval[:2], casa_wcs.wcs.crval[:2])
assert np.all(cube.wcs.wcs.cdelt[:2] == casa_wcs.wcs.cdelt[:2])
assert np.all(list(cube.wcs.wcs.cunit)[:2] == list(casa_wcs.wcs.cunit)[:2])
assert np.all(list(cube.wcs.wcs.ctype)[:2] == list(casa_wcs.wcs.ctype)[:2])
# Reference pixels in CASA are 1 pixel off.
assert_allclose(cube.wcs.wcs.crpix, casa_wcs.wcs.crpix, atol=1.0)
def _write_stacked_image(imagename, pixels, template_image, stampsize):
import os
import shutil
import numpy as np
from taskinit import ia
# global stampsize
if os.access(imagename, os.F_OK): shutil.rmtree(imagename)
ia.open(template_image)
beam = ia.restoringbeam()
cs = ia.coordsys()
ia.done()
csnew = cs.copy()
csnew.setreferencevalue([0.]*2, 'dir')
csnew.setreferencepixel([int(stampsize/2+0.5)]*2, 'dir')
ia.fromarray(imagename, pixels=pixels, csys = csnew.torecord())
ia.open(imagename)
ia.setrestoringbeam(beam=beam)
ia.done()
def _calculate_flux_weights(coords):
from taskinit import ia
import re
fluxmap = []
r = re.compile('(.*)\.image/*')
for imagename in coords.imagenames:
match = r.match(imagename)
if match:
fluximage = match.group(1)+'.flux'
if ia.open(fluximage):
fluxmap.append(ia.getregion())
ia.done()
for i,coord in enumerate(coords):
coord.weight = (fluxmap[coord.image][int(coord.x+0.5), int(coord.y+0.5), 0, 0])**2
return coords
def _getPixelCoords1Im(coords, imagename):
from interval import interval
import math
try:
from taskinit import ia
ia.open(imagename)
cs = ia.coordsys()
Nx = ia.shape()[0]
Ny = ia.shape()[1]
ia.done()
x0 = cs.referencevalue()['numeric'][0]
y0 = cs.referencevalue()['numeric'][1]
x_pix_ref = cs.referencepixel()['numeric'][0]
y_pix_ref = cs.referencepixel()['numeric'][1]
x_pix_inc = cs.increment()['numeric'][0]
y_pix_inc = cs.increment()['numeric'][1]
# If we fail to load ia, we will use pyrap instead.
# This probably means stacker was loaded from outside casapy.
except ImportError:
from pyrap.images import image
im = image(imagename)
cs = im.coordinates().get_coordinate('direction')
dir_axis_index = im.coordinates().get_axes().index(cs.get_axes())
imshape = im.shape()
try:
x_axis_index = cs.get_axes().index('Right Ascension')
except ValueError:
raise ValueError('Could not find direction coordinate: '\
'RightAscension')
try:
y_axis_index = cs.get_axes().index('Declination')
except ValueError:
raise ValueError('Could not find direction coordinate: '\
'Declination')
Nx = im.shape()[dir_axis_index + x_axis_index]
Ny = im.shape()[dir_axis_index + y_axis_index]
x0 = cs.get_referencevalue()[x_axis_index]
y0 = cs.get_referencevalue()[y_axis_index]
x_pix_ref = cs.get_referencepixel()[x_axis_index]
y_pix_ref = cs.get_referencepixel()[y_axis_index]
x_pix_inc = cs.get_increment()[x_axis_index]
y_pix_inc = cs.get_increment()[y_axis_index]
pixcoords = []
for coord in coords:
dx = (coord.x - x0) * math.cos(coord.y)
dy = math.asin(math.sin(coord.y) / math.cos(dx)) - y0
x = dx / x_pix_inc + x_pix_ref
y = dy / y_pix_inc + y_pix_ref
if x in interval[0, Nx - 1] and y in interval[0., Ny - 1]:
# pixcoords.append(Coord(x,y, coord.weight))
c = Coord(x, y, coord.weight)
try:
c.index = coord.index
except AttributeError:
pass
pixcoords.append(c)
return pixcoords
def make_casa_mask(SpecCube, outname, append_to_image=True,
img=None, add_stokes=True, stokes_posn=None,
overwrite=False
):
'''
Outputs the mask attached to the SpectralCube object as a CASA image, or
optionally appends the mask to a preexisting CASA image.
Parameters
----------
SpecCube : SpectralCube
SpectralCube object containing mask.
outname : str
Name of the outputted mask file.
append_to_image : bool, optional
Appends the mask to a given image.
img : str, optional
Image to be appended to. Must be specified if append_to_image is
enabled.
add_stokes: bool, optional
Adds a Stokes axis onto the wcs from SpecCube.
stokes_posn : int, optional
Sets the position of the new Stokes axis. Defaults to the last axis.
overwrite : bool, optional
Overwrite the image and mask files if they exist?
'''
try:
from casatools import image
ia = image()
except ImportError:
try:
from taskinit import ia
except ImportError:
raise ImportError("Cannot import casa. Must be run in a CASA environment.")
# the 'mask name' is distinct from the mask _path_
maskname = os.path.split(outname)[1]
maskpath = outname
# Get the header info from the image
# There's not wcs_astropy2casa (yet), so create a temporary file for
# CASA to open.
temp = tempfile.NamedTemporaryFile()
# CASA is closing this file at some point so set it to manual delete.
temp2 = tempfile.NamedTemporaryFile(delete=False)
# Grab wcs
# Optionally re-add on the Stokes axis
if add_stokes:
my_wcs = SpecCube.wcs
if stokes_posn is None:
stokes_posn = my_wcs.wcs.naxis
new_wcs = add_stokes_axis_to_wcs(my_wcs, stokes_posn)
header = new_wcs.to_header()
# Transpose the shape so we're adding the axis at the place CASA will
# recognize. Then transpose back.
shape = SpecCube.shape[::-1]
shape = shape[:stokes_posn] + (1,) + shape[stokes_posn:]
shape = shape[::-1]
else:
# Just grab the header from SpecCube
header = SpecCube.header
shape = SpecCube.shape
hdu = fits.PrimaryHDU(header=header,
data=np.empty(shape, dtype='int16'))
hdu.writeto(temp.name)
ia.fromfits(infile=temp.name, outfile=temp2.name, overwrite=overwrite)
temp.close()
cs = ia.coordsys()
ia.done()
ia.close()
temp2.close()
mask_arr = SpecCube.mask.include()
# Reshape mask with possible Stokes axis
mask_arr = mask_arr.reshape(shape)
# Transpose to match CASA axes
mask_arr = mask_arr.T
ia.fromarray(outfile=maskpath,
pixels=mask_arr.astype('int16'),
overwrite=overwrite)
ia.done()
ia.close()
ia.open(maskpath, cache=False)
ia.setcoordsys(cs.torecord())
ia.done()
ia.close()
if append_to_image:
if img is None:
raise TypeError("img argument must be specified to append the mask.")
ia.open(maskpath, cache=False)
ia.calcmask(maskname+">0.5")
ia.done()
ia.close()
ia.open(img, cache=False)
ia.maskhandler('copy', [maskpath+":mask0", maskname])
ia.maskhandler('set', maskname)
ia.done()
ia.close()
result_dict['component0']['spectrum']['frequency']['m0'][
'value'] = float(freqs[ll])
results[itpp]['results'][comp] = result_dict['component0']
results[itpp]['converged'].append(True)
except:
results[itpp]['converged'].append(False)
results[itpp]['results']['nelements'] = results[itpp]['results'].keys()
# update timestamp
timstr = hdr['date-obs']
return [True, timstr, img, results]
except Exception, instance:
casalog.post(str('*** Error in imfit ***') + str(instance))
# raise instance
return [False, timstr, img, {}]
finally:
ia.done()
def pmaxfit(imagefiles, ncpu, box, width):
from functools import partial
# check if imagefiles is a single file or a list of files
if isinstance(imagefiles, str):
imagefiles = [imagefiles]
if (not isinstance(imagefiles, list)):
print 'input "imagefiles" is not a list. Abort...'
# check file existence
imgfiles = []
for i, img in enumerate(imagefiles):
if os.path.exists(img):
def stack(coords, outfile, stampsize = 32, imagenames= [], method = 'mean',
weighting = None, maxmaskradius=None, psfmode = 'point', primarybeam = None):
"""
Performs stacking in the image domain.
coords -- A coordList object of all target coordinates.
outfile -- Target name for stacked image.
stampsize -- size of target image in pixels
imagenames -- Name of images to extract flux from.
method -- 'mean' or 'median', will determined how pixels are calculated
weighting -- only for method 'mean', if set to None will use weights in coords.
maxmaskradius -- allows blanking of centre pixels in weight calculation
psfmode -- Allows application of filters to stacking, currently not supported.
primarybeam -- only applies if weighting='pb'
returns: Estimate of stacked flux assuming point source.
"""
from ..interval import interval
import os
import shutil
import numpy as np
from taskinit import ia, casalog
casalog.origin('stacker')
casalog.post('#'*42,'INFO')
casalog.post('#'*5 + ' {0: <31}'.format("Begin Task: Stacker")+'#'*5, 'INFO')
global skymap
global data
global oldimagenames
if coords.coord_type == 'physical':
coords = stacker.getPixelCoords(coords, imagenames)
# Important that len(coords) here is for the pixel coordinates, not physical!
_allocate_buffers(coords.imagenames, stampsize, len(coords))
ia.open(coords.imagenames[0])
cs = ia.coordsys()
outnchans = ia.boundingbox()['trc'][2]+1
outnstokes = ia.boundingbox()['trc'][3]+1
ia.done()
for imagename in coords.imagenames:
ia.open(imagename)
if ia.shape()[2] != outnchans or ia.shape()[3] != outnstokes:
print('Channels/polarisations do not match in all images! You probably want to stacking do stacking on continuum data and not on spectral cube.')
return
ia.done()
_load_stack(coords, psfmode)
if method == 'mean' and weighting == 'sigma2':
coords = _calculate_sigma2_weights(coords, maxmaskradius)
elif method == 'mean' and weighting == 'sigma':
coords = _calculate_sigma_weights(coords, maxmaskradius)
elif method == 'mean' and weighting == 'pb':
coords = calculate_pb_weights(coords, primarybeam, imagenames)
npos = len([c.weight for c in coords if c.weight > 1e-6])
casalog.post('Number of stacking positions: {0}'.format(npos),
priority='INFO')
stacked_im = _stack_stack(method, coords)
_write_stacked_image(outfile, stacked_im,
coords.imagenames[0], stampsize)
casalog.post('#'*5 + ' {0: <31}'.format("End Task: stacker")+'#'*5)
casalog.post('#'*42)
return stacked_im[int(stampsize/2), int(stampsize/2),0,0]
def _allocate_buffers( imagenames, new_stampsize, nstackpos):
import numpy as np
try:
from taskinit import ia
dataread = 'casa'
except ImportError:
from pyrap.images import image
dataread = 'pyrap'
global skymap
global data
global oldimagenames
global stampsize
global imagesizes
if dataread == 'casa':
ia.open(imagenames[0])
cs = ia.coordsys()
outnchans = ia.boundingbox()['trc'][2]+1
outnstokes = ia.boundingbox()['trc'][3]+1
ia.done()
elif dataread == 'pyrap':
im = image(imagenames[0])
cs = im.coordinates()
outnchans = im.shape()[cs.get_axes().index(cs.get_coordinate('spectral').get_axes())]
outnstokes = im.shape()[cs.get_axes().index(cs.get_coordinate('stokes').get_axes())]
# To improve performance this module will keep buffers between run.
# This following code resets these buffers if they have grown obsolete.
if oldimagenames == []:
oldimagenames = imagenames
if oldimagenames != imagenames:
oldimagenames = imagenames
skymap = []
data = []
if stampsize == 0:
stampsize = new_stampsize
elif stampsize != new_stampsize:
stampsize = new_stampsize
data = []
skymap = []
if not(data == []) and nstackpos != data.shape[0]:
data = []
# If there is no data buffer create one.
# The data buffer is used to save the right stacking positions before stacking them.
# During stacking this is where the full stack will actually be saved.
if data == []:
data = np.zeros((nstackpos, new_stampsize, new_stampsize, outnstokes, outnchans))
else:
data = 0.*data
# If there is no skymap buffer create one.
# This is the data that is most important to buffer.
# Reading a skymap from disk is a time consuming task and we don't want to do this too much.
if skymap == []:
for imagename in imagenames:
if dataread == 'casa':
ia.open(imagename)
skymap.append(ia.getregion())
ia.done()
elif dataread == 'pyrap':
buff = im.getdata()
dir_axis = cs.get_axes().index(cs.get_coordinate('direction').get_axes())
x_axis = dir_axis+cs.get_coordinate('direction').get_axes().index('Right Ascension')
y_axis = dir_axis+cs.get_coordinate('direction').get_axes().index('Declination')
specax = cs.get_axes().index(cs.get_coordinate('spectral').get_axes())
stokesax = cs.get_axes().index(cs.get_coordinate('stokes').get_axes())
axis_order = [x_axis, y_axis, stokesax, specax]
for i in range(len(axis_order)-1):
if axis_order[i] != i:
target = axis_order.index(i)
origin = i
buff = buff.swapaxes(axis_order[origin], axis_order[target])
axis_order[origin], axis_order[target] =\
axis_order[target], axis_order[origin]
skymap.append(buff)
imagesizes = []
for imagename in imagenames:
if dataread == 'casa':
ia.open(imagename)
imagesizes.append((ia.shape()[0], ia.shape()[1]))
ia.done()
elif dataread == 'pyrap':
dir_axis = cs.get_axes().index(cs.get_coordinate('direction').get_axes())
x_axis_index = dir_axis+cs.get_coordinate('direction').get_axes().index('Right Ascension')
y_axis_index = dir_axis+cs.get_coordinate('direction').get_axes().index('Declination')
imagesizes.append((im.shape()[x_axis_index], im.shape()[y_axis_index]))
