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()