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