def process(self, mg):
fig_prefix = self.params['fig_name']
# distributed along feed
mg.dataset_common_to_distributed('eigval', distributed_axis=1)
mg.dataset_common_to_distributed('gain', distributed_axis=1)
feed = mpiutil.scatter_array(mg.attrs['feed'])
for idx, fd in enumerate(feed):
plt.figure()
plt.subplot(411)
plt.imshow(mg['gain'].local_data[:, idx, 0, :].T.real,
origin='lower') # xx
plt.colorbar()
plt.subplot(412)
plt.imshow(mg['gain'].local_data[:, idx, 0, :].T.imag,
origin='lower') # xx
plt.colorbar()
plt.subplot(413)
plt.imshow(mg['gain'].local_data[:, idx, 1, :].T.real,
origin='lower') # yy
plt.colorbar()
plt.subplot(414)
plt.imshow(mg['gain'].local_data[:, idx, 1, :].T.imag,
origin='lower') # yy
plt.colorbar()
fig_name = '%s_%d.png' % (fig_prefix, fd)
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.clf()
return mg
def __init__(self, parameter_file_or_dict=None, feedback=2):
super(OneAndOne, self).__init__(parameter_file_or_dict, feedback)
self.input_files = input_path(format_list(self.params['input_files']))
self.output_files = output_path(format_list(self.params['output_files']), mkdir=False)
self.iterable = self.params['iterable']
if self.iterable:
self.iter_start = self.params['iter_start']
self.iter_step = self.params['iter_step']
self.iter_num = self.params['iter_num']
self._iter_cnt = 0 # inner iter counter
# Inspect the `process` method to see how many arguments it takes.
pro_argspec = inspect.getargspec(self.process)
n_args = len(pro_argspec.args) - 1
if n_args > 1:
msg = ("`process` method takes more than 1 argument, which is not"
" allowed")
raise PipelineConfigError(msg)
if pro_argspec.varargs or pro_argspec.keywords or pro_argspec.defaults:
msg = ("`process` method may not have variable length or optional"
" arguments")
raise PipelineConfigError(msg)
if n_args == 0:
self._no_input = True
else: # n_args == 1
self._no_input = False
if len(self._in) != n_args and len(self.input_files) == 0:
msg = ("No data to iterate over. There are no 'in' keys and no 'input_files'")
raise PipelineConfigError(msg)
def write_output(self, output):
"""Method for writing time ordered data output. """
exclude = self.params['exclude']
check_status = self.params['check_status']
write_hints = self.params['write_hints']
libver = self.params['libver']
chunk_vis = self.params['chunk_vis']
chunk_shape = self.params['chunk_shape']
chunk_size = self.params['chunk_size']
output_failed_continue = self.params['output_failed_continue']
tag_output_iter = self.params['tag_output_iter']
if self.iterable and tag_output_iter:
output_files = output_path(self.output_files,
relative=False,
iteration=self.iteration)
else:
output_files = self.output_files
try:
output.to_files(output_files, exclude, check_status, write_hints,
libver, chunk_vis, chunk_shape, chunk_size)
except Exception as e:
if output_failed_continue:
msg = 'Process %d writing output to files failed...' % mpiutil.rank
logger.warning(msg)
traceback.print_exc(file=sys.stdout)
else:
raise e
def __init__(self, pipefile=None, feedback=2):
# Read in the parameters.
self.params, self.task_params = parse_ini.parse(pipefile, self.params_init, prefix=self.prefix, return_undeclared=True, feedback=feedback)
self.tasks = self.params['tasks']
# timing the running
if self.params['timing']:
self.start_time = datetime.datetime.now()
if mpiutil.rank0:
print 'Start the pipeline at %s...' % self.start_time
# set environment var
os.environ['TL_OUTPUT'] = self.params['outdir'] + '/'
# copy pipefile to outdir if required
if self.params['copy']:
base_name = path.basename(pipefile)
dst_file = '%s/%s' % (self.params['outdir'], base_name)
outdir = output_path(dst_file, relative=False, mkdir=True)
if mpiutil.rank0:
if self.params['overwrite']:
shutil.copy2(pipefile, dst_file)
else:
if not path.exists(dst_file):
shutil.copy2(pipefile, dst_file)
else:
base, ext = path.splitext(dst_file)
for cnt in itertools.count(1):
dst = '%s_%d%s' % (base, cnt, ext) # add cnt to file name
if not path.exists(dst):
shutil.copy2(pipefile, dst)
break
mpiutil.barrier()
def combine_results(self):
output_combined = self.params['output_combined']
output_combined = output_path(
output_combined, relative=not output_combined.startswith('/'))
self.allocate_output(output_combined, 'w')
self.df_out[-1]['mode_list'] = self.mode_list
for _m in self.mode_list:
mode_key = 'cleaned_%02dmode' % _m
self.create_dataset(-1,
mode_key,
dset_shp=self.dset_shp,
dset_info=self.map_info)
self.create_dataset(-1,
mode_key + '_weight',
dset_shp=self.dset_shp,
dset_info=self.map_info)
mask = np.ones(self.dset_shp[:1]).astype('bool')
for ii in range(len(self.output_files)):
for key in self.df_out[ii][mode_key].keys():
self.df_out[-1][mode_key][:]\
+= self.df_out[ii]['weight'][:]\
* self.df_out[ii]['%s/%s'%(mode_key, key)][:]
self.df_out[-1][mode_key + '_weight'][:]\
+= self.df_out[ii]['weight'][:]
mask *= ~(self.df_out[ii]['mask'][:].astype('bool'))
weight = self.df_out[-1][mode_key + '_weight'][:]
weight[weight == 0] = np.inf
self.df_out[-1][mode_key][:] /= weight
self.df_out[-1][mode_key + '_mask'] = (~mask).astype('int')
def __init__(self, pipefile=None, feedback=2):
# Read in the parameters.
self.params, self.task_params = parse_ini.parse(pipefile, self.params_init, prefix=self.prefix, return_undeclared=True, feedback=feedback)
self.tasks = self.params['tasks']
# set environment var
os.environ['TL_OUTPUT'] = self.params['outdir'] + '/'
# copy pipefile to outdir if required
if self.params['copy']:
base_name = path.basename(pipefile)
dst_file = '%s/%s' % (self.params['outdir'], base_name)
outdir = output_path(dst_file, relative=False, mkdir=True)
if mpiutil.rank0:
if self.params['overwrite']:
shutil.copy2(pipefile, dst_file)
else:
if not path.exists(dst_file):
shutil.copy2(pipefile, dst_file)
else:
base, ext = path.splitext(dst_file)
for cnt in itertools.count(1):
dst = '%s_%d%s' % (base, cnt, ext) # add cnt to file name
if not path.exists(dst):
shutil.copy2(pipefile, dst)
break
mpiutil.barrier()
def process(self, mg):
fig_prefix = self.params['fig_name']
# distributed along feed
mg.dataset_common_to_distributed('eigval', distributed_axis=1)
mg.dataset_common_to_distributed('gain', distributed_axis=1)
feed = mpiutil.scatter_array(mg.attrs['feed'])
for idx, fd in enumerate(feed):
plt.figure()
plt.subplot(411)
plt.imshow(mg['gain'].local_data[:, idx, 0, :].T.real, origin='lower') # xx
plt.colorbar()
plt.subplot(412)
plt.imshow(mg['gain'].local_data[:, idx, 0, :].T.imag, origin='lower') # xx
plt.colorbar()
plt.subplot(413)
plt.imshow(mg['gain'].local_data[:, idx, 1, :].T.real, origin='lower') # yy
plt.colorbar()
plt.subplot(414)
plt.imshow(mg['gain'].local_data[:, idx, 1, :].T.imag, origin='lower') # yy
plt.colorbar()
fig_name = '%s_%d.png' % (fig_prefix, fd)
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.clf()
return mg
def write_to_file_splitmock(self):
output_prefix = '/%s_cube_%s' % (self.params['prefix'],
self.params['scenario'])
for outfile in self.outfiles_split:
mock_idx = self.iter_list[self.iter]
output_file = output_prefix + '_%03d_%s.h5' % (mock_idx, outfile)
output_file = output_path(output_file, relative=True)
#self.allocate_output(output_file, 'w')
if outfile is 'optsim':
map_key = 'delta'
weight_key = 'separable'
map_name = 'sim_map_optsim'
weight_name = 'sel'
else:
map_key = self.params['map_tmp_key']
weight_key = self.params['map_tmp_weight']
map_name = 'sim_map_' + outfile
weight_name = 'weight'
map_pad = self.map_pad
map_shp = self.map_tmp.shape
_map_shp = (map_shp[0], map_shp[1] - 2 * map_pad,
map_shp[2] - 2 * map_pad)
map_slice = (slice(0, None), slice(map_pad, map_shp[1] - map_pad),
slice(map_pad, map_shp[2] - map_pad))
with h5py.File(output_file, 'w') as f:
d = f.create_dataset(map_key,
_map_shp,
dtype=self.map_tmp.dtype)
_d = getattr(self, map_name)
d[:] = _d[map_slice]
for key, value in self.map_tmp.info.iteritems():
d.attrs[key] = repr(value)
d.attrs['freq_delta'] = repr(self.map_tmp.info['freq_delta'] /
1.e6)
d.attrs['freq_centre'] = repr(
self.map_tmp.info['freq_centre'] / 1.e6)
#d.attrs['ra_delta'] = repr(self.map_tmp.info['ra_delta'])
#d.attrs['ra_centre'] = repr(self.map_tmp.info['ra_centre'])
#d.attrs['dec_delta'] = repr(self.map_tmp.info['dec_delta'])
#d.attrs['dec_centre'] = repr(self.map_tmp.info['dec_centre'])
if hasattr(self, weight_name):
#mask =
#mask[mask!=0] = 1.
#d[:] *= mask
n = f.create_dataset(weight_key,
_map_shp,
dtype=self.map_tmp.dtype)
n[:] = getattr(self, weight_name)[map_slice]
for key, value in self.map_tmp.info.iteritems():
n.attrs[key] = repr(value)
n.attrs['freq_delta'] = repr(
self.map_tmp.info['freq_delta'] / 1.e6)
n.attrs['freq_centre'] = repr(
self.map_tmp.info['freq_centre'] / 1.e6)
def init_output(self):
suffix = '_%s.h5'%self.params['data_sets']
output_file = self.output_files[0]
output_file = output_path(output_file + suffix,
relative = not output_file.startswith('/'))
#if mpiutil.rank0:
# self.df = h5py.File(output_file, mode='w')
self.allocate_output(output_file, 'w')
def read_process_write(self, tstream):
"""Overwrite the method of superclass."""
if isinstance(tstream, timestream.Timestream):
return self.process(tstream)
else:
ts_dir = output_path(self.params['ts_dir'])
ts_name = self.params['ts_name']
if mpiutil.rank0:
print 'Try to load tstream from %s/%s' % (ts_dir, ts_name)
tstream = timestream.Timestream.load(ts_dir, ts_name)
return self.process(tstream)
def write_output(self, output):
"""Method for writing time ordered data output. """
exclude = self.params['exclude']
check_status = self.params['check_status']
libver = self.params['libver']
tag_output_iter = self.params['tag_output_iter']
if self.iterable and tag_output_iter:
output_files = output_path(self.output_files,
relative=False,
iteration=self.iteration)
else:
output_files = self.output_files
output.to_files(output_files, exclude, check_status, libver)
def read_input(self):
for input_file in self.input_files:
print input_file
self.open(input_file)
self.map_tmp = al.make_vect(al.load_h5(self.df_in[0], 'dirty_map'))
self.map_shp = self.map_tmp.shape
for output_file in self.output_files:
output_file = output_path(output_file,
relative=not output_file.startswith('/'))
self.allocate_output(output_file, 'w')
self.create_dataset_like(-1, 'clean_map', self.map_tmp)
self.create_dataset_like(-1, 'noise_diag', self.map_tmp)
return 1
def write_output(self, output):
norm = self.hitmap
norm[norm==0] = np.inf
self.avgmap /= norm
nside = self.params['nside']
map_name = 'avgmap_nside%d_f%6.2fMHz.fits'%(nside, self.freq)
map_name = output_path(map_name)
hp.write_map(map_name, self.avgmap, coord='C', overwrite=True)
self.avgmap[self.avgmap==0] = hp.UNSEEN
hp.mollview(self.avgmap, title='Map @ %6.2f MHz'%self.freq, badcolor='none')
hp.graticule(coord='C', dpar=30, dmer=30)
plt.show()
def write_output(self, output):
#super(PlotSpectrum, self).write_output(output)
fig_prefix = self.params['output_files'][0]
main_data = self.params['main_data']
ymin = self.params['ymin']
ymax = self.params['ymax']
axhh = self.axhh
axvv = self.axvv
fig = self.fig
x_label = self.x_label
axhh.set_xticklabels([])
#axhh.set_ylabel(r'${\rm frequency\, [GHz]\, HH}$')
#axhh.set_ylabel('HH Polarization')
#if xmin is None: xmin = x_axis[0]
#if xmax is None: xmax = x_axis[-1]
xmin = self.xmin
xmax = self.xmax
axhh.set_xlim(xmin=self.xmin, xmax=self.xmax)
axhh.set_ylim(ymin=ymin, ymax=ymax)
axhh.minorticks_on()
axhh.tick_params(length=4, width=1, direction='in')
axhh.tick_params(which='minor', length=2, width=1, direction='in')
axhh.legend(title=self.params['legend_title'])
axhh.set_ylabel('HH Polarization')
#axvv.set_xlabel(r'$({\rm time} - {\rm UT}\quad %s\,) [{\rm hour}]$'%t_start)
axvv.set_xlabel(x_label)
#axvv.set_ylabel(r'${\rm frequency\, [GHz]\, VV}$')
#axvv.set_ylabel('VV Polarization')
axvv.set_xlim(xmin=self.xmin, xmax=self.xmax)
axvv.set_ylim(ymin=ymin, ymax=ymax)
axvv.minorticks_on()
axvv.tick_params(length=4, width=1, direction='in')
axvv.tick_params(which='minor', length=2, width=1, direction='in')
axvv.set_ylabel('VV Polarization')
if fig_prefix is not None:
fig_name = '%s_%s_Spec.png' % (fig_prefix, main_data)
fig_name = output_path(fig_name)
plt.savefig(fig_name, formate='png') #, dpi=100)
def write_output(self, output):
super(PlotPointingvsTime, self).write_output(output)
fig_prefix = self.params['output_files'][0]
main_data = self.params['main_data']
axhh = self.axhh
axvv = self.axvv
fig = self.fig
axhh.set_ylabel('Azimuth [arcmin]')
axvv.set_ylabel('Elevation [arcmin]')
if fig_prefix is not None:
fig_name = '%s_%s_AzEl.png' % (fig_prefix, main_data)
fig_name = output_path(fig_name)
plt.savefig(fig_name, formate='png') #, dpi=100)
def write_output(self, output):
super(PlotVvsTime, self).write_output(output)
fig_prefix = self.params['output_files'][0]
main_data = self.params['main_data']
axhh = self.axhh
axvv = self.axvv
fig = self.fig
axhh.set_ylabel('HH Polarization')
axvv.set_ylabel('VV Polarization')
if fig_prefix is not None:
fig_name = '%s_%s_TS.png' % (fig_prefix, main_data)
fig_name = output_path(fig_name)
plt.savefig(fig_name, formate='png') #, dpi=100)
def read_input(self):
for input_file in self.input_files:
if mpiutil.rank0:
logger.info('%s' % input_file)
self.open(input_file)
map_tmp = al.load_h5(self.df_in[0], 'dirty_map')
self.map_tmp = al.make_vect(map_tmp, axis_names=map_tmp.info['axes'])
self.map_shp = self.map_tmp.shape
for output_file in self.output_files:
output_file = output_path(output_file,
relative=not output_file.startswith('/'))
self.allocate_output(output_file, 'w')
self.create_dataset_like(-1, 'clean_map', self.map_tmp)
self.create_dataset_like(-1, 'noise_diag', self.map_tmp)
self.create_dataset_like(-1, 'dirty_map', self.map_tmp)
return 1
def init_output(self):
output_file = self.output_files[0]
output_file = output_path(output_file,
relative= not output_file.startswith('/'))
self.allocate_output(output_file, 'w')
# load one file to get the ant_n and pol_n
self.create_dataset('binavg_1d', self.dset_shp + (self.knum,))
self.create_dataset('counts_1d', self.dset_shp + (self.knum,))
self.create_dataset('binavg_2d', self.dset_shp + (self.knum_x, self.knum_y))
self.create_dataset('counts_2d', self.dset_shp + (self.knum_x, self.knum_y))
self.df['kbin'] = self.kbin
self.df['kbin_x'] = self.kbin_x
self.df['kbin_y'] = self.kbin_y
self.df['kbin_edges'] = self.kbin_edges
self.df['kbin_x_edges'] = self.kbin_x_edges
self.df['kbin_y_edges'] = self.kbin_y_edges
self.df['sim'] = self.params['sim']
def write_to_file(self):
output_prefix = '/%s_cube_%s' % (self.params['prefix'],
self.params['scenario'])
for outfile in self.outfiles:
mock_idx = self.iter_list[self.iter]
output_file = output_prefix + '_%03d_%s.h5' % (mock_idx, outfile)
output_file = output_path(output_file, relative=True)
#self.allocate_output(output_file, 'w')
if outfile is 'optsim':
map_key = 'delta'
weight_key = 'separable'
map_name = 'sim_map_optsim'
weight_name = 'sel'
else:
map_key = self.params['map_tmp_key']
weight_key = self.params['map_tmp_weight']
map_name = 'sim_map_' + outfile
weight_name = 'weight'
with h5py.File(output_file, 'w') as f:
d = f.create_dataset(map_key,
self.map_tmp.shape,
dtype=self.map_tmp.dtype)
d[:] = getattr(self, map_name)
for key, value in self.map_tmp.info.iteritems():
d.attrs[key] = repr(value)
if hasattr(self, weight_name):
#mask =
#mask[mask!=0] = 1.
#d[:] *= mask
n = f.create_dataset(weight_key,
self.map_tmp.shape,
dtype=self.map_tmp.dtype)
n[:] = getattr(self, weight_name)
for key, value in self.map_tmp.info.iteritems():
n.attrs[key] = repr(value)
def open_outputfiles(self):
output_prefix = '/%s_cube_%s' % (self.params['prefix'],
self.params['scenario'])
output_file = output_prefix + '.h5'
output_file = output_path(output_file, relative=True)
self.allocate_output(output_file, 'w')
dset_shp = (self.params['mock_n'], ) + self.map_tmp.shape
dset_info = {}
dset_info['axes'] = ('mock', 'freq', 'ra', 'dec')
dset_info['type'] = 'vect'
dset_info['mock_delta'] = 1
dset_info['mock_centre'] = self.params['mock_n'] // 2
dset_info['freq_delta'] = self.map_tmp.info['freq_delta'] / 1.e6
dset_info['freq_centre'] = self.map_tmp.info['freq_centre'] / 1.e6
dset_info['ra_delta'] = self.map_tmp.info['ra_delta']
dset_info['ra_centre'] = self.map_tmp.info['ra_centre']
dset_info['dec_delta'] = self.map_tmp.info['dec_delta']
dset_info['dec_centre'] = self.map_tmp.info['dec_centre']
#if self.params['outfile_raw']:
for outfile in self.outfiles:
self.create_dataset(outfile, dset_shp, dset_info)
def setup(self):
super(SurveySimToMap, self).setup()
params = self.params
freq = params['freq']
self.n_freq = freq.shape[0]
self.freq_spacing = freq[1] - freq[0]
self.n_ra, self.n_dec = params['map_shape']
self.map_shp = (self.n_freq, self.n_ra, self.n_dec)
self.spacing = params['pixel_spacing']
self.dec_spacing = self.spacing
self.ra_spacing = self.spacing / sp.cos(
params['field_centre'][1] * sp.pi / 180.)
axis_names = ('freq', 'ra', 'dec')
map_tmp = np.zeros(self.map_shp, dtype=__dtype__)
map_tmp = al.make_vect(map_tmp, axis_names=axis_names)
map_tmp.set_axis_info('freq', freq[self.n_freq // 2],
self.freq_spacing)
map_tmp.set_axis_info('ra', params['field_centre'][0], self.ra_spacing)
map_tmp.set_axis_info('dec', params['field_centre'][1],
self.dec_spacing)
self.map_tmp = map_tmp
mock_n = self.mock_n
#for ii in range(mock_n):
for ii in self.HI_mock_ids:
output_file = 'sim_mock%03d_%s_%s_%s_%s.h5' % (
ii, self.params['prefix'], self.params['survey_mode'],
self.params['HI_scenario'], self.params['HI_model_type'])
output_file = output_path(output_file, relative=True)
self.allocate_output(output_file, 'w')
self.create_dataset_like(-1, 'dirty_map', map_tmp)
self.create_dataset_like(-1, 'clean_map', map_tmp)
self.create_dataset_like(-1, 'count_map', map_tmp)
def flag(self, vis, vis_mask, li, gi, bl, ts, **kwargs):
"""Function that does the actual flag."""
freq_window = self.params['freq_window']
time_window = self.params['time_window']
freq_sigma = self.params['freq_sigma']
time_sigma = self.params['time_sigma']
plot_fit = self.params['plot_fit']
freq_fig_prefix = self.params['freq_fig_name']
time_fig_prefix = self.params['time_fig_name']
tag_output_iter = self.params['tag_output_iter']
iteration = self.iteration
freq_flag = kwargs.get('freq_flag')
time_flag = kwargs.get('time_flag')
time = ts.time[:]
nt = len(time)
freq = ts.freq[:]
nfreq = len(freq)
if isinstance(ts, Timestream): # for Timestream
pol = bl[0]
bl = tuple(bl[1])
elif isinstance(ts, RawTimestream): # for RawTimestream
pol = None
bl = tuple(bl)
else:
raise ValueError('Need either a RawTimestream or Timestream')
if freq_flag:
# time integration
tm_vis = np.ma.mean(np.ma.array(vis, mask=vis_mask), axis=0)
abs_vis = np.abs(tm_vis)
if np.ma.count_masked(tm_vis) > 0: # has masked value
abs_vis_valid = abs_vis[~abs_vis.mask]
inds_valid = np.arange(nfreq)[~abs_vis.mask]
itp = InterpolatedUnivariateSpline(inds_valid, abs_vis_valid)
abs_vis_itp = itp(np.arange(nfreq))
abs_vis1 = abs_vis_itp.copy()
else:
abs_vis1 = abs_vis.copy()
for cnt in xrange(10):
if cnt != 0:
abs_vis1[inds] = smooth[inds]
smooth = savitzky_golay(abs_vis1, freq_window, 3)
# flage RFI
diff = abs_vis1 - smooth
mean = np.mean(diff)
std = np.std(diff)
inds = np.where(np.abs(diff - mean) > freq_sigma*std)[0]
if len(inds) == 0:
break
diff = abs_vis - smooth
mean = np.mean(diff)
std = np.std(diff)
inds = np.where(np.abs(diff - mean) > freq_sigma*std)[0] # masked inds
vis_mask[:, inds] = True # set mask
if plot_fit:
plt.figure()
plt.plot(freq, abs_vis, label='data')
plt.plot(freq[inds], abs_vis[inds], 'ro', label='flag')
plt.plot(freq, smooth, label='smooth')
plt.xlabel(r'$\nu$ / MHz')
plt.legend(loc='best')
if pol is None:
fig_name = '%s_%d_%d.png' % (freq_fig_prefix, bl[0], bl[1])
else:
fig_name = '%s_%d_%d_%s.png' % (freq_fig_prefix, bl[0], bl[1], pol)
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
if time_flag:
# freq integration
fm_vis = np.ma.mean(np.ma.array(vis, mask=vis_mask), axis=1)
abs_vis = np.abs(fm_vis)
if np.ma.count_masked(fm_vis) > 0: # has masked value
abs_vis_valid = abs_vis[~abs_vis.mask]
inds_valid = np.arange(nt)[~abs_vis.mask]
itp = InterpolatedUnivariateSpline(inds_valid, abs_vis_valid)
abs_vis_itp = itp(np.arange(nt))
abs_vis1 = abs_vis_itp.copy()
else:
abs_vis1 = abs_vis.copy()
for cnt in xrange(10):
if cnt != 0:
abs_vis1[inds] = smooth[inds]
smooth = savitzky_golay(abs_vis1, time_window, 3)
# flage RFI
diff = abs_vis1 - smooth
mean = np.mean(diff)
std = np.std(diff)
inds = np.where(np.abs(diff - mean) > time_sigma*std)[0]
if len(inds) == 0:
break
diff = abs_vis - smooth
mean = np.mean(diff)
std = np.std(diff)
inds = np.where(np.abs(diff - mean) > time_sigma*std)[0] # masked inds
# Addtional threshold
# inds1 = np.where(np.abs(diff[inds]) > 1.0e-2*np.abs(smooth[inds]))[0]
# inds = inds[inds1]
vis_mask[inds] = True # set mask
if plot_fit:
plt.figure()
plt.plot(time, abs_vis, label='data')
plt.plot(time[inds], abs_vis[inds], 'ro', label='flag')
plt.plot(time, smooth, label='smooth')
plt.xlabel(r'$t$ / Julian Date')
plt.legend(loc='best')
if pol is None:
fig_name = '%s_%d_%d.png' % (time_fig_prefix, bl[0], bl[1])
else:
fig_name = '%s_%d_%d_%s.png' % (time_fig_prefix, bl[0], bl[1], pol)
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
return vis, vis_mask
def plot(self, vis, vis_mask, li, gi, bl, ts, **kwargs):
"""Function that does the actual plot work."""
plot_type = self.params['plot_type']
bl_incl = self.params['bl_incl']
bl_excl = self.params['bl_excl']
flag_mask = self.params['flag_mask']
flag_ns = self.params['flag_ns']
slices = self.params['slices']
fig_prefix = self.params['fig_name']
rotate_xdate = self.params['rotate_xdate']
feed_no = self.params['feed_no']
tag_output_iter = self.params['tag_output_iter']
iteration = self.iteration
if isinstance(ts, Timestream): # for Timestream
pol = bl[0]
bl = tuple(bl[1])
feed_no = True
elif isinstance(ts, RawTimestream): # for RawTimestream
pol = None
bl = tuple(bl)
if feed_no:
pol = ts['bl_pol'].local_data[li]
bl = tuple(ts['true_blorder'].local_data[li])
else:
raise ValueError('Need either a RawTimestream or Timestream')
if bl_incl != 'all':
bl1 = set(bl)
bl_incl = [{f1, f2} for (f1, f2) in bl_incl]
bl_excl = [{f1, f2} for (f1, f2) in bl_excl]
if (not bl1 in bl_incl) or (bl1 in bl_excl):
return
if plot_type == 'time':
nt = vis.shape[0]
c = nt / 2
s = max(0, c - slices / 2)
e = min(nt, s + slices)
if flag_mask:
vis1 = np.ma.array(vis[s:e], mask=vis_mask[s:e])
elif flag_ns:
vis1 = vis[s:e].copy()
if 'ns_on' in ts.iterkeys():
ns_on = ts['ns_on'][s:e]
on = np.where(ns_on)[0]
vis1[on] = complex(np.nan, np.nan)
else:
vis1 = vis[s:e]
o = c - s
shift = 0.1 * np.ma.max(np.abs(vis1[o]))
ax_val = ts.freq[:]
xlabel = r'$\nu$ / MHz'
elif plot_type == 'freq':
nfreq = vis.shape[1]
c = nfreq / 2
s = max(0, c - slices / 2)
e = min(nfreq, s + slices)
if flag_mask:
vis1 = np.ma.array(vis[:, s:e], mask=vis_mask[:, s:e])
elif flag_ns:
vis1 = vis[:, s:e].copy()
if 'ns_on' in ts.iterkeys():
ns_on = ts['ns_on'][:]
on = np.where(ns_on)[0]
vis1[on] = complex(np.nan, np.nan)
else:
vis1 = vis[:, s:e]
o = c - s
shift = 0.1 * np.ma.max(np.abs(vis1[:, o]))
# ax_val = ts.time[:]
# xlabel = r'$t$ / Julian Date'
ax_val = np.array(
[datetime.fromtimestamp(sec) for sec in ts['sec1970'][:]])
xlabel = '%s' % ax_val[0].date()
ax_val = mdates.date2num(ax_val)
else:
raise ValueError(
'Unknown plot_type %s, must be either time or freq' %
plot_type)
plt.figure()
f, axarr = plt.subplots(3, sharex=True)
for i in range(e - s):
if plot_type == 'time':
axarr[0].plot(ax_val,
vis1[i].real + (i - o) * shift,
label='real')
elif plot_type == 'freq':
axarr[0].plot(ax_val,
vis1[:, i].real + (i - o) * shift,
label='real')
if i == 0:
axarr[0].legend()
if plot_type == 'time':
axarr[1].plot(ax_val,
vis1[i].imag + (i - o) * shift,
label='imag')
elif plot_type == 'freq':
axarr[1].plot(ax_val,
vis1[:, i].imag + (i - o) * shift,
label='imag')
if i == 0:
axarr[1].legend()
if plot_type == 'time':
axarr[2].plot(ax_val,
np.abs(vis1[i]) + (i - o) * shift,
label='abs')
elif plot_type == 'freq':
axarr[2].plot(ax_val,
np.abs(vis1[:, i]) + (i - o) * shift,
label='abs')
if i == 0:
axarr[2].legend()
if plot_type == 'freq':
axarr[2].xaxis_date()
date_format = mdates.DateFormatter('%H:%M')
axarr[2].xaxis.set_major_formatter(date_format)
if rotate_xdate:
# set the x-axis tick labels to diagonal so it fits better
f.autofmt_xdate()
else:
# half the number of date ticks so they do not overlap
# axarr[2].set_xticks(axarr[2].get_xticks()[::2])
# reduce the number of tick locators
locator = MaxNLocator(nbins=6)
axarr[2].xaxis.set_major_locator(locator)
axarr[2].xaxis.set_minor_locator(AutoMinorLocator(2))
axarr[2].set_xlabel(xlabel)
if feed_no:
fig_name = '%s_%s_%d_%d_%s.png' % (fig_prefix, plot_type, bl[0],
bl[1], ts.pol_dict[pol])
else:
fig_name = '%s_%s_%d_%d.png' % (fig_prefix, plot_type, bl[0],
bl[1])
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
def process(self, ts):
mask_daytime = self.params['mask_daytime']
mask_time_range = self.params['mask_time_range']
tsys = self.params['tsys']
accuracy_boost = self.params['accuracy_boost']
l_boost = self.params['l_boost']
bl_range = self.params['bl_range']
auto_correlations = self.params['auto_correlations']
time_avg = self.params['time_avg']
pol = self.params['pol']
interp = self.params['interp']
beam_dir = output_path(self.params['beam_dir'])
use_existed_beam = self.params['use_existed_beam']
gen_inv = self.params['gen_invbeam']
noise_weight = self.params['noise_weight']
ts_dir = output_path(self.params['ts_dir'])
ts_name = self.params['ts_name']
no_m_zero = self.params['no_m_zero']
simulate = self.params['simulate']
input_maps = self.params['input_maps']
prior_map = self.params['prior_map']
add_noise = self.params['add_noise']
dirty_map = self.params['dirty_map']
nbin = self.params['nbin']
method = self.params['method']
normalize = self.params['normalize']
threshold = self.params['threshold']
eps = self.params['epsilon']
correct_order = self.params['correct_order']
if use_existed_beam:
# load the saved telescope from disk
tel = None
else:
assert isinstance(ts, Timestream), '%s only works for Timestream object' % self.__class__.__name__
ts.redistribute('baseline')
lat = ts.attrs['sitelat']
# lon = ts.attrs['sitelon']
lon = 0.0
# lon = np.degrees(ts['ra_dec'][0, 0]) # the first ra
local_origin = False
freqs = ts.freq[:] # MHz
nfreq = freqs.shape[0]
band_width = ts.attrs['freqstep'] # MHz
try:
ndays = ts.attrs['ndays']
except KeyError:
ndays = 1
feeds = ts['feedno'][:]
bl_order = mpiutil.gather_array(ts.local_bl, axis=0, root=None, comm=ts.comm)
bls = [ tuple(bl) for bl in bl_order ]
az, alt = ts['az_alt'][0]
az = np.degrees(az)
alt = np.degrees(alt)
pointing = [az, alt, 0.0]
feedpos = ts['feedpos'][:]
if ts.is_dish:
from tlpipe.map.drift.telescope import tl_dish
dish_width = ts.attrs['dishdiam']
tel = tl_dish.TlUnpolarisedDishArray(lat, lon, freqs, band_width, tsys, ndays, accuracy_boost, l_boost, bl_range, auto_correlations, local_origin, dish_width, feedpos, pointing)
elif ts.is_cylinder:
from tlpipe.map.drift.telescope import tl_cylinder
# factor = 1.2 # suppose an illumination efficiency, keep same with that in timestream_common
factor = 0.79 # for xx
# factor = 0.88 # for yy
cyl_width = factor * ts.attrs['cywid']
tel = tl_cylinder.TlUnpolarisedCylinder(lat, lon, freqs, band_width, tsys, ndays, accuracy_boost, l_boost, bl_range, auto_correlations, local_origin, cyl_width, feedpos)
else:
raise RuntimeError('Unknown array type %s' % ts.attrs['telescope'])
if not simulate:
# select the corresponding vis and vis_mask
if pol == 'xx':
local_vis = ts.local_vis[:, :, 0, :]
local_vis_mask = ts.local_vis_mask[:, :, 0, :]
elif pol == 'yy':
local_vis = ts.local_vis[:, :, 1, :]
local_vis_mask = ts.local_vis_mask[:, :, 1, :]
elif pol == 'I':
xx_vis = ts.local_vis[:, :, 0, :]
xx_vis_mask = ts.local_vis_mask[:, :, 0, :]
yy_vis = ts.local_vis[:, :, 1, :]
yy_vis_mask = ts.local_vis_mask[:, :, 1, :]
local_vis = np.zeros_like(xx_vis)
for ti in xrange(local_vis.shape[0]):
for fi in xrange(local_vis.shape[1]):
for bi in xrange(local_vis.shape[2]):
if xx_vis_mask[ti, fi, bi] != yy_vis_mask[ti, fi, bi]:
if xx_vis_mask[ti, fi, bi]:
local_vis[ti, fi, bi] = yy_vis[ti, fi, bi]
else:
local_vis[ti, fi, bi] = xx_vis[ti, fi, bi]
else:
local_vis[ti, fi, bi] = 0.5 * (xx_vis[ti, fi, bi] + yy_vis[ti, fi, bi])
local_vis_mask = xx_vis_mask | yy_vis_mask
else:
raise ValueError('Invalid pol: %s' % pol)
if interp != 'none':
for fi in xrange(local_vis.shape[1]):
for bi in xrange(local_vis.shape[2]):
# interpolate for local_vis
true_inds = np.where(local_vis_mask[:, fi, bi])[0] # masked inds
if len(true_inds) > 0:
false_inds = np.where(~local_vis_mask[:, fi, bi])[0] # un-masked inds
if len(false_inds) > 0.1 * local_vis.shape[0]:
# nearest interpolate for local_vis
if interp in ('linear', 'nearest'):
itp_real = interp1d(false_inds, local_vis[false_inds, fi, bi].real, kind=interp, fill_value='extrapolate', assume_sorted=True)
itp_imag = interp1d(false_inds, local_vis[false_inds, fi, bi].imag, kind=interp, fill_value='extrapolate', assume_sorted=True)
elif interp == 'rbf':
itp_real = Rbf(false_inds, local_vis[false_inds, fi, bi].real, smooth=10)
itp_imag = Rbf(false_inds, local_vis[false_inds, fi, bi].imag, smooth=10)
else:
raise ValueError('Unknown interpolation method: %s' % interp)
local_vis[true_inds, fi, bi] = itp_real(true_inds) + 1.0J * itp_imag(true_inds) # the interpolated vis
else:
local_vis[:, fi, bi] = 0 # TODO: may need to take special care
# average data
nt = ts['sec1970'].shape[0]
phi_size = 2*tel.mmax + 1
# phi = np.zeros((phi_size,), dtype=ts['ra_dec'].dtype)
phi = np.linspace(0, 2*np.pi, phi_size, endpoint=False)
vis = np.zeros((phi_size,)+local_vis.shape[1:], dtype=local_vis.dtype)
if time_avg == 'avg':
nt_m = float(nt) / phi_size
# roll data to have phi=0 near the first
roll_len = np.int(np.around(0.5*nt_m))
local_vis[:] = np.roll(local_vis[:], roll_len, axis=0)
if interp == 'none':
local_vis_mask[:] = np.roll(local_vis_mask[:], roll_len, axis=0)
# ts['ra_dec'][:] = np.roll(ts['ra_dec'][:], roll_len, axis=0)
repeat_inds = np.repeat(np.arange(nt), phi_size)
num, start, end = mpiutil.split_m(nt*phi_size, phi_size)
# average over time
for idx in xrange(phi_size):
inds, weight = unique(repeat_inds[start[idx]:end[idx]], return_counts=True)
if interp == 'none':
vis[idx] = average(np.ma.array(local_vis[inds], mask=local_vis_mask[inds]), axis=0, weights=weight) # time mean
else:
vis[idx] = average(local_vis[inds], axis=0, weights=weight) # time mean
# phi[idx] = np.average(ts['ra_dec'][:, 0][inds], axis=0, weights=weight)
elif time_avg == 'fft':
if interp == 'none':
raise ValueError('Can not do fft average without first interpolation')
Vm = np.fft.fftshift(np.fft.fft(local_vis, axis=0), axes=0)
vis[:] = np.fft.ifft(np.fft.ifftshift(Vm[nt/2-tel.mmax:nt/2+tel.mmax+1], axes=0), axis=0) / (1.0 * nt / phi_size)
# for fi in xrange(vis.shape[1]):
# for bi in xrange(vis.shape[2]):
# # plot local_vis and vis
# import matplotlib
# matplotlib.use('Agg')
# import matplotlib.pyplot as plt
# phi0 = np.linspace(0, 2*np.pi, nt, endpoint=False)
# phi1 = np.linspace(0, 2*np.pi, phi_size, endpoint=False)
# plt.figure()
# plt.subplot(211)
# plt.plot(phi0, local_vis[:, fi, bi].real, label='v0.real')
# plt.plot(phi1, vis[:, fi, bi].real, label='v1.real')
# plt.legend()
# plt.subplot(212)
# plt.plot(phi0, local_vis[:, fi, bi].imag, label='v0.imag')
# plt.plot(phi1, vis[:, fi, bi].imag, label='v1.imag')
# plt.legend()
# plt.savefig('vis_fft/vis_%d_%d.png' % (fi, bi))
# plt.close()
else:
raise ValueError('Unknown time_avg: %s' % time_avg)
del local_vis
del local_vis_mask
# mask daytime data
if mask_daytime:
day_or_night = np.where(ts['local_hour'][:]>=mask_time_range[0] & ts['local_hour'][:]<=mask_time_range[1], True, False)
day_inds = np.where(np.repeat(day_or_night, phi_size).reshape(nt, phi_size).astype(np.int).sum(axis=1).astype(bool))[0]
vis[day_inds] = 0
del ts # no longer need ts
# redistribute vis to time axis
vis = mpiarray.MPIArray.wrap(vis, axis=2).redistribute(0).local_array
allpairs = tel.allpairs
redundancy = tel.redundancy
nrd = len(redundancy)
# reorder bls according to allpairs
vis_tmp = np.zeros_like(vis)
for ind, (a1, a2) in enumerate(allpairs):
try:
b_ind = bls.index((feeds[a1], feeds[a2]))
vis_tmp[:, :, ind] = vis[:, :, b_ind]
except ValueError:
b_ind = bls.index((feeds[a2], feeds[a1]))
vis_tmp[:, :, ind] = vis[:, :, b_ind].conj()
del vis
# average over redundancy
vis_stream = np.zeros(vis_tmp.shape[:-1]+(nrd,), dtype=vis_tmp.dtype)
red_bin = np.cumsum(np.insert(redundancy, 0, 0)) # redundancy bin
# average over redundancy
for ind in xrange(nrd):
vis_stream[:, :, ind] = np.sum(vis_tmp[:, :, red_bin[ind]:red_bin[ind+1]], axis=2) / redundancy[ind]
del vis_tmp
# beamtransfer
bt = beamtransfer.BeamTransfer(beam_dir, tel, noise_weight, True)
if not use_existed_beam:
bt.generate()
if tel is None:
tel = bt.telescope
if simulate:
ndays = 733
tstream = timestream.simulate(bt, ts_dir, ts_name, input_maps, ndays, add_noise=add_noise)
else:
# timestream and map-making
tstream = timestream.Timestream(ts_dir, ts_name, bt, no_m_zero)
parent_path = os.path.dirname(tstream._fdir(0))
if os.path.exists(parent_path + '/COMPLETED'):
if mpiutil.rank0:
print 'Use existed timestream_f files in %s' % parent_path
else:
for fi in mpiutil.mpirange(nfreq):
# Make directory if required
if not os.path.exists(tstream._fdir(fi)):
os.makedirs(tstream._fdir(fi))
# create memh5 object and write data to temporary file
vis_h5 = memh5.MemGroup(distributed=True)
vis_h5.create_dataset('/timestream', data=mpiarray.MPIArray.wrap(vis_stream, axis=0))
tmp_file = parent_path +'/vis_stream_temp.hdf5'
vis_h5.to_hdf5(tmp_file, hints=False)
del vis_h5
# re-organize data as need for tstream
# make load even among nodes
for fi in mpiutil.mpirange(nfreq, method='rand'):
# read the needed data from the temporary file
with h5py.File(tmp_file, 'r') as f:
vis_fi = f['/timestream'][:, fi, :]
# Write file contents
with h5py.File(tstream._ffile(fi), 'w') as f:
# Timestream data
# allocate space for vis_stream
shp = (nrd, phi_size)
f.create_dataset('/timestream', data=vis_fi.T)
f.create_dataset('/phi', data=phi)
# Telescope layout data
f.create_dataset('/feedmap', data=tel.feedmap)
f.create_dataset('/feedconj', data=tel.feedconj)
f.create_dataset('/feedmask', data=tel.feedmask)
f.create_dataset('/uniquepairs', data=tel.uniquepairs)
f.create_dataset('/baselines', data=tel.baselines)
# Telescope frequencies
f.create_dataset('/frequencies', data=freqs)
# Write metadata
f.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
f.attrs['ntime'] = phi_size
mpiutil.barrier()
# remove temp file
if mpiutil.rank0:
os.remove(tmp_file)
# mark all frequencies tstream files are saved correctly
open(parent_path + '/COMPLETED', 'a').close()
tstream.generate_mmodes()
nside = hputil.nside_for_lmax(tel.lmax, accuracy_boost=tel.accuracy_boost)
if dirty_map:
tstream.mapmake_full(nside, 'map_full_dirty.hdf5', nbin, dirty=True, method=method, normalize=normalize, threshold=threshold)
else:
tstream.mapmake_full(nside, 'map_full.hdf5', nbin, dirty=False, method=method, normalize=normalize, threshold=threshold, eps=eps, correct_order=correct_order, prior_map_file=prior_map)
# ts.add_history(self.history)
return tstream
def process(self, ts):
assert isinstance(ts, Timestream), '%s only works for Timestream object' % self.__class__.__name__
mask_daytime = self.params['mask_daytime']
mask_time_range = self.params['mask_time_range']
beam_theta_range = self.params['beam_theta_range']
tsys = self.params['tsys']
accuracy_boost = self.params['accuracy_boost']
l_boost = self.params['l_boost']
bl_range = self.params['bl_range']
auto_correlations = self.params['auto_correlations']
pol = self.params['pol']
beam_dir = output_path(self.params['beam_dir'])
gen_inv = self.params['gen_invbeam']
noise_weight = self.params['noise_weight']
ts_dir = output_path(self.params['ts_dir'])
ts_name = self.params['ts_name']
simulate = self.params['simulate']
input_maps = self.params['input_maps']
add_noise = self.params['add_noise']
dirty_map = self.params['dirty_map']
nbin = self.params['nbin']
method = self.params['method']
normalize = self.params['normalize']
threshold = self.params['threshold']
ts.redistribute('frequency')
lat = ts.attrs['sitelat']
# lon = ts.attrs['sitelon']
lon = 0.0
# lon = np.degrees(ts['ra_dec'][0, 0]) # the first ra
local_origin = False
freq = ts.freq
freqs = ts.freq.data.to_numpy_array(root=None) # MHz
band_width = ts.attrs['freqstep'] # MHz
ndays = 1
feeds = ts['feedno'][:]
az, alt = ts['az_alt'][0]
az = np.degrees(az)
alt = np.degrees(alt)
pointing = [az, alt, 0.0]
feedpos = ts['feedpos'][:]
if ts.is_dish:
dish_width = ts.attrs['dishdiam']
tel = tl_dish.TlUnpolarisedDishArray(lat, lon, freqs, band_width, tsys, ndays, accuracy_boost, l_boost, bl_range, auto_correlations, local_origin, dish_width, feedpos, pointing)
elif ts.is_cylinder:
factor = 1.2 # suppose an illumination efficiency, keep same with that in timestream_common
cyl_width = factor * ts.attrs['cywid']
tel = tl_cylinder.TlUnpolarisedCylinder(lat, lon, freqs, band_width, tsys, ndays, accuracy_boost, l_boost, bl_range, auto_correlations, local_origin, cyl_width, feedpos)
else:
raise RuntimeError('Unknown array type %s' % ts.attrs['telescope'])
if not simulate:
# mask daytime data
if mask_daytime:
day_inds = np.where(np.logical_and(ts['local_hour'][:]>=mask_time_range[0], ts['local_hour'][:]<=mask_time_range[1]))[0]
ts.local_vis_mask[day_inds] = True # do not change vis directly
# average data
nt = ts['sec1970'].shape[0]
phi_size = 2*tel.mmax + 1
nt_m = float(nt) / phi_size
# roll data to have phi=0 near the first
roll_len = np.int(np.around(0.5*nt_m))
ts.local_vis[:] = np.roll(ts.local_vis[:], roll_len, axis=0)
ts.local_vis_mask[:] = np.roll(ts.local_vis_mask[:], roll_len, axis=0)
ts['ra_dec'][:] = np.roll(ts['ra_dec'][:], roll_len, axis=0)
repeat_inds = np.repeat(np.arange(nt), phi_size)
num, start, end = mpiutil.split_m(nt*phi_size, phi_size)
# phi = np.zeros((phi_size,), dtype=ts['ra_dec'].dtype)
phi = np.linspace(0, 2*np.pi, phi_size, endpoint=False)
vis = np.zeros((phi_size,)+ts.local_vis.shape[1:], dtype=ts.vis.dtype)
# average over time
for idx in xrange(phi_size):
inds, weight = unique(repeat_inds[start[idx]:end[idx]], return_counts=True)
vis[idx] = average(np.ma.array(ts.local_vis[inds], mask=ts.local_vis_mask[inds]), axis=0, weights=weight) # time mean
# phi[idx] = np.average(ts['ra_dec'][:, 0][inds], axis=0, weights=weight)
if pol == 'xx':
vis = vis[:, :, 0, :]
elif pol == 'yy':
vis = vis[:, :, 1, :]
elif pol == 'I':
vis = 0.5 * (vis[:, :, 0, :] + vis[:, :, 1, :])
elif pol == 'all':
vis = np.sum(vis, axis=2) # sum over all pol
else:
raise ValueError('Invalid pol: %s' % pol)
allpairs = tel.allpairs
redundancy = tel.redundancy
# reorder bls according to allpairs
vis_tmp = np.zeros_like(vis)
bls = [ tuple(bl) for bl in ts['blorder'][:] ]
for ind, (a1, a2) in enumerate(allpairs):
try:
b_ind = bls.index((feeds[a1], feeds[a2]))
vis_tmp[:, :, ind] = vis[:, :, b_ind]
except ValueError:
b_ind = bls.index((feeds[a2], feeds[a1]))
vis_tmp[:, :, ind] = vis[:, :, b_ind].conj()
# average over redundancy
vis_stream = np.zeros(vis.shape[:-1]+(len(redundancy),), dtype=vis_tmp.dtype)
red_bin = np.cumsum(np.insert(redundancy, 0, 0)) # redundancy bin
# average over redundancy
for ind in xrange(len(redundancy)):
vis_stream[:, :, ind] = np.sum(vis_tmp[:, :, red_bin[ind]:red_bin[ind+1]], axis=2) / redundancy[ind]
del vis
del vis_tmp
# beamtransfer
bt = beamtransfer.BeamTransfer(beam_dir, tel, noise_weight, True)
bt.generate()
if simulate:
ndays = 733
print ndays
tstream = timestream.simulate(bt, ts_dir, ts_name, input_maps, ndays, add_noise=add_noise)
else:
# timestream and map-making
tstream = timestream.Timestream(ts_dir, ts_name, bt)
for lfi, fi in freq.data.enumerate(axis=0):
# Make directory if required
if not os.path.exists(tstream._fdir(fi)):
os.makedirs(tstream._fdir(fi))
# Write file contents
with h5py.File(tstream._ffile(fi), 'w') as f:
# Timestream data
f.create_dataset('/timestream', data=vis_stream[:, lfi].T)
f.create_dataset('/phi', data=phi)
# Telescope layout data
f.create_dataset('/feedmap', data=tel.feedmap)
f.create_dataset('/feedconj', data=tel.feedconj)
f.create_dataset('/feedmask', data=tel.feedmask)
f.create_dataset('/uniquepairs', data=tel.uniquepairs)
f.create_dataset('/baselines', data=tel.baselines)
# Telescope frequencies
f.create_dataset('/frequencies', data=freqs)
# Write metadata
f.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
f.attrs['ntime'] = phi_size
mpiutil.barrier()
tstream.generate_mmodes()
nside = hputil.nside_for_lmax(tel.lmax, accuracy_boost=tel.accuracy_boost)
if dirty_map:
tstream.mapmake_full(nside, 'map_full_dirty.hdf5', nbin, dirty=True, method=method, normalize=normalize, threshold=threshold)
else:
tstream.mapmake_full(nside, 'map_full.hdf5', nbin, dirty=False, method=method, normalize=normalize, threshold=threshold)
# ts.add_history(self.history)
return tstream
def process(self, ts):
calibrator = self.params['calibrator']
catalog = self.params['catalog']
file_prefix = self.params['file_name']
plot_closure = self.params['plot_closure']
fig_prefix = self.params['fig_name']
bins = self.params['bins']
gauss_fit = self.params['gauss_fit']
tag_output_iter = self.params['tag_output_iter']
ts.redistribute('frequency')
nfreq = len(ts.local_freq[:]) # local nfreq
feedno = ts['feedno'][:].tolist()
pol = ts['pol'][:].tolist()
bl = ts.local_bl[:] # local bls
bls = [ tuple(b) for b in bl ]
# calibrator
srclist, cutoff, catalogs = a.scripting.parse_srcs(calibrator, catalog)
cat = a.src.get_catalog(srclist, cutoff, catalogs)
assert(len(cat) == 1), 'Allow only one calibrator'
s = cat.values()[0]
if mpiutil.rank0:
print 'Calibrating for source %s with' % calibrator,
print 'strength', s._jys, 'Jy',
print 'measured at', s.mfreq, 'GHz',
print 'with index', s.index
ra = ts['ra_dec'][:, 0]
ra = np.unwrap(ra)
ra = ra[np.logical_not(ts['ns_on'][:])] # get only ns_off values
abs_diff = np.abs(np.diff(s._ra - ra))
ind1 = np.argmin(abs_diff)
print 'ind1:', ind1
for pi in [ pol.index('xx'), pol.index('yy') ]: # xx and yy
if nfreq > 0: # skip empty processes
# find the ind that not be all masked
for i in xrange(20):
if not ts.local_vis_mask[ind1+i, :, pi].all():
ind = ind1 + i
break
if not ts.local_vis_mask[ind1-i, :, pi].all():
ind = ind1 - i
break
else:
raise RuntimeError('vis is masked during this period for pol %s' % pol[pi])
print 'ind:', ind
for fi in xrange(nfreq):
gfi = fi + ts.freq.local_offset[0] # global freq index
vis = ts.local_vis[ind, fi, pi, :] # only I
vis_mask = ts.local_vis_mask[ind, fi, pi, :] # only I
closure = []
cache = dict()
for i, j, k in itertools.combinations(feedno, 3):
bi, vij = mk_cache(cache, (i, j), bls, vis)
if vis_mask[bi]:
continue
bi, vjk = mk_cache(cache, (j, k), bls, vis)
if vis_mask[bi]:
continue
bi, vki = mk_cache(cache, (k, i), bls, vis)
if vis_mask[bi]:
continue
# closure.append(np.angle(vij, True) + np.angle(vjk, True) + np.angle(vki, True)) # in degree, have 360 deg problem
c = lambda x: np.complex128(x) # complex128 to avoid overflow in the product
ang = np.angle(c(vij) * c(vjk) * c(vki), True) # in degree
closure.append(ang) # in degree
# save closure phase to file
file_name = '%s_%d_%s.hdf5' % (file_prefix, gfi, pol[pi])
if tag_output_iter:
file_name = output_path(file_name, iteration=self.iteration)
else:
file_name = output_path(file_name)
with h5py.File(file_name, 'w') as f:
f.create_dataset('closure_phase', data=np.array(closure))
if plot_closure:
# plot all closure phase
plt.figure()
plt.plot(closure, 'o')
fig_name = '%s_all_%d_%s.png' % (fig_prefix, gfi, pol[pi])
if tag_output_iter:
fig_name = output_path(fig_name, iteration=self.iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
# plot histogram of closure phase
# histogram
plt.figure()
data = plt.hist(closure, bins=bins)
plt.xlabel('Closure phase / degree')
if gauss_fit:
# Generate data from bins as a set of points
x = [0.5 * (data[1][i] + data[1][i+1]) for i in xrange(len(data[1])-1)]
y = data[0]
popt, pcov = optimize.curve_fit(f, x, y)
A, b, c = popt
xmax = max(abs(x[0]), abs(x[-1]))
x_fit = np.linspace(-xmax, xmax, bins)
y_fit = f(x_fit, *popt)
lable = r'$a \, \exp{(- \frac{(x - \mu)^2} {2 \sigma^2})}$' + '\n\n' + r'$a = %f$' % A + '\n' + r'$\mu = %f$' % b + '\n' + r'$\sigma = %f$' % np.abs(c)
plt.plot(x_fit, y_fit, lw=2, color="r", label=lable)
plt.xlim(-xmax, xmax)
plt.legend()
fig_name = '%s_hist_%d_%s.png' % (fig_prefix, gfi, pol[pi])
if tag_output_iter:
fig_name = output_path(fig_name, iteration=self.iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
mpiutil.barrier()
ts.add_history(self.history)
return ts
def process(self, ts):
calibrator = self.params['calibrator']
catalog = self.params['catalog']
span = self.params['span']
save_gain = self.params['save_gain']
gain_file = self.params['gain_file']
ts.redistribute('frequency')
lfreq = ts.local_freq[:] # local freq
feedno = ts['feedno'][:].tolist()
pol = ts['pol'][:].tolist()
bl = ts.bl[:]
bls = [ tuple(b) for b in bl ]
# # antpointing = np.radians(ts['antpointing'][-1, :, :]) # radians
# transitsource = ts['transitsource'][:]
# transit_time = transitsource[-1, 0] # second, sec1970
# int_time = ts.attrs['inttime'] # second
# calibrator
srclist, cutoff, catalogs = a.scripting.parse_srcs(calibrator, catalog)
cat = a.src.get_catalog(srclist, cutoff, catalogs)
assert(len(cat) == 1), 'Allow only one calibrator'
s = cat.values()[0]
if mpiutil.rank0:
print 'Calibrating for source %s with' % calibrator,
print 'strength', s._jys, 'Jy',
print 'measured at', s.mfreq, 'GHz',
print 'with index', s.index
# get transit time of calibrator
# array
aa = ts.array
aa.set_jultime(ts['jul_date'][0]) # the first obs time point
next_transit = aa.next_transit(s)
transit_time = a.phs.ephem2juldate(next_transit) # Julian date
if transit_time > ts['jul_date'][-1]:
local_next_transit = ephem.Date(next_transit + 8.0 * ephem.hour)
raise RuntimeError('Data does not contain local transit time %s of source %s' % (local_next_transit, calibrator))
# the first transit index
transit_inds = [ np.searchsorted(ts['jul_date'][:], transit_time) ]
# find all other transit indices
aa.set_jultime(ts['jul_date'][0] + 1.0)
transit_time = a.phs.ephem2juldate(aa.next_transit(s)) # Julian date
cnt = 2
while(transit_time <= ts['jul_date'][-1]):
transit_inds.append(np.searchsorted(ts['jul_date'][:], transit_time))
aa.set_jultime(ts['jul_date'][0] + 1.0*cnt)
transit_time = a.phs.ephem2juldate(aa.next_transit(s)) # Julian date
cnt += 1
print transit_inds
### now only use the first transit point to do the cal
### may need to improve in the future
transit_ind = transit_inds[0]
int_time = ts.attrs['inttime'] # second
start_ind = transit_ind - np.int(span / int_time)
end_ind = transit_ind + np.int(span / int_time)
nt = end_ind - start_ind
nfeed = len(feedno)
eigval = np.empty((nt, nfeed, 2, len(lfreq)), dtype=np.float64)
eigval[:] = np.nan
gain = np.empty((nt, nfeed, 2, len(lfreq)), dtype=np.complex128)
gain[:] = complex(np.nan, np.nan)
# construct visibility matrix for a single time, pol, freq
Vmat = np.zeros((nfeed, nfeed), dtype=ts.main_data.dtype)
for ind, ti in enumerate(range(start_ind, end_ind)):
# when noise on, just pass
if ts['ns_on'][ti]:
continue
aa.set_jultime(ts['jul_date'][ti])
s.compute(aa)
# get fluxes vs. freq of the calibrator
Sc = s.get_jys()
# get the topocentric coordinate of the calibrator at the current time
s_top = s.get_crds('top', ncrd=3)
aa.sim_cache(cat.get_crds('eq', ncrd=3)) # for compute bm_response and sim
for pi in [pol.index('xx'), pol.index('yy')]: # xx, yy
aa.set_active_pol(pol[pi])
for fi, freq in enumerate(lfreq): # mpi among freq
for i, ai in enumerate(feedno):
for j, aj in enumerate(feedno):
# uij = aa.gen_uvw(i, j, src='z').squeeze() # (rj - ri)/lambda
uij = aa.gen_uvw(i, j, src='z')[:, 0, :] # (rj - ri)/lambda
# bmij = aa.bm_response(i, j).squeeze() # will get error for only one local freq
# import pdb
# pdb.set_trace()
bmij = aa.bm_response(i, j).reshape(-1)
try:
bi = bls.index((ai, aj))
# Vmat[i, j] = ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(-2.0J * np.pi * np.dot(s_top, uij[:, fi]))) # xx, yy
Vmat[i, j] = ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(2.0J * np.pi * np.dot(s_top, uij[:, fi]))) # xx, yy
except ValueError:
bi = bls.index((aj, ai))
# Vmat[i, j] = np.conj(ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(-2.0J * np.pi * np.dot(s_top, uij[:, fi])))) # xx, yy
Vmat[i, j] = np.conj(ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(2.0J * np.pi * np.dot(s_top, uij[:, fi])))) # xx, yy
# Eigen decomposition
Vmat = np.where(np.isfinite(Vmat), Vmat, 0)
e, U = eigh(Vmat)
eigval[ind, :, pi, fi] = e[::-1] # descending order
# max eigen-val
lbd = e[-1] # lambda
# the gain vector for this freq
gvec = np.sqrt(lbd) * U[:, -1] # only eigen-vector corresponding to the maximum eigen-val
gain[ind, :, pi, fi] = gvec
# apply gain to vis
# get the time mean gain
tgain = np.ma.mean(np.ma.masked_invalid(gain), axis=0) # time mean
tgain = mpiutil.gather_array(tgain, axis=-1, root=None)
ts.redistribute('baseline')
ts.pol_and_bl_data_operate(cal, tgain=tgain)
# save gain if required:
if save_gain:
gain_file = output_path(gain_file)
eigval = mpiarray.MPIArray.wrap(eigval, axis=3)
gain = mpiarray.MPIArray.wrap(gain, axis=3)
mem_gain = memh5.MemGroup(distributed=True)
mem_gain.create_dataset('eigval', data=eigval)
mem_gain.create_dataset('gain', data=gain)
# add attris
mem_gain.attrs['jul_data'] = ts['jul_date'][start_ind:end_ind]
mem_gain.attrs['feed'] = np.array(feedno)
mem_gain.attrs['pol'] = np.array(['xx', 'yy'])
mem_gain.attrs['freq'] = ts.freq[:] # freq should be common
# save to file
mem_gain.to_hdf5(gain_file, hints=False)
ts.add_history(self.history)
return ts
def process(self, ts):
mask_daytime = self.params["mask_daytime"]
mask_time_range = self.params["mask_time_range"]
beam_theta_range = self.params["beam_theta_range"]
tsys = self.params["tsys"]
accuracy_boost = self.params["accuracy_boost"]
l_boost = self.params["l_boost"]
bl_range = self.params["bl_range"]
auto_correlations = self.params["auto_correlations"]
pol = self.params["pol"]
beam_dir = output_path(self.params["beam_dir"])
gen_inv = self.params["gen_invbeam"]
noise_weight = self.params["noise_weight"]
ts_dir = output_path(self.params["ts_dir"])
ts_name = self.params["ts_name"]
simulate = self.params["simulate"]
input_maps = self.params["input_maps"]
add_noise = self.params["add_noise"]
ts.redistribute("frequency")
lat = ts.attrs["sitelat"]
# lon = ts.attrs['sitelon']
lon = 0.0
# lon = np.degrees(ts['ra_dec'][0, 0]) # the first ra
freqs = ts.freq.data.to_numpy_array(root=None)
ndays = 1
feeds = ts["feedno"][:]
az, alt = ts["az_alt"][0]
az = np.degrees(az)
alt = np.degrees(alt)
pointing = [az, alt, 0.0]
feedpos = ts["feedpos"][:]
if ts.is_dish:
dish_width = ts.attrs["dishdiam"]
tel = tl_dish.TlUnpolarisedDishArray(
lat,
lon,
freqs,
beam_theta_range,
tsys,
ndays,
accuracy_boost,
l_boost,
bl_range,
auto_correlations,
dish_width,
feedpos,
pointing,
)
elif ts.is_cylinder:
cyl_width = ts.attrs["cywid"]
tel = tl_cylinder.TlUnpolarisedCylinder(
lat,
lon,
freqs,
beam_theta_range,
tsys,
ndays,
accuracy_boost,
l_boost,
bl_range,
auto_correlations,
cyl_width,
feedpos,
)
else:
raise RuntimeError("Unknown array type %s" % ts.attrs["telescope"])
# import matplotlib
# matplotlib.use('Agg')
# import matplotlib.pyplot as plt
# plt.figure()
# plt.plot(ts['ra_dec'][:])
# # plt.plot(ts['az_alt'][:])
# plt.savefig('ra_dec1.png')
if not simulate:
# mask daytime data
if mask_daytime:
day_inds = np.where(
np.logical_and(ts["local_hour"][:] >= mask_time_range[0], ts["local_hour"][:] <= mask_time_range[1])
)[0]
ts.local_vis_mask[day_inds] = True # do not change vis directly
# average data
nt = ts["sec1970"].shape[0]
phi_size = tel.phi_size
nt_m = float(nt) / phi_size
# roll data to have phi=0 near the first
roll_len = np.int(np.around(0.5 * nt_m))
ts.local_vis[:] = np.roll(ts.local_vis[:], roll_len, axis=0)
ts.local_vis_mask[:] = np.roll(ts.local_vis_mask[:], roll_len, axis=0)
ts["ra_dec"][:] = np.roll(ts["ra_dec"][:], roll_len, axis=0)
repeat_inds = np.repeat(np.arange(nt), phi_size)
num, start, end = mpiutil.split_m(nt * phi_size, phi_size)
# phi = np.zeros((phi_size,), dtype=ts['ra_dec'].dtype)
phi = np.linspace(0, 2 * np.pi, phi_size, endpoint=False)
vis = np.zeros((phi_size,) + ts.local_vis.shape[1:], dtype=ts.vis.dtype)
# average over time
for idx in xrange(phi_size):
inds, weight = unique(repeat_inds[start[idx] : end[idx]], return_counts=True)
vis[idx] = average(
np.ma.array(ts.local_vis[inds], mask=ts.local_vis_mask[inds]), axis=0, weights=weight
) # time mean
# phi[idx] = np.average(ts['ra_dec'][:, 0][inds], axis=0, weights=weight)
if pol == "xx":
vis = vis[:, :, 0, :]
elif pol == "yy":
vis = vis[:, :, 1, :]
elif pol == "I":
vis = 0.5 * (vis[:, :, 0, :] + vis[:, :, 1, :])
elif pol == "all":
vis = np.sum(vis, axis=2) # sum over all pol
else:
raise ValueError("Invalid pol: %s" % pol)
allpairs = tel.allpairs
redundancy = tel.redundancy
# reorder bls according to allpairs
vis_tmp = np.zeros_like(vis)
bls = [tuple(bl) for bl in ts["blorder"][:]]
for ind, (a1, a2) in enumerate(allpairs):
try:
b_ind = bls.index((feeds[a1], feeds[a2]))
vis_tmp[:, :, ind] = vis[:, :, b_ind]
except ValueError:
b_ind = bls.index((feeds[a2], feeds[a1]))
vis_tmp[:, :, ind] = vis[:, :, b_ind].conj()
# average over redundancy
vis_stream = np.zeros(vis.shape[:-1] + (len(redundancy),), dtype=vis_tmp.dtype)
red_bin = np.cumsum(np.insert(redundancy, 0, 0)) # redundancy bin
# average over redundancy
for ind in xrange(len(redundancy)):
vis_stream[:, :, ind] = np.sum(vis_tmp[:, :, red_bin[ind] : red_bin[ind + 1]], axis=2) / redundancy[ind]
del vis
del vis_tmp
vis_stream = mpiarray.MPIArray.wrap(vis_stream, axis=1)
vis_h5 = memh5.MemGroup(distributed=True)
vis_h5.create_dataset("/timestream", data=vis_stream)
vis_h5.create_dataset("/phi", data=phi)
# Telescope layout data
vis_h5.create_dataset("/feedmap", data=tel.feedmap)
vis_h5.create_dataset("/feedconj", data=tel.feedconj)
vis_h5.create_dataset("/feedmask", data=tel.feedmask)
vis_h5.create_dataset("/uniquepairs", data=tel.uniquepairs)
vis_h5.create_dataset("/baselines", data=tel.baselines)
# Telescope frequencies
vis_h5.create_dataset("/frequencies", data=freqs)
# Write metadata
# vis_h5.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
vis_h5.attrs["ntime"] = phi_size
# beamtransfer
bt = beamtransfer.BeamTransfer(beam_dir, tel, gen_inv, noise_weight)
bt.generate()
if simulate:
ndays = 733
print ndays
ts = timestream.simulate(bt, ts_dir, ts_name, input_maps, ndays, add_noise=add_noise)
else:
# timestream and map-making
ts = timestream.Timestream(ts_dir, ts_name, bt)
# Make directory if required
try:
os.makedirs(ts._tsdir)
except OSError:
# directory exists
pass
vis_h5.to_hdf5(ts._tsfile)
# ts.generate_mmodes(vis_stream.to_numpy_array(root=None))
ts.generate_mmodes()
nside = hputil.nside_for_lmax(tel.lmax, accuracy_boost=tel.accuracy_boost)
ts.mapmake_full(nside, "full")
# ts.add_history(self.history)
return ts
def _init_output_files(self):
self.output_files = output_path(format_list(
self.params['output_files']),
mkdir=False)
def flag(self, vis, vis_mask, li, gi, bl, ts, **kwargs):
"""Function that does the actual flag."""
freq_window = self.params['freq_window']
time_window = self.params['time_window']
freq_sigma = self.params['freq_sigma']
time_sigma = self.params['time_sigma']
plot_fit = self.params['plot_fit']
freq_fig_prefix = self.params['freq_fig_name']
time_fig_prefix = self.params['time_fig_name']
tag_output_iter = self.params['tag_output_iter']
iteration = self.iteration
freq_flag = kwargs.get('freq_flag')
time_flag = kwargs.get('time_flag')
time = ts.time[:]
nt = len(time)
freq = ts.freq[:]
nfreq = len(freq)
if isinstance(ts, Timestream): # for Timestream
pol = bl[0]
bl = tuple(bl[1])
elif isinstance(ts, RawTimestream): # for RawTimestream
pol = None
bl = tuple(bl)
else:
raise ValueError('Need either a RawTimestream or Timestream')
if freq_flag:
# time integration
tm_vis = np.ma.mean(np.ma.array(vis, mask=vis_mask), axis=0)
abs_vis = np.abs(tm_vis)
if np.ma.count_masked(tm_vis) > 0: # has masked value
abs_vis_valid = abs_vis[~abs_vis.mask]
inds_valid = np.arange(nfreq)[~abs_vis.mask]
itp = InterpolatedUnivariateSpline(inds_valid, abs_vis_valid)
abs_vis_itp = itp(np.arange(nfreq))
abs_vis1 = abs_vis_itp.copy()
else:
abs_vis1 = abs_vis.copy()
for cnt in xrange(10):
if cnt != 0:
abs_vis1[inds] = smooth[inds]
smooth = savitzky_golay(abs_vis1, freq_window, 3)
# flage RFI
diff = abs_vis1 - smooth
mean = np.mean(diff)
std = np.std(diff)
inds = np.where(np.abs(diff - mean) > freq_sigma * std)[0]
if len(inds) == 0:
break
diff = abs_vis - smooth
mean = np.mean(diff)
std = np.std(diff)
inds = np.where(
np.abs(diff - mean) > freq_sigma * std)[0] # masked inds
vis_mask[:, inds] = True # set mask
if plot_fit:
plt.figure()
plt.plot(freq, abs_vis, label='data')
plt.plot(freq[inds], abs_vis[inds], 'ro', label='flag')
plt.plot(freq, smooth, label='smooth')
plt.xlabel(r'$\nu$ / MHz')
plt.legend(loc='best')
if pol is None:
fig_name = '%s_%d_%d.png' % (freq_fig_prefix, bl[0], bl[1])
else:
fig_name = '%s_%d_%d_%s.png' % (freq_fig_prefix, bl[0],
bl[1], pol)
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
if time_flag:
# freq integration
fm_vis = np.ma.mean(np.ma.array(vis, mask=vis_mask), axis=1)
abs_vis = np.abs(fm_vis)
if np.ma.count_masked(fm_vis) > 0: # has masked value
abs_vis_valid = abs_vis[~abs_vis.mask]
inds_valid = np.arange(nt)[~abs_vis.mask]
itp = InterpolatedUnivariateSpline(inds_valid, abs_vis_valid)
abs_vis_itp = itp(np.arange(nt))
abs_vis1 = abs_vis_itp.copy()
else:
abs_vis1 = abs_vis.copy()
for cnt in xrange(10):
if cnt != 0:
abs_vis1[inds] = smooth[inds]
smooth = savitzky_golay(abs_vis1, time_window, 3)
# flage RFI
diff = abs_vis1 - smooth
mean = np.mean(diff)
std = np.std(diff)
inds = np.where(np.abs(diff - mean) > time_sigma * std)[0]
if len(inds) == 0:
break
diff = abs_vis - smooth
mean = np.mean(diff)
std = np.std(diff)
inds = np.where(
np.abs(diff - mean) > time_sigma * std)[0] # masked inds
# Addtional threshold
# inds1 = np.where(np.abs(diff[inds]) > 1.0e-2*np.abs(smooth[inds]))[0]
# inds = inds[inds1]
vis_mask[inds] = True # set mask
if plot_fit:
plt.figure()
plt.plot(time, abs_vis, label='data')
plt.plot(time[inds], abs_vis[inds], 'ro', label='flag')
plt.plot(time, smooth, label='smooth')
plt.xlabel(r'$t$ / Julian Date')
plt.legend(loc='best')
if pol is None:
fig_name = '%s_%d_%d.png' % (time_fig_prefix, bl[0], bl[1])
else:
fig_name = '%s_%d_%d_%s.png' % (time_fig_prefix, bl[0],
bl[1], pol)
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
def plot(self, vis, vis_mask, li, gi, bl, ts, **kwargs):
"""Function that does the actual plot work."""
bl_incl = self.params['bl_incl']
bl_excl = self.params['bl_excl']
flag_mask = self.params['flag_mask']
flag_ns = self.params['flag_ns']
interpolate_ns = self.params['interpolate_ns']
y_axis = self.params['y_axis']
plot_abs = self.params['plot_abs']
fig_prefix = self.params['fig_name']
tag_output_iter = self.params['tag_output_iter']
iteration = self.iteration
if isinstance(ts, Timestream): # for Timestream
pol = bl[0]
bl = tuple(bl[1])
elif isinstance(ts, RawTimestream): # for RawTimestream
pol = None
bl = tuple(bl)
else:
raise ValueError('Need either a RawTimestream or Timestream')
if bl_incl != 'all':
bl1 = set(bl)
bl_incl = [ {f1, f2} for (f1, f2) in bl_incl ]
bl_excl = [ {f1, f2} for (f1, f2) in bl_excl ]
if (not bl1 in bl_incl) or (bl1 in bl_excl):
return vis, vis_mask
if flag_mask:
vis1 = np.ma.array(vis, mask=vis_mask)
elif flag_ns:
vis1 = vis.copy()
on = np.where(ts['ns_on'][:])[0]
if not interpolate_ns:
vis1[on] = complex(np.nan, np.nan)
else:
off = np.where(np.logical_not(ts['ns_on'][:]))[0]
for fi in xrange(vis1.shape[1]):
itp_real = InterpolatedUnivariateSpline(off, vis1[off, fi].real)
itp_imag= InterpolatedUnivariateSpline(off, vis1[off, fi].imag)
vis1[on, fi] = itp_real(on) + 1.0J * itp_imag(on)
else:
vis1 = vis
freq = ts.freq[:]
if y_axis == 'jul_date':
y_aixs = ts.time[:]
y_label = r'$t$ / Julian Date'
elif y_axis == 'ra':
y_aixs = ts['ra_dec'][:, 0]
y_label = r'RA / radian'
extent = [freq[0], freq[-1], y_aixs[0], y_aixs[-1]]
plt.figure()
if plot_abs:
fig, axarr = plt.subplots(1, 3, sharey=True)
else:
fig, axarr = plt.subplots(1, 2, sharey=True)
im = axarr[0].imshow(vis1.real, extent=extent, origin='lower', aspect='auto')
axarr[0].set_xlabel(r'$\nu$ / MHz')
axarr[0].set_ylabel(y_label)
plt.colorbar(im, ax=axarr[0])
im = axarr[1].imshow(vis1.imag, extent=extent, origin='lower', aspect='auto')
axarr[1].set_xlabel(r'$\nu$ / MHz')
plt.colorbar(im, ax=axarr[1])
if plot_abs:
im = axarr[2].imshow(np.abs(vis1), extent=extent, origin='lower', aspect='auto')
axarr[2].set_xlabel(r'$\nu$ / MHz')
plt.colorbar(im, ax=axarr[2])
if pol is None:
fig_name = '%s_%d_%d.png' % (fig_prefix, bl[0], bl[1])
else:
fig_name = '%s_%d_%d_%s.png' % (fig_prefix, bl[0], bl[1], pol)
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
return vis, vis_mask
def plot(self, vis, vis_mask, li, gi, bl, ts, **kwargs):
"""Function that does the actual plot work."""
bl_incl = self.params['bl_incl']
bl_excl = self.params['bl_excl']
flag_mask = self.params['flag_mask']
flag_ns = self.params['flag_ns']
interpolate_ns = self.params['interpolate_ns']
y_axis = self.params['y_axis']
plot_abs = self.params['plot_abs']
abs_only = self.params['abs_only']
gray_color = self.params['gray_color']
color_flag = self.params['color_flag']
flag_color = self.params['flag_color']
transpose = self.params['transpose']
hist_equal = self.params['hist_equal']
fig_prefix = self.params['fig_name']
rotate_xdate = self.params['rotate_xdate']
feed_no = self.params['feed_no']
order_bl = self.params['order_bl']
tag_output_iter = self.params['tag_output_iter']
iteration = self.iteration
if isinstance(ts, Timestream): # for Timestream
pol = bl[0]
bl = tuple(bl[1])
feed_no = True
elif isinstance(ts, RawTimestream): # for RawTimestream
pol = None
bl = tuple(bl)
if feed_no:
pol = ts['bl_pol'].local_data[li]
bl = tuple(ts['true_blorder'].local_data[li])
if order_bl and (bl[0] > bl[1]):
bl = (bl[1], bl[0])
vis = vis.conj()
else:
raise ValueError('Need either a RawTimestream or Timestream')
if bl_incl != 'all':
bl1 = set(bl)
bl_incl = [{f1, f2} for (f1, f2) in bl_incl]
bl_excl = [{f1, f2} for (f1, f2) in bl_excl]
if (not bl1 in bl_incl) or (bl1 in bl_excl):
return
if flag_mask:
vis1 = np.ma.array(vis, mask=vis_mask)
elif flag_ns:
if 'ns_on' in ts.iterkeys():
vis1 = vis.copy()
on = np.where(ts['ns_on'][:])[0]
if not interpolate_ns:
vis1[on] = complex(np.nan, np.nan)
else:
off = np.where(np.logical_not(ts['ns_on'][:]))[0]
for fi in xrange(vis1.shape[1]):
itp_real = InterpolatedUnivariateSpline(
off, vis1[off, fi].real)
itp_imag = InterpolatedUnivariateSpline(
off, vis1[off, fi].imag)
vis1[on, fi] = itp_real(on) + 1.0J * itp_imag(on)
else:
vis1 = vis
else:
vis1 = vis
freq = ts.freq[:]
x_label = r'$\nu$ / MHz'
if y_axis == 'jul_date':
y_aixs = ts.time[:]
y_label = r'$t$ / Julian Date'
elif y_axis == 'ra':
y_aixs = ts['ra_dec'][:, 0]
y_label = r'RA / radian'
elif y_axis == 'time':
y_aixs = [
datetime.fromtimestamp(s)
for s in (ts['sec1970'][0], ts['sec1970'][-1])
]
y_label = '%s' % y_aixs[0].date()
# convert datetime objects to the correct format for matplotlib to work with
y_aixs = mdates.date2num(y_aixs)
else:
raise ValueError(
'Invalid y_axis %s, can only be "time", "jul_data" or "ra"' %
y_axis)
freq_extent = [freq[0], freq[-1]]
time_extent = [y_aixs[0], y_aixs[-1]]
extent = freq_extent + time_extent
plt.figure()
if gray_color:
# cmap = 'gray'
cmap = plt.cm.gray
if color_flag:
cmap.set_bad(flag_color)
else:
cmap = None
if abs_only:
if transpose:
vis1 = vis1.T
x_label, y_label = y_label, x_label
extent = time_extent + freq_extent
fig, ax = plt.subplots()
vis_abs = np.abs(vis1)
if hist_equal:
if isinstance(vis_abs, np.ma.MaskedArray):
vis_hist = hist_eq.hist_eq(vis_abs.filled(0))
vis_abs = np.ma.array(vis_hist,
mask=np.ma.getmask(vis_abs))
else:
vis_hist = hist_eq.hist_eq(
np.where(np.isfinite(vis_abs), vis_abs, 0))
mask = np.where(np.isfinite(vis_abs), False, True)
vis_abs = np.ma.array(vis_hist, mask=mask)
im = ax.imshow(vis_abs,
extent=extent,
origin='lower',
aspect='auto',
cmap=cmap)
# convert axis to datetime string
if transpose:
ax.xaxis_date()
else:
ax.yaxis_date()
# format datetime string
# date_format = mdates.DateFormatter('%y/%m/%d %H:%M')
date_format = mdates.DateFormatter('%H:%M')
if transpose:
ax.xaxis.set_major_formatter(date_format)
else:
ax.yaxis.set_major_formatter(date_format)
if transpose:
if rotate_xdate:
# set the x-axis tick labels to diagonal so it fits better
fig.autofmt_xdate()
else:
# reduce the number of tick locators
locator = MaxNLocator(nbins=6)
ax.xaxis.set_major_locator(locator)
ax.xaxis.set_minor_locator(AutoMinorLocator(2))
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
plt.colorbar(im)
else:
if plot_abs:
fig, axarr = plt.subplots(1, 3, sharey=True)
else:
fig, axarr = plt.subplots(1, 2, sharey=True)
im = axarr[0].imshow(vis1.real,
extent=extent,
origin='lower',
aspect='auto',
cmap=cmap)
axarr[0].set_xlabel(x_label)
axarr[0].yaxis_date()
# format datetime string
date_format = mdates.DateFormatter('%H:%M')
axarr[0].yaxis.set_major_formatter(date_format)
axarr[0].set_ylabel(y_label)
plt.colorbar(im, ax=axarr[0])
im = axarr[1].imshow(vis1.imag,
extent=extent,
origin='lower',
aspect='auto',
cmap=cmap)
axarr[1].set_xlabel(x_label)
plt.colorbar(im, ax=axarr[1])
if plot_abs:
im = axarr[2].imshow(np.abs(vis1),
extent=extent,
origin='lower',
aspect='auto',
cmap=cmap)
axarr[2].set_xlabel(x_label)
plt.colorbar(im, ax=axarr[2])
if feed_no:
fig_name = '%s_%d_%d_%s.png' % (fig_prefix, bl[0], bl[1],
ts.pol_dict[pol])
else:
fig_name = '%s_%d_%d.png' % (fig_prefix, bl[0], bl[1])
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
def cal(self, vis, vis_mask, li, gi, fbl, rt, **kwargs):
"""Function that does the actual cal."""
num_mean = self.params['num_mean']
plot_phs = self.params['plot_phs']
fig_prefix = self.params['fig_name']
rotate_xdate = self.params['rotate_xdate']
feed_no = self.params['feed_no']
tag_output_iter = self.params['tag_output_iter']
iteration = self.iteration
bls_plt = kwargs['bls_plt']
freq_plt = kwargs['freq_plt']
if np.prod(vis.shape) == 0 :
return
fi = gi[0] # freq idx for this cal
bl = tuple(fbl[1]) # bl for this cal
nt = vis.shape[0]
on_time = rt['ns_on'].attrs['on_time']
num_mean = min(num_mean, on_time-2)
if num_mean <= 0:
raise RuntimeError('Do not have enough noise on time samples to do the ns_cal')
ns_on = rt['ns_on'][:]
ns_on = np.where(ns_on, 1, 0)
diff_ns = np.diff(ns_on)
inds = np.where(diff_ns==1)[0]
# if inds[0]-num_mean < 0:
if inds[0]-1 < 0: # no off data in the beginning to use
inds = inds[1:]
# if inds[-1]+num_mean+1 > len(ns_on)-1:
if inds[-1]+2 > nt-1: # no on data in the end to use
inds = inds[:-1]
valid_inds = []
phase = []
for ind in inds:
if ind == inds[0]: # the first ind
lower = max(0, ind-num_mean)
else:
lower = ind - num_mean
off_sec = np.ma.array(vis[lower:ind], mask=vis_mask[lower:ind])
if off_sec.count() > 0: # not all data in this section are masked
valid_inds.append(ind)
if ind == inds[-1]: # the last ind
upper = min(nt, ind+2+num_mean)
else:
upper = ind + 2 + num_mean
phase.append( np.angle(np.mean(vis[ind+2:upper]) - np.ma.mean(off_sec)) ) # in radians
# not enough valid data to do the ns_cal
if len(phase) <= 3:
vis_mask[:] = True # mask the vis as no ns_cal has done
return
phase = np.unwrap(phase) # unwrap 2pi discontinuity
f = InterpolatedUnivariateSpline(valid_inds, phase)
all_phase = f(np.arange(nt))
if plot_phs and (bl in bls_plt and fi in freq_plt):
plt.figure()
fig, ax = plt.subplots()
ax_val = np.array([ datetime.fromtimestamp(sec) for sec in rt['sec1970'][:] ])
xlabel = '%s' % ax_val[0].date()
ax_val = mdates.date2num(ax_val)
ax.plot(ax_val, all_phase)
ax.plot(ax_val[valid_inds], phase, 'ro')
ax.xaxis_date()
date_format = mdates.DateFormatter('%H:%M')
ax.xaxis.set_major_formatter(date_format)
if rotate_xdate:
# set the x-axis tick labels to diagonal so it fits better
fig.autofmt_xdate()
else:
# reduce the number of tick locators
locator = MaxNLocator(nbins=6)
ax.xaxis.set_major_locator(locator)
ax.xaxis.set_minor_locator(AutoMinorLocator(2))
ax.set_xlabel(xlabel)
ax.set_ylabel(r'$\Delta \phi$ / radian')
if feed_no:
pol = rt['bl_pol'].local_data[li[1]]
bl = tuple(rt['true_blorder'].local_data[li[1]])
fig_name = '%s_%f_%d_%d_%s.png' % (fig_prefix, fbl[0], bl[0], bl[1], rt.pol_dict[pol])
else:
fig_name = '%s_%f_%d_%d.png' % (fig_prefix, fbl[0], fbl[1][0], fbl[1][1])
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
vis[:] = vis * np.exp(-1.0J * all_phase)
def process(self, ts):
assert isinstance(
ts, Timestream
), '%s only works for Timestream object' % self.__class__.__name__
mask_daytime = self.params['mask_daytime']
mask_time_range = self.params['mask_time_range']
beam_theta_range = self.params['beam_theta_range']
tsys = self.params['tsys']
accuracy_boost = self.params['accuracy_boost']
l_boost = self.params['l_boost']
bl_range = self.params['bl_range']
auto_correlations = self.params['auto_correlations']
pol = self.params['pol']
beam_dir = output_path(self.params['beam_dir'])
gen_inv = self.params['gen_invbeam']
noise_weight = self.params['noise_weight']
ts_dir = output_path(self.params['ts_dir'])
ts_name = self.params['ts_name']
simulate = self.params['simulate']
input_maps = self.params['input_maps']
add_noise = self.params['add_noise']
dirty_map = self.params['dirty_map']
nbin = self.params['nbin']
method = self.params['method']
normalize = self.params['normalize']
threshold = self.params['threshold']
ts.redistribute('frequency')
lat = ts.attrs['sitelat']
# lon = ts.attrs['sitelon']
lon = 0.0
# lon = np.degrees(ts['ra_dec'][0, 0]) # the first ra
local_origin = False
freq = ts.freq
freqs = ts.freq.data.to_numpy_array(root=None) # MHz
band_width = ts.attrs['freqstep'] # MHz
ndays = ts.attrs['ndays']
feeds = ts['feedno'][:]
az, alt = ts['az_alt'][0]
az = np.degrees(az)
alt = np.degrees(alt)
pointing = [az, alt, 0.0]
feedpos = ts['feedpos'][:]
if ts.is_dish:
dish_width = ts.attrs['dishdiam']
tel = tl_dish.TlUnpolarisedDishArray(lat, lon, freqs, band_width,
tsys, ndays, accuracy_boost,
l_boost, bl_range,
auto_correlations,
local_origin, dish_width,
feedpos, pointing)
elif ts.is_cylinder:
# factor = 1.2 # suppose an illumination efficiency, keep same with that in timestream_common
factor = 0.79 # for xx
# factor = 0.88 # for yy
cyl_width = factor * ts.attrs['cywid']
tel = tl_cylinder.TlUnpolarisedCylinder(
lat, lon, freqs, band_width, tsys, ndays, accuracy_boost,
l_boost, bl_range, auto_correlations, local_origin, cyl_width,
feedpos)
else:
raise RuntimeError('Unknown array type %s' % ts.attrs['telescope'])
if not simulate:
# mask daytime data
if mask_daytime:
day_inds = np.where(
np.logical_and(
ts['local_hour'][:] >= mask_time_range[0],
ts['local_hour'][:] <= mask_time_range[1]))[0]
ts.local_vis_mask[
day_inds] = True # do not change vis directly
# average data
nt = ts['sec1970'].shape[0]
phi_size = 2 * tel.mmax + 1
nt_m = float(nt) / phi_size
# roll data to have phi=0 near the first
roll_len = np.int(np.around(0.5 * nt_m))
ts.local_vis[:] = np.roll(ts.local_vis[:], roll_len, axis=0)
ts.local_vis_mask[:] = np.roll(ts.local_vis_mask[:],
roll_len,
axis=0)
ts['ra_dec'][:] = np.roll(ts['ra_dec'][:], roll_len, axis=0)
repeat_inds = np.repeat(np.arange(nt), phi_size)
num, start, end = mpiutil.split_m(nt * phi_size, phi_size)
# phi = np.zeros((phi_size,), dtype=ts['ra_dec'].dtype)
phi = np.linspace(0, 2 * np.pi, phi_size, endpoint=False)
vis = np.zeros((phi_size, ) + ts.local_vis.shape[1:],
dtype=ts.vis.dtype)
# average over time
for idx in xrange(phi_size):
inds, weight = unique(repeat_inds[start[idx]:end[idx]],
return_counts=True)
vis[idx] = average(np.ma.array(ts.local_vis[inds],
mask=ts.local_vis_mask[inds]),
axis=0,
weights=weight) # time mean
# phi[idx] = np.average(ts['ra_dec'][:, 0][inds], axis=0, weights=weight)
if pol == 'xx':
vis = vis[:, :, 0, :]
elif pol == 'yy':
vis = vis[:, :, 1, :]
elif pol == 'I':
vis = 0.5 * (vis[:, :, 0, :] + vis[:, :, 1, :])
elif pol == 'all':
vis = np.sum(vis, axis=2) # sum over all pol
else:
raise ValueError('Invalid pol: %s' % pol)
allpairs = tel.allpairs
redundancy = tel.redundancy
# reorder bls according to allpairs
vis_tmp = np.zeros_like(vis)
bls = [tuple(bl) for bl in ts['blorder'][:]]
for ind, (a1, a2) in enumerate(allpairs):
try:
b_ind = bls.index((feeds[a1], feeds[a2]))
vis_tmp[:, :, ind] = vis[:, :, b_ind]
except ValueError:
b_ind = bls.index((feeds[a2], feeds[a1]))
vis_tmp[:, :, ind] = vis[:, :, b_ind].conj()
# average over redundancy
vis_stream = np.zeros(vis.shape[:-1] + (len(redundancy), ),
dtype=vis_tmp.dtype)
red_bin = np.cumsum(np.insert(redundancy, 0, 0)) # redundancy bin
# average over redundancy
for ind in xrange(len(redundancy)):
vis_stream[:, :,
ind] = np.sum(
vis_tmp[:, :, red_bin[ind]:red_bin[ind + 1]],
axis=2) / redundancy[ind]
del vis
del vis_tmp
# beamtransfer
bt = beamtransfer.BeamTransfer(beam_dir, tel, noise_weight, True)
bt.generate()
if simulate:
ndays = 733
print ndays
tstream = timestream.simulate(bt,
ts_dir,
ts_name,
input_maps,
ndays,
add_noise=add_noise)
else:
# timestream and map-making
tstream = timestream.Timestream(ts_dir, ts_name, bt)
for lfi, fi in freq.data.enumerate(axis=0):
# Make directory if required
if not os.path.exists(tstream._fdir(fi)):
os.makedirs(tstream._fdir(fi))
# Write file contents
with h5py.File(tstream._ffile(fi), 'w') as f:
# Timestream data
f.create_dataset('/timestream', data=vis_stream[:, lfi].T)
f.create_dataset('/phi', data=phi)
# Telescope layout data
f.create_dataset('/feedmap', data=tel.feedmap)
f.create_dataset('/feedconj', data=tel.feedconj)
f.create_dataset('/feedmask', data=tel.feedmask)
f.create_dataset('/uniquepairs', data=tel.uniquepairs)
f.create_dataset('/baselines', data=tel.baselines)
# Telescope frequencies
f.create_dataset('/frequencies', data=freqs)
# Write metadata
f.attrs['beamtransfer_path'] = os.path.abspath(
bt.directory)
f.attrs['ntime'] = phi_size
mpiutil.barrier()
tstream.generate_mmodes()
nside = hputil.nside_for_lmax(tel.lmax,
accuracy_boost=tel.accuracy_boost)
if dirty_map:
tstream.mapmake_full(nside,
'map_full_dirty.hdf5',
nbin,
dirty=True,
method=method,
normalize=normalize,
threshold=threshold)
else:
tstream.mapmake_full(nside,
'map_full.hdf5',
nbin,
dirty=False,
method=method,
normalize=normalize,
threshold=threshold)
# ts.add_history(self.history)
return tstream
def fit(vis_obs, vis_mask, vis_sim, start_ind, end_ind, num_shift, idx,
plot_fit, fig_prefix, iteration, tag_output_iter, bls_plt, freq_plt):
vis_obs = np.ma.array(vis_obs, mask=vis_mask)
num_nomask = vis_obs.count()
if num_nomask == 0: # no valid vis data
return 1.0, 0
fi, pi, (i, j) = idx
gains = []
chi2s = []
shifts = xrange(-num_shift / 2, num_shift / 2 + 1)
for si in shifts:
# vis = vis_obs[start_ind+si:end_ind+si].astype(np.complex128) # improve precision
vis = vis_obs[start_ind + si:end_ind + si]
num_nomask = vis.count()
if num_nomask == 0: # no valid vis data
continue
else:
xx = np.ma.dot(vis_sim.conj(), vis_sim)
xy = np.ma.dot(vis_sim.conj(), vis)
gain = xy / xx
vis_cal = gain * vis_sim
err = vis - vis_cal
chi2 = np.ma.dot(err.conj(), err).real
gains.append(gain)
chi2s.append(chi2 / num_nomask)
if len(gains) == 0: # no valid vis data
return 1.0, 0
chi2s = np.array(chi2s)
if np.allclose(chi2s, np.sort(chi2s)):
if mpiutil.rank0:
print 'Warn: chi2 increasing for %s...' % (idx, )
if np.allclose(chi2s, np.sort(chi2s)[::-1]):
if mpiutil.rank0:
print 'Warn: chi2 decreasing for %s...' % (idx, )
ind = np.argmin(chi2s)
gain = gains[ind]
# chi2 = chi2s[ind]
si = shifts[ind]
obs_data = np.ma.array(vis_obs[start_ind:end_ind],
mask=vis_mask[start_ind:end_ind])
factor = np.ma.max(np.ma.abs(obs_data)) / np.max(np.abs(vis_sim))
obs_data = obs_data / factor # make amp close to each other
vis_cal = np.ma.array(vis_obs[start_ind + si:end_ind + si],
mask=vis_mask[start_ind + si:end_ind + si]) / gain
if si != 0 and mpiutil.rank0:
print 'shift %d for %s...' % (si, idx)
if plot_fit and (fi in freq_plt and (i, j) in bls_plt):
# plot the fit
plt.figure()
plt.subplot(311)
plt.plot(obs_data.real, label='obs, real')
if not vis_cal is np.ma.masked: # in case gain is --
plt.plot(vis_cal.real, label='cal, real')
plt.plot(vis_sim.real, label='sim, real')
plt.legend(loc='best')
plt.subplot(312)
plt.plot(obs_data.imag, label='obs, imag')
if not vis_cal is np.ma.masked: # in case gain is --
plt.plot(vis_cal.imag, label='cal, imag')
plt.plot(vis_sim.imag, label='sim, imag')
plt.legend(loc='best')
plt.subplot(313)
plt.plot(np.abs(obs_data), label='obs, abs')
if not vis_cal is np.ma.masked: # in case gain is --
plt.plot(np.abs(vis_cal), label='cal, abs')
plt.plot(np.abs(vis_sim), label='sim, abs')
plt.legend(loc='best')
fig_name = '%s_%d_%d_%d_%d.png' % (fig_prefix, fi, pi, i, j)
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
return gain, si
def plot(self, vis, vis_mask, li, gi, bl, ts, **kwargs):
"""Function that does the actual plot work."""
integral = self.params['integral']
bl_incl = self.params['bl_incl']
bl_excl = self.params['bl_excl']
flag_mask = self.params['flag_mask']
flag_ns = self.params['flag_ns']
fig_prefix = self.params['fig_name']
tag_output_iter = self.params['tag_output_iter']
iteration= self.iteration
if isinstance(ts, Timestream): # for Timestream
pol = bl[0]
bl = tuple(bl[1])
elif isinstance(ts, RawTimestream): # for RawTimestream
pol = None
bl = tuple(bl)
else:
raise ValueError('Need either a RawTimestream or Timestream')
if bl_incl != 'all':
bl1 = set(bl)
bl_incl = [ {f1, f2} for (f1, f2) in bl_incl ]
bl_excl = [ {f1, f2} for (f1, f2) in bl_excl ]
if (not bl1 in bl_incl) or (bl1 in bl_excl):
return vis, vis_mask
if flag_mask:
vis1 = np.ma.array(vis, mask=vis_mask)
elif flag_ns:
vis1 = vis.copy()
on = np.where(ts['ns_on'][:])[0]
vis1[on] = complex(np.nan, np.nan)
else:
vis1 = vis
if integral == 'time':
ax_val = ts.freq[:]
vis1 = np.ma.mean(np.ma.masked_invalid(vis1), axis=0)
xlabel = r'$\nu$ / MHz'
elif integral == 'freq':
ax_val = ts.time[:]
vis1 = np.ma.mean(np.ma.masked_invalid(vis1), axis=1)
xlabel = r'$t$ / Julian Date'
else:
raise ValueError('Unknown integral type %s' % integral)
plt.figure()
f, axarr = plt.subplots(3, sharex=True)
axarr[0].plot(ax_val, vis1.real, label='real')
axarr[0].legend()
axarr[1].plot(ax_val, vis1.imag, label='imag')
axarr[1].legend()
axarr[2].plot(ax_val, np.abs(vis1), label='abs')
axarr[2].legend()
axarr[2].set_xlabel(xlabel)
if pol is None:
fig_name = '%s_%s_%d_%d.png' % (fig_prefix, integral, bl[0], bl[1])
else:
fig_name = '%s_%s_%d_%d_%s.png' % (fig_prefix, integral, bl[0], bl[1], pol)
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
return vis, vis_mask
def process(self, ts):
assert isinstance(
ts, Timestream
), '%s only works for Timestream object' % self.__class__.__name__
calibrator = self.params['calibrator']
catalog = self.params['catalog']
vis_conj = self.params['vis_conj']
zero_diag = self.params['zero_diag']
span = self.params['span']
reserve_high_gain = self.params['reserve_high_gain']
plot_figs = self.params['plot_figs']
fig_prefix = self.params['fig_name']
tag_output_iter = self.params['tag_output_iter']
save_src_vis = self.params['save_src_vis']
src_vis_file = self.params['src_vis_file']
subtract_src = self.params['subtract_src']
replace_with_src = self.params['replace_with_src']
apply_gain = self.params['apply_gain']
save_gain = self.params['save_gain']
save_phs_change = self.params['save_phs_change']
gain_file = self.params['gain_file']
# temperature_convert = self.params['temperature_convert']
show_progress = self.params['show_progress']
progress_step = self.params['progress_step']
if save_src_vis or subtract_src or apply_gain or save_gain:
pol_type = ts['pol'].attrs['pol_type']
if pol_type != 'linear':
raise RuntimeError('Can not do ps_cal for pol_type: %s' %
pol_type)
ts.redistribute('baseline')
feedno = ts['feedno'][:].tolist()
pol = [ts.pol_dict[p] for p in ts['pol'][:]] # as string
gain_pd = {
'xx': 0,
'yy': 1,
0: 'xx',
1: 'yy'
} # for gain related op
bls = mpiutil.gather_array(ts.local_bl[:], root=None, comm=ts.comm)
# # antpointing = np.radians(ts['antpointing'][-1, :, :]) # radians
# transitsource = ts['transitsource'][:]
# transit_time = transitsource[-1, 0] # second, sec1970
# int_time = ts.attrs['inttime'] # second
# get the calibrator
try:
s = calibrators.get_src(calibrator)
except KeyError:
if mpiutil.rank0:
print 'Calibrator %s is unavailable, available calibrators are:'
for key, d in calibrators.src_data.items():
print '%8s -> %12s' % (key, d[0])
raise RuntimeError('Calibrator %s is unavailable')
if mpiutil.rank0:
print 'Try to calibrate with %s...' % s.src_name
# get transit time of calibrator
# array
aa = ts.array
aa.set_jultime(ts['jul_date'][0]) # the first obs time point
next_transit = aa.next_transit(s)
transit_time = a.phs.ephem2juldate(next_transit) # Julian date
# get time zone
pattern = '[-+]?\d+'
tz = re.search(pattern, ts.attrs['timezone']).group()
tz = int(tz)
local_next_transit = ephem.Date(
next_transit + tz * ephem.hour) # plus 8h to get Beijing time
# if transit_time > ts['jul_date'][-1]:
if transit_time > max(ts['jul_date'][-1], ts['jul_date'][:].max()):
raise RuntimeError(
'Data does not contain local transit time %s of source %s'
% (local_next_transit, calibrator))
# the first transit index
transit_inds = [np.searchsorted(ts['jul_date'][:], transit_time)]
# find all other transit indices
aa.set_jultime(ts['jul_date'][0] + 1.0)
transit_time = a.phs.ephem2juldate(
aa.next_transit(s)) # Julian date
cnt = 2
while (transit_time <= ts['jul_date'][-1]):
transit_inds.append(
np.searchsorted(ts['jul_date'][:], transit_time))
aa.set_jultime(ts['jul_date'][0] + 1.0 * cnt)
transit_time = a.phs.ephem2juldate(
aa.next_transit(s)) # Julian date
cnt += 1
if mpiutil.rank0:
print 'transit ind of %s: %s, time: %s' % (
s.src_name, transit_inds, local_next_transit)
### now only use the first transit point to do the cal
### may need to improve in the future
transit_ind = transit_inds[0]
int_time = ts.attrs['inttime'] # second
start_ind = transit_ind - np.int(span / int_time)
end_ind = transit_ind + np.int(
span /
int_time) + 1 # plus 1 to make transit_ind is at the center
start_ind = max(0, start_ind)
end_ind = min(end_ind, ts.vis.shape[0])
if vis_conj:
ts.local_vis[:] = ts.local_vis.conj()
nt = end_ind - start_ind
t_inds = range(start_ind, end_ind)
freq = ts.freq[:] # MHz
nf = len(freq)
nlb = len(ts.local_bl[:])
nfeed = len(feedno)
tfp_inds = list(
itertools.product(
t_inds, range(nf),
[pol.index('xx'), pol.index('yy')])) # only for xx and yy
ns, ss, es = mpiutil.split_all(len(tfp_inds), comm=ts.comm)
# gather data to make each process to have its own data which has all bls
for ri, (ni, si, ei) in enumerate(zip(ns, ss, es)):
lvis = np.zeros((ni, nlb), dtype=ts.vis.dtype)
lvis_mask = np.zeros((ni, nlb), dtype=ts.vis_mask.dtype)
for ii, (ti, fi, pi) in enumerate(tfp_inds[si:ei]):
lvis[ii] = ts.local_vis[ti, fi, pi]
lvis_mask[ii] = ts.local_vis_mask[ti, fi, pi]
# gather vis from all process for separate bls
gvis = mpiutil.gather_array(lvis,
axis=1,
root=ri,
comm=ts.comm)
gvis_mask = mpiutil.gather_array(lvis_mask,
axis=1,
root=ri,
comm=ts.comm)
if ri == mpiutil.rank:
tfp_linds = tfp_inds[si:ei] # inds for this process
this_vis = gvis
this_vis_mask = gvis_mask
del tfp_inds
del lvis
del lvis_mask
tfp_len = len(tfp_linds)
# lotl_mask = np.zeros((tfp_len, nfeed, nfeed), dtype=bool)
cnan = complex(np.nan, np.nan) # complex nan
if save_src_vis or subtract_src:
# save calibrator src vis
lsrc_vis = np.full((tfp_len, nfeed, nfeed),
cnan,
dtype=ts.vis.dtype)
if save_src_vis:
# save sky vis
lsky_vis = np.full((tfp_len, nfeed, nfeed),
cnan,
dtype=ts.vis.dtype)
# save outlier vis
lotl_vis = np.full((tfp_len, nfeed, nfeed),
cnan,
dtype=ts.vis.dtype)
if apply_gain or save_gain:
lGain = np.full((tfp_len, nfeed), cnan, dtype=ts.vis.dtype)
# find indices mapping between Vmat and vis
# bis = range(nbl)
bis_conj = [] # indices that shold be conj
mis = [
] # indices in the nfeed x nfeed matrix by flatten it to a vector
mis_conj = [
] # indices (of conj vis) in the nfeed x nfeed matrix by flatten it to a vector
for bi, (fdi, fdj) in enumerate(bls):
ai, aj = feedno.index(fdi), feedno.index(fdj)
mis.append(ai * nfeed + aj)
if ai != aj:
bis_conj.append(bi)
mis_conj.append(aj * nfeed + ai)
# construct visibility matrix for a single time, freq, pol
Vmat = np.full((nfeed, nfeed), cnan, dtype=ts.vis.dtype)
# get flus of the calibrator in the observing frequencies
if show_progress and mpiutil.rank0:
pg = progress.Progress(tfp_len, step=progress_step)
for ii, (ti, fi, pi) in enumerate(tfp_linds):
if show_progress and mpiutil.rank0:
pg.show(ii)
# when noise on, just pass
if 'ns_on' in ts.iterkeys() and ts['ns_on'][ti]:
continue
# aa.set_jultime(ts['jul_date'][ti])
# s.compute(aa)
# get the topocentric coordinate of the calibrator at the current time
# s_top = s.get_crds('top', ncrd=3)
# aa.sim_cache(cat.get_crds('eq', ncrd=3)) # for compute bm_response and sim
Vmat.flat[mis] = np.ma.array(
this_vis[ii], mask=this_vis_mask[ii]).filled(cnan)
Vmat.flat[mis_conj] = np.ma.array(
this_vis[ii, bis_conj],
mask=this_vis_mask[ii, bis_conj]).conj().filled(cnan)
if save_src_vis:
lsky_vis[ii] = Vmat
# set invalid val to 0
invalid = ~np.isfinite(Vmat) # a bool array
# if too many masks
if np.where(invalid)[0].shape[0] > 0.3 * nfeed**2:
continue
Vmat[invalid] = 0
# if all are zeros
if np.allclose(Vmat, 0.0):
continue
# fill diagonal of Vmat to 0
if zero_diag:
np.fill_diagonal(Vmat, 0)
# initialize the outliers
med = np.median(Vmat.real) + 1.0J * np.median(Vmat.imag)
diff = Vmat - med
S0 = np.where(
np.abs(diff) > 3.0 * rpca_decomp.MAD(Vmat), diff, 0)
# stable PCA decomposition
V0, S = rpca_decomp.decompose(Vmat,
rank=1,
S=S0,
max_iter=200,
threshold='hard',
tol=1.0e-6,
debug=False)
# # find abnormal values in S
# # first check diagonal elements
# import pdb; pdb.set_trace()
# svals = np.diag(S)
# smed = np.median(svals.real) + 1.0J * np.median(svals.imag)
# smad = rpca_decomp.MAD(svals)
# # abnormal indices
# abis = np.where(np.abs(svals - smed) > 3.0 * smad)[0]
# for abi in abis:
# lotl_mask[ii, abi, abi] = True
# # then check non-diagonal elements
# for rii in range(nfeed):
# for cii in range(nfeed):
# if rii == cii:
# continue
# rli = max(0, rii-2)
# rhi = min(nfeed, rii+3)
# cli = max(0, cii-2)
# chi = min(nfeed, cii+3)
# svals = np.array([ S[xi, yi] for xi in range(rli, rhi) for yi in range(cli, chi) if xi != yi ])
# smed = np.median(svals.real) + 1.0J * np.median(svals.imag)
# smad = rpca_decomp.MAD(svals)
# if np.abs(S[rii, cii] - smed) > 3.0 * smad:
# lotl_mask[ii, rii, cii] = True
if save_src_vis or subtract_src:
lsrc_vis[ii] = V0
if save_src_vis:
lotl_vis[ii] = S
# plot
if plot_figs:
ind = ti - start_ind
# plot Vmat
plt.figure(figsize=(13, 5))
plt.subplot(121)
plt.imshow(Vmat.real,
aspect='equal',
origin='lower',
interpolation='nearest')
plt.colorbar(shrink=1.0)
plt.subplot(122)
plt.imshow(Vmat.imag,
aspect='equal',
origin='lower',
interpolation='nearest')
plt.colorbar(shrink=1.0)
fig_name = '%s_V_%d_%d_%s.png' % (fig_prefix, ind, fi,
pol[pi])
if tag_output_iter:
fig_name = output_path(fig_name,
iteration=self.iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
# plot V0
plt.figure(figsize=(13, 5))
plt.subplot(121)
plt.imshow(V0.real,
aspect='equal',
origin='lower',
interpolation='nearest')
plt.colorbar(shrink=1.0)
plt.subplot(122)
plt.imshow(V0.imag,
aspect='equal',
origin='lower',
interpolation='nearest')
plt.colorbar(shrink=1.0)
fig_name = '%s_V0_%d_%d_%s.png' % (fig_prefix, ind, fi,
pol[pi])
if tag_output_iter:
fig_name = output_path(fig_name,
iteration=self.iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
# plot S
plt.figure(figsize=(13, 5))
plt.subplot(121)
plt.imshow(S.real,
aspect='equal',
origin='lower',
interpolation='nearest')
plt.colorbar(shrink=1.0)
plt.subplot(122)
plt.imshow(S.imag,
aspect='equal',
origin='lower',
interpolation='nearest')
plt.colorbar(shrink=1.0)
fig_name = '%s_S_%d_%d_%s.png' % (fig_prefix, ind, fi,
pol[pi])
if tag_output_iter:
fig_name = output_path(fig_name,
iteration=self.iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
# plot N
N = Vmat - V0 - S
plt.figure(figsize=(13, 5))
plt.subplot(121)
plt.imshow(N.real,
aspect='equal',
origin='lower',
interpolation='nearest')
plt.colorbar(shrink=1.0)
plt.subplot(122)
plt.imshow(N.imag,
aspect='equal',
origin='lower',
interpolation='nearest')
plt.colorbar(shrink=1.0)
fig_name = '%s_N_%d_%d_%s.png' % (fig_prefix, ind, fi,
pol[pi])
if tag_output_iter:
fig_name = output_path(fig_name,
iteration=self.iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
if apply_gain or save_gain:
# use v_ij = gi gj^* \int Ai Aj^* e^(2\pi i n \cdot uij) T(x) d^2n
# precisely, we shold have
# V0 = (lambda^2 * Sc / (2 k_B)) * gi gj^* Ai Aj^* e^(2\pi i n0 \cdot uij)
e, U = la.eigh(V0, eigvals=(nfeed - 1, nfeed - 1))
g = U[:, -1] * e[
-1]**0.5 # = \sqrt(lambda^2 * Sc / (2 k_B)) * gi Ai * e^(2\pi i n0 \cdot ui)
if g[0].real < 0:
g *= -1.0 # make all g[0] phase 0, instead of pi
lGain[ii] = g
# plot Gain
if plot_figs:
plt.figure()
plt.plot(feedno, g.real, 'b-', label='real')
plt.plot(feedno, g.real, 'bo')
plt.plot(feedno, g.imag, 'g-', label='imag')
plt.plot(feedno, g.imag, 'go')
plt.plot(feedno, np.abs(g), 'r-', label='abs')
plt.plot(feedno, np.abs(g), 'ro')
plt.xlim(feedno[0] - 1, feedno[-1] + 1)
yl, yh = plt.ylim()
plt.ylim(yl, yh + (yh - yl) / 5)
plt.xlabel('Feed number')
plt.legend()
fig_name = '%s_ants_%d_%d_%s.png' % (fig_prefix, ind,
fi, pol[pi])
if tag_output_iter:
fig_name = output_path(fig_name,
iteration=self.iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
# # apply outlier mask
# nbl = len(bls)
# lom = np.zeros((lotl_mask.shape[0], nbl), dtype=lotl_mask.dtype)
# for bi, (fd1, fd2) in enumerate(bls):
# b1, b2 = feedno.index(fd1), feedno.index(fd2)
# lom[:, bi] = lotl_mask[:, b1, b2]
# lom = mpiarray.MPIArray.wrap(lom, axis=0, comm=ts.comm)
# lom = lom.redistribute(axis=1).local_array.reshape(nt, nf, 2, -1)
# ts.local_vis_mask[start_ind:end_ind, :, pol.index('xx')] |= lom[:, :, 0]
# ts.local_vis_mask[start_ind:end_ind, :, pol.index('yy')] |= lom[:, :, 1]
# subtract the vis of calibrator from self.vis
if subtract_src:
nbl = len(bls)
lv = np.zeros((lsrc_vis.shape[0], nbl), dtype=lsrc_vis.dtype)
for bi, (fd1, fd2) in enumerate(bls):
b1, b2 = feedno.index(fd1), feedno.index(fd2)
lv[:, bi] = lsrc_vis[:, b1, b2]
lv = mpiarray.MPIArray.wrap(lv, axis=0, comm=ts.comm)
lv = lv.redistribute(axis=1).local_array.reshape(nt, nf, 2, -1)
if replace_with_src:
ts.local_vis[start_ind:end_ind, :,
pol.index('xx')] = lv[:, :, 0]
ts.local_vis[start_ind:end_ind, :,
pol.index('yy')] = lv[:, :, 1]
else:
if 'ns_on' in ts.iterkeys():
lv[ts['ns_on']
[start_ind:
end_ind]] = 0 # avoid ns_on signal to become nan
ts.local_vis[start_ind:end_ind, :,
pol.index('xx')] -= lv[:, :, 0]
ts.local_vis[start_ind:end_ind, :,
pol.index('yy')] -= lv[:, :, 1]
del lv
if not save_src_vis:
if subtract_src:
del lsrc_vis
else:
if tag_output_iter:
src_vis_file = output_path(src_vis_file,
iteration=self.iteration)
else:
src_vis_file = output_path(src_vis_file)
# create file and allocate space first by rank0
if mpiutil.rank0:
with h5py.File(src_vis_file, 'w') as f:
# allocate space
shp = (nt, nf, 2, nfeed, nfeed)
f.create_dataset('sky_vis', shp, dtype=lsky_vis.dtype)
f.create_dataset('src_vis', shp, dtype=lsrc_vis.dtype)
f.create_dataset('outlier_vis',
shp,
dtype=lotl_vis.dtype)
# f.create_dataset('outlier_mask', shp, dtype=lotl_mask.dtype)
f.attrs['calibrator'] = calibrator
f.attrs['dim'] = 'time, freq, pol, feed, feed'
try:
f.attrs['time'] = ts.time[start_ind:end_ind]
except RuntimeError:
f.create_dataset('time',
data=ts.time[start_ind:end_ind])
f.attrs['time'] = '/time'
f.attrs['freq'] = freq
f.attrs['pol'] = np.array(['xx', 'yy'])
f.attrs['feed'] = np.array(feedno)
mpiutil.barrier()
# write data to file
for i in range(10):
try:
# NOTE: if write simultaneously, will loss data with processes distributed in several nodes
for ri in xrange(mpiutil.size):
if ri == mpiutil.rank:
with h5py.File(src_vis_file, 'r+') as f:
for ii, (ti, fi,
pi) in enumerate(tfp_linds):
ti_ = ti - start_ind
pi_ = gain_pd[pol[pi]]
f['sky_vis'][ti_, fi,
pi_] = lsky_vis[ii]
f['src_vis'][ti_, fi,
pi_] = lsrc_vis[ii]
f['outlier_vis'][ti_, fi,
pi_] = lotl_vis[ii]
# f['outlier_mask'][ti_, fi, pi_] = lotl_mask[ii]
mpiutil.barrier()
break
except IOError:
time.sleep(0.5)
continue
else:
raise RuntimeError('Could not open file: %s...' %
src_vis_file)
del lsrc_vis
del lsky_vis
del lotl_vis
# del lotl_mask
mpiutil.barrier()
if apply_gain or save_gain:
# flag outliers in lGain along each feed
lG_abs = np.full_like(lGain, np.nan, dtype=lGain.real.dtype)
for i in range(lGain.shape[0]):
valid_inds = np.where(np.isfinite(lGain[i]))[0]
if len(valid_inds) > 3:
vabs = np.abs(lGain[i, valid_inds])
vmed = np.median(vabs)
vabs_diff = np.abs(vabs - vmed)
vmad = np.median(vabs_diff) / 0.6745
if reserve_high_gain:
# reserve significantly higher ones, flag only significantly lower ones
lG_abs[i, valid_inds] = np.where(
vmed - vabs > 3.0 * vmad, np.nan, vabs)
else:
# flag both significantly higher and lower ones
lG_abs[i, valid_inds] = np.where(
vabs_diff > 3.0 * vmad, np.nan, vabs)
# choose data slice near the transit time
li = max(start_ind, transit_ind - 10) - start_ind
hi = min(end_ind, transit_ind + 10 + 1) - start_ind
ci = transit_ind - start_ind # center index for transit_ind
# compute s_top for this time range
n0 = np.zeros(((hi - li), 3))
for ti, jt in enumerate(ts.time[start_ind:end_ind][li:hi]):
aa.set_jultime(jt)
s.compute(aa)
n0[ti] = s.get_crds('top', ncrd=3)
if save_phs_change:
n0t = np.zeros((nt, 3))
for ti, jt in enumerate(ts.time[start_ind:end_ind]):
aa.set_jultime(jt)
s.compute(aa)
n0t[ti] = s.get_crds('top', ncrd=3)
# get the positions of feeds
feedpos = ts['feedpos'][:]
# wrap and redistribute Gain and flagged G_abs
Gain = mpiarray.MPIArray.wrap(lGain, axis=0, comm=ts.comm)
Gain = Gain.redistribute(axis=1).reshape(
nt, nf, 2, None).redistribute(axis=0).reshape(
None, nf * 2 * nfeed).redistribute(axis=1)
G_abs = mpiarray.MPIArray.wrap(lG_abs, axis=0, comm=ts.comm)
G_abs = G_abs.redistribute(axis=1).reshape(
nt, nf, 2, None).redistribute(axis=0).reshape(
None, nf * 2 * nfeed).redistribute(axis=1)
fpd_inds = list(
itertools.product(range(nf), range(2),
range(nfeed))) # only for xx and yy
fpd_linds = mpiutil.mpilist(fpd_inds,
method='con',
comm=ts.comm)
del fpd_inds
# create data to save the solved gain for each feed
lgain = np.full((len(fpd_linds), ), cnan,
dtype=Gain.dtype) # gain for each feed
if save_phs_change:
lphs = np.full((nt, len(fpd_linds)),
np.nan,
dtype=Gain.real.dtype
) # phase change with time for each feed
# check for conj
num_conj = 0
for ii, (fi, pi, di) in enumerate(fpd_linds):
y = G_abs.local_array[li:hi, ii]
inds = np.where(np.isfinite(y))[0]
if len(inds) >= max(4, 0.5 * len(y)):
# get the approximate magnitude by averaging the central G_abs
# solve phase by least square fit
ui = (feedpos[di] - feedpos[0]) * (
1.0e6 * freq[fi]
) / const.c # position of this feed (relative to the first feed) in unit of wavelength
exp_factor = np.exp(2.0J * np.pi * np.dot(n0, ui))
ef = exp_factor
Gi = Gain.local_array[li:hi, ii]
e_phs = np.dot(ef[inds].conj(),
Gi[inds] / y[inds]) / len(inds)
ea = np.abs(e_phs)
e_phs_conj = np.dot(ef[inds],
Gi[inds] / y[inds]) / len(inds)
eac = np.abs(e_phs_conj)
if eac > ea:
num_conj += 1
# reduce num_conj from all processes
num_conj = mpiutil.allreduce(num_conj, comm=ts.comm)
if num_conj > 0.5 * (nf * 2 * nfeed): # 2 for 2 pols
if mpiutil.rank0:
print '!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
print '!!! Detect data should be their conjugate... !!!'
print '!!! Correct it automatically... !!!'
print '!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
mpiutil.barrier()
# correct vis
ts.local_vis[:] = ts.local_vis.conj()
# correct G
Gain.local_array[:] = Gain.local_array.conj()
# solve for gain
for ii, (fi, pi, di) in enumerate(fpd_linds):
y = G_abs.local_array[li:hi, ii]
inds = np.where(np.isfinite(y))[0]
if len(inds) >= max(4, 0.5 * len(y)):
# get the approximate magnitude by averaging the central G_abs
mag = np.mean(
y[inds]
) # = \sqrt(lambda^2 * Sc / (2 k_B)) * |gi| Ai
# solve phase by least square fit
ui = (feedpos[di] - feedpos[0]) * (
1.0e6 * freq[fi]
) / const.c # position of this feed (relative to the first feed) in unit of wavelength
exp_factor = np.exp(2.0J * np.pi * np.dot(n0, ui))
ef = exp_factor
Gi = Gain.local_array[li:hi, ii]
e_phs = np.dot(ef[inds].conj(), Gi[inds] /
y[inds]) / len(inds) # the phase of gi
ea = np.abs(e_phs)
if np.abs(ea - 1.0) < 0.1:
# compute gain for this feed
lgain[
ii] = mag * e_phs # \sqrt(lambda^2 * Sc / (2 k_B)) * gi Ai
if save_phs_change:
lphs[:, ii] = np.angle(
np.exp(-2.0J * np.pi * np.dot(n0t, ui)) *
Gain.local_array[:, ii])
else:
e_phs_conj = np.dot(ef[inds],
Gi[inds] / y[inds]) / len(inds)
eac = np.abs(e_phs_conj)
if eac > ea:
if np.abs(eac - 1.0) < 0.01:
print 'feedno = %d, fi = %d, pol = %s: may need to be conjugated' % (
feedno[di], fi, gain_pd[pi])
else:
print 'feedno = %d, fi = %d, pol = %s: maybe wrong abs(e_phs): %s' % (
feedno[di], fi, gain_pd[pi], ea)
# gather local gain
gain = mpiutil.gather_array(lgain,
axis=0,
root=None,
comm=ts.comm)
del lgain
gain = gain.reshape(nf, 2, nfeed)
if save_phs_change:
phs = mpiutil.gather_array(lphs,
axis=1,
root=0,
comm=ts.comm)
del lphs
if mpiutil.rank0:
phs = phs.reshape(nt, nf, 2, nfeed)
# normalize to get the exact gain
Sc = s.get_jys(1.0e-3 * freq)
# Omega = aa.ants[0].beam.Omega ### TODO: implement Omega for dish
Ai = aa.ants[0].beam.response(n0[ci - li])
lmd = const.c / (1.0e6 * freq)
factor = np.sqrt(
(lmd**2 * 1.0e-26 * Sc) /
(2 * const.k_B)) * Ai # NOTE: 1Jy = 1.0e-26 W m^-2 Hz^-1
gain /= factor[:, np.newaxis, np.newaxis]
# apply gain to vis
if apply_gain:
for fi in range(nf):
for pi in [pol.index('xx'), pol.index('yy')]:
pi_ = gain_pd[pol[pi]]
for bi, (fd1, fd2) in enumerate(
ts['blorder'].local_data):
g1 = gain[fi, pi_, feedno.index(fd1)]
g2 = gain[fi, pi_, feedno.index(fd2)]
if np.isfinite(g1) and np.isfinite(g2):
if fd1 == fd2:
# auto-correlation should be real
ts.local_vis[:, fi, pi,
bi] /= (g1 *
np.conj(g2)).real
else:
ts.local_vis[:, fi, pi,
bi] /= (g1 * np.conj(g2))
else:
# mask the un-calibrated vis
ts.local_vis_mask[:, fi, pi, bi] = True
# in unit K after the calibration
ts.vis.attrs['unit'] = 'K'
# save gain to file
if save_gain:
if tag_output_iter:
gain_file = output_path(gain_file,
iteration=self.iteration)
else:
gain_file = output_path(gain_file)
if mpiutil.rank0:
with h5py.File(gain_file, 'w') as f:
# allocate space for Gain
dset = f.create_dataset('Gain', (nt, nf, 2, nfeed),
dtype=Gain.dtype)
dset.attrs['calibrator'] = calibrator
dset.attrs['dim'] = 'time, freq, pol, feed'
try:
dset.attrs['time'] = ts.time[start_ind:end_ind]
except RuntimeError:
f.create_dataset(
'time', data=ts.time[start_ind:end_ind])
dset.attrs['time'] = '/time'
dset.attrs['freq'] = freq
dset.attrs['pol'] = np.array(['xx', 'yy'])
dset.attrs['feed'] = np.array(feedno)
dset.attrs['transit_ind'] = transit_ind
# save gain
dset = f.create_dataset('gain', data=gain)
dset.attrs['calibrator'] = calibrator
dset.attrs['dim'] = 'freq, pol, feed'
dset.attrs['freq'] = freq
dset.attrs['pol'] = np.array(['xx', 'yy'])
dset.attrs['feed'] = np.array(feedno)
# save phs
if save_phs_change:
f.create_dataset('phs', data=phs)
mpiutil.barrier()
# save Gain
for i in range(10):
try:
# NOTE: if write simultaneously, will loss data with processes distributed in several nodes
for ri in xrange(mpiutil.size):
if ri == mpiutil.rank:
with h5py.File(gain_file, 'r+') as f:
for ii, (ti, fi,
pi) in enumerate(tfp_linds):
ti_ = ti - start_ind
pi_ = gain_pd[pol[pi]]
f['Gain'][ti_, fi, pi_] = lGain[ii]
mpiutil.barrier()
break
except IOError:
time.sleep(0.5)
continue
else:
raise RuntimeError('Could not open file: %s...' %
gain_file)
mpiutil.barrier()
return super(PsCal, self).process(ts)
def plot(self, vis, vis_mask, li, gi, bl, ts, **kwargs):
"""Function that does the actual plot work."""
integral = self.params['integral']
bl_incl = self.params['bl_incl']
bl_excl = self.params['bl_excl']
flag_mask = self.params['flag_mask']
flag_ns = self.params['flag_ns']
fig_prefix = self.params['fig_name']
rotate_xdate = self.params['rotate_xdate']
feed_no = self.params['feed_no']
tag_output_iter = self.params['tag_output_iter']
iteration = self.iteration
if isinstance(ts, Timestream): # for Timestream
pol = bl[0]
bl = tuple(bl[1])
feed_no = True
elif isinstance(ts, RawTimestream): # for RawTimestream
pol = None
bl = tuple(bl)
if feed_no:
pol = ts['bl_pol'].local_data[li]
bl = tuple(ts['true_blorder'].local_data[li])
else:
raise ValueError('Need either a RawTimestream or Timestream')
if bl_incl != 'all':
bl1 = set(bl)
bl_incl = [{f1, f2} for (f1, f2) in bl_incl]
bl_excl = [{f1, f2} for (f1, f2) in bl_excl]
if (not bl1 in bl_incl) or (bl1 in bl_excl):
return
if flag_mask:
vis1 = np.ma.array(vis, mask=vis_mask)
elif flag_ns:
if 'ns_on' in ts.iterkeys():
vis1 = vis.copy()
on = np.where(ts['ns_on'][:])[0]
vis1[on] = complex(np.nan, np.nan)
else:
vis1 = vis
else:
vis1 = vis
if integral == 'time':
vis1 = np.ma.mean(np.ma.masked_invalid(vis1), axis=0)
ax_val = ts.freq[:]
xlabel = r'$\nu$ / MHz'
elif integral == 'freq':
vis1 = np.ma.mean(np.ma.masked_invalid(vis1), axis=1)
# ax_val = ts.time[:]
# xlabel = r'$t$ / Julian Date'
ax_val = np.array(
[datetime.fromtimestamp(sec) for sec in ts['sec1970'][:]])
xlabel = '%s' % ax_val[0].date()
ax_val = mdates.date2num(ax_val)
else:
raise ValueError('Unknown integral type %s' % integral)
plt.figure()
f, axarr = plt.subplots(3, sharex=True)
axarr[0].plot(ax_val, vis1.real, label='real')
axarr[0].legend()
axarr[1].plot(ax_val, vis1.imag, label='imag')
axarr[1].legend()
axarr[2].plot(ax_val, np.abs(vis1), label='abs')
axarr[2].legend()
axarr[2].xaxis_date()
date_format = mdates.DateFormatter('%H:%M')
axarr[2].xaxis.set_major_formatter(date_format)
if rotate_xdate:
# set the x-axis tick labels to diagonal so it fits better
f.autofmt_xdate()
else:
# reduce the number of tick locators
locator = MaxNLocator(nbins=6)
axarr[2].xaxis.set_major_locator(locator)
axarr[2].xaxis.set_minor_locator(AutoMinorLocator(2))
axarr[2].set_xlabel(xlabel)
if feed_no:
fig_name = '%s_%s_%d_%d_%s.png' % (fig_prefix, integral, bl[0],
bl[1], ts.pol_dict[pol])
else:
fig_name = '%s_%s_%d_%d.png' % (fig_prefix, integral, bl[0], bl[1])
if tag_output_iter:
fig_name = output_path(fig_name, iteration=iteration)
else:
fig_name = output_path(fig_name)
plt.savefig(fig_name)
plt.close()
def process(self, ts):
calibrator = self.params['calibrator']
catalog = self.params['catalog']
span = self.params['span']
save_gain = self.params['save_gain']
gain_file = self.params['gain_file']
ts.redistribute('frequency')
lfreq = ts.local_freq[:] # local freq
feedno = ts['feedno'][:].tolist()
pol = ts['pol'][:].tolist()
bl = ts.bl[:]
bls = [tuple(b) for b in bl]
# # antpointing = np.radians(ts['antpointing'][-1, :, :]) # radians
# transitsource = ts['transitsource'][:]
# transit_time = transitsource[-1, 0] # second, sec1970
# int_time = ts.attrs['inttime'] # second
# calibrator
srclist, cutoff, catalogs = a.scripting.parse_srcs(calibrator, catalog)
cat = a.src.get_catalog(srclist, cutoff, catalogs)
assert (len(cat) == 1), 'Allow only one calibrator'
s = cat.values()[0]
if mpiutil.rank0:
print 'Calibrating for source %s with' % calibrator,
print 'strength', s._jys, 'Jy',
print 'measured at', s.mfreq, 'GHz',
print 'with index', s.index
# get transit time of calibrator
# array
aa = ts.array
aa.set_jultime(ts['jul_date'][0]) # the first obs time point
next_transit = aa.next_transit(s)
transit_time = a.phs.ephem2juldate(next_transit) # Julian date
if transit_time > ts['jul_date'][-1]:
local_next_transit = ephem.Date(next_transit + 8.0 * ephem.hour)
raise RuntimeError(
'Data does not contain local transit time %s of source %s' %
(local_next_transit, calibrator))
# the first transit index
transit_inds = [np.searchsorted(ts['jul_date'][:], transit_time)]
# find all other transit indices
aa.set_jultime(ts['jul_date'][0] + 1.0)
transit_time = a.phs.ephem2juldate(aa.next_transit(s)) # Julian date
cnt = 2
while (transit_time <= ts['jul_date'][-1]):
transit_inds.append(
np.searchsorted(ts['jul_date'][:], transit_time))
aa.set_jultime(ts['jul_date'][0] + 1.0 * cnt)
transit_time = a.phs.ephem2juldate(
aa.next_transit(s)) # Julian date
cnt += 1
print transit_inds
### now only use the first transit point to do the cal
### may need to improve in the future
transit_ind = transit_inds[0]
int_time = ts.attrs['inttime'] # second
start_ind = transit_ind - np.int(span / int_time)
end_ind = transit_ind + np.int(span / int_time)
nt = end_ind - start_ind
nfeed = len(feedno)
eigval = np.empty((nt, nfeed, 2, len(lfreq)), dtype=np.float64)
eigval[:] = np.nan
gain = np.empty((nt, nfeed, 2, len(lfreq)), dtype=np.complex128)
gain[:] = complex(np.nan, np.nan)
# construct visibility matrix for a single time, pol, freq
Vmat = np.zeros((nfeed, nfeed), dtype=ts.main_data.dtype)
for ind, ti in enumerate(range(start_ind, end_ind)):
# when noise on, just pass
if 'ns_on' in ts.iterkeys() and ts['ns_on'][ti]:
continue
aa.set_jultime(ts['jul_date'][ti])
s.compute(aa)
# get fluxes vs. freq of the calibrator
Sc = s.get_jys()
# get the topocentric coordinate of the calibrator at the current time
s_top = s.get_crds('top', ncrd=3)
aa.sim_cache(cat.get_crds(
'eq', ncrd=3)) # for compute bm_response and sim
for pi in [pol.index('xx'), pol.index('yy')]: # xx, yy
aa.set_active_pol(pol[pi])
for fi, freq in enumerate(lfreq): # mpi among freq
for i, ai in enumerate(feedno):
for j, aj in enumerate(feedno):
# uij = aa.gen_uvw(i, j, src='z').squeeze() # (rj - ri)/lambda
uij = aa.gen_uvw(i, j,
src='z')[:,
0, :] # (rj - ri)/lambda
# bmij = aa.bm_response(i, j).squeeze() # will get error for only one local freq
# import pdb
# pdb.set_trace()
bmij = aa.bm_response(i, j).reshape(-1)
try:
bi = bls.index((ai, aj))
# Vmat[i, j] = ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(-2.0J * np.pi * np.dot(s_top, uij[:, fi]))) # xx, yy
Vmat[i, j] = ts.local_vis[ti, fi, pi, bi] / (
Sc[fi] * bmij[fi] *
np.exp(2.0J * np.pi * np.dot(
s_top, uij[:, fi]))) # xx, yy
except ValueError:
bi = bls.index((aj, ai))
# Vmat[i, j] = np.conj(ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(-2.0J * np.pi * np.dot(s_top, uij[:, fi])))) # xx, yy
Vmat[i, j] = np.conj(
ts.local_vis[ti, fi, pi, bi] /
(Sc[fi] * bmij[fi] *
np.exp(2.0J * np.pi * np.dot(
s_top, uij[:, fi])))) # xx, yy
# Eigen decomposition
Vmat = np.where(np.isfinite(Vmat), Vmat, 0)
e, U = eigh(Vmat)
eigval[ind, :, pi, fi] = e[::-1] # descending order
# max eigen-val
lbd = e[-1] # lambda
# the gain vector for this freq
gvec = np.sqrt(
lbd
) * U[:,
-1] # only eigen-vector corresponding to the maximum eigen-val
gain[ind, :, pi, fi] = gvec
# apply gain to vis
# get the time mean gain
tgain = np.ma.mean(np.ma.masked_invalid(gain), axis=0) # time mean
tgain = mpiutil.gather_array(tgain, axis=-1, root=None)
ts.redistribute('baseline')
ts.pol_and_bl_data_operate(cal, tgain=tgain)
# save gain if required:
if save_gain:
gain_file = output_path(gain_file)
eigval = mpiarray.MPIArray.wrap(eigval, axis=3)
gain = mpiarray.MPIArray.wrap(gain, axis=3)
mem_gain = memh5.MemGroup(distributed=True)
mem_gain.create_dataset('eigval', data=eigval)
mem_gain.create_dataset('gain', data=gain)
# add attris
mem_gain.attrs['jul_data'] = ts['jul_date'][start_ind:end_ind]
mem_gain.attrs['feed'] = np.array(feedno)
mem_gain.attrs['pol'] = np.array(['xx', 'yy'])
mem_gain.attrs['freq'] = ts.freq[:] # freq should be common
# save to file
mem_gain.to_hdf5(gain_file, hints=False)
return super(PsCal, self).process(ts)
def process(self, ts):
excl_auto = self.params['excl_auto']
plot_stats = self.params['plot_stats']
fig_prefix = self.params['fig_name']
rotate_xdate = self.params['rotate_xdate']
tag_output_iter = self.params['tag_output_iter']
ts.redistribute('baseline')
if ts.local_vis_mask.ndim == 3: # RawTimestream
if excl_auto:
bl = ts.local_bl
vis_mask = ts.local_vis_mask[:, :, bl[:, 0] != bl[:, 1]].copy()
else:
vis_mask = ts.local_vis_mask.copy()
nt, nf, lnb = vis_mask.shape
elif ts.local_vis_mask.ndim == 4: # Timestream
# suppose masks are the same for all 4 pols
if excl_auto:
bl = ts.local_bl
vis_mask = ts.local_vis_mask[:, :, 0,
bl[:, 0] != bl[:, 1]].copy()
else:
vis_mask = ts.local_vis_mask[:, :, 0].copy()
nt, nf, lnb = vis_mask.shape
else:
raise RuntimeError('Incorrect vis_mask shape %s' %
ts.local_vis_mask.shape)
# total number of bl
nb = mpiutil.allreduce(lnb, comm=ts.comm)
# un-mask ns-on positions
if 'ns_on' in ts.iterkeys():
vis_mask[ts['ns_on'][:]] = False
# statistics along time axis
time_mask = np.sum(vis_mask, axis=(1, 2)).reshape(-1, 1)
# gather local array to rank0
time_mask = mpiutil.gather_array(time_mask,
axis=1,
root=0,
comm=ts.comm)
if mpiutil.rank0:
time_mask = np.sum(time_mask, axis=1)
# statistics along time axis
freq_mask = np.sum(vis_mask, axis=(0, 2)).reshape(-1, 1)
# gather local array to rank0
freq_mask = mpiutil.gather_array(freq_mask,
axis=1,
root=0,
comm=ts.comm)
if mpiutil.rank0:
freq_mask = np.sum(freq_mask, axis=1)
if plot_stats and mpiutil.rank0:
time_fig_name = '%s_%s.png' % (fig_prefix, 'time')
if tag_output_iter:
time_fig_name = output_path(time_fig_name,
iteration=self.iteration)
else:
time_fig_name = output_path(time_fig_name)
# plot time_mask
plt.figure()
fig, ax = plt.subplots()
x_vals = np.array([
(datetime.utcfromtimestamp(s) + timedelta(hours=8))
for s in ts['sec1970'][:]
])
xlabel = '%s' % x_vals[0].date()
x_vals = mdates.date2num(x_vals)
ax.plot(x_vals, 100 * time_mask / np.float(nf * nb))
ax.xaxis_date()
date_format = mdates.DateFormatter('%H:%M')
ax.xaxis.set_major_formatter(date_format)
if rotate_xdate:
# set the x-axis tick labels to diagonal so it fits better
fig.autofmt_xdate()
else:
# reduce the number of tick locators
locator = MaxNLocator(nbins=6)
ax.xaxis.set_major_locator(locator)
ax.xaxis.set_minor_locator(AutoMinorLocator(2))
ax.set_xlabel(xlabel)
ax.set_ylabel(r'RFI (%)')
plt.savefig(time_fig_name)
plt.close()
freq_fig_name = '%s_%s.png' % (fig_prefix, 'freq')
if tag_output_iter:
freq_fig_name = output_path(freq_fig_name,
iteration=self.iteration)
else:
freq_fig_name = output_path(freq_fig_name)
# plot freq_mask
plt.figure()
plt.plot(ts.freq[:], 100 * freq_mask / np.float(nt * nb))
plt.grid(True)
plt.xlabel(r'$\nu$ / MHz')
plt.ylabel(r'RFI (%)')
plt.savefig(freq_fig_name)
plt.close()
return super(Stats, self).process(ts)