def process(self, tstream):
bt = tstream.beamtransfer
tel = bt.telescope
mmode_dir = tstream.output_directory + '/mmodes'
count_file = mmode_dir + '/count.hdf5'
if os.path.exists(count_file):
# get count
with h5py.File(count_file, 'r') as f:
N = f['count'][:]
# normalize mmode
for mi in mpiutil.mpirange(tel.mmax + 1):
with h5py.File(tstream._mfile(mi), 'r+') as f:
f['/mmode'][:] /= N[:, np.newaxis, :]
mpiutil.barrier()
# delete the count file
if mpiutil.rank0:
os.remove(count_file)
else:
if mpiutil.rank0:
print 'Count file %s does not exist, do noting...' % count_file
# mpiutil.barrier()
return tstream
def generate(self, regen=False):
"""Perform the KL-transform for all m-modes and save the result.
Uses MPI to distribute the work (if available).
Parameters
----------
mlist : array_like, optional
Set of m's to calculate KL-modes for By default do all m-modes.
"""
if mpiutil.rank0:
st = time.time()
print "======== Starting KL calculation ========"
# Iterate list over MPI processes.
for mi in mpiutil.mpirange(self.telescope.mmax+1):
if os.path.exists(self._evfile % mi) and not regen:
print "m index %i. File: %s exists. Skipping..." % (mi, (self._evfile % mi))
continue
self.transform_save(mi)
# If we're part of an MPI run, synchronise here.
mpiutil.barrier()
if mpiutil.rank0:
et = time.time()
print "======== Ending KL calculation (time=%f) ========" % (et - st)
# Collect together the eigenvalues
self._collect()
def generate_mmodes_kl(self):
"""Generate the KL modes for the Timestream.
"""
kl = self.manager.kltransforms[self.klname]
# Iterate over local m's, project mode and save to disk.
for mi in mpiutil.mpirange(self.telescope.mmax + 1):
if os.path.exists(self._klfile(mi)):
print "File %s exists. Skipping..." % self._klfile(mi)
continue
svdm = self.mmode_svd(
mi) #.reshape(self.telescope.nfreq, 2*self.telescope.npairs)
#svdm = self.beamtransfer.project_vector_telescope_to_svd(mi, tm)
klm = kl.project_vector_svd_to_kl(mi,
svdm,
threshold=self.klthreshold)
with h5py.File(self._klfile(mi), 'w') as f:
f.create_dataset('mmode_kl', data=klm)
f.attrs['m'] = mi
mpiutil.barrier()
def generate(self, regen=False):
"""Perform the KL-transform for all m-modes and save the result.
Uses MPI to distribute the work (if available).
Parameters
----------
mlist : array_like, optional
Set of m's to calculate KL-modes for By default do all m-modes.
"""
if mpiutil.rank0:
st = time.time()
print("======== Starting KL calculation ========")
# Iterate list over MPI processes.
for mi in mpiutil.mpirange(self.telescope.mmax + 1):
if os.path.exists(self._evfile % mi) and not regen:
print(
"m index %i. File: %s exists. Skipping..."
% (mi, (self._evfile % mi))
)
continue
self.transform_save(mi)
# If we're part of an MPI run, synchronise here.
mpiutil.barrier()
if mpiutil.rank0:
et = time.time()
print("======== Ending KL calculation (time=%f) ========" % (et - st))
# Collect together the eigenvalues
self._collect()
def compute_Pkp(sky_map, kp):
N = kp.shape[0]
Nf, Np = sky_map.shape
Pkp = np.zeros(N)
for pi in mpiutil.mpirange(Np):
gkp = (cd_span / N) * ndft(cd, sky_map[:, pi], kp) # K (h Mpc^-1)^-1
# according to the definition <\delta(k) \delta(k')^*> = (2\pi)^3 \delta(k - k') P(k), but for 1D, <\delta(k) \delta(k')^*> = (2\pi) \delta(k - k') P(k)
Pkp += np.abs(gkp)**2 / (2.0 * np.pi) # K^2 (Mpc / h)
Pkp /= Np # K^2 (Mpc / h)
Pkp *= 1.0e6 # mK^2 (Mpc / h)
tmp = mpiutil.gather_array(Pkp.reshape(1, -1), axis=0, root=0)
if mpiutil.rank0:
print 'tmp shape: ', tmp.shape
Pkp = np.sum(tmp, axis=0)
print 'Pkp shape: ', Pkp.shape
return Pkp # only rank0 has the correct Pkp
def generate_mmodes_svd(self):
"""Generate the SVD modes for the Timestream."""
# Iterate over local m's, project mode and save to disk.
for mi in mpiutil.mpirange(self.telescope.mmax + 1):
if os.path.exists(self._svdfile(mi)):
print("File %s exists. Skipping..." % self._svdfile(mi))
continue
tm = self.mmode(mi).reshape(self.telescope.nfreq, 2 * self.telescope.npairs)
svdm = self.beamtransfer.project_vector_telescope_to_svd(mi, tm)
with h5py.File(self._svdfile(mi), "w") as f:
f.create_dataset("mmode_svd", data=svdm)
f.attrs["m"] = mi
mpiutil.barrier()
def process(self, input):
def _indx_f(x, shp):
if x >= np.prod(shp): return
_i = [
int(x / np.prod(shp[1:])),
]
for i in range(1, len(shp)):
x -= _i[i - 1] * np.prod(shp[i:])
_i += [
int(x / np.prod(shp[i + 1:])),
]
return tuple(_i)
diag_cov = self.params['diag_cov']
threshold = self.params['threshold']
task_n = np.prod(self.map_shp[:-2])
for task_ind in mpiutil.mpirange(task_n):
indx = _indx_f(task_ind, self.map_shp[:-2])
#print mpiutil.rank, indx
print "RANK%03d: (" % mpiutil.rank + ("%04d, " *
len(indx)) % indx + ")"
map_shp = self.map_shp[-2:]
_dirty_map = np.zeros(map_shp, dtype=__dtype__)
if diag_cov:
_cov_inv = np.zeros(map_shp, dtype=__dtype__)
else:
_cov_inv = np.zeros(map_shp * 2, dtype=__dtype__)
for ii, df in enumerate(self.df_in):
_dirty_map += df['dirty_map'][indx + (slice(None), )]
self.read_block_from_dset(ii, 'cov_inv', indx, _cov_inv)
#_cov_inv += df['cov_inv'][indx + (slice(None), )]
self.df_out[-1]['dirty_map'][indx + (slice(None), )] = _dirty_map
clean_map, noise_diag = make_cleanmap(_dirty_map, _cov_inv,
diag_cov, threshold)
self.df_out[-1]['clean_map'][indx + (slice(None), )] = clean_map
self.df_out[-1]['noise_diag'][indx + (slice(None), )] = noise_diag
del _cov_inv
gc.collect()
mpiutil.barrier()
def generate_mmodes_svd(self):
"""Generate the SVD modes for the Timestream.
"""
# Iterate over local m's, project mode and save to disk.
for mi in mpiutil.mpirange(self.telescope.mmax + 1):
if os.path.exists(self._svdfile(mi)):
print "File %s exists. Skipping..." % self._svdfile(mi)
continue
tm = self.mmode(mi).reshape(self.telescope.nfreq, 2*self.telescope.npairs)
svdm = self.beamtransfer.project_vector_telescope_to_svd(mi, tm)
with h5py.File(self._svdfile(mi), 'w') as f:
f.create_dataset('mmode_svd', data=svdm)
f.attrs['m'] = mi
mpiutil.barrier()
def fake_kl_data(self):
kl = self.manager.kltransforms[self.klname]
# Iterate over local m's, project mode and save to disk.
for mi in mpiutil.mpirange(self.telescope.mmax + 1):
evals = kl.evals_m(mi)
if evals is None:
klmode = np.array([], dtype=np.complex128)
else:
modeamp = ((evals + 1.0) / 2.0)**0.5
klmode = modeamp * (np.array([1.0, 1.0J]) * np.random.standard_normal((modeamp.shape[0], 2))).sum(axis=1)
with h5py.File(self._klfile(mi), 'w') as f:
f.create_dataset('mmode_kl', data=klmode)
f.attrs['m'] = mi
mpiutil.barrier()
def fake_kl_data(self):
kl = self.manager.kltransforms[self.klname]
# Iterate over local m's, project mode and save to disk.
for mi in mpiutil.mpirange(self.telescope.mmax + 1, method='rand'):
evals = kl.evals_m(mi)
if evals is None:
klmode = np.array([], dtype=np.complex128)
else:
modeamp = ((evals + 1.0) / 2.0)**0.5
klmode = modeamp * (np.array([1.0, 1.0J]) * np.random.standard_normal((modeamp.shape[0], 2))).sum(axis=1)
with h5py.File(self._klfile(mi), 'w') as f:
f.create_dataset('mmode_kl', data=klmode)
f.attrs['m'] = mi
mpiutil.barrier()
def generate_mmodes_kl(self):
"""Generate the KL modes for the Timestream.
"""
kl = self.manager.kltransforms[self.klname]
# Iterate over local m's, project mode and save to disk.
for mi in mpiutil.mpirange(self.telescope.mmax + 1):
if os.path.exists(self._klfile(mi)):
print "File %s exists. Skipping..." % self._klfile(mi)
continue
svdm = self.mmode_svd(mi) #.reshape(self.telescope.nfreq, 2*self.telescope.npairs)
#svdm = self.beamtransfer.project_vector_telescope_to_svd(mi, tm)
klm = kl.project_vector_svd_to_kl(mi, svdm, threshold=self.klthreshold)
with h5py.File(self._klfile(mi), 'w') as f:
f.create_dataset('mmode_kl', data=klm)
f.attrs['m'] = mi
mpiutil.barrier()
def mpi(self):
rank = mpiutil.rank
size = mpiutil.size
if rank == 0:
#g = free_free(v = self.v, nside = self.nside)
#delt_m, params = g.delta_m()
delt_m = np.ones(np.int(12 * self.nside**2))
params = [1, 2, 3, 4, 5, 6, 7]
else:
delt_m = None
#local_delt_m = mpiutil.mpilist(delt_m, method = 'con',comm = MPI.COMM_WORLD)
local_range = mpiutil.mpirange(0, hp.nside2npix(self.nside))
delt_m = mpiutil.bcast(delt_m, root=0)
params = mpiutil.bcast(params, root=0)
result = []
for pix_number in local_range:
a = time.time()
l, b = hp.pix2ang(self.nside, pix_number, nest=False, lonlat=True)
#pix_value = delt_m[pix_number] + 100.
pix_value = self.integrate_by_hand(self._new,
0.01,
dist,
args=(l, b, delt_m[pix_number],
params))
b = time.time()
print 'pix_number', pix_number, 'delta_time', b - a
result.append([pix_number, pix_value])
result = mpiutil.gather_list(result, root=None)
if rank == 0:
with h5py.File('text_0to1.hdf5', 'w') as f:
f.create_dataset('result', data=result)
print 'end'
def setup(self):
ant_file = self.params['ant_file']
#ant_dat = np.genfromtxt(ant_file,
# dtype=[('name', 'S4'), ('X', 'f8'), ('Y', 'f8'), ('Z', 'f8')])
ant_dat = pd.read_fwf(ant_file,
header=None,
names=['name', 'X', 'Y', 'Z', 'px', 'py'])
self.ants = np.array(ant_dat['name'], dtype='str')
ants_pos = [
np.array(ant_dat['X'])[:, None],
np.array(ant_dat['Y'])[:, None],
np.array(ant_dat['Z'])[:, None]
]
self.ants_pos = np.concatenate(ants_pos, axis=1)
freq = self.params['freq']
dfreq = freq[1] - freq[0]
freq_n = freq.shape[0]
self.SM = globals()[self.params['survey_mode']](
self.params['schedule_file'])
#self.SM.generate_altaz(startalt, startaz, starttime, obs_len, obs_speed, obs_int)
self.SM.generate_altaz()
self.SM.radec_list([ant_dat['px'], ant_dat['py']])
#starttime = Time(self.params['starttime'])
#startalt, startaz = self.params['startpointing']
#startalt *= u.deg
#startaz *= u.deg
#obs_speed = self.params['obs_speed']
obs_int = self.SM.obs_int #self.params['obs_int']
self.obs_int = obs_int
samplerate = ((1. / obs_int).to(u.Hz)).value
#obs_tot = self.SM.obs_tot # self.params['obs_tot']
#obs_len = int((obs_tot / obs_int).decompose().value)
#self.block_time = self.SM.sche['block_time'] #self.params['block_time']
self.block_time = np.array(self.SM.sche['block_time'])
#self.block_len = int((block_time / obs_int).decompose().value)
block_num = self.block_time.shape[0]
_obs_int = (obs_int.to(u.second)).value
self._RMS = self.params['T_rec'] / np.sqrt(_obs_int * dfreq * 1.e6)
if self.params['fg_syn_model'] is not None:
self.syn_model = hp.read_map(self.params['fg_syn_model'],
range(freq.shape[0]))
self.syn_model = self.syn_model.T
if self.params['HI_model'] is not None:
with h5py.File(self.params['HI_model'], 'r') as fhi:
#_HI_model = al.make_vect(
_HI_model = al.load_h5(fhi, self.params['HI_model_type'])
logger.info('HI bias %3.2f' % self.params['HI_bias'])
_HI_model *= self.params['HI_bias']
if self.params['HI_mock_ids'] is not None:
self.HI_mock_ids = list(self.params['HI_mock_ids'])
_HI_model = _HI_model[self.HI_mock_ids]
else:
self.HI_mock_ids = range(_HI_model.shape[0])
self.mock_n = _HI_model.shape[0]
self.HI_model = al.make_vect(_HI_model)
else:
self.mock_n = self.params['mock_n']
if self.params['fnoise']:
self.FN = fnoise.FNoise(dtime=obs_int.value,
dfreq=dfreq,
alpha=self.params['alpha'],
f0=self.params['f0'],
beta=self.params['beta'])
self.get_blorder()
#self.iter_list = mpiutil.mpirange(0, obs_len, self.block_len)
self.iter_list = mpiutil.mpirange(0, block_num)
self.iter = 0
self.iter_num = len(self.iter_list)
def iterpstasks(self, input):
refinement = self.params['refinement']
task_list = self.task_list
for task_ind in mpiutil.mpirange(len(task_list)):
tind_l, tind_r, tind_o = task_list[task_ind]
tind_l = tuple(tind_l)
tind_r = tuple(tind_r)
tind_o = tuple(tind_o)
msg = ("RANK %03d est. ps.(" + "%03d,"*len(tind_l) + ") x ("\
+ "%03d,"*len(tind_r) + ")")%((mpiutil.rank, ) + tind_l + tind_r)
logger.info(msg)
cube = []
cube_w = []
tind_list = [tind_l, tind_r]
for i in range(2):
tind = tind_list[i]
map_key = self.params['map_key'][i]
input_map = input[tind[0]][map_key][tind[1:] + (slice(None), )]
input_map_mask = ~np.isfinite(input_map)
if (map_key is not 'delta') and (len(self.params['freq_mask']) != 0):
# ignore freqency mask for optical data
logger.info('apply freq_mask')
input_map_mask[self.params['freq_mask'], ...] = True
#input_map_mask += input_map == 0.
input_map[input_map_mask] = 0.
if self.params['prewhite']:
input_map_mask = input_map == 0.
_mean = np.ma.mean(np.ma.masked_equal(input_map, 0), axis=(1, 2))
input_map -= _mean[:, None, None]
input_map[input_map_mask] = 0.
input_map = al.make_vect(input_map, axis_names = ['freq', 'ra', 'dec'])
for key in input_map.info['axes']:
input_map.set_axis_info(key,
self.map_info[key+'_centre'],
self.map_info[key+'_delta'])
weight_key = self.params['weight_key'][i]
if weight_key is not None:
weight = input[tind[0]][weight_key][tind[1:] + (slice(None), )]
weight[input_map_mask] = 0.
if weight_key == 'noise_diag':
weight = fgrm.make_noise_factorizable(weight)
if weight_key == 'separable':
logger.debug('apply FKP weight')
weight = weight / (1. + weight * self.params['FKP'])
weight = al.make_vect(weight, axis_names = ['freq', 'ra', 'dec'])
weight.info = input_map.info
if not self.params['cube_input'][i]:
c, c_info = physical_grid(input_map, refinement=refinement,order=0)
else:
logger.debug('cube input')
c = input_map
c_info = None
cube.append(c)
if weight_key is not None:
if not self.params['cube_input'][i]:
cw, cw_info = physical_grid(weight, refinement=refinement,order=0)
else:
cw = weight
cw_info = None
#cw[c==0] = 0.
cube_w.append(cw)
del weight
else:
cw = al.ones_like(c)
cw[c==0] = 0.
cube_w.append(cw)
del c, c_info, cw, input_map
if tind_l == tind_r:
cube.append(cube[0])
cube_w.append(cube_w[0])
break
yield tind_o, cube, cube_w
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
g_abs = np.random.standard_normal(m.telescope.nfeed) * sigma_abs
g = (1.0 + g_abs) * np.exp(1.0j * g_phase)
gmat = np.outer(g, g.conj())
wmask = np.where(np.logical_and(t.feedmask, np.logical_not(t.feedconj)))
bf0 = np.bincount(t.feedmap[wmask])
bf1r = np.bincount(t.feedmap[wmask], weights=gmat[wmask].real)
bf1i = np.bincount(t.feedmap[wmask], weights=gmat[wmask].imag)
ug = (bf1r + 1.0j * bf1i) / bf0
return ug
# Distribute over frequencies to apply gain fluctuations.
for fi in mpiutil.mpirange(t.nfreq):
tsf = h5py.File(ts._ffile(fi))
# Generate gain fluctuations for this frequency
gain_fluc = mk_gain_fluctuation(sigma_g_abs, sigma_g_phase)
# Apply to each visibility
for ti in range(ts.ntime):
tsf["timestream"][:, ti] *= gain_fluc
tsf.close()
def setup(self):
self.refinement = self.params['refinement']
self.scenario = self.params['scenario']
map_pad = self.params['map_pad']
if self.params['map_tmp'] is None:
freq = self.params['freq'] * 1.e6
freq_d = freq[1] - freq[0]
freq_n = freq.shape[0]
freq_c = freq[freq_n // 2]
field_centre = self.params['field_centre']
spacing = self.params['pixel_spacing']
dec_spacing = spacing
ra_spacing = -spacing / np.cos(field_centre[1] * np.pi / 180.)
axis_names = ['freq', 'ra', 'dec']
map_shp = [x + map_pad for x in self.params['map_shape']]
map_tmp = np.zeros([
freq_n,
] + map_shp)
map_tmp = al.make_vect(map_tmp, axis_names=axis_names)
map_tmp.set_axis_info('freq', freq_c, freq_d)
map_tmp.set_axis_info('ra', field_centre[0], ra_spacing)
map_tmp.set_axis_info('dec', field_centre[1], dec_spacing)
self.map_tmp = map_tmp
else:
pad_shp = ((0, 0), (map_pad, map_pad), (map_pad, map_pad))
with h5py.File(self.params['map_tmp'], 'r') as f:
_map_tmp = al.load_h5(f, self.params['map_tmp_key'])
_axis_names = _map_tmp.info['axes']
_info = _map_tmp.info
_map_tmp = np.pad(_map_tmp, pad_shp, 'constant')
_map_tmp = al.make_vect(_map_tmp, axis_names=_axis_names)
_map_tmp.info.update(_info)
_weight = al.load_h5(f, self.params['map_tmp_weight'])
_weight = np.pad(_weight, pad_shp, 'constant')
_weight = al.make_vect(_weight, axis_names=_axis_names)
#self.map_tmp = al.zeros_like(_map_tmp)
self.map_tmp = _map_tmp
self.weight = _weight
# here we use 300 h km/s from WiggleZ for streaming dispersion
self.streaming_dispersion = 300. * 0.72
self.map_pad = map_pad
#self.beam_data = np.array([1., 1., 1.])
#self.beam_freq = np.array([900, 1100, 1400]) #* 1.e6
if self.params['beam_file'] is not None:
_bd = np.loadtxt(self.params['beam_file'])
self.beam_freq = _bd[:, 0] * 1.e6
self.beam_data = _bd[:, 1]
else:
fwhm1400 = 0.9
self.beam_freq = np.linspace(800., 1600., 500).astype('float')
self.beam_data = 1.2 * fwhm1400 * 1400. / self.beam_freq
self.beam_freq *= 1.e6
random.seed(3936650408)
seeds = random.random_integers(100000000, 1000000000, mpiutil.size)
self.seed = seeds[mpiutil.rank]
print "RANK: %02d with random seed [%d]" % (mpiutil.rank, self.seed)
random.seed(self.seed)
self.outfiles = self.params['outfiles']
self.outfiles_split = self.params['outfiles_split']
self.open_outputfiles()
self.iter_list = mpiutil.mpirange(self.params['mock_n'])
self.iter = 0
self.iter_num = len(self.iter_list)
def process(self, input):
task_list = self.task_list
for task_ind in mpiutil.mpirange(len(task_list)):
tind_l, tind_r, tind_o = task_list[task_ind]
tind_l = tuple(tind_l)
tind_r = tuple(tind_r)
tind_o = tind_o
print ("RANK %03d fgrm.\n(" + "%03d,"*len(tind_l) + ") x ("\
+ "%03d,"*len(tind_r) + ")\n")%((mpiutil.rank, ) + tind_l + tind_r)
tind_list = [tind_l, tind_r]
maps = []
weights = []
freq_good = np.ones(self.dset_shp[0]).astype('bool')
if len(self.params['freq_mask']) != 0:
freq_good[self.params['freq_mask']] = False
for i in range(2):
tind = tind_list[i]
map_key = self.params['map_key'] #'clean_map'
input_map = al.load_h5(input[tind[0]], map_key)
input_map = al.make_vect(input_map,
axis_names=['freq', 'ra', 'dec'])
maps.append(input_map)
weight_key = self.params['weight_key'] #'noise_diag'
if weight_key is not None:
weight = al.load_h5(input[tind[0]], weight_key)
if weight_key is 'noise_diag':
weight_prior = self.params['weight_prior']
logger.info('using wp %e' % weight_prior)
weight = make_noise_factorizable(weight, weight_prior)
else:
weight = np.ones_like(input_map)
weight[input_map == 0] = 0.
weight = al.make_vect(weight, axis_names=['freq', 'ra', 'dec'])
weight.info = input_map.info
try:
freq_good *= ~(input[tind[0]]['mask'][:]).astype('bool')
except KeyError:
logger.info('mask doesn\' exist')
pass
weights.append(weight)
maps[0][~freq_good] = 0.
maps[1][~freq_good] = 0.
weights[0][~freq_good] = 0.
weights[1][~freq_good] = 0.
if self.params['conv_factor'] != 0:
maps, weights = degrade_resolution(
maps,
weights,
conv_factor=self.params['conv_factor'],
mode='constant',
beam_file=self.params['beam_file'],
fwhm1400=self.params['fwhm1400'])
else:
logger.info('common reso. conv. ignored')
if self.params['add_map'] is not None:
_maps = self.params['add_map']
_map_A_path, _map_A_name = os.path.split(
os.path.splitext(_maps[0])[0])
_map_B_path, _map_B_name = os.path.split(
os.path.splitext(_maps[1])[0])
logger.info('add real map pair (%s %s)' %
(_map_A_name, _map_B_name))
with h5.File(os.path.join(_map_A_path, _map_A_name + '.h5'),
'r') as f:
maps[0][:] += al.load_h5(f,
'cleaned_00mode/%s' % _map_B_name)
with h5.File(os.path.join(_map_B_path, _map_B_name + '.h5'),
'r') as f:
maps[1][:] += al.load_h5(f,
'cleaned_00mode/%s' % _map_A_name)
svd_info = self.svd_info
if svd_info is None:
freq_cov, counts = find_modes.freq_covariance(
maps[0], maps[1], weights[0], weights[1], freq_good,
freq_good)
svd_info = find_modes.get_freq_svd_modes(
freq_cov, np.sum(freq_good))
mode_list = self.mode_list
mode_list_ed = copy.deepcopy(mode_list)
mode_list_st = copy.deepcopy(mode_list)
mode_list_st[1:] = mode_list_st[:-1]
dset_key = tind_o[0] + '_sigvalu'
self.df_out[tind_l[0]][dset_key] = svd_info[0]
dset_key = tind_o[0] + '_sigvect'
self.df_out[tind_l[0]][dset_key] = svd_info[1]
self.df_out[tind_l[0]]['weight'][:] = weights[0]
self.df_out[tind_l[0]]['mask'][:] = (~freq_good).astype('int')
if tind_o[1] != tind_o[0]:
dset_key = tind_o[1] + '_sigvalu'
self.df_out[tind_r[0]][dset_key] = svd_info[0]
dset_key = tind_o[1] + '_sigvect'
self.df_out[tind_r[0]][dset_key] = svd_info[2]
self.df_out[tind_r[0]]['weight'][:] = weights[1]
self.df_out[tind_r[0]]['mask'][:] = (~freq_good).astype('int')
for (n_modes_st, n_modes_ed) in zip(mode_list_st, mode_list_ed):
svd_modes = svd_info[1][n_modes_st:n_modes_ed]
group_name = 'cleaned_%02dmode/' % n_modes_ed
maps[0], amp = find_modes.subtract_frequency_modes(
maps[0], svd_modes, weights[0], freq_good)
dset_key = group_name + tind_o[0]
self.df_out[tind_l[0]][dset_key][:] = copy.deepcopy(maps[0])
if tind_o[0] != tind_o[1]:
svd_modes = svd_info[2][n_modes_st:n_modes_ed]
maps[1], amp = find_modes.subtract_frequency_modes(
maps[1], svd_modes, weights[1], freq_good)
dset_key = group_name + tind_o[1]
self.df_out[tind_r[0]][dset_key][:] = copy.deepcopy(
maps[1])
# for the case of auto with different svd svd modes
if 'Combined' in self.df_out[tind_r[0]][group_name].keys():
dset_key = group_name + 'Combined'
_map = maps[0].copy() * weights[0].copy()\
+ maps[1].copy() * weights[1].copy()
_wet = weights[0].copy() + weights[1].copy()
_wet[_wet == 0] = np.inf
_map /= _wet
self.df_out[tind_r[0]][dset_key][:] = copy.deepcopy(_map)
if self.params['output_combined'] is not None:
self.combine_results()
for ii in range(self.input_files_num):
input[ii].close()
print '-' * 35 + ' ' + tag + ' ' + '-' * 35
# mpilist
separator(sec, 'mpilist')
full_list = [1, 2.5, 'a', True, (3, 4), {'x': 1}]
local_list = mpiutil.mpilist(full_list)
print "rank %d has %s with method = 'con'" % (rank, local_list)
local_list = mpiutil.mpilist(full_list, method='alt')
print "rank %d has %s with method = 'alt'" % (rank, local_list)
local_list = mpiutil.mpilist(full_list, method='rand')
print "rank %d has %s with method = 'rand'" % (rank, local_list)
# mpirange
separator(sec, 'mpirange')
local_ary = mpiutil.mpirange(1, 7)
print "rank %d has %s with method = 'con'" % (rank, local_ary)
local_ary = mpiutil.mpirange(1, 7, method='alt')
print "rank %d has %s with method = 'alt'" % (rank, local_ary)
local_ary = mpiutil.mpirange(1, 7, method='rand')
print "rank %d has %s with method = 'rand'" % (rank, local_ary)
# bcast
separator(sec, 'bcast')
if rank == 0:
sendobj = 'obj'
else:
sendobj = None
sendobj = mpiutil.bcast(sendobj, root=0)
print 'rank %d has sendobj = %s after bcast' % (rank, sendobj)
g = (1.0 + g_abs) * np.exp(1.0J * g_phase) gmat = np.outer(g, g.conj()) wmask = np.where(np.logical_and(t.feedmask, np.logical_not(t.feedconj))) bf0 = np.bincount(t.feedmap[wmask]) bf1r = np.bincount(t.feedmap[wmask], weights=gmat[wmask].real) bf1i = np.bincount(t.feedmap[wmask], weights=gmat[wmask].imag) ug = (bf1r + 1.0J * bf1i) / bf0 return ug # Distribute over frequencies to apply gain fluctuations. for fi in mpiutil.mpirange(t.nfreq): tsf = h5py.File(ts._ffile(fi)) # Generate gain fluctuations for this frequency gain_fluc = mk_gain_fluctuation(sigma_g_abs, sigma_g_phase) # Apply to each visibility for ti in range(ts.ntime): tsf['timestream'][:, ti] *= gain_fluc tsf.close()
def mpi(self):
rank = mpiutil.rank
size = mpiutil.size
if rank == 0:
g = free_free(v = self.v, nside = self.nside,index_type = self.index_type,dist = self.dist,emi_form = self.emi_form,I_E_form = self.I_E_form,R0_R1_equal = self.R0_R1_equal,using_raw_diffuse = self.using_raw_diffuse,only_fit_Anu = self.only_fit_Anu)
delt_m, params = g.delta_m()
else:
delt_m = None
params = None
#local_delt_m = mpiutil.mpilist(delt_m, method = 'con',comm = MPI.COMM_WORLD)
local_range = mpiutil.mpirange(0,hp.nside2npix(self.nside))
delt_m = mpiutil.bcast(delt_m, root = 0)
params = mpiutil.bcast(params, root = 0)
result_absorb = []
for pix_number in local_range:
a = time.time()
l, b = hp.pix2ang(self.nside, pix_number, nest = False, lonlat = True)
if self.test == True:
pix_value =self.integrate_by_hand(self._new, 0.1, dist, args=(l, b, delt_m[pix_number], params))
else:
pix_value =self.integrate_by_hand(self._new, 0.01, dist, args=(l, b, delt_m[pix_number], params))
distance = self.critical_distance(l,b,delt_m[pix_number],params)
l, b = hp.pix2ang(self.nside, pix_number, nest = False, lonlat = True)
b = time.time()
if self.test == True:
result_absorb.append([pix_number, pix_value])
else:
result_absorb.append([pix_number, pix_value, distance])
if self.test == True:
result_absorb = mpiutil.gather_list(result_absorb, root = None)
else:
result_absorb = mpiutil.gather_list(result_absorb, root = None)
if rank == 0:
if self.test == True:
with h5py.File('./' + str(self.emi_form)+str(self.v) + 'F2py_absorb.hdf5', 'w') as f:
f.create_dataset('F2py_absorb', data = result_absorb)
else:
with h5py.File('./' + 'exp'+str(self.v)+'Mhz_delt_m_and_unabsorb_and_delt_m_percentage.hdf5','r') as f:
#print f.keys()
unabsorb = f['integrated_temperature_total_m'][:]
diffuse_raw = f['diffuse_raw'][:]
result_absorb = np.array(result_absorb)
absorb = result_absorb[:,1]
I_E = self.I_E(self.v)
result = []
print ('in the beginning')
for pix_number in range(unabsorb.size):
print ('left number of pixel',unabsorb.size - pix_number)
X = unabsorb[pix_number] - self.I_E(self.v)
l, b = hp.pix2ang(self.nside, pix_number, nest = False, lonlat = True)
tao = self.Fortran2Py_optical_deepth(l, b)
Y = absorb[pix_number] - I_E * np.exp(-tao[-1])
mean_exptao = Y / X
pix_value = diffuse_raw[pix_number] * mean_exptao + I_E*np.exp(-tao[-1])
result.append([pix_number,pix_value])
#print 'pixel_number',pix_number
with h5py.File('./' + str(self.emi_form)+str(self.v)+'MHz_global_spectrum_with_perterbation.hdf5','w') as h:
h.create_dataset('result',data = result)
#h.create_dataset('smooth_result',data = result_absorb)
print ('end, good job!, you are the best')
return result
def process(self, input):
def _indx_f(x, shp):
if x >= np.prod(shp): return
_i = [
int(x / np.prod(shp[1:])),
]
for i in range(1, len(shp)):
x -= _i[i - 1] * np.prod(shp[i:])
_i += [
int(x / np.prod(shp[i + 1:])),
]
return tuple(_i)
threshold = self.params['threshold']
loop_n = np.prod(self.map_shp[:-2])
block_length = self.params['block_length']
block_olap = self.params['block_overlap']
ra_length_tot = self.map_shp[-2]
dec_length_tot = self.map_shp[-1]
task_n = int(ra_length_tot / block_length) + 1
if mpiutil.rank0:
logger.debug('RA split into %4d task with block length %4d' %
(task_n, block_length))
for task_ind in mpiutil.mpirange(task_n):
ra_st = task_ind * block_length
ra_ed = min((task_ind + 1) * block_length, ra_length_tot)
ra_st_olap = max(ra_st - block_olap, 0)
ra_ed_olap = min(ra_ed + block_olap, ra_length_tot)
ra_length = ra_ed - ra_st
ra_length_olap = ra_ed_olap - ra_st_olap
olap_lower = ra_st - ra_st_olap
logger.debug('RANK %03d: RA %4d - %4d' %
(mpiutil.rank, ra_st, ra_ed))
logger.debug('RANK %03d: RA olap %4d - %4d' %
(mpiutil.rank, ra_st_olap, ra_ed_olap))
radec_slice = (slice(ra_st, ra_ed), slice(None))
radec_slice_olap = (slice(ra_st_olap, ra_ed_olap), slice(None))
for loop_ind in xrange(loop_n):
indx = _indx_f(loop_ind, self.map_shp[:-2])
logger.debug('RANK %03d: Loop idx '%mpiutil.rank\
+ '%3d'*len(indx)%indx)
map_shp = (ra_length_olap, dec_length_tot)
_dirty_map = np.zeros(map_shp)
_cov_inv = np.zeros(map_shp * 2, dtype=float)
for df in self.df_in:
_dirty_map += df['dirty_map'][indx + radec_slice_olap]
_cov_inv += df['cov_inv'][indx + radec_slice_olap * 2]
clean_map, noise_diag = make_cleanmap(_dirty_map, _cov_inv,
threshold)
clean_map = clean_map[olap_lower:olap_lower + ra_length]
noise_diag = noise_diag[olap_lower:olap_lower + ra_length]
self.df_out[-1]['clean_map'][indx + radec_slice] = clean_map
self.df_out[-1]['noise_diag'][indx + radec_slice] = noise_diag
mpiutil.barrier()
