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