def svd_clean(data):
I = data.data[:, 0, :]
freq = np.ones(I.shape[0]).astype("bool")
corr, weight = find_modes.freq_covariance(I, I, None, None, freq, freq, no_weight=True, normalize=False)
svd_result = find_modes.get_freq_svd_modes(corr, corr.shape[0])
return svd_result
def test():
#sim_map = algebra.make_vect(algebra.load( './maps/pks/sim_raw_000.npy'))
sim_map_raw = algebra.make_vect(algebra.load( './maps/gbt/sim_raw_000.npy'))
#sim_map_wb = convolve_beam(sim_map, apply_bandpass=False,
# frequency_depended=False, plot_beam=True)
#exit()
sim_map = copy.deepcopy(sim_map_raw)
freq = sim_map.get_axis('freq')
ra = sim_map.get_axis('ra')
dec = sim_map.get_axis('dec')
print freq.max(), freq.min()
#plot_map(sim_map, filename='sim_raw.png')
ps_2d, kn_2d, ps_1d, kn_1d, k_bin = est_power(sim_map)
#plot_power([ps_1d,], [k_bin,], label_list=['raw',])
#exit()
pts_map, spec_list = add_pointsources(sim_map.shape, freq/1.e6, number=50)
sim_map += pts_map
pts_map_unresolved = add_pointsources_unresolved(sim_map.shape, freq/1.e6 )
sim_map += pts_map_unresolved
#plot_map(pts_map, filename='pts_map.png', freq=freq, ra=ra, dec=dec)
#plot_map(pts_map_unresolved, filename='pts_map_unresolved.png', freq=freq, ra=ra, dec=dec)
#plot_map(sim_map, filename='sim_plus_ps.png', freq=freq, ra=ra, dec=dec)
corr, weight = find_modes.freq_covariance( sim_map, sim_map, None, None,
range(freq.shape[0]), range(freq.shape[0]), no_weight=True)
svd = find_modes.get_freq_svd_modes(corr, 30)
#corr_norm, weight_norm = find_modes.freq_covariance( sim_map, sim_map, None, None,
# range(freq.shape[0]), range(freq.shape[0]), no_weight=True, normalize=True)
#svd_norm = find_modes.get_freq_svd_modes(corr_norm, 30)
# convolve beam
#sim_map_wb = convolve_beam(sim_map, apply_bandpass=True, frequency_depended=True)
#sim_map_wb = convolve_beam(sim_map, apply_bandpass=False, frequency_depended=True)
sim_map_wb = convolve_beam(copy.deepcopy(sim_map), apply_bandpass=True, frequency_depended=False)
#sim_map_wb = convolve_beam(sim_map, apply_bandpass=False, frequency_depended=False)
#plot_map(sim_map_wb, filename='sim_plus_ps_cov.png')
#log_sim_map_wb = np.log(np.abs(copy.deepcopy(sim_map_wb)))
#log_sim_map_wb[np.isnan(log_sim_map_wb)] = 0
#log_corr_wb, log_weight_wb = find_modes.freq_covariance(
# log_sim_map_wb, log_sim_map_wb, None, None,
# range(freq.shape[0]), range(freq.shape[0]), no_weight=True )
#log_svd_wb = find_modes.get_freq_svd_modes(log_corr_wb, 30)
#for i in range(5):
# sub_modes(log_sim_map_wb, log_svd_wb[1][i:i+1])
#lin_sim_map_wb = np.exp(log_sim_map_wb) * np.abs(sim_map_wb)/sim_map_wb
#lin_sim_map_wb[np.isnan(lin_sim_map_wb)] = 0.
#corr_wb, weight_wb = find_modes.freq_covariance( lin_sim_map_wb, lin_sim_map_wb, None, None,
corr_wb, weight_wb = find_modes.freq_covariance( sim_map_wb, sim_map_wb, None, None,
range(freq.shape[0]), range(freq.shape[0]), no_weight=True )
log_svd_wb = find_modes.get_freq_svd_modes(np.log(corr_wb), 30)
svd_wb = find_modes.get_freq_svd_modes(corr_wb, 30)
for i in range(3):
sub_modes(sim_map, svd[1][i:i+1])
#plot_map(sim_map, filename='sim_plot_ps_sub%02dmodes.png'%(i+1),
# freq=freq, ra=ra, dec=dec)
for i in range(3):
sub_modes(sim_map_wb, log_svd_wb[1][i:i+1])
#plot_map(sim_map_wb, filename='sim_plot_ps_cov_sub%02dmodes.png'%(i+1),
# freq=freq, ra=ra, dec=dec)
#plot_eig([log_svd_wb, svd_wb], mode_n=10)
#exit()
ps_2d_cln, kn_2d_cln, ps_1d_cln, kn_1d_cln, k_bin_cln = est_power(sim_map)
ps_2d_cln_wb, kn_2d_cln_wb, ps_1d_cln_wb, kn_1d_cln_wb, k_bin_cln_wb = est_power(sim_map_wb, sim_map_raw)
plot_power([ps_1d, ps_1d_cln, ps_1d_cln_wb],
[k_bin, k_bin_cln, k_bin_cln_wb],
label_list=['raw', 'cleaned', 'cleaned with beam'])
plt.show()
def svd_simulation(input_root, output_root='./', sim_num=100, fg_num=1, mode_num=4):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
for ii in range(sim_num):
# initialize the hdf5 output
sr = SimResult(output_root + 'RANK%02d_sim_raw_%03d.h5'%(rank, ii))
print "using simulation map: " + input_root + "sim_raw_%03d.npy"%ii
print "output simulation map: " + output_root + "sim_raw_%03d.h5"%ii
sys.stdout.flush()
# load the raw simulation maps, and init the freq, ra, and dec
sim_map_raw = algebra.make_vect(algebra.load(input_root + 'sim_raw_%03d.npy'%ii))
freq = sim_map_raw.get_axis('freq')
ra = sim_map_raw.get_axis('ra')
dec = sim_map_raw.get_axis('dec')
# estimate the ps for raw simulation map
ps_result_raw = est_power(sim_map_raw)
#[ps_2d_raw, kn_2d_raw, ps_1d_raw, kn_1d_raw, k_bin_raw] = ps_result_raw
# initialize the map information
sr.init_info(sim_map_raw)
# save the raw simulation map
if rank == 0:
print "output sim_map_raw: map and ps_result"
sr.add_map('/raw/sim/', sim_map_raw)
sr.add_ps('/raw/sim/', ps_result_raw)
# convolve the raw simulation map with common beam
sim_map_raw_beam = convolve_beam(copy.deepcopy(sim_map_raw),
apply_bandpass=False, frequency_depended=False)
ps_result_raw_beam_auto = est_power(sim_map_raw_beam)
ps_result_raw_beam_cros = est_power(sim_map_raw_beam, sim_map_raw)
if rank == 0:
print "output sim_map_raw_beam: map and ps_result"
sr.add_map('/beam/sim/', sim_map_raw_beam)
sr.add_ps('/beam/sim/', ps_result_raw_beam_auto, suffix='_auto')
sr.add_ps('/beam/sim/', ps_result_raw_beam_cros, suffix='_cros')
# convolve the raw simulation map with frequency depended beam
sim_map_raw_beamf = convolve_beam(copy.deepcopy(sim_map_raw),
apply_bandpass=False, frequency_depended=True)
ps_result_raw_beamf_auto = est_power(sim_map_raw_beamf)
ps_result_raw_beamf_cros = est_power(sim_map_raw_beamf, sim_map_raw)
if rank == 0:
print "output sim_map_raw_beamf: map and ps_result"
sr.add_map('/beamf/sim/', sim_map_raw_beamf)
sr.add_ps('/beamf/sim/', ps_result_raw_beamf_auto, suffix='_auto')
sr.add_ps('/beamf/sim/', ps_result_raw_beamf_cros, suffix='_cros')
# convolve the raw simulation map with beam * bandpass
sim_map_raw_beamb = convolve_beam(copy.deepcopy(sim_map_raw),
apply_bandpass=True, frequency_depended=False)
ps_result_raw_beamb_auto = est_power(sim_map_raw_beamb)
ps_result_raw_beamb_cros = est_power(sim_map_raw_beamb, sim_map_raw)
if rank == 0:
print "output sim_map_raw_beamb: map and ps_result"
sr.add_map('/beamb/sim/', sim_map_raw_beamb)
sr.add_ps('/beamb/sim/', ps_result_raw_beamb_auto, suffix='_auto')
sr.add_ps('/beamb/sim/', ps_result_raw_beamb_cros, suffix='_cros')
if rank == 0:
print 50*'='
sys.stdout.flush()
# add foreground , both resolved and unresolved
for jj in range(fg_num):
# make a copy of the simulation map
sim_map = copy.deepcopy(sim_map_raw)
pts_map, spec_list = add_pointsources(sim_map.shape, freq/1.e6, number=50)
sim_map += pts_map
pts_map_unresolved = add_pointsources_unresolved(sim_map.shape, freq/1.e6 )
sim_map += pts_map_unresolved
# save the simulatin maps with foreground added
if rank == 0:
print "output sim_map + fg: map"
sr.add_map('/raw/sim_fg/', sim_map, name_id=jj)
# convolved a beam of 1.1 times largest beam, which DO NOT have
# frequency evolution.
# this case is similar to our common beam convolved result.
sim_map_beam = convolve_beam(copy.deepcopy(sim_map),
apply_bandpass=False, frequency_depended=False)
if rank == 0:
print "output sim_map_beam + fg: map"
sr.add_map('/beam/sim_fg/', sim_map_beam, name_id=jj)
# convoled with a beam having frequency evolution
sim_map_beamf = convolve_beam(copy.deepcopy(sim_map),
apply_bandpass=False, frequency_depended=True)
if rank == 0:
print "output sim_map_beamf + fg: map"
sr.add_map('/beamf/sim_fg/', sim_map_beamf, name_id=jj)
# convoled with a beam, no frequency evolution, but have strange bandpass
sim_map_beamb = convolve_beam(copy.deepcopy(sim_map),
apply_bandpass=True, frequency_depended=False)
if rank == 0:
print "output sim_map_beamb + fg: map"
sr.add_map('/beamb/sim_fg/', sim_map_beamb, name_id=jj)
# remove the foreground
corr, weight = find_modes.freq_covariance( sim_map, sim_map, None, None,
range(freq.shape[0]), range(freq.shape[0]), no_weight=True)
svd = find_modes.get_freq_svd_modes(corr, 30)
if rank == 0:
print "output sim_map + fg: svd"
sr.add_svd('/raw/sim_fg_cln/', svd, name_id=jj)
for i in range(mode_num):
sub_modes(sim_map, svd[1][i:i+1])
ps_result_auto = est_power(sim_map)
ps_result_cros = est_power(sim_map, sim_map_raw)
if rank == 0:
print "output sim_map + fg: %02d modes map and ps_result"%i
sr.add_map('/raw/sim_fg_cln/', sim_map, name_id=jj, suffix='_%02dmode'%i)
sr.add_ps('/raw/sim_fg_cln/', ps_result_auto, name_id=jj, suffix='_auto_%02dmode'%i)
sr.add_ps('/raw/sim_fg_cln/', ps_result_cros, name_id=jj, suffix='_cros_%02dmode'%i)
corr_beam, weight_beam = find_modes.freq_covariance( sim_map_beam, sim_map_beam,
None, None, range(freq.shape[0]), range(freq.shape[0]), no_weight=True)
if rank == 0:
print "output sim_map_beam + fg: svd"
svd_beam = find_modes.get_freq_svd_modes(corr_beam, 30)
sr.add_svd('/beam/sim_fg_cln/', svd_beam, name_id=jj)
for i in range(mode_num):
sub_modes(sim_map_beam, svd_beam[1][i:i+1])
ps_result_beam_auto = est_power(sim_map_beam)
ps_result_beam_cros = est_power(sim_map_beam, sim_map_raw)
if rank == 0:
print "output sim_map_beam + fg: %02d modes map and ps_result"%i
sr.add_map('/beam/sim_fg_cln/', sim_map_beam, name_id=jj, suffix='_%02dmode'%i)
sr.add_ps('/beam/sim_fg_cln/', ps_result_beam_auto, name_id=jj, suffix='_auto_%02dmode'%i)
sr.add_ps('/beam/sim_fg_cln/', ps_result_beam_cros, name_id=jj, suffix='_cros_%02dmode'%i)
corr_beamf, weight_beamf = find_modes.freq_covariance( sim_map_beamf, sim_map_beamf,
None, None, range(freq.shape[0]), range(freq.shape[0]), no_weight=True)
if rank == 0:
print "output sim_map_beamf + fg: svd"
svd_beamf = find_modes.get_freq_svd_modes(corr_beamf, 30)
sr.add_svd('/beamf/sim_fg_cln/', svd_beamf, name_id=jj)
for i in range(mode_num):
sub_modes(sim_map_beamf, svd_beamf[1][i:i+1])
ps_result_beamf_auto = est_power(sim_map_beamf)
ps_result_beamf_cros = est_power(sim_map_beamf, sim_map_raw)
if rank == 0:
print "output sim_map_beamf + fg: %02d modes map and ps_result"%i
sr.add_map('/beamf/sim_fg_cln/', sim_map_beamf, name_id=jj, suffix='_%02dmode'%i)
sr.add_ps('/beamf/sim_fg_cln/', ps_result_beamf_auto, name_id=jj, suffix='_auto_%02dmode'%i)
sr.add_ps('/beamf/sim_fg_cln/', ps_result_beamf_cros, name_id=jj, suffix='_cros_%02dmode'%i)
corr_beamb, weight_beamb = find_modes.freq_covariance( sim_map_beamb, sim_map_beamb,
None, None, range(freq.shape[0]), range(freq.shape[0]), no_weight=True)
if rank == 0:
print "output sim_map_beamb + fg: svd"
svd_beamb = find_modes.get_freq_svd_modes(corr_beamb, 30)
sr.add_svd('/beamb/sim_fg_cln/', svd_beamb, name_id=jj)
for i in range(mode_num):
sub_modes(sim_map_beamb, svd_beamb[1][i:i+1])
ps_result_beamb_auto = est_power(sim_map_beamb)
ps_result_beamb_cros = est_power(sim_map_beamb, sim_map_raw)
if rank == 0:
print "output sim_map_beamb + fg: %02d modes map and ps_result"%i
sr.add_map('/beamb/sim_fg_cln/', sim_map_beamb, name_id=jj, suffix='_%02dmode'%i)
sr.add_ps('/beamb/sim_fg_cln/', ps_result_beamb_auto, name_id=jj, suffix='_auto_%02dmode'%i)
sr.add_ps('/beamb/sim_fg_cln/', ps_result_beamb_cros, name_id=jj, suffix='_cros_%02dmode'%i)
if rank == 0:
print 50*'-'
sys.stdout.flush()
comm.barrier()
if rank == 0:
print 50*'='
def corr_svd(data, params, file_ind, freq=None, time=None):
if not hasattr(data, "mask"):
data = np.ma.array(data)
data.mask = np.logical_or(data.mask, np.logical_not(np.isfinite(data)))
if np.all(data.mask[:, ::3, ...]):
msg = "WARNING: all data are masked, no svd performed"
warnings.warn(msg)
return data
# flag out the bad dots
# sigma = np.ma.var(data)
# mean = np.ma.mean(data)
# data.mask = np.logical_or(data.mask, data < mean - 5*sigma)
threshold = 4
bad_mask = np.any(data[:, ::3, ...] < threshold, axis=(1, 2))
data.mask = np.logical_or(data.mask, bad_mask[:, None, None, :])
mask_perc = np.ma.count_masked(data[:, ::3, ...]) / float(data[:, ::3, ...].size)
if mask_perc > 0.8:
msg = ("WARNING: %f%% data are masked, no svd performed") % (mask_perc * 100)
warnings.warn(msg)
return data
data[..., :100] = np.ma.masked
data[..., -100:] = np.ma.masked
data[..., 1640:1740] = np.ma.masked
data[..., 2066:2166] = np.ma.masked
# freq_mask = np.logical_not(np.any(np.all(data.mask, axis=0), axis=1))
time_mask = np.logical_not(np.any(np.all(data.mask, axis=-1), axis=(1, 2)))
freq_mask = np.logical_not(np.any(data.mask[time_mask, ...], axis=(0, 2)))
# data[:,0,:,np.logical_not(freq_mask[0,:])] = np.ma.masked
# data[:,1,:,np.logical_not(freq_mask[1,:])] = np.ma.masked
# data[:,2,:,np.logical_not(freq_mask[2,:])] = np.ma.masked
# data[:,3,:,np.logical_not(freq_mask[3,:])] = np.ma.masked
data.mask = np.logical_or(data.mask, np.logical_not(freq_mask[None, :, None, :]))
if not np.any(freq_mask[::3, :]):
msg = "WARNING: all freq channels are masked, no svd performed"
warnings.warn(msg)
return data
weights = np.ma.ones(data.shape)
weights[data.mask] = np.ma.masked
# subtract mean of each frequency
# data_mean = np.sum(data, axis=0) / np.sum(weights, axis=0)
# data_mean[data_mean==0] = np.inf
# data /= data_mean[None, ...]
# for XX
if params["save_svd"]:
map1_raw = copy.deepcopy(data[:, 0, 0, :].T)
map2_raw = copy.deepcopy(data[:, 0, 1, :].T)
if params["save_plot"]:
f_name = params["output_root"] + params["file_middles"][file_ind] + "_raw_XX.hdf5"
utils.mkparents(f_name)
check_map(f_name, data[:, 0, :, :], time, freq)
f_name = params["output_root"] + params["file_middles"][file_ind] + "_raw_YY.hdf5"
utils.mkparents(f_name)
check_map(f_name, data[:, 3, :, :], time, freq)
corr, weight = find_modes.freq_covariance(
data[:, 0, 0, :].T,
data[:, 0, 1, :].T,
weights[:, 0, 0, :].T,
weights[:, 0, 1, :].T,
freq_mask[0, :],
freq_mask[0, :],
no_weight=False,
)
svd_result = find_modes.get_freq_svd_modes(corr, corr.shape[0])
map1, map2, outmap_left, outmap_right = subtract_foregrounds(
svd_result,
data[:, 0, 0, :].T,
data[:, 0, 1, :].T,
weights[:, 0, 0, :].T,
weights[:, 0, 1, :].T,
freq_mask[0, :],
0,
params["modes"],
)
data[:, 0, 0, :] = map1.T
data[:, 0, 1, :] = map2.T
if params["save_svd"]:
f_name = params["output_root"] + params["file_middles"][file_ind] + "_svd_XX.hdf5"
utils.mkparents(f_name)
f = h5py.File(f_name, "w")
f["singular_values"] = svd_result[0]
f["left_vectors"] = svd_result[1]
f["right_vectors"] = svd_result[2]
f["outmap_left"] = outmap_left
f["outmap_right"] = outmap_right
f["map_left"] = map1
f["map_right"] = map2
f["raw_left"] = map1_raw
f["raw_right"] = map2_raw
f["freq_mask"] = freq_mask[0, :]
f["freq"] = freq
f["time"] = time
f.close()
if params["save_plot"]:
f_name = params["output_root"] + params["file_middles"][file_ind] + "_svd_XX.hdf5"
utils.mkparents(f_name)
check_svd(f_name, svd_result, freq_mask[0, :], freq)
check_map(f_name, np.ma.array(data[:, 0, :, :]), time, freq)
f_name = params["output_root"] + params["file_middles"][file_ind] + "_spec_XX.hdf5"
check_spec(f_name, np.ma.array(data[:, 0, :, :]), freq)
# f_name = params['output_root'] + \
# params['file_middles'][file_ind] + '_corr_XX.hdf5'
# check_corr(f_name, corr, weight)
del corr, weight, svd_result, map1, map2, outmap_left, outmap_right
gc.collect()
# for YY
if params["save_svd"]:
map1_raw = copy.deepcopy(data[:, 3, 0, :].T)
map2_raw = copy.deepcopy(data[:, 3, 1, :].T)
corr, weight = find_modes.freq_covariance(
data[:, 3, 0, :].T,
data[:, 3, 1, :].T,
weights[:, 3, 0, :].T,
weights[:, 3, 1, :].T,
freq_mask[3, :],
freq_mask[3, :],
no_weight=False,
)
svd_result = find_modes.get_freq_svd_modes(corr, corr.shape[0])
map1, map2, outmap_left, outmap_right = subtract_foregrounds(
svd_result,
data[:, 3, 0, :].T,
data[:, 3, 1, :].T,
weights[:, 3, 0, :].T,
weights[:, 3, 1, :].T,
freq_mask[3, :],
0,
params["modes"],
)
data[:, 3, 0, :] = map1.T
data[:, 3, 1, :] = map2.T
if params["save_svd"]:
f_name = params["output_root"] + params["file_middles"][file_ind] + "_svd_YY.hdf5"
utils.mkparents(f_name)
f = h5py.File(f_name, "w")
f["singular_values"] = svd_result[0]
f["left_vectors"] = svd_result[1]
f["right_vectors"] = svd_result[2]
f["outmap_left"] = outmap_left
f["outmap_right"] = outmap_right
f["map_left"] = map1
f["map_right"] = map2
f["raw_left"] = map1_raw
f["raw_right"] = map2_raw
f["freq_mask"] = freq_mask[3, :]
f["freq"] = freq
f["time"] = time
f.close()
if params["save_plot"]:
f_name = params["output_root"] + params["file_middles"][file_ind] + "_svd_YY.hdf5"
utils.mkparents(f_name)
check_svd(f_name, svd_result, freq_mask[3, :], freq)
check_map(f_name, np.ma.array(data[:, 3, :, :]), time, freq)
f_name = params["output_root"] + params["file_middles"][file_ind] + "_spec_YY.hdf5"
check_spec(f_name, np.ma.array(data[:, 3, :, :]), freq)
# f_name = params['output_root'] + \
# params['file_middles'][file_ind] + '_corr_XX.hdf5'
# check_corr(f_name, corr, weight)
del corr, weight, svd_result, map1, map2, outmap_left, outmap_right
gc.collect()
# data[data_mask] = np.inf
# data = np.ma.array(data)
# data[data_mask] = np.ma.masked
return data
