def _generate_ws_template(self, nrow=1, nchan=1, npol=1):
u = sakura.empty_aligned((nrow, ), dtype=numpy.float64)
v = sakura.empty_like_aligned(u)
rdata = sakura.empty_aligned((
nrow,
npol,
nchan,
),
dtype=numpy.float32)
idata = sakura.empty_like_aligned(rdata)
flag = sakura.empty_aligned(rdata.shape, dtype=numpy.bool)
weight = sakura.empty_aligned((
nrow,
nchan,
), dtype=numpy.float32)
row_flag = sakura.empty_aligned(u.shape, dtype=numpy.bool)
channel_map = sakura.empty_aligned((nchan, ), dtype=numpy.int32)
ws = almagridder.GridderWorkingSet(data_id=0,
u=u,
v=v,
rdata=rdata,
idata=idata,
flag=flag,
weight=weight,
row_flag=row_flag,
channel_map=channel_map)
# set defualt value for some attributes
ws.flag[:] = True
ws.row_flag[:] = True
ws.weight[:] = 1.0
ws.channel_map[:] = list(range(nchan))
return ws
def grid2ws(grid_real, grid_imag, wgrid_real, wgrid_imag):
"""convert gridder result into workingset instance
Arguments:
grid_real {numpy.ndarray} -- real part of the visibility
grid_imag {numpy.ndarray} -- imaginary part of the visibility
wgrid_real {numpy.ndarray} -- weight of the visibility
wgrid_imag {numpy.ndarray} -- weight of the visibility
Returns:
VisibilityWorkingSet -- visibility working set
"""
gridshape = grid_real.shape
nonzero_real = numpy.where(wgrid_real != 0.0)
nonzero_imag = numpy.where(wgrid_imag != 0.0)
# uv location
# assumption here is that the first index corresponds to v while
# the second one corresponds u so that [i,j] represents
# the value at uv location (j,i).
# since the array is C-contiguous, memory layout is contiguous
# along u-axis.
#
# | 9|10|11|
# | 6| 7| 8|
# | 3| 4| 5|
# | 0| 1| 2|
npol = gridshape[2]
nchan = gridshape[3]
data_id = 0
num_vis = len(nonzero_real[0])
u = sakura.empty_aligned((num_vis, ), dtype=numpy.int32)
v = sakura.empty_like_aligned(u)
rdata = sakura.empty_aligned((num_vis, ), dtype=numpy.float64)
idata = sakura.empty_like_aligned(rdata)
wdata = sakura.empty_like_aligned(rdata)
xpos = 0
for ipol in range(npol):
for ichan in range(nchan):
ir = numpy.where(
numpy.logical_and(nonzero_real[2] == ipol,
nonzero_real[3] == ichan))
#ii = numpy.where(numpy.logical_and(nonzero_imag[2] == ipol,
# nonzero_imag[3] == ichan))
xlen = len(v)
nextpos = xpos + xlen
v[xpos:nextpos] = nonzero_real[0][ir]
u[xpos:nextpos] = nonzero_real[1][ir]
rdata[xpos:nextpos] = grid_real[v, u, ipol, ichan]
idata[xpos:nextpos] = grid_imag[v, u, ipol, ichan]
wdata[xpos:nextpos] = wgrid_real[v, u, ipol, ichan]
xpos = nextpos
visibility_data = VisibilityWorkingSet(data_id=data_id,
u=u,
v=v,
rdata=rdata,
idata=idata,
weight=wdata)
return visibility_data
def flatten(self, working_set):
"""
Generator yielding list of working_sets divided by spectral channels
"""
nchan = working_set.nchan
nrow_ws = working_set.nrow
chan_image = numpy.unique(working_set.channel_map)
num_ws = len(chan_image)
for imchan in chan_image:
vischans = numpy.where(working_set.channel_map == imchan)[0]
nchan = len(vischans)
nrow = nrow_ws * nchan
ws_shape = (
nrow,
1,
1,
)
ws_shape2 = (
nrow,
1,
)
real = sakura.empty_aligned(ws_shape,
dtype=working_set.rdata.dtype)
imag = sakura.empty_aligned(ws_shape,
dtype=working_set.idata.dtype)
flag = sakura.empty_aligned(ws_shape, dtype=working_set.flag.dtype)
weight = sakura.empty_aligned(ws_shape2,
dtype=working_set.weight.dtype)
row_flag = sakura.empty_aligned((nrow, ),
dtype=working_set.row_flag.dtype)
channel_map = sakura.empty_aligned(
(1, ), dtype=working_set.channel_map.dtype)
u = sakura.empty_aligned((nrow, ), dtype=working_set.u.dtype)
v = sakura.empty_like_aligned(u)
row_start = 0
for ichan in vischans:
row_end = row_start + nrow_ws
real[row_start:row_end, 0, 0] = working_set.rdata[:, 0, ichan]
imag[row_start:row_end, 0, 0] = working_set.idata[:, 0, ichan]
flag[row_start:row_end, 0, 0] = working_set.flag[:, 0, ichan]
weight[row_start:row_end, 0] = working_set.weight[:, ichan]
row_flag[row_start:row_end] = working_set.row_flag[:]
channel_map[0] = imchan
u[row_start:row_end] = working_set.u[:, ichan]
v[row_start:row_end] = working_set.v[:, ichan]
row_start = row_end
ws = gridder.GridderWorkingSet(data_id=working_set.data_id,
u=u,
v=v,
rdata=real,
idata=imag,
flag=flag,
weight=weight,
row_flag=row_flag,
channel_map=channel_map)
#print('yielding channelized working set from channels {}'.format(vischans))
yield ws
def generate_subset(self, subset_id):
self.subset_id = subset_id
# grid data shape: (nv, nu, npol, nchan)
rdata = self.visibility.rdata
idata = self.visibility.idata
weight = self.visibility.weight
u = self.visibility.u
v = self.visibility.v
for local_id in range(self.num_fold):
self.subset_id = local_id
# random index
random_index = self.index_generator.get_subset_index(self.subset_id)
#print('DEBUG_TN: subset ID {0} random_index = {1}'.format(self.subset_id, list(random_index)))
# mask array for active visibility
num_vis = len(self.visibility)
mask = numpy.zeros(num_vis, dtype=numpy.bool)
mask[:] = True
mask[random_index] = False
# uv location
# assumption here is that the first index corresponds to v while
# the second one corresponds u so that [i,j] represents
# the value at uv location (j,i).
# since the array is C-contiguous, memory layout is contiguous
# along u-axis.
#
# | 9|10|11|
# | 6| 7| 8|
# | 3| 4| 5|
# | 0| 1| 2|
# visibility data to be cached
# here, we assume npol == 1 (Stokes visibility I_v) and nchan == 1
num_subvis = len(random_index)
rcache = sakura.empty_aligned((num_subvis,), dtype=rdata.dtype)
icache = sakura.empty_like_aligned(rcache)
wcache = sakura.empty_like_aligned(rcache)
ucache = sakura.empty_aligned((num_subvis,), dtype=u.dtype)
vcache = sakura.empty_like_aligned(ucache)
rcache[:] = rdata[random_index]
icache[:] = idata[random_index]
wcache[:] = weight[random_index]
ucache[:] = u[random_index]
vcache[:] = v[random_index]
# generate subset
assert num_subvis < num_vis
num_active = num_vis - num_subvis
ractive = sakura.empty_aligned((num_active,), dtype=rdata.dtype)
iactive = sakura.empty_like_aligned(ractive)
wactive = sakura.empty_like_aligned(ractive)
uactive = sakura.empty_aligned((num_active,), dtype=u.dtype)
vactive = sakura.empty_like_aligned(uactive)
ractive[:] = rdata[mask]
iactive[:] = idata[mask]
wactive[:] = weight[mask]
uactive[:] = u[mask]
vactive[:] = v[mask]
self.visibility_active = datacontainer.VisibilityWorkingSet(data_id=0,
rdata=ractive,
idata=iactive,
weight=wactive,
u=uactive,
v=vactive)
self.visibility_cache = [datacontainer.VisibilityWorkingSet(data_id=0, # nominal data ID
rdata=rcache,
idata=icache,
weight=wcache,
u=ucache,
v=vcache)]
try:
yield self
finally:
self._clear()
def generate_subset(self, subset_id):
self.subset_id = subset_id
# grid data shape: (nv, nu, npol, nchan)
grid_real = self.visibility.real
grid_imag = self.visibility.imag
wgrid_real = self.visibility.wreal
gdata_shape = grid_real.shape
# random index
random_index = self.index_generator.get_subset_index(self.subset_id)
#print('DEBUG_TN: subset ID {0} random_index = {1}'.format(self.subset_id, list(random_index)))
# uv location
# assumption here is that the first index corresponds to v while
# the second one corresponds u so that [i,j] represents
# the value at uv location (j,i).
# since the array is C-contiguous, memory layout is contiguous
# along u-axis.
#
# | 9|10|11|
# | 6| 7| 8|
# | 3| 4| 5|
# | 0| 1| 2|
num_subvis = len(random_index)
u = sakura.empty_aligned((num_subvis,), dtype=numpy.float64)
v = sakura.empty_like_aligned(u)
uid = self.active_index[1][random_index]
vid = self.active_index[0][random_index]
cellu = self.uvgrid.cellu
cellv = self.uvgrid.cellv
offsetu = self.uvgrid.offsetu
offsetv = self.uvgrid.offsetv
nu = self.uvgrid.nu
nv = self.uvgrid.nv
assert gdata_shape[0] == nv
assert gdata_shape[1] == nu
#print 'subset ID {0}: uid{0}={1}; vid{0}={2}'.format(self.subset_id, uid.tolist(), vid.tolist())
u[:] = (uid - offsetu) * cellu
v[:] = (vid - offsetv) * cellv
# visibility data to be cached
# here, we assume npol == 1 (Stokes visibility I_v) and nchan == 1
assert len(gdata_shape) == 4
assert gdata_shape[2] == 1 # npol should be 1
assert gdata_shape[3] == 1 # nchan should be 1
real = sakura.empty_aligned((num_subvis,), dtype=numpy.float32)
imag = sakura.empty_like_aligned(real)
wreal = sakura.empty_like_aligned(real)
real[:] = grid_real[self.active_index][random_index]
imag[:] = grid_imag[self.active_index][random_index]
wreal[:] = wgrid_real[self.active_index][random_index]
# generate subset
self.visibility_active = self.visibility
self.__replace_with(self.visibility_active.real, random_index, 0.0)
self.__replace_with(self.visibility_active.imag, random_index, 0.0)
self.__replace_with(self.visibility_active.wreal, random_index, 0.0)
self.visibility_cache = [datacontainer.VisibilityWorkingSet(data_id=0, # nominal data ID
rdata=real,
idata=imag,
weight=wreal,
u=u,
v=v)]
try:
yield self
finally:
self.restore_visibility()
def fill_uvw(self, ws, chunk, lsr_edge_frequency):
"""
Fill UV coordinate
ws -- working set to be filled
chunk -- input data chunk
lsr_edge_frequency -- channel edge frequency (LSRK)
"""
phasecenter = self.imageparam.phasecenter
self._check_phasecenter(phasecenter)
#self._warn_refocus()
qa = casa.CreateCasaQuantity()
speed_of_light = qa.constants('c')
c = qa.convert(speed_of_light, 'm/s')['value']
data_shape = chunk['data'].shape
nrow = data_shape[2]
nchan = data_shape[1]
uvw = chunk['uvw']
data_desc_ids = chunk['data_desc_id']
uvgrid = self.imageparam.uvgridconfig
delta_u = uvgrid.cellu
delta_v = uvgrid.cellv
offset_u = uvgrid.offsetu
offset_v = uvgrid.offsetv
# UVW conversion
u = sakura.empty_aligned((nrow, nchan), dtype=uvw.dtype)
v = sakura.empty_like_aligned(u)
for irow in range(nrow):
# TODO: phase rotation if image phasecenter is different from
# the reference direction of the observation
pass
# conversion from physical baseline length to the value
# normalized by observing wavelength
u0 = uvw[0, irow]
v0 = uvw[1, irow]
spw_id = data_desc_ids[irow]
#chan_freq = lsr_frequency
#chan_width = (lsr_edge_frequency[1:] - lsr_edge_frequency[:-1]).mean()
#freq_start = chan_freq[0] - chan_width / 2
#freq_end = chan_freq[-1] + chan_width / 2
#freq_start = lsr_edge_frequency[0]
#freq_end = lsr_edge_frequency[-1]
#center_freq = (freq_start + freq_end) / 2
center_freq = numpy.fromiter(
(numpy.mean(lsr_edge_frequency[i:i + 2])
for i in range(nchan)),
dtype=numpy.float64)
u[irow] = u0 * center_freq / c # divided by wavelength
v[irow] = v0 * center_freq / c # divided by wavelength
# TODO?: refocus UVW if distance to the source is known
pass
# project uv-coordinate value onto gridding pixel plane
# pixel size is determined by an inverse of image extent
# along x (RA) and y (DEC) direction
#
# v
# nv-1|(nmin,vmax) (umax,vmax)
# .|
# .|
# .|
# .|
# nv/2| (0,0)
# .|
# .|
# 2|
# 1|
# 0|(umin,vmin) (umax,vmin)
# |__________________________
# 0 1 2 ......nu/2...nu-1 u
u[irow] = u[irow] / delta_u + offset_u
# Sign of v must be inverted so that MFISTA routine generates
# proper image. Otherwise, image will be flipped in the vertical
# axis.
v[irow] = -v[irow] / delta_v + offset_v
ws.u = u
ws.v = v