def get_cstat(data, statistic, comm=None): """ Compute a collective statistic across all ranks and return as float. Must be called by all ranks. """ #if isinstance(data, MeshSource): # data = data.compute().value if isinstance(data, RealField) or isinstance(data, ComplexField): data = data.value else: assert type(data) == np.ndarray if comm is None: from nbodykit import CurrentMPIComm comm = CurrentMPIComm.get() if statistic == 'min': return comm.allreduce(data.min(), op=MPI.MIN) elif statistic == 'max': return comm.allreduce(data.max(), op=MPI.MAX) elif statistic == 'mean': # compute the mean csum = comm.allreduce(data.sum(), op=MPI.SUM) csize = comm.allreduce(data.size, op=MPI.SUM) return csum / float(csize) elif statistic == 'rms': mean = get_cmean(data, comm=comm) rsum = comm.allreduce(((data - mean)**2).sum()) csize = comm.allreduce(data.size) rms = (rsum / float(csize))**0.5 return rms else: raise Exception("Invalid statistic %s" % statistic)
def get_galaxy_catalog_from_source_catalog( self, source_cat, rand_seed_for_galaxy_sampling=123): assert self.log10M_column in source_cat.columns #cat = deepcopy(source_cat) cat = copy(source_cat) comm = CurrentMPIComm.get() # For each halo draw a random number RAND b/w 0 and 1. # For each halo, compute prob to be a galaxy. # Keep only halos where RAND<=prob_gal, remove rest from catalog. # This is our galaxy catalog. # Draw random number b/w 0 and 1 rng = MPIRandomState(comm, seed=rand_seed_for_galaxy_sampling, size=cat.size, chunksize=100000) cat['RAND'] = rng.uniform(low=0.0, high=1.0, dtype='f8') #print(cat[self.log10M_column]) #cat['PROB_GAL'] = 0.0 #cat[self.log10M_column] cat['PROB_GAL'] = 0.5 * (1.0 + erf( (cat[self.log10M_column].compute() - self.log10Mmin) / self.sigma_log10M)) print('Nhalos:', cat.csize) cat = cat[cat['RAND'] <= cat['PROB_GAL']] print('Ngalaxies:', cat.csize) print('Galaxy mass: ', get_cstats_string(cat[self.log10M_column].compute())) return cat
def get_mesh(self): """ Get the model as a MeshSource object. """ comm = CurrentMPIComm.get() if comm.rank == 0: print('Read %s' % self.in_fname) # read mesh from disk if self.read_mode == 'velocity': # get rho (don't divide by mean) mesh = read_vel_from_bigfile(self.in_fname) elif self.read_mode == 'delta from rho': # compute fractional delta (taking out mean), assuming file has rho mesh = read_delta_from_rho_bigfile(self.in_fname) elif self.read_mode == 'delta from 1+delta': # compute delta, assuming file has 1+delta mesh = read_delta_from_1plusdelta_bigfile(self.in_fname) elif self.read_mode == 'delta from 2+delta': # compute delta, assuming file has 1+delta mesh = read_delta_from_2plusdelta_bigfile(self.in_fname) else: raise Exception('Invalid read_mode: %s' % str(self.read_mode)) # multiply mesh by rescale factor if self.rescale_factor not in [None, 1.0]: if comm.rank == 0: print('Apply rescale fac to %s: %s' % (self.name, str(self.rescale_factor))) if type(self.rescale_factor) in [float, np.float, np.float64]: mesh = FieldMesh( mesh.compute(mode='real') * self.rescale_factor) else: print(type(self.rescale_factor)) raise Exception("Invalid rescale factor: %s" % str(self.rescale_factor)) # apply filters to mesh for filter_fcn in self.filters: # make copy mesh2 = FieldMesh(mesh.compute()) # apply filter mesh2 = mesh2.apply(filter_fcn, kind='wavenumber', mode='complex') # compute and save in mesh mesh = FieldMesh(mesh2.compute(mode='real')) cstats = get_cstats_string(mesh.compute()) if comm.rank == 0: print('MESH %s: %s\n' % (self.name, cstats)) return mesh
def print_cstats(data, prefix="", logger=None, comm=None): """ Must be called by all ranks. """ if comm is None: comm = CurrentMPIComm.get() if logger is None: logger = logging.getLogger("utils") cstats = get_cstats_string(data, comm) if comm.rank == 0: logger.info('%s%s' % (prefix, cstats)) print('%s%s' % (prefix, cstats))
def __new__(cls, *args, **kwargs): obj = object.__new__(cls) # ensure self.comm is set, though usually already set by the child. obj.comm = kwargs.get('comm', CurrentMPIComm.get()) # user-provided overrides and defaults for columns obj._overrides = {} # stores memory owner obj.base = None return obj
def read_ms_hdf5_catalog(fname, root='Subsample/'): from nbodykit import CurrentMPIComm comm = CurrentMPIComm.get() logger = logging.getLogger('paint_utils') if comm.rank == 0: logger.info("Try reading %s" % fname) cat_nbk = HDFCatalog(fname, root=root) # Positions in Marcel hdf5 catalog files go from 0 to 1 but nbkit requires 0 to boxsize maxpos = get_cstat(cat_nbk['Position'].compute(), 'max') if (maxpos < 0.1 * np.min(cat_nbk.attrs['BoxSize'][:])) or maxpos < 1.5: cat_nbk['Position'] = cat_nbk['Position'] * cat_nbk.attrs['BoxSize'] if comm.rank == 0: logger.info("cat: %s" % str(cat_nbk)) logger.info("columns: %s" % str(cat_nbk.columns)) return cat_nbk
def __enter__(self): """ Split the base communicator such that each task gets allocated the specified number of cpus to perform the task with """ chain_ranks = [] color = 0 total_ranks = 0 nworkers = 0 # split the ranks for i, ranks in split_ranks(self.size, self.cpus_per_task, include_all=self.use_all_cpus): chain_ranks.append(ranks[0]) if self.rank in ranks: color = i+1 total_ranks += len(ranks) nworkers = nworkers + 1 self.workers = nworkers # store the total number of workers # check for no workers! if self.workers == 0: raise ValueError("no pool workers available; try setting `use_all_cpus` = True") leftover = (self.size - 1) - total_ranks if leftover and self.rank == 0: args = (self.cpus_per_task, self.size-1, leftover) self.logger.warning("with `cpus_per_task` = %d and %d available rank(s), %d rank(s) will do no work" %args) self.logger.warning("set `use_all_cpus=True` to use all available cpus") # crash if we only have one process or one worker if self.size <= self.workers: args = (self.size, self.workers+1, self.workers) raise ValueError("only have %d ranks; need at least %d to use the desired %d workers" %args) # ranks that will do work have a nonzero color now self._valid_worker = color > 0 # split the comm between the workers self.comm = self.basecomm.Split(color, 0) self.original_comm = CurrentMPIComm.get() CurrentMPIComm.set(self.comm) return self
def to_nbodykit(self, fields=None): from nbodykit.base.catalog import CatalogSource from nbodykit import CurrentMPIComm comm = CurrentMPIComm.get() if comm.rank == 0: source = self else: source = None source = comm.bcast(source) # compute the size start = comm.rank * source.size // comm.size end = (comm.rank + 1) * source.size // comm.size new = object.__new__(CatalogSource) new._size = end - start CatalogSource.__init__(new, comm=comm) for key in source.fields: new[key] = new.make_column(source[key])[start:end] new.attrs.update(source.attrs) return new
def get_displacement_from_density_rfield(in_density_rfield, in_density_cfield=None, component=None, Psi_type=None, smoothing=None, smoothing_Psi3LPT=None, prefac_Psi_1storder=1.0, prefac_Psi_2ndorder=1.0, prefac_Psi_3rdorder=1.0, RSD=False, RSD_line_of_sight=None, RSD_f_log_growth=None): """ Given density delta(x) in real space, compute Zeldovich displacemnt Psi_component(x) given by Psi_component(\vk) = k_component / k^2 * W(k) * delta(\vk), where W(k) is smoothing window. For Psi_type='Zeldovich' compute 1st order displacement. For Psi_type='2LPT' compute 1st plus 2nd order displacement. etc Supply either in_density_rfield or in_density_cfield. Multiply 1st order displacement by prefac_Psi_1storder, 2nd order by prefac_Psi_2ndorder, etc. Use this for getting time derivative of Psi. Follow http://rainwoodman.github.io/pmesh/intro.html. Parameters ---------- RSD : boolean If True, include RSD by displacing by \vecPsi(q)+f (\e_LOS.\vecPsi(q)) \e_LOS, where \ve_LOS is unit vector in line of sight direction. RSD_line_of_sight : array_like, (3,) Line of sight direction, e.g. [0,0,1] for z axis. """ assert (component in [0, 1, 2]) assert Psi_type in ['Zeldovich', '2LPT', '-2LPT'] comm = CurrentMPIComm.get() if in_density_cfield is None: # copy so we don't do any in-place changes by accident density_rfield = in_density_rfield.copy() density_cfield = density_rfield.r2c() else: assert in_density_rfield is None density_cfield = in_density_cfield.copy() if Psi_type in ['Zeldovich', '2LPT', '-2LPT']: # get zeldovich displacement in direction given by component def potential_transfer_function(k, v): k2 = sum(ki**2 for ki in k) return np.where(k2 == 0.0, 0 * v, v / (k2)) # get potential pot = delta/k^2 #pot_k = density_rfield.r2c().apply(potential_transfer_function) pot_k = density_cfield.apply(potential_transfer_function) #print("pot_k head:\n", pot_k[:2,:2,:2]) # apply smoothing if smoothing is not None: pot_k = smoothen_cfield(pot_k, **smoothing) #print("pot_k head2:\n", pot_k[:2,:2,:2]) # get zeldovich displacement def force_transfer_function(k, v, d=component): # MS: not sure if we want a factor of -1 here. return k[d] * 1j * v Psi_component_rfield = pot_k.apply(force_transfer_function).c2r() if RSD: # Add linear RSD displacement f (\e_LOS.\vecPsi^(1)(q)) \e_LOS. assert RSD_f_log_growth is not None if RSD_line_of_sight in [[0, 0, 1], [0, 1, 0], [1, 0, 0]]: # If [0,0,1] simply shift by Psi_z along z axis. Similarly in the other cases. if RSD_line_of_sight[component] == 0: # nothing to do in this direction pass elif RSD_line_of_sight[component] == 1: # add f Psi_component(q) Psi_component_rfield += RSD_f_log_growth * Psi_component_rfield if comm.rank == 0: print('%d: Added RSD in direction %d' % (comm.rank, component)) else: # Need to compute (\e_LOS.\vecPsi(q)) which requires all Psi components. raise Exception('RSD_line_of_sight %s not implemented' % str(RSD_line_of_sight)) Psi_component_rfield *= prefac_Psi_1storder # if comm.rank == 0: # print('mean, rms, max Psi^{1}_%d: %g, %g, %g' % ( # component, np.mean(Psi_component_rfield), # np.mean(Psi_component_rfield**2)**0.5, # np.max(Psi_component_rfield))) if Psi_type in ['2LPT', '-2LPT']: # add 2nd order Psi on top of Zeldovich if in_density_rfield is not None: in_density_fieldmesh = FieldMesh(in_density_rfield) else: in_density_fieldmesh = FieldMesh(in_density_cfield) # compute G2 G2_cfield = calc_quadratic_field( base_field_mesh=in_density_fieldmesh, quadfield='tidal_G2', smoothing_of_base_field=smoothing).compute(mode='complex') # compute Psi_2ndorder = -3/14 ik/k^2 G2(k). checked sign: improves rcc with deltaNL # if we use -3/14, but get worse rcc when using +3/14. Psi_2ndorder_rfield = -3. / 14. * ( G2_cfield.apply(potential_transfer_function).apply( force_transfer_function).c2r()) del G2_cfield if Psi_type == '-2LPT': # this is just to test sign Psi_2ndorder_rfield *= -1.0 if RSD: # Add 2nd order RSD displacement 2*f*(\e_LOS.\vecPsi^(2)(q)) \e_LOS. # Notice factor of 2 b/c \dot\psi enters for RSD. if RSD_line_of_sight in [[0, 0, 1], [0, 1, 0], [1, 0, 0]]: # If [0,0,1] simply shift by Psi_z along z axis. Similarly in the other cases. if RSD_line_of_sight[component] == 0: # nothing to do in this direction pass elif RSD_line_of_sight[component] == 1: # add 2 f Psi^{(2)}_component(q) Psi_2ndorder_rfield += (2.0 * RSD_f_log_growth * Psi_2ndorder_rfield) if comm.rank == 0: print('%d: Added 2nd order RSD in direction %d' % (comm.rank, component)) else: # Need to compute (\e_LOS.\vecPsi(q)) which requires all Psi components. raise Exception('RSD_line_of_sight %s not implemented' % str(RSD_line_of_sight)) # if comm.rank == 0: # print('mean, rms, max Psi^{2}_%d: %g, %g, %g' % ( # component, np.mean(Psi_2ndorder_rfield), # np.mean(Psi_2ndorder_rfield**2)**0.5, # np.max(Psi_2ndorder_rfield))) Psi_2ndorder_rfield *= prefac_Psi_2ndorder # add 2nd order to Zeldoivhc displacement Psi_component_rfield += Psi_2ndorder_rfield del Psi_2ndorder_rfield return Psi_component_rfield
def main(): """ Calculate BOSS window function multipoles. See https://nbodykit.readthedocs.io/en/latest/api/_autosummary/nbodykit.algorithms.html#nbodykit.algorithms.SurveyDataPairCount and https://arxiv.org/pdf/1607.03150.pdf eq 22 """ # Parse command line arguments ap = ArgumentParser() ap.add_argument('--rmin', default=100.0, type=float, help='Number of bins for separation r between pairs.') ap.add_argument('--rmax', default=500.0, type=float, help='Number of bins for separation r between pairs.') ap.add_argument('--Nr', default=20, type=int, help='Number of bins for separation r between pairs.') ap.add_argument('--Nmu', default=10, type=int, help='Number of bins for angle mu w.r.t. line of sight.') ap.add_argument('--download_dir', default='$SCRATCH/lss/sdss_data/', help='Where to store/read downloaded data.') default_sample = 'DR12v5_CMASSLOWZTOT_South' #default_sample = 'DR12v5_LOWZ_South' ap.add_argument( '--boss_sample', default=default_sample, help= 'Which BOSS sample to use. See https://data.sdss.org/sas/dr12/boss/lss/' ) ap.add_argument('--out_dir', default='window/', help='Folder where to store measured window function.') ap.add_argument('--out_base', default='paircount', help='Prefix for where to store measured window function.') ap.add_argument('--FKP', default=0, type=int, help='Include FKP weight.') ap.add_argument('--randoms1_catalog_id', default=0, type=int, help='ID for randoms1 catalog') ap.add_argument('--randoms2_catalog_id', default=1, type=int, help='ID for randoms2 catalog') ap.add_argument('--subsample_fraction', default=1e-4, type=float, help='If less than 1, use random subsample of randoms.') cmd_args = ap.parse_args() setup_logging() comm = CurrentMPIComm.get() # download the data to the current directory download_dir = os.path.expandvars(cmd_args.download_dir) if comm.rank == 0: print('download_dir:', download_dir) boss_sample = cmd_args.boss_sample if comm.rank == 0: download_data(download_dir, boss_sample=boss_sample, random_catalog_id=cmd_args.randoms1_catalog_id) download_data(download_dir, boss_sample=boss_sample, random_catalog_id=cmd_args.randoms2_catalog_id) # NOTE: change this path if you downloaded the data somewhere else! randoms1_path = os.path.join( download_dir, 'random%d_%s.fits' % (cmd_args.randoms1_catalog_id, boss_sample)) randoms2_path = os.path.join( download_dir, 'random%d_%s.fits' % (cmd_args.randoms2_catalog_id, boss_sample)) # initialize the FITS catalog objects for data and randoms randoms1 = FITSCatalog(randoms1_path) randoms2 = FITSCatalog(randoms2_path) if comm.rank == 0: print('randoms1 columns = ', randoms1.columns) print('randoms2 columns = ', randoms2.columns) # Select redshift range if boss_sample in ['DR12v5_LOWZ_South', 'DR12v5_LOWZ_North']: ZMIN = 0.15 ZMAX = 0.43 elif boss_sample in ['DR12v5_CMASS_South', 'DR12v5_CMASS_North']: ZMIN = 0.43 ZMAX = 0.7 elif boss_sample in [ 'DR12v5_CMASSLOWZTOT_South', 'DR12v5_CMASSLOWZTOT_North' ]: ZMIN = 0.5 ZMAX = 0.75 else: raise Exception('Must specify ZMIN and ZMAX for boss_sample=%s' % str(boss_sample)) # slice the randoms valid1 = (randoms1['Z'] > ZMIN) & (randoms1['Z'] < ZMAX) randoms1 = randoms1[valid1] valid2 = (randoms2['Z'] > ZMIN) & (randoms2['Z'] < ZMAX) randoms2 = randoms2[valid2] if cmd_args.subsample_fraction < 1.0: # Create random subsamples rng1 = MPIRandomState(randoms1.comm, seed=123, size=randoms1.size) rr1 = rng1.uniform(0.0, 1.0, itemshape=(1, )) randoms1 = randoms1[rr1[:, 0] < cmd_args.subsample_fraction] rng2 = MPIRandomState(randoms2.comm, seed=456, size=randoms2.size) rr2 = rng2.uniform(0.0, 1.0, itemshape=(1, )) randoms2 = randoms2[rr2[:, 0] < cmd_args.subsample_fraction] Nrandoms1 = randoms1.csize Nrandoms2 = randoms2.csize if comm.rank == 0: print('Nrandoms1:', Nrandoms1) print('Nrandoms2:', Nrandoms2) # weights if cmd_args.FKP == 0: randoms1['Weight'] = 1.0 randoms2['Weight'] = 1.0 else: randoms1['Weight'] = randoms1['WEIGHT_FKP'] randoms2['Weight'] = randoms2['WEIGHT_FKP'] # the fiducial BOSS DR12 cosmology cosmo = cosmology.Cosmology(h=0.676).match(Omega0_m=0.31) # bins for separation edges_r = np.logspace(np.log10(cmd_args.rmin), np.log10(cmd_args.rmax), num=cmd_args.Nr + 1) print('edges_r', edges_r) if comm.rank == 0: print("Start pair count...") paircount = SurveyDataPairCount(mode='2d', first=randoms1, second=randoms2, edges=edges_r, cosmo=cosmo, Nmu=cmd_args.Nmu, pimax=None, ra='RA', dec='DEC', redshift='Z', weight='Weight', show_progress=True, domain_factor=4) if comm.rank == 0: print("Done pair count") if comm.rank == 0: print('paircount', paircount) for key in paircount.attrs: print("%s = %s" % (key, str(paircount.attrs[key]))) # save results if comm.rank == 0: if not os.path.exists(cmd_args.out_dir): os.makedirs(cmd_args.out_dir) # save window to file, in nbodykit format out_file_base = os.path.join( cmd_args.out_dir, '%s_%s_rmin%.1f_rmax%.1f_Nr%d_Nmu%d_randID1%d_randID2%d_SUB%g_FKP%d' % (cmd_args.out_base, boss_sample, cmd_args.rmin, cmd_args.rmax, cmd_args.Nr, cmd_args.Nmu, cmd_args.randoms1_catalog_id, cmd_args.randoms2_catalog_id, cmd_args.subsample_fraction, cmd_args.FKP)) fname = '%s.nbk.dat' % out_file_base paircount.save(fname) print('Wrote %s' % fname)
def main(): """ Measure BOSS power spectrum. Based on https://nbodykit.readthedocs.io/en/latest/cookbook/boss-dr12-data.html. """ # Parse command line arguments ap = ArgumentParser() ap.add_argument('--Nmesh', default=64, type=int, help='Nmesh used to compute power spectrum.') ap.add_argument('--download_dir', default='$SCRATCH/lss/sdss_data/', help='Where to store/read downloaded data.') default_sample = 'DR12v5_CMASSLOWZTOT_South' #default_sample = 'DR12v5_LOWZ_South' ap.add_argument('--boss_sample', default=default_sample, help='Which BOSS sample to use. See https://data.sdss.org/sas/dr12/boss/lss/') ap.add_argument('--out_dir', default='power/', help='Folder where to store measured power spectrum.') ap.add_argument('--out_base', default='power', help='Prefix for where to store measured power spectrum.') ap.add_argument('--plot', dest='plot', action='store_true', help='Plot power spectra.') ap.add_argument('--subtract_shot', dest='subtract_shot', action='store_true', help='Subtract shot noise from monopole power spectrum.') ap.add_argument('--random_catalog_id', default=0, type=int, help='Which random catalog to use.') cmd_args = ap.parse_args() # Setup things if cmd_args.plot: from nbodykit import style import matplotlib.pyplot as plt plt.style.use(style.notebook) setup_logging() comm = CurrentMPIComm.get() # download the data to the current directory download_dir = os.path.expandvars(cmd_args.download_dir) print('download_dir:', download_dir) boss_sample = cmd_args.boss_sample if comm.rank == 0: download_data(download_dir, boss_sample=boss_sample, random_catalog_id=cmd_args.random_catalog_id) # NOTE: change this path if you downloaded the data somewhere else! data_path = os.path.join(download_dir, 'galaxy_%s.fits' % boss_sample) randoms_path = os.path.join(download_dir, 'random%d_%s.fits' % ( cmd_args.random_catalog_id, boss_sample)) # initialize the FITS catalog objects for data and randoms data = FITSCatalog(data_path) randoms = FITSCatalog(randoms_path) print('data columns = ', data.columns) print('randoms columns = ', randoms.columns) # Select redshift range if boss_sample in ['DR12v5_LOWZ_South', 'DR12v5_LOWZ_North']: ZMIN = 0.15 ZMAX = 0.43 elif boss_sample in ['DR12v5_CMASS_South', 'DR12v5_CMASS_North']: ZMIN = 0.43 ZMAX = 0.7 elif boss_sample in ['DR12v5_CMASSLOWZTOT_South', 'DR12v5_CMASSLOWZTOT_North']: ZMIN = 0.5 ZMAX = 0.75 else: raise Exception('Must specify ZMIN and ZMAX for boss_sample=%s' % str(boss_sample)) # slice the randoms valid = (randoms['Z'] > ZMIN)&(randoms['Z'] < ZMAX) randoms = randoms[valid] # slice the data valid = (data['Z'] > ZMIN)&(data['Z'] < ZMAX) data = data[valid] Ngalaxies = data.csize print('Ngalaxies:', Ngalaxies) Nrandoms = randoms.csize print('Nrandoms:', Nrandoms) # the fiducial BOSS DR12 cosmology cosmo = cosmology.Cosmology(h=0.676).match(Omega0_m=0.31) # add Cartesian position column data['Position'] = transform.SkyToCartesian(data['RA'], data['DEC'], data['Z'], cosmo=cosmo) randoms['Position'] = transform.SkyToCartesian(randoms['RA'], randoms['DEC'], randoms['Z'], cosmo=cosmo) randoms['WEIGHT'] = 1.0 data['WEIGHT'] = data['WEIGHT_SYSTOT'] * (data['WEIGHT_NOZ'] + data['WEIGHT_CP'] - 1.0) # combine the data and randoms into a single catalog fkp = FKPCatalog(data, randoms) mesh = fkp.to_mesh(Nmesh=cmd_args.Nmesh, nbar='NZ', fkp_weight='WEIGHT_FKP', comp_weight='WEIGHT', window='tsc') # compute the multipoles r = ConvolvedFFTPower(mesh, poles=[0,2,4], dk=0.005, kmin=0.) for key in r.attrs: print("%s = %s" % (key, str(r.attrs[key]))) # save results if not os.path.exists(cmd_args.out_dir): os.makedirs(cmd_args.out_dir) # save power to file, in nbodykit format fname = os.path.join(cmd_args.out_dir, '%s_%s_Nmesh%d.nbk.dat' % (cmd_args.out_base, boss_sample, cmd_args.Nmesh)) r.save(fname) print('Wrote %s' % fname) # Also save plain txt file poles = r.poles Nk = poles['k'].shape[0] mat = np.zeros((Nk, 4)) + np.nan header = 'boss_sample=%s\n' % boss_sample header += 'ZMIN=%g\n' % ZMIN header += 'ZMAX=%g\n' % ZMAX header += 'Nmesh=%d\n' % cmd_args.Nmesh header += 'Ngalaxies=%d\n' % Ngalaxies if cmd_args.subtract_shot: header += 'subtract_shot=True\n' else: header += 'subtract_shot=False\n' header += 'random_catalog_id=%d\n' % cmd_args.random_catalog_id header += 'Columns: k, P_0, P_2, P_4' mat[:,0] = poles['k'] if cmd_args.subtract_shot: mat[:,1] = poles['power_0'].real - r.attrs['shotnoise'] else: mat[:,1] = poles['power_0'].real mat[:,2] = poles['power_2'].real mat[:,3] = poles['power_4'].real # save out_file_base = os.path.join(cmd_args.out_dir, '%s_%s_Nmesh%d_subtrShot%d_randID%d' % ( cmd_args.out_base, boss_sample, cmd_args.Nmesh, int(cmd_args.subtract_shot==True), cmd_args.random_catalog_id )) fname = out_file_base + '.txt' np.savetxt(fname, mat, header=header) print('Wrote %s' % fname) if cmd_args.plot: # run code with --plot to plot for ell in [0, 2, 4]: label = r'$\ell=%d$' % (ell) P = poles['power_%d' %ell].real if cmd_args.subtract_shot: if ell == 0: P = P - r.attrs['shotnoise'] plt.plot(poles['k'], poles['k']*P, label=label) # format the axes plt.legend(loc=0) plt.xlabel(r"$k$ [$h \ \mathrm{Mpc}^{-1}$]") plt.ylabel(r"$k \ P_\ell$ [$h^{-2} \ \mathrm{Mpc}^2$]") plt.xlim(0.01, 0.25) fname = out_file_base + '.pdf' plt.savefig(fname) print('Made %s' % fname)
def get_catalog(self, keep_all_columns=False): """ Get the target catalog, with columns 'Position' and 'val'. """ comm = CurrentMPIComm.get() if self.in_fname is None: # use copy of in_catalog in memory cat = catalog_persist(self.in_catalog, columns=self.in_catalog.columns) else: assert self.in_catalog is None if comm.rank == 0: print('Read %s' % self.in_fname) # Read rho, not dividing or subtracting mean, to get velocity correctly cat = BigFileCatalog(self.in_fname, dataset='./', header='Header') # Set Position column if self.position_column != 'Position': print('position_column: ', self.position_column) cat['Position'] = cat[self.position_column] print('Catalog columns:', cat.columns) # add RSD if self.apply_RSD_to_position: print('Apply RSD') cat['Position'] += (self.RSDFactor * cat[self.velocity_column] * self.RSD_los) # cuts Ngal_before_cuts = cat.csize for cut in self.cuts: if type(cut) == tuple: assert len(cut) == 3 cut_column, cut_op, cut_value = cut print('cut:', cut_column, cut_op, cut_value) if type(cut_value) in (list, tuple): # cut value is a list or tuple, e.g. to cut 'Position' for i, cv in enumerate(cut_value): if cut_op == 'min': cat = cat[cat[cut_column][:, i] >= cv] elif cut_op == 'max': cat = cat[cat[cut_column][:, i] < cv] else: raise Exception('Invalid cut operation %s' % str(cut_op)) else: # cut value is a scalar, e.g. to cut mass if cut_op == 'min': cat = cat[cat[cut_column] >= cut_value] elif cut_op == 'max': cat = cat[cat[cut_column] < cut_value] else: raise Exception('Invalid cut operation %s' % str(cut_op)) elif isinstance(cut, SimGalaxyCatalogCreator): cat = cut.get_galaxy_catalog_from_source_catalog(cat) else: raise Exception('Invalid cut %s' % str(cut)) Ngal_after_cuts = cat.csize print('Cuts removed %g%% of objects' % (100. * float(Ngal_after_cuts - Ngal_before_cuts) / float(Ngal_before_cuts))) # compute the value we are interested in, save in 'val' column component = self.val_component if self.val_column is not None: if component is None: cat['val'] = cat[self.val_column][:] else: cat['val'] = cat[self.val_column][:, component] # catalog rescale factor to be applied to value column if self.rescale_factor is not None: if self.rescale_factor == 'RSDFactor': fac = cat.attrs['RSDFactor'][0] elif type(self.rescale_factor) == float: fac = self.rescale_factor else: raise Exception('Invalid rescale_factor %s' % self.rescale_factor) if cat.comm.rank == 0: print('Apply rescale factor: %g' % fac) cat['val'] *= fac if keep_all_columns: pass else: # keep only Position and val columns, delete all other columns cat = catalog_persist(cat, columns=['Position', 'val']) if 'val' in cat.columns: cstats = get_cstats_string(cat['val'].compute()) if comm.rank == 0: print('TARGET CATALOG %s: %s\n' % (self.name, cstats)) return cat
def main(): """ Combine source fields to get proxy of a target field. This is stage 0 of reconstruction, and can be used to quantify Perror of a bias model. For example: - Combine delta_m, delta_m^2, s^2 to get proxy of target=delta_halo. - Combine different mass-weighted delta_h fields to get proxy of target=delta_m. Usage examples: ./run.sh python main_calc_Perr_test.py or ./run.sh mpiexec -n 4 python main_calc_Perr_test.py --SimSeed 403 """ ##################################### # PARSE COMMAND LINE ARGS ##################################### ap = ArgumentParser() ap.add_argument('--SimSeed', type=int, default=403, help='Simulation seed to load.') ap.add_argument('--HaloMassString', default='13.8_15.1', help="Halo mass string, for example '13.8_15.1'.") cmd_args = ap.parse_args() ##################################### # OPTIONS ##################################### opts = OrderedDict() # Bump this when changing code without changing options. Otherwise pickle # loading might wrongly read old pickles. opts['main_calc_Perr_test_version'] = '1.5' # Simulation options. Will be used by path_utils to get input path, and # to compute deltalin at the right redshift. seed = cmd_args.SimSeed opts['sim_opts'] = parameters_ms_gadget.MSGadgetSimOpts.load_default_opts( sim_name='ms_gadget_test_data', sim_seed=seed, halo_mass_string=cmd_args.HaloMassString) # Grid options. Ngrid = 64 opts['grid_opts'] = parameters.GridOpts( Ngrid=Ngrid, kmax=2.0*np.pi/opts['sim_opts'].boxsize * float(Ngrid)/2.0, grid_ptcle2grid_deconvolution=None ) # Options for measuring power spectrum. Use defaults. opts['power_opts'] = parameters.PowerOpts() # Transfer function options. See lsstools.parameters.py for details. opts['trf_fcn_opts'] = parameters.TrfFcnOpts( Rsmooth_for_quadratic_sources=0.1, Rsmooth_for_quadratic_sources2=0.1, N_ortho_iter=1, orth_method='CholeskyDecomp', interp_kind='manual_Pk_k_bins' ) # External grids to load: deltalin, delta_m, shifted grids opts['ext_grids_to_load'] = opts['sim_opts'].get_default_ext_grids_to_load( Ngrid=opts['grid_opts'].Ngrid) # Catalogs to read opts['cats'] = opts['sim_opts'].get_default_catalogs() # Specify bias expansions to test opts['trf_specs'] = [] # Quadratic Lagrangian bias: delta_Z + b1 deltalin(q+Psi) + b2 # [deltalin^2-<deltalin^2>](q+Psi) + bG2 [G2](q+Psi) opts['trf_specs'].append( TrfSpec(linear_sources=[ 'deltalin_SHIFTEDBY_deltalin', 'deltalin_growth-mean_SHIFTEDBY_deltalin', 'deltalin_G2_SHIFTEDBY_deltalin' ], fixed_linear_sources=['1_SHIFTEDBY_deltalin'], field_to_smoothen_and_square=None, quadratic_sources=[], target_field='delta_h', save_bestfit_field= 'hat_delta_h_from_1_Tdeltalin2G2_SHIFTEDBY_PsiZ')) # Save results opts['keep_pickle'] = False opts['pickle_file_format'] = 'dill' opts['pickle_path'] = '$SCRATCH/perr/pickle/' # Save some additional power spectra that are useful for plotting later opts['Pkmeas_helper_columns'] = [ 'delta_h', 'delta_m', '1_SHIFTEDBY_deltalin', 'deltalin' ] # Save grids for 2d slice plots and histograms opts['save_grids4plots'] = False opts['grids4plots_base_path'] = '$SCRATCH/perr/grids4plots/' opts['grids4plots_R'] = 0.0 # Gaussian smoothing applied to grids4plots # Cache path opts['cache_base_path'] = '$SCRATCH/perr/cache/' # Run the program given the above opts. outdict = model_error.calculate_model_error(opts) # Compare vs expected result. residual_key = '[hat_delta_h_from_1_Tdeltalin2G2_SHIFTEDBY_PsiZ]_MINUS_[delta_h]' Perr = outdict['Pkmeas'][(residual_key, residual_key)].P Perr_expected = np.array([ 9965.6, 17175.8, 22744.4, 19472.3, 19081.2, 19503.4, 19564.9, 18582.9, 19200.1, 16911.3, 16587.4, 16931.9, 15051.0, 13835.1, 13683.8, 13109.9, 12353.5, 11900.2, 11085.1, 11018.4, 10154.0, 9840.7, 8960.6, 8484.1, 7942.2, 7426.7, 6987.8, 6578.1, 6269.1, 5810.7, 5511.7 ]) setup_logging() comm = CurrentMPIComm.get() logger = logging.getLogger('TestPerrCalc') if comm.rank == 0: Perr_lst = ['%.1f' % a for a in list(Perr)] Perr_expected_lst = ['%.1f' % a for a in list(Perr)] logger.info('Perr:\n%s' % str(','.join(Perr_lst))) logger.info('Expected Perr:\n%s' % str(','.join(Perr_expected_lst))) if np.allclose(Perr, Perr_expected, rtol=1e-3): logger.info('TEST Perr: OK') else: logger.info('TEST Perr: FAILED') raise Exception('Test failed')
def main(): """ Convert grids4plots 3d grids to 2d slices. """ ##################################### # PARSE COMMAND LINE ARGS ##################################### ap = ArgumentParser() ap.add_argument('--inbasepath', type=str, default='$SCRATCH/perr/grids4plots/', help='Input base path.') ap.add_argument('--outbasepath', type=str, default='$SCRATCH/perr/grids4plots/', help='Output base path.') ap.add_argument( '--inpath', type=str, default= 'main_calc_Perr_2020_Sep_22_18:44:31_time1600800271.dill', # laptop #default='main_calc_Perr_2020_Aug_26_02:49:57_time1598410197.dill', # cluster help='Input path.') ap.add_argument('--Rsmooth', type=float, default=2.0, help='3D Gaussian smoothing applied to field.') # min and max index included in output. inclusive. ap.add_argument('--ixmin', type=int, default=5, help='xmin of output. must be between 0 and Ngrid.') ap.add_argument('--ixmax', type=int, default=5, help='xmax of output') ap.add_argument('--iymin', type=int, default=0, help='ymin of output') ap.add_argument('--iymax', type=int, default=-1, help='ymax of output') ap.add_argument('--izmin', type=int, default=0, help='zmin of output') ap.add_argument('--izmax', type=int, default=-1, help='zmax of output') cmd_args = ap.parse_args() verbose = True ##################################### # START PROGRAM ##################################### comm = CurrentMPIComm.get() rank = comm.rank path = os.path.join(os.path.expandvars(cmd_args.inbasepath), cmd_args.inpath) if rank == 0: print('path: ', path) initialized_slicecat = False for fname in os.listdir(path): if fname.startswith('SLICE'): continue full_fname = os.path.join(path, fname) print('%d Reading %s' % (rank, full_fname)) inmesh = BigFileMesh(full_fname, dataset='tmp4storage', header='header') Ngrid = inmesh.attrs['Ngrid'] boxsize = inmesh.attrs['boxsize'] # apply smoothing mesh = apply_smoothing(inmesh, mode='Gaussian', R=cmd_args.Rsmooth) del inmesh # convert indices to modulo ngrid ixmin = cmd_args.ixmin % Ngrid ixmax = cmd_args.ixmax % Ngrid iymin = cmd_args.iymin % Ngrid iymax = cmd_args.iymax % Ngrid izmin = cmd_args.izmin % Ngrid izmax = cmd_args.izmax % Ngrid # convert to boxsize units (Mpc/h) xmin = float(ixmin) / float(Ngrid) * boxsize xmax = float(ixmax + 1) / float(Ngrid) * boxsize ymin = float(iymin) / float(Ngrid) * boxsize ymax = float(iymax + 1) / float(Ngrid) * boxsize zmin = float(izmin) / float(Ngrid) * boxsize zmax = float(izmax + 1) / float(Ngrid) * boxsize if not initialized_slicecat: # Generate catalog with positions of slice points, to readout mesh there. # First generate all 3D points. Then keep only subset in slice. # THen readout mesh at those points. # use pmesh generate_uniform_particle_grid # http://rainwoodman.github.io/pmesh/pmesh.pm.html?highlight= # readout#pmesh.pm.ParticleMesh.generate_uniform_particle_grid partmesh = ParticleMesh(BoxSize=boxsize, Nmesh=[Ngrid, Ngrid, Ngrid]) ptcles = partmesh.generate_uniform_particle_grid(shift=0.0, dtype='f8') #print("type ptcles", type(ptcles), ptcles.shape) #print("head ptcles:", ptcles[:5,:]) dtype = np.dtype([('Position', ('f8', 3))]) # number of rows is given by number of ptcles on this rank uni_cat_array = np.empty((ptcles.shape[0], ), dtype=dtype) uni_cat_array['Position'] = ptcles uni_cat = ArrayCatalog(uni_cat_array, comm=None, BoxSize=boxsize * np.ones(3), Nmesh=[Ngrid, Ngrid, Ngrid]) del ptcles del uni_cat_array print("%d: Before cut: local Nptcles=%d, global Nptcles=%d" % (comm.rank, uni_cat.size, uni_cat.csize)) # only keep points in the slice uni_cat = uni_cat[(uni_cat['Position'][:, 0] >= xmin) & (uni_cat['Position'][:, 0] < xmax) & (uni_cat['Position'][:, 1] >= ymin) & (uni_cat['Position'][:, 1] < ymax) & (uni_cat['Position'][:, 2] >= zmin) & (uni_cat['Position'][:, 2] < zmax)] print("%d: After cut: local Nptcles=%d, global Nptcles=%d" % (comm.rank, uni_cat.size, uni_cat.csize)) initialized_slicecat = True # read out full 3D mesh at catalog positions. this is a numpy array slicecat = readout_mesh_at_cat_pos(mesh=mesh, cat=uni_cat, readout_window='nearest') if rank == 0: print('slicecat type:', type(slicecat)) slicecat = GatherArray(slicecat, comm, root=0) if rank == 0: if not slicecat.shape == ((ixmax - ixmin + 1) * (iymax - iymin + 1) * (izmax - izmin + 1), ): raise Exception( 'Unexpected shape of particles read out on slice: %s' % str(slicecat.shape)) slicecat = slicecat.reshape( (ixmax - ixmin + 1, iymax - iymin + 1, izmax - izmin + 1)) print('slicecat shape:', slicecat.shape) if verbose: print('slicecat:', slicecat) # convert to a mesh. assume full numpy array sits on rank 0. Lx = xmax - xmin Ly = ymax - ymin Lz = zmax - zmin if Lx == 0.: Lx = boxsize / float(Ngrid) if Ly == 0.: Ly = boxsize / float(Ngrid) if Lz == 0.: Lz = boxsize / float(Ngrid) BoxSize_slice = np.array([Lx, Ly, Lz]) slicemesh = ArrayMesh(slicecat, BoxSize=BoxSize_slice, root=0) outshape = slicemesh.compute(mode='real').shape if verbose: print('slicemesh: ', slicemesh.compute(mode='real')) # write to disk outpath = os.path.join( os.path.expandvars(cmd_args.outbasepath), cmd_args.inpath, 'SLICE_R%g_%d-%d_%d-%d_%d-%d/' % (cmd_args.Rsmooth, ixmin, ixmax, iymin, iymax, izmin, izmax)) if rank == 0: if not os.path.exists(outpath): os.makedirs(outpath) full_outfname = os.path.join(outpath, fname) if rank == 0: print('Writing %s' % full_outfname) slicemesh.save(full_outfname) if rank == 0: print('Wrote %s' % full_outfname)
def weigh_and_shift_uni_cat(delta_for_weights, displacements, Nptcles_per_dim, out_Ngrid, boxsize, internal_scale_factor_for_weights=None, out_scale_factor=None, cosmo_params=None, weighted_CIC_mode=None, uni_cat_generator='pmesh', plot_slices=False, verbose=False, return_catalog=False): """ Make uniform catalog, weigh by delta_for_weights, displace by displacements and interpolate to grid which we return. Parameters ---------- delta_for_weights : None or pmesh.pm.RealField object Particles are weighted by 1+delta_for_weights. If None, use weight=1. displacements : list [Psi_x, Psi_y, Psi_z] where Psi_i are pmesh.pm.RealField objects, holding the displacement field in different directions (on the grid). If None, do not shift. uni_cat_generator : string If 'pmesh', use pmesh for generating uniform catalog. If 'manual', use an old serial code. Returns ------- delta_shifted : FieldMesh object Density delta_shifted of shifted weighted particles (normalized to mean of 1 i.e. returning 1+delta). Returned if return_catalog=False. attrs : meshsource attrs Attrs of delta_shifted. Returned if return_catalog=False. catalog_shifted : ArrayCatalog object Shifted catalog, returned if return_catalog=True. """ comm = CurrentMPIComm.get() # ###################################################################### # Generate uniform catalog with Nptcles_per_dim^3 particles on regular # grid # ###################################################################### if uni_cat_generator == 'pmesh': # use pmesh generate_uniform_particle_grid # http://rainwoodman.github.io/pmesh/pmesh.pm.html?highlight= # readout#pmesh.pm.ParticleMesh.generate_uniform_particle_grid pmesh = ParticleMesh( BoxSize=boxsize, Nmesh=[Nptcles_per_dim, Nptcles_per_dim, Nptcles_per_dim]) ptcles = pmesh.generate_uniform_particle_grid(shift=0.0, dtype='f8') #print("type ptcles", type(ptcles), ptcles.shape) #print("head ptcles:", ptcles[:5,:]) dtype = np.dtype([('Position', ('f8', 3))]) # number of rows is given by number of ptcles on this rank uni_cat_array = np.empty((ptcles.shape[0], ), dtype=dtype) uni_cat_array['Position'] = ptcles uni_cat = ArrayCatalog( uni_cat_array, comm=None, BoxSize=boxsize * np.ones(3), Nmesh=[Nptcles_per_dim, Nptcles_per_dim, Nptcles_per_dim]) print("%d: local Nptcles=%d, global Nptcles=%d" % (comm.rank, uni_cat.size, uni_cat.csize)) del ptcles del uni_cat_array elif uni_cat_generator == 'manual': # Old serial code to generate regular grid and scatter across ranks. if comm.rank == 0: # Code copied from do_rec_v1.py and adopted # Note that nbkit UniformCatalog is random catalog, but we want a catalog # where each ptcle sits at grid points of a regular grid. # This is what we call 'regular uniform' catalog. Np = Nptcles_per_dim dtype = np.dtype([('Position', ('f8', 3))]) # Have Np**3 particles, and each particle has position x,y,z and weight 'Weight' uni_cat_array = np.empty((Np**3, ), dtype=dtype) # x components in units such that box ranges from 0 to 1. Note dx=1/Np. #x_components_1d = np.linspace(0.0, (Np-1)*(L/float(Np)), num=Np, endpoint=True)/L x_components_1d = np.linspace(0.0, (Np - 1) / float(Np), num=Np, endpoint=True) ones_1d = np.ones(x_components_1d.shape) # Put particles on the regular grid print("%d: Fill regular uniform catalog" % comm.rank) uni_cat_array['Position'][:, 0] = np.einsum('a,b,c->abc', x_components_1d, ones_1d, ones_1d).reshape( (Np**3, )) uni_cat_array['Position'][:, 1] = np.einsum('a,b,c->abc', ones_1d, x_components_1d, ones_1d).reshape( (Np**3, )) uni_cat_array['Position'][:, 2] = np.einsum('a,b,c->abc', ones_1d, ones_1d, x_components_1d).reshape( (Np**3, )) print("%d: Done filling regular uniform catalog" % comm.rank) # in nbkit0.3 units must be in Mpc/h uni_cat_array['Position'] *= boxsize else: uni_cat_array = None # Scatter across all ranks print("%d: Scatter array" % comm.rank) from nbodykit.utils import ScatterArray uni_cat_array = ScatterArray(uni_cat_array, comm, root=0, counts=None) print("%d: Scatter array done. Shape: %s" % (comm.rank, str(uni_cat_array.shape))) # Save in ArrayCatalog object uni_cat = ArrayCatalog(uni_cat_array) uni_cat.attrs['BoxSize'] = np.ones(3) * boxsize uni_cat.attrs['Nmesh'] = np.ones(3) * Nptcles_per_dim uni_cat.attrs['Nmesh_internal'] = np.ones(3) * Nmesh_orig else: raise Exception('Invalid uni_cat_generator %s' % uni_cat_generator) ######################################################################## # Set weight of particles in uni_cat to delta (interpolated to ptcle # positions) ######################################################################## if delta_for_weights is None: # set all weights to 1 uni_cat['Mass'] = np.ones(uni_cat['Position'].shape[0]) else: # weight by delta_for_weights nbkit03_utils.interpolate_pm_rfield_to_catalog( delta_for_weights, uni_cat, catalog_column_to_save_to='Mass') print("%d: rms Mass: %g" % (comm.rank, np.sqrt(np.mean(np.array(uni_cat['Mass'])**2)))) # optionally plot weighted uniform cat before shifting if plot_slices: # paint the original uni_cat to a grid and plot slice import matplotlib.pyplot as plt tmp_meshsource = uni_cat.to_mesh(Nmesh=out_Ngrid, value='Mass', window='cic', compensated=False, interlaced=False) # paint to get delta(a_internal) tmp_outfield = tmp_meshsource.paint(mode='real') # linear rescale factor from internal_scale_factor_for_weights to # out_scale_factor rescalefac = nbkit03_utils.linear_rescale_fac( internal_scale_factor_for_weights, out_scale_factor, cosmo_params=cosmo_params) tmp_outfield = 1.0 + rescalefac * (tmp_outfield - 1.0) tmp_mesh = FieldMesh(tmp_outfield) plt.imshow(tmp_mesh.preview(Nmesh=32, axes=(0, 1))) if comm.rank == 0: plt_fname = 'inmesh_Np%d_Nm%d_Ng%d.pdf' % (Nptcles_per_dim, Nmesh_orig, out_Ngrid) plt.savefig(plt_fname) print("Made %s" % plt_fname) del tmp_meshsource, rescalefac, tmp_outfield, tmp_mesh # ###################################################################### # Shift uniform catalog particles by Psi (changes uni_cat) # ###################################################################### nbkit03_utils.shift_catalog_by_psi_grid( cat=uni_cat, in_displacement_rfields=displacements, pos_column='Position', pos_units='Mpc/h', displacement_units='Mpc/h', boxsize=boxsize, verbose=verbose) #del Psi_rfields if return_catalog: # return shifted catalog return uni_cat else: # return density of shifted catalog, delta_shifted # ###################################################################### # paint shifted catalog to grid, using field_to_shift as weights # ###################################################################### print("%d: paint shifted catalog to grid using mass weights" % comm.rank) # this gets 1+delta if weighted_CIC_mode == 'sum': delta_shifted, attrs = paint_utils.weighted_paint_cat_to_delta( uni_cat, weight='Mass', Nmesh=out_Ngrid, weighted_paint_mode=weighted_CIC_mode, normalize=True, # compute 1+delta verbose=verbose, to_mesh_kwargs={ 'window': 'cic', 'compensated': False, 'interlaced': False }) # this get rho elif weighted_CIC_mode == 'avg': delta_shifted, attrs = paint_utils.mass_avg_weighted_paint_cat_to_rho( uni_cat, weight='Mass', Nmesh=out_Ngrid, verbose=verbose, to_mesh_kwargs={ 'window': 'cic', 'compensated': False, 'interlaced': False }) else: raise Exception('Invalid weighted_CIC_mode %s' % weighted_CIC_mode) # ###################################################################### # rescale to output redshift # ###################################################################### if internal_scale_factor_for_weights != out_scale_factor: # linear rescale factor from internal_scale_factor_for_weights to # out_scale_factor rescalefac = nbkit03_utils.linear_rescale_fac( internal_scale_factor_for_weights, out_scale_factor, cosmo_params=cosmo_params) delta_shifted *= rescalefac # print some info: if comm.rank == 0: print( "%d: Linear rescalefac from a=%g to a=%g, rescalefac=%g" % (comm.rank, internal_scale_factor_for_weights, out_scale_factor, rescalefac)) raise Exception( 'Check if rescaling of delta_shifted is correct. Looks like 1+delta.' ) if verbose: print("%d: delta_shifted: min, mean, max, rms(x-1):" % comm.rank, np.min(delta_shifted), np.mean(delta_shifted), np.max(delta_shifted), np.mean((delta_shifted - 1.)**2)**0.5) # get 1+deta mesh from field #outmesh = FieldMesh(1 + out_delta) # print some info: this makes code never finish (race condition maybe?) #nbkit03_utils.rfield_print_info(outfield, comm, 'outfield: ') return delta_shifted, attrs
def avg_value_mass_weighted_paint_cat_to_rho( cat=None, value_column=None, weight_ptcles_by=None, Ngrid=None, fill_empty_cells='RandNeighb', RandNeighbSeed=1234, raise_exception_if_too_many_empty_cells=True, to_mesh_kwargs=None, verbose=False ): """ Helper function that paints cat[value_column] to grid, averaging over values of all particles belonging to a cell, and allowing for additional particle mass weights. Also has several methods to fill empty cells. """ # In the code 'value' is called 'chi', because value is chi in reconstruction # code. if to_mesh_kwargs is None: to_mesh_kwargs = { 'window': 'cic', 'compensated': False, 'interlaced': False} comm = CurrentMPIComm.get() logger = logging.getLogger('paint_utils') ## Get mass density rho so we can normalize chi later. Assume mass=1, or given by # weight_ptcles_by. # This is to get avg chi if multiple ptcles are in same cell. # 1 Sep 2017: Want chi_avg = sum_i m_i chi_i / sum_j m_i where m_i is particle mass, # because particle mass says how much the average should be dominated by a single ptcle # that can represent many original no-mass particles. # Compute rho4chi = sum_i m_i rho4chi, rho4chi_attrs = weighted_paint_cat_to_delta( cat, weight=weight_ptcles_by, weighted_paint_mode='sum', to_mesh_kwargs=to_mesh_kwargs, normalize=False, # want rho not 1+delta Nmesh=Ngrid, set_mean=None, verbose=verbose) # compute chi weighted by ptcle mass chi(x)m(x) weighted_col = 'TMP weighted %s' % value_column if weight_ptcles_by is not None: cat[weighted_col] = cat[weight_ptcles_by] * cat[value_column] else: # weight 1 for each ptcle cat[weighted_col] = cat[value_column] thisChi, thisChi_attrs = weighted_paint_cat_to_delta( cat, weight=weighted_col, # chi weighted by ptcle mass weighted_paint_mode='sum', to_mesh_kwargs=to_mesh_kwargs, normalize=False, # want rho not 1+delta (TODO: check) Nmesh=Ngrid, set_mean=None, verbose=verbose) # Normalize Chi by dividing by rho: So far, our chi will get larger if there are # more particles, because it sums up displacements over all particles. # To normalize, divide by rho (=mass density on grid if all ptcles have mass m=1, # or mass given by weight_ptcles_by). # (i.e. divide by number of contributions to a cell) if fill_empty_cells in [None, 'SetZero']: # Set chi=0 if there are not ptcles in grid cell. Used until 7 April 2017. # Seems ok for correl coeff and BAO, but gives large-scale bias in transfer # function or broad-band power because violates mass conservation. raise Exception('Possible bug: converting to np array only uses root rank?') thisChi = FieldMesh( np.where( rho4chi.compute(mode='real') == 0, rho4chi.compute(mode='real') * 0, thisChi.compute(mode='real') / rho4chi.compute(mode='real'))) #thisChi = np.where(gridx.G['rho4chi']==0, thisChi*0, thisChi/gridx.G['rho4chi']) elif fill_empty_cells in [ 'RandNeighb', 'RandNeighbReadout', 'AvgAndRandNeighb']: # Set chi in empty cells equal to a random neighbor cell. Do this until all empty # cells are filled. # First set all empty cells to nan. #thisChi = np.where(gridx.G['rho4chi']==0, thisChi*0+np.nan, thisChi/gridx.G['rho4chi']) thisChi = thisChi / rho4chi # get nan when rho4chi=0 if True: # test if nan ok ww1 = np.where(rho4chi == 0) #ww2 = np.where(np.isnan(thisChi.compute(mode='real'))) ww2 = np.where(np.isnan(thisChi)) assert np.allclose(ww1, ww2) del ww1, ww2 # Progressively replace nan by random neighbors: Ng = Ngrid #thisChi = thisChi.reshape((Ng,Ng,Ng)) logger.info('thisChi.shape: %s' % str(thisChi.shape)) #assert thisChi.shape == (Ng,Ng,Ng) # indices of empty cells on this rank ww = np.where(np.isnan(thisChi)) # number of empty cells across all ranks Nfill = comm.allreduce(ww[0].shape[0], op=MPI.SUM) have_empty_cells = (Nfill > 0) if fill_empty_cells in ['RandNeighb', 'RandNeighbReadout']: i_iter = -1 while have_empty_cells: i_iter += 1 if comm.rank == 0: logger.info( "Fill %d empty chi cells (%g percent) using random neighbors" % (Nfill, Nfill / float(Ng)**3 * 100.)) if Nfill / float(Ng)**3 >= 0.999: if raise_exception_if_too_many_empty_cells: raise Exception( "Stop because too many empty chi cells") else: logger.warning( "More than 99.9 percent of cells are empty") # draw -1,0,+1 for each empty cell, in 3 directions # r = np.random.randint(-1,2, size=(ww[0].shape[0],3), dtype='int') rng = MPIRandomState(comm, seed=RandNeighbSeed + i_iter * 100, size=ww[0].shape[0], chunksize=100000) r = rng.uniform(low=-2, high=2, dtype='int', itemshape=(3,)) assert np.all(r >= -1) assert np.all(r <= 1) # Old serial code to replace nan by random neighbors. # thisChi[ww[0],ww[1],ww[2]] = thisChi[(ww[0]+r[:,0])%Ng, (ww[1]+r[:,1])%Ng, (ww[2]+r[:,2])%Ng] if fill_empty_cells == 'RandNeighbReadout': # New parallel code, 1st implementation. # Use readout to get field at positions [(ww+rank_offset+r)%Ng] dx. BoxSize = cat.attrs['BoxSize'] dx = BoxSize / (float(Ng)) #pos_wanted = ((np.array(ww).transpose() + r) % Ng) * dx # ranges from 0 to BoxSize # more carefully: pos_wanted = np.zeros((ww[0].shape[0], 3)) + np.nan for idir in [0, 1, 2]: pos_wanted[:, idir] = ( (np.array(ww[idir] + thisChi.start[idir]) + r[:, idir]) % Ng) * dx[idir] # ranges from 0..BoxSize # use readout to get neighbors readout_window = 'nnb' layout = thisChi.pm.decompose(pos_wanted, smoothing=readout_window) # interpolate field to particle positions (use pmesh 'readout' function) thisChi_neighbors = thisChi.readout( pos_wanted, resampler=readout_window, layout=layout) if False: # print dbg info for ii in range(10000, 10004): if comm.rank == 1: logger.info( 'chi manual neighbor: %g' % thisChi[(ww[0][ii] + r[ii, 0]) % Ng, (ww[1][ii] + r[ii, 1]) % Ng, (ww[2][ii] + r[ii, 2]) % Ng]) logger.info('chi readout neighbor: %g' % thisChi_neighbors[ii]) thisChi[ww] = thisChi_neighbors elif fill_empty_cells == 'RandNeighb': # New parallel code, 2nd implementation. # Use collective getitem and only work with indices. # http://rainwoodman.github.io/pmesh/pmesh.pm.html#pmesh.pm.Field.cgetitem. # Note ww are indices of local slab, need to convert to global indices. thisChi_neighbors = None my_cindex_wanted = None for root in range(comm.size): # bcast to all ranks b/c must call cgetitem collectively with same args on each rank if comm.rank == root: # convert local index to collective index using ltoc which gives 3 tuple assert len(ww) == 3 wwarr = np.array(ww).transpose() #cww = np.array([ # ltoc(field=thisChi, index=[ww[0][i],ww[1][i],ww[2][i]]) # for i in range(ww[0].shape[0]) ]) cww = ltoc_index_arr(field=thisChi, lindex_arr=wwarr) #logger.info('cww: %s' % str(cww)) #my_cindex_wanted = [(cww[:,0]+r[:,0])%Ng, (cww[1][:]+r[:,1])%Ng, (cww[2][:]+r[:,2])%Ng] my_cindex_wanted = (cww + r) % Ng #logger.info('my_cindex_wanted: %s' % str(my_cindex_wanted)) cindex_wanted = comm.bcast(my_cindex_wanted, root=root) glob_thisChi_neighbors = cgetitem_index_arr( thisChi, cindex_wanted) # slower version doing the same # glob_thisChi_neighbors = [ # thisChi.cgetitem([cindex_wanted[i,0], cindex_wanted[i,1], cindex_wanted[i,2]]) # for i in range(cindex_wanted.shape[0]) ] if comm.rank == root: thisChi_neighbors = np.array( glob_thisChi_neighbors) #thisChi_neighbors = thisChi.cgetitem([40,42,52]) #print('thisChi_neighbors:', thisChi_neighbors) if False: # print dbg info (rank 0 ok, rank 1 fails to print) for ii in range(11000, 11004): if comm.rank == 1: logger.info( 'ww: %s' % str([ww[0][ii], ww[1][ii], ww[2][ii]])) logger.info( 'chi[ww]: %g' % thisChi[ww[0][ii], ww[1][ii], ww[2][ii]] ) logger.info( 'chi manual neighbor: %g' % thisChi[(ww[0][ii] + r[ii, 0]) % Ng, (ww[1][ii] + r[ii, 1]) % Ng, (ww[2][ii] + r[ii, 2]) % Ng]) logger.info('chi bcast neighbor: %g' % thisChi_neighbors[ii]) raise Exception('just dbg') thisChi[ww] = thisChi_neighbors ww = np.where(np.isnan(thisChi)) Nfill = comm.allreduce(ww[0].shape[0], op=MPI.SUM) have_empty_cells = (Nfill > 0) comm.barrier() elif fill_empty_cells == 'AvgAndRandNeighb': raise NotImplementedError # while have_empty_cells: # print("Fill %d empty chi cells (%g percent) using avg and random neighbors" % ( # ww[0].shape[0],ww[0].shape[0]/float(Ng)**3*100.)) # # first take average (only helps empty cells surrounded by filled cells) # thisChi[ww[0],ww[1],ww[2]] = 0.0 # for r0 in range(-1,2): # for r1 in range(-1,2): # for r2 in range(-1,2): # if (r0==0) and (r1==0) and (r2==0): # # do not include center point in avg b/c this is nan # continue # else: # # average over 27-1 neighbor points # thisChi[ww[0],ww[1],ww[2]] += thisChi[(ww[0]+r0)%Ng, (ww[1]+r1)%Ng, (ww[2]+r2)%Ng]/26.0 # # get indices of cells that are still empty (happens if a neighbor was nan above) # ww = np.where(np.isnan(thisChi)) # have_empty_cells = (ww[0].shape[0] > 0) # if have_empty_cells: # # draw -1,0,+1 for each empty cell, in 3 directions # r = np.random.randint(-1,2, size=(ww[0].shape[0],3), dtype='int') # # replace nan by random neighbors # thisChi[ww[0],ww[1],ww[2]] = thisChi[(ww[0]+r[:,0])%Ng, (ww[1]+r[:,1])%Ng, (ww[2]+r[:,2])%Ng] # # recompute indices of nan cells # ww = np.where(np.isnan(thisChi)) # have_empty_cells = (ww[0].shape[0] > 0) else: raise Exception("Invalid fill_empty_cells option: %s" % str(fill_empty_cells)) return thisChi, thisChi_attrs
def main(): """ Script to calculate skew spectra of an input density. """ ##################################### # PARSE COMMAND LINE ARGS ##################################### ap = ArgumentParser() ap.add_argument('--SimSeed', type=int, default=400, help='Simulation seed to load.') ap.add_argument('--boxsize', default=1500.0, type=float, help='Boxsize in Mpc/h.') ap.add_argument('--ApplyRSD', type=int, default=0, help='0: No RSD. 1: Include RSD in catalog.') ap.add_argument('--Rsmooth', default=20.0, type=float, help='Smoothing of quad field.') ap.add_argument('--Ngrid', default=64, type=int, help='Ngrid used to compute skew spectra.') ap.add_argument('--SubsampleRatio', default=0.0015, type=float, help='Subsample ratio of DM snapshot to use as input.') ap.add_argument('--MaxDisplacement', default=100.0, type=float, help='Maximum RSD displacement in Mpc/h.') ap.add_argument('--DensitySource', default='catalog', type=str, help='Source from which to compute the density. catalog or delta_2SPT') ap.add_argument('--b1', default=1.0, type=float, help='b1 bias. Only used if DensitySource=delta_2SPT.') ap.add_argument('--b2', default=0.0, type=float, help='b2 bias. Only used if DensitySource=delta_2SPT.') ap.add_argument('--bG2', default=0.0, type=float, help='bG2 bias. Only used if DensitySource=delta_2SPT.') ap.add_argument('--fLogGrowth', default=0.786295, type=float, help='Logarithmic growth factor f. Only used if DensitySource=delta_2SPT.') cmd_args = ap.parse_args() ##################################### # OPTIONS ##################################### opts = OrderedDict() opts['code_version_for_outputs'] = '0.2' # Parse options opts['sim_seed'] = cmd_args.SimSeed opts['boxsize'] = cmd_args.boxsize basedir = os.path.expandvars('$SCRATCH/lss/ms_gadget/run4/00000%d-01536-%.1f-wig/' % ( opts['sim_seed'], opts['boxsize'])) opts['sim_scale_factor'] = 0.625 opts['Rsmooth'] = cmd_args.Rsmooth opts['Ngrid'] = cmd_args.Ngrid opts['LOS'] = np.array([0,0,1]) opts['APPLY_RSD'] = bool(cmd_args.ApplyRSD) opts['subsample_ratio'] = cmd_args.SubsampleRatio opts['max_displacement'] = cmd_args.MaxDisplacement # which multipoles (ell) to compute opts['poles'] = [0,2] # Source from which to compute density: 'catalog' or 'delta_2SPT' opts['density_source'] = cmd_args.DensitySource # more options if source is catalog if opts['density_source'] == 'catalog': # Catalog with particle positions: 'DM_subsample' or 'gal_ptchall_with_RSD' opts['positions_catalog'] = 'DM_subsample' # Velocity source: DM_sim, deltalin_D2, deltalin_D2_2SPT opts['velocity_source'] = 'DM_sim' elif opts['density_source'] == 'delta_2SPT': opts['positions_catalog'] = None opts['velocity_source'] = None else: raise Exception('Invalid density_source %s' % opts['density_source']) # cosmology of ms_gadget sims (to compute D_lin(z)) # omega_m = 0.307494 # omega_bh2 = 0.022300 # omega_ch2 = 0.118800 # h = math.sqrt((omega_bh2 + omega_ch2) / omega_m) = 0.6774 opts['cosmo_params'] = dict(Om_m=0.307494, Om_L=1.0 - 0.307494, Om_K=0.0, Om_r=0.0, h0=0.6774) #opts['f_log_growth'] = 0.7862951 # np.sqrt(0.61826) opts['f_log_growth'] = cmd_args.fLogGrowth # where to save output (todo: move code to a function ) DS_string = '_DS%s' % opts['density_source'] if opts['density_source'] == 'catalog': sr_string = '_sr%g' % opts['subsample_ratio'] vel_string = '_v%s' % opts['velocity_source'] MD_string = '_MD%g' % opts['max_displacement'] else: sr_string, vel_string, MD_string = '', '', '' b1, b2, bG2 = cmd_args.b1, cmd_args.b2, cmd_args.bG2 if opts['density_source'] == 'catalog': b_string, f_string = '', '' elif opts['density_source'] == 'delta_2SPT': b_string = '' if (b1==1. and b2==0. and bG2==0.) else ('_b%g_%g_%g' % (b1,b2,bG2)) f_string = '' if opts['f_log_growth']==0.786295 else ('_f%g' % opts['f_log_growth']) opts['outdir'] = 'data/Pskew_sims/00000%d-01536-%.1f-wig/R%.1f_Ng%d_RSD%d%s%s%s%s%s%s/' % ( opts['sim_seed'], opts['boxsize'], opts['Rsmooth'], opts['Ngrid'], int(opts['APPLY_RSD']), DS_string, sr_string, vel_string, MD_string, b_string, f_string) # ########################### # START PROGRAM # ########################### Nmesh = opts['Ngrid'] BoxSize = np.array([opts['boxsize'], opts['boxsize'], opts['boxsize']]) comm = CurrentMPIComm.get() LOS = opts['LOS'] LOS_string = 'LOS%d%d%d' % (LOS[0], LOS[1], LOS[2]) # Below, 'D' stands for RSD displacement in Mpc/h: D=v/(aH)=f*PsiDot. ### Catalogs # DM subsample if opts['subsample_ratio'] != 1.0: DM_subsample = Target( name='DM_subsample', in_fname=os.path.join(basedir, 'snap_%.4f_sub_sr%g_ssseed40%d.bigfile' % ( opts['sim_scale_factor'], opts['subsample_ratio'], opts['sim_seed'])), position_column='Position' ) else: # full DM sample DM_subsample = Target( name='DM_subsample', in_fname=os.path.join(basedir, 'snap_%.4f' % ( opts['sim_scale_factor'])), position_column='Position' ) # DM_D2 = Target( # name='DM_D2', # in_fname=os.path.join(basedir, 'snap_%.4f_sub_sr0.0015_ssseed40%d.bigfile' % ( # opts['sim_scale_factor'], opts['sim_seed'])), # position_column='Position', # val_column='Velocity', # val_component=2, # rescale_factor='RSDFactor' # ) if False: # PT Challenge galaxies from rockstar halos. Rockstar gives core positions and velocities. # Units: 1/(aH) = 1./(a * H0*np.sqrt(Om_m/a**3+Om_L)) * (H0/100.) in Mpc/h / (km/s). # For ms_gadget, get 1/(aH) = 0.01145196 Mpc/h/(km/s) = 0.0183231*0.6250 Mpc/h/(km/s). # Note that MP-Gadget files have RSDFactor=1/(a^2H)=0.0183231 for a=0.6250 b/c they use a^2\dot x for Velocity. from lsstools.sim_galaxy_catalog_creator import PTChallengeGalaxiesFromRockstarHalos assert opts['sim_scale_factor'] == 0.625 # PT challenge galaxies, apply RSD to position (TEST) assert opts['sim_scale_factor'] == 0.625 gal_ptchall_with_RSD = Target( name='gal_ptchall_with_RSD', in_fname=os.path.join(basedir, 'snap_%.4f.gadget3/rockstar_out_0.list.parents.bigfile' % opts['sim_scale_factor']), position_column='Position', velocity_column='Velocity', apply_RSD_to_position=True, RSD_los=LOS, RSDFactor=0.01145196, #val_column='Velocity', # This is rockstar velocity, which is v=a\dot x in km/s ("Velocities in km / s (physical, peculiar)") #val_component=0, #rescale_factor=0.01145196, # RSD displacement in Mpc/h is D=v/(aH)=0.01145196*v. cuts=[PTChallengeGalaxiesFromRockstarHalos( log10M_column='log10Mvir', log10Mmin=12.97, sigma_log10M=0.35, RSD=False), #('Position', 'max', [100.,100.,20.]) ] ) ### Densities # Linear density z_rescalefac = linear_rescale_fac(current_scale_factor=1.0, desired_scale_factor=opts['sim_scale_factor'], cosmo_params=opts['cosmo_params']) print('z_rescalefac:', z_rescalefac) deltalin = Model( name='deltalin', in_fname=os.path.join(basedir, 'IC_LinearMesh_z0_Ng%d' % opts['Ngrid']), rescale_factor=z_rescalefac, read_mode='delta from 1+delta', filters=None, readout_window='cic') # Linear RSD displacement def k2ovksq_filter_fcn(k, v, d=2): ksq = sum(ki**2 for ki in k) return np.where(ksq == 0.0, 0*v, 1j*k[d] * v / (ksq)) deltalin_D2 = Model( name='deltalin_D2', in_fname=os.path.join(basedir, 'IC_LinearMesh_z0_Ng%d' % opts['Ngrid']), rescale_factor=opts['f_log_growth']*z_rescalefac, read_mode='delta from 1+delta', filters=[k2ovksq_filter_fcn], readout_window='cic') # 2nd order SPT displacement: ik/k^2*G2 deltalin_D2_2ndorderSPTcontri = Model( name='deltalin_D2_2ndorderSPTcontri', in_fname=os.path.join(basedir, 'IC_LinearMesh_z0_Ng%d' % opts['Ngrid']), rescale_factor=opts['f_log_growth']*z_rescalefac, read_mode='delta from 1+delta', filters=[k2ovksq_filter_fcn], readout_window='cic') if opts['density_source'] == 'delta_2SPT': # compute 2nd order density from linear mesh if comm.rank == 0: print('Warning: Not smoothing when generating 2nd order field') if not opts['APPLY_RSD']: # Compute delta1 + F2[delta1] delta_mesh = FieldMesh( deltalin.get_mesh().compute() + QuadField(composite='F2').compute_from_mesh(deltalin.get_mesh()).compute(mode='real') ) else: # compute 2nd order density in redshift space f = opts['f_log_growth'] delta_mesh = FieldMesh( b1 * deltalin.get_mesh().compute() + f * LinField(m=2*LOS, n=-2).compute_from_mesh(deltalin.get_mesh()).compute() + b1 * QuadField(composite='F2').compute_from_mesh(deltalin.get_mesh()).compute() + f * QuadField(composite=('velocity_G2_par_%s' % LOS_string)).compute_from_mesh(deltalin.get_mesh()).compute(mode='real') + b1 * f * QuadField(nprime=-2, mprime=LOS, mprimeprime=LOS).compute_from_mesh(deltalin.get_mesh()).compute() + f**2 * QuadField(n=-2, m=2*LOS, nprime=-2, mprime=LOS, mprimeprime=LOS).compute_from_mesh(deltalin.get_mesh()).compute() ) if b2 != 0.: delta_mesh = FieldMesh( delta_mesh.compute() + b2/2.0 * QuadField().compute_from_mesh(deltalin.get_mesh()).compute()) if bG2 != 0.0: delta_mesh = FieldMesh( delta_mesh.compute() + bG2 * QuadField(composite='tidal_G2').compute_from_mesh(deltalin.get_mesh()).compute()) elif opts['density_source'] == 'catalog': ########################################################################## # Get DM catalog in redshift space (if APPLY_RSD==True) ########################################################################## # get the catalog if opts['positions_catalog'] == 'gal_ptchall_with_RSD': target = gal_ptchall_with_RSD elif opts['positions_catalog'] == 'DM_subsample': target = DM_subsample cat = target.get_catalog(keep_all_columns=True) if comm.rank == 0: print('Positions catalog:') print(cat.attrs) # add redshift space positions, assuming LOS is in z direction if opts['velocity_source'] == 'DM_sim': # use DM velocity displacement = cat['Velocity']*cat.attrs['RSDFactor'] * opts['LOS'] displacement = displacement.compute() if opts['max_displacement'] is not None: ww = np.where(np.linalg.norm(displacement, axis=1)>opts['max_displacement'])[0] print(r'%d: max_displacement=%g. Set velocity to 0 for %g percent of particles' % ( comm.rank, opts['max_displacement'], 100.*float(len(ww)/displacement.shape[0]))) displacement[ww,0] *= 0 displacement[ww,1] *= 0 displacement[ww,2] *= 0 cat['RSDPosition'] = cat['Position'] + displacement elif opts['velocity_source'] in ['deltalin_D2', 'deltalin_D2_2SPT']: if opts['velocity_source'] == 'deltalin_D2': print('Warning: linear velocity does not have 2nd order G2 velocity, so expect wrong bispectrum') assert np.all(opts['LOS'] == np.array([0,0,1])) mtp = ModelTargetPair(model=deltalin_D2, target=DM_subsample) cat['RSDPosition'] = cat['Position'].compute() # read out model at each catalog position on current rank. model_at_target_pos = mtp.readout_model_at_target_pos() mat = np.zeros((model_at_target_pos.size,3)) assert np.all(opts['LOS'] == np.array([0,0,1])) mat[:,2] = model_at_target_pos cat['RSDPosition'] += mat if opts['velocity_source'] == 'deltalin_D2_2SPT': # add 2nd order G2 contribution to velocity raise Exception('todo') else: raise Exception('Invalid velocity_source: %s' % str(opts['velocity_source'])) if comm.rank == 0: print('rms RSD displacement: %g Mpc/h' % np.mean((cat['Position'].compute()-cat['RSDPosition'].compute())**2)**0.5) print('max RSD displacement: %g Mpc/h' % np.max(np.abs(cat['Position'].compute()-cat['RSDPosition'].compute()))) # Get redshift space catalog RSDcat = catalog_persist(cat, columns=['ID','PID','Position','RSDPosition','Velocity', 'log10Mvir']) del cat if opts['APPLY_RSD']: if comm.rank == 0: print('Applying RSD') RSDcat['Position'] = RSDcat['RSDPosition'] else: if comm.rank == 0: print('Not applying RSD') # paint to mesh delta_mesh = FieldMesh(RSDcat.to_mesh(Nmesh=Nmesh, BoxSize=BoxSize, window='cic', interlaced=False, compensated=False).compute()-1) if comm.rank == 0: print('# objects: Original: %d' % (RSDcat.csize)) else: raise Exception('Invalid density_source %s' % opts['density_source']) ########################################################################## # Calculate power spectrum multipoles ########################################################################## def calc_power(mesh, second=None, mode='1d', k_bin_width=1.0, verbose=False, los=None, poles=None): BoxSize = mesh.attrs['BoxSize'] assert BoxSize[0] == BoxSize[1] assert BoxSize[0] == BoxSize[2] boxsize = BoxSize[0] dk = 2.0 * np.pi / boxsize * k_bin_width kmin = 2.0 * np.pi / boxsize / 2.0 if mode == '1d': res = FFTPower(first=mesh, second=second, mode=mode, dk=dk, kmin=kmin) elif mode == '2d': if poles is None: poles = [0,2,4] res = FFTPower(first=mesh, second=second, mode=mode, dk=dk, kmin=kmin, poles=poles, Nmu=5, los=los) else: raise Exception("Mode not implemented: %s" % mode) return res ## Compute density power spectrum Pdd = calc_power(delta_mesh, los=opts['LOS'], mode='2d', poles=opts['poles']) ########################################################################## # Get RSD skew spectra ########################################################################## power_kwargs={'mode': '2d', 'poles': opts['poles']} # Apply smoothing smoothers = [smoothing.GaussianSmoother(R=opts['Rsmooth'])] delta_mesh_smoothed = FieldMesh(delta_mesh.compute(mode='real')) for smoother in smoothers: delta_mesh_smoothed = smoother.apply_smoothing(delta_mesh_smoothed) # Compute skew spectra skew_spectra = SkewSpectrum.get_list_of_standard_skew_spectra( LOS=LOS, redshift_space_spectra=True) for skew_spec in skew_spectra: # compute and store in skew_spec.Pskew skew_spec.compute_from_mesh(mesh=delta_mesh_smoothed, third_mesh=delta_mesh, power_kwargs=power_kwargs, store_key='default_key') # Store results in files if comm.rank == 0: if not os.path.exists(opts['outdir']): os.makedirs(opts['outdir']) if False: for skew_spec in skew_spectra: # store as json fname = os.path.join(opts['outdir'], skew_spec.name+'.json') skew_spec.Pskew.save(fname) print('Wrote %s' % fname) # store as plain text fname = os.path.join(opts['outdir'], skew_spec.name+'.txt') skew_spec.save_plaintext(fname) print('Wrote %s' % fname) # store all in one file for each multipole if skew_spectra != []: for ell in skew_spec.Pskew['default_key'].attrs['poles']: mydtype = [('k', 'f8')] for skew_spec in skew_spectra: mydtype.append((skew_spec.name, 'f8')) arr = np.empty(shape=skew_spec.Pskew['default_key'].poles['k'].shape, dtype=mydtype) arr['k'] = skew_spec.Pskew['default_key'].poles['k'] for skew_spec in skew_spectra: arr[skew_spec.name] = skew_spec.Pskew['default_key'].poles['power_%d'%ell].real fname = os.path.join(opts['outdir'], 'Sn_ell%d.txt'%ell) header = 'Columns: ' + str(arr.dtype.names) np.savetxt(fname, arr, header=header) print('Wrote %s' % fname) # also store density power for ell in Pdd.attrs['poles']: mydtype = [('k', 'f8'), ('P', 'f8')] arr = np.empty(shape=Pdd.poles['k'].shape, dtype=mydtype) arr['k'] = Pdd.poles['k'] arr['P'] = Pdd.poles['power_%d'%ell].real fname = os.path.join(opts['outdir'], 'P_ell%d.txt'%ell) header = 'Columns: ' + str(arr.dtype.names) np.savetxt(fname, arr, header=header) print('Wrote %s' % fname)
def calc_model_errors_at_cat_pos(trf_specs=None, paths=None, cat_specs=None, ext_grids_to_load=None, trf_fcn_opts=None, grid_opts=None, sim_opts=None, power_opts=None, Pkmeas_helper_columns=None, Pkmeas_helper_columns_calc_crosses=False, f_log_growth=None, debug=True): """ Compute models at catalog positions, and residual to target, for each trf_spec. """ comm = CurrentMPIComm.get() result_of_trf_spec = OrderedDict() for trf_spec in trf_specs: cats = OrderedDict() if True: ## Load target catalog cat_id = trf_spec.target_field cat_opts = cat_specs[cat_id] # read file fname = os.path.join(paths['in_path'], cat_opts['in_fname']) if comm.rank == 0: print('Read %s' % fname) # Read rho, not dividing or subtracting mean, to get velocity correctly cat = BigFileCatalog(fname, dataset='./', header='Header') # cuts for cut_column, cut_instruction in cat_opts['cuts'].items(): cut_op, cut_value = cut_instruction if cut_op == 'min': cat = cat[cat[cut_column] >= cut_value] elif cut_op == 'max': cat = cat[cat[cut_column] < cut_value] else: raise Exception('Invalid cut operation %s' % str(cut_op)) # Set Position column cat['Position'] = cat[cat_opts['position_column']] # compute the value we are interested in, save in 'val' column component = cat_opts['val_component'] if component is None: cat['val'] = cat[cat_opts['val_column']][:] else: cat['val'] = cat[cat_opts['val_column']][:, component] # catalog rescale factor if cat_opts.get('rescale_factor', None) is not None: if cat_opts['rescale_factor'] == 'RSDFactor': cat['val'] *= cat.attrs['RSDFactor'][0] else: raise Exception('Invalid rescale_factor %s' % cat_opts['rescale_factor']) # additional rescale factors or post processing options for catalog if hasattr(trf_spec, 'field_opts'): if cat_id in trf_spec.field_opts: this_field_opts = trf_spec.field_opts[cat_id] if this_field_opts.get('additional_rescale_factor', 1.0) != 1.0: resc_fac = this_field_opts['additional_rescale_factor'] if type(resc_fac) == float: cat['val'] *= resc_fac else: raise Exception("Invalid rescale factor: %s" % resc_fac) # keep only Position and val columns, delete all other columns cat2 = catalog_persist(cat, columns=['Position', 'val']) del cat cstats = get_cstats_string(cat2['val'].compute()) if comm.rank == 0: print('CATALOG %s: %s\n' % (cat_id, cstats)) # Store in cats dict cats[cat_id] = cat2 del cat2 target_cat = cats[trf_spec.target_field] # Make sure we have no linear source (trf fcns not implemented here) if trf_spec.linear_sources is not None: raise Exception( 'Trf fcns not implemented for model error at cat pos') ## Load remaining fields, assuming they are bigfiles meshs, and get # residual at target positions. residual_cat = target_cat.copy() for mesh_id in trf_spec.fixed_linear_sources: fname = os.path.join(paths['in_path'], mesh_id) if comm.rank == 0: print('Read %s' % fname) if hasattr(trf_spec, 'field_opts'): fopts = trf_spec.field_opts[mesh_id] else: # default options for meshs fopts = {'read_mode': 'velocity', 'readout_window': 'cic'} # read mesh if fopts['read_mode'] == 'velocity': # get rho (don't divide by mean) mesh = read_vel_from_bigfile(fname) elif fopts['read_mode'] == 'density': # compute fractional delta (taking out mean) mesh = read_delta_from_bigfile(fname) else: raise Exception('Invalid read_mode: %s' % fopts['read_mode']) # more post processing of mesh if fopts.get('additional_rescale_factor', 1.0) != 1.0: resc_fac = fopts['additional_rescale_factor'] if comm.rank == 0: print('Apply rescale fac to %s: %s' % (mesh_id, str(resc_fac))) if type(resc_fac) == float: mesh = FieldMesh(mesh.compute(mode='real') * resc_fac) elif resc_fac == 'f_log_growth': mesh = FieldMesh(mesh.compute(mode='real') * f_log_growth) else: raise Exception("Invalid rescale factor: %s" % resc_fac) cstats = get_cstats_string(mesh.compute()) if comm.rank == 0: print('MESH %s: %s\n' % (mesh_id, cstats)) # Read out mesh at Position of target catalog if comm.rank == 0: print('Read out mesh at target pos: %s' % mesh_id) mesh_at_target_pos = readout_mesh_at_cat_pos( mesh=mesh, cat=target_cat, readout_window=fopts['readout_window']) ## Compute residual = target - fixed_linear_sources at target position residual_cat['val'] -= mesh_at_target_pos cstats = get_cstats_string(residual_cat['val'].compute()) if comm.rank == 0: print('RESIDUAL %s - %s: %s\n' % (trf_spec.target_field, trf_spec.save_bestfit_field, cstats)) # plot histogram if debug: import matplotlib.pyplot as plt plt.hist(residual_cat['val'].compute(), bins=50, range=(-10, 10), alpha=0.5) plt.hist(target_cat['val'].compute(), bins=50, range=(-10, 10), alpha=0.5) plt.hist(mesh_at_target_pos, bins=50, range=(-10, 10), alpha=0.5) plt.show() ## paint to mesh to_mesh_kwargs = { #'Nmesh': grid_opts.Ngrid, 'window': 'cic', 'compensated': False, 'interlaced': False, #'BoxSize': residual_cat.attrs['BoxSize'], 'dtype': 'f8' } # rho mesh of target, taking avg value instead of summing target_mesh = FieldMesh( mass_avg_weighted_paint_cat_to_rho(target_cat, weight='val', Nmesh=grid_opts.Ngrid, to_mesh_kwargs=to_mesh_kwargs, rho_of_empty_cells=0.0, verbose=True)[0]) # rho mesh of residual, taking avg value instead of summing residual_mesh = FieldMesh( mass_avg_weighted_paint_cat_to_rho(residual_cat, weight='val', Nmesh=grid_opts.Ngrid, to_mesh_kwargs=to_mesh_kwargs, rho_of_empty_cells=0.0, verbose=True)[0]) ## Compute power spectra pow_kwargs = { 'mode': power_opts.Pk_1d_2d_mode, 'k_bin_width': power_opts.k_bin_width } ## Save results of this trf_spec in dict trf_res = OrderedDict() # Power spectra trf_res['power'] = OrderedDict() trf_res['power']['Target-Model'] = calc_power(residual_mesh, **pow_kwargs) trf_res['power']['Target'] = calc_power(target_mesh, **pow_kwargs) if debug: # plot power spectra import matplotlib.pyplot as plt print('plot power') k = trf_res['power']['Target'].power['k'] plt.loglog(k, k**2 * trf_res['power']['Target'].power['power']) k = trf_res['power']['Target-Model'].power['k'] plt.loglog(k, k**2 * trf_res['power']['Target-Model'].power['power']) plt.show() # target catalog results trf_res['target_cat_attrs'] = target_cat.attrs trf_res['target_cat_csize'] = target_cat.csize # todo: save rms of fields and target result_of_trf_spec[trf_spec.save_bestfit_field] = trf_res # save all results in a big pickle_dict pickle_dict = OrderedDict() pickle_dict['result_of_trf_spec'] = result_of_trf_spec print('pickle_dict:', pickle_dict) return pickle_dict
def calculate_model_error(sim_opts=None, grid_opts=None, power_opts=None, trf_fcn_opts=None, ext_grids_to_load=None, xgrids_in_memory=None, kgrids_in_memory=None, cats=None, trf_specs=None, keep_pickle=False, pickle_file_format='dill', pickle_path='$SCRATCH/perr/pickle/', Pkmeas_helper_columns=None, Pkmeas_helper_columns_calc_crosses=False, store_Pkmeas_in_trf_results=False, save_grids4plots=False, grids4plots_base_path=None, grids4plots_R=None, cache_base_path=None, RSDstrings=None, code_version_for_pickles=None, return_fields=None, shifted_fields_Np=None, shifted_fields_Nmesh=None, shifted_fields_RPsi=None): """ Calculate the model error for all models specified by trf_specs. Use return_fields=['bestfit'] or ['residual'] or ['bestfit','residual'] to return the fields as well, as lists in same order as trf_specs. """ # store opts in dict so we can save in pickle later opts = dict( sim_opts=sim_opts, grid_opts=grid_opts, power_opts=power_opts, trf_fcn_opts=trf_fcn_opts, ext_grids_to_load=ext_grids_to_load, #xgrids_in_memory=xgrids_in_memory, #kgrids_in_memory=kgrids_in_memory, cats=cats, trf_specs=trf_specs, keep_pickle=keep_pickle, pickle_file_format=pickle_file_format, pickle_path=pickle_path, Pkmeas_helper_columns=Pkmeas_helper_columns, Pkmeas_helper_columns_calc_crosses=Pkmeas_helper_columns_calc_crosses, store_Pkmeas_in_trf_results=store_Pkmeas_in_trf_results, save_grids4plots=save_grids4plots, grids4plots_base_path=grids4plots_base_path, grids4plots_R=grids4plots_R, cache_base_path=cache_base_path, code_version_for_pickles=code_version_for_pickles) ##################################### # Initialize ##################################### # make sure we keep the pickle if it is a big run and do not plot if grid_opts.Ngrid > 256: keep_pickle = True # load defaults if not set if ext_grids_to_load is None: ext_grids_to_load = sim_opts.get_default_ext_grids_to_load( Ngrid=grid_opts.Ngrid) if cats is None: cats = sim_opts.get_default_catalogs() ### derived options (do not move above b/c command line args might ### overwrite some options!) opts['in_path'] = path_utils.get_in_path(opts) # for output densities opts['out_rho_path'] = os.path.join(opts['in_path'], 'out_rho_Ng%d' % grid_opts.Ngrid) # expand environment names in paths paths = {} for key in [ 'in_path', 'in_fname', 'in_fname_PTsim_psi_calibration', 'in_fname_halos_to_displace_by_mchi', 'pickle_path', 'cache_base_path', 'grids4plots_base_path', 'out_rho_path' ]: if opts.has_key(key): if opts[key] is None: paths[key] = None else: paths[key] = os.path.expandvars(opts[key]) setup_logging() comm = CurrentMPIComm.get() logger = logging.getLogger('PerrCalc') model_spec.check_trf_specs_consistency(trf_specs) # Init Pickler instance to save pickle later (this will init pickle fname) pickler = None if comm.rank == 0: pickler = Pickler(path=paths['pickle_path'], base_fname='main_calc_Perr', file_format=pickle_file_format, rand_sleep=(grid_opts.Ngrid > 128)) print("Pickler: ", pickler.full_fname) pickler = comm.bcast(pickler, root=0) # where to save grids for slice and scatter plots if save_grids4plots: paths['grids4plots_path'] = os.path.join( paths['grids4plots_base_path'], os.path.basename(pickler.full_fname)) if comm.rank == 0: if not os.path.exists(paths['grids4plots_path']): os.makedirs(paths['grids4plots_path']) print("grids4plots_path:", paths['grids4plots_path']) paths['cache_path'] = utils.make_cache_path(paths['cache_base_path'], comm) # Get list of all densities actually needed for trf fcns. densities_needed_for_trf_fcns = utils.get_densities_needed_for_trf_fcns( trf_specs) #opts['densities_needed_for_trf_fcns'] = densities_needed_for_trf_fcns # ########################################################################## # Run program. # ########################################################################## #if opts.get('RSDstrings', ['']) != ['']: if True or RSDstrings not in [None, ['']]: # calculate D and f cosmo = CosmoModel(**sim_opts.cosmo_params) calc_Da = generate_calc_Da(cosmo=cosmo) f_log_growth = calc_f_log_growth_rate(a=sim_opts.sim_scale_factor, calc_Da=calc_Da, cosmo=cosmo, do_test=True) # save in opts so we can easily access it throughout code (although strictly # speaking it is not a free option but derived from cosmo_params) opts['f_log_growth'] = f_log_growth else: opts['f_log_growth'] = None # Compute best-fit model and power spectra # TODO: maybe split into method computing field and separate method to compute power spectra. # For now, load fields from cache as workaround (see below) pickle_dict = combine_fields.paint_combine_and_calc_power( trf_specs=trf_specs, paths=paths, catalogs=cats, needed_densities=densities_needed_for_trf_fcns, ext_grids_to_load=ext_grids_to_load, xgrids_in_memory=xgrids_in_memory, kgrids_in_memory=kgrids_in_memory, trf_fcn_opts=trf_fcn_opts, grid_opts=grid_opts, sim_opts=sim_opts, power_opts=power_opts, save_grids4plots=save_grids4plots, grids4plots_R=grids4plots_R, Pkmeas_helper_columns=Pkmeas_helper_columns, Pkmeas_helper_columns_calc_crosses=Pkmeas_helper_columns_calc_crosses, store_Pkmeas_in_trf_results=store_Pkmeas_in_trf_results, f_log_growth=opts['f_log_growth']) # Load fields from cache if they shall be returned # (actually not used anywhere, could delete) if return_fields is not None: if 'bestfit' in return_fields: bestfit_fields = [] for trf_spec in trf_specs: # load bestfit fields from cache gridk = ComplexGrid(fname=pickle_dict['gridk_cache_fname'], read_columns=[trf_spec.save_bestfit_field]) bestfit_fields.append(gridk.G[trf_spec.save_bestfit_field]) del gridk if 'residual' in return_fields: residual_fields = [] for trf_spec in trf_specs: # load residual field from cache residual_key = '[%s]_MINUS_[%s]' % ( trf_spec.save_bestfit_field, trf_spec.target_field) gridk = ComplexGrid(fname=pickle_dict['gridk_cache_fname'], read_columns=[residual_key]) residual_fields.append(gridk.G[residual_key]) del gridk # copy over opts so they are saved assert not pickle_dict.has_key('opts') pickle_dict['opts'] = opts.copy() # save all resutls to pickle if comm.rank == 0: pickler.write_pickle(pickle_dict) # print path with grids for slice and scatter plotting if save_grids4plots: print("grids4plots_path: %s" % paths['grids4plots_path']) # print save_bestfit_fields save_bestfit_fields = [t.save_bestfit_field for t in opts['trf_specs']] print('\nsave_bestfit_fields:\n' + '\n'.join(save_bestfit_fields)) # delete pickle if not wanted any more if comm.rank == 0: if keep_pickle: print("Pickle: %s" % pickler.full_fname) else: pickler.delete_pickle_file() # delete cache dir from shutil import rmtree rmtree(paths['cache_path']) if return_fields in [False, None]: return pickle_dict elif return_fields == ['bestfit']: return bestfit_fields, pickle_dict elif return_fields == ['residual']: return residual_fields, pickle_dict elif return_fields == ['bestfit', 'residual']: return bestfit_fields, residual_fields, pickle_dict
def main(): """ Read illustris ICs """ #### # OPTIONS #### Nmesh_lst = [2500] if False: # Illustris-3 # https://www.illustris-project.org/data/downloads/Illustris-3/ fname = '/data/mschmittfull/lss/IllustrisTNG/Illustris-3/output/snap_ics.hdf5' else: # TNG300-1 / L205n2500TNG fname = '/data/mschmittfull/lss/IllustrisTNG/L205n2500TNG/output/snap_ics.hdf5' #### # Run code #### comm = CurrentMPIComm.get() print('Greetings from rank %d' % comm.rank) for Nmesh in Nmesh_lst: if comm.rank == 0: print('Reading header from %s' % fname) # Read header f = h5py.File(fname, 'r') if comm.rank == 0: print('header: ', f['Header'].attrs.keys()) attrs = dict( Redshift=f['Header'].attrs[u'Redshift'], Time=f['Header'].attrs[u'Time'], BoxSize=f['Header'].attrs[u'BoxSize'] / 1e3 * np.ones(3), #convert to Mpc/h NumFilesPerSnapshot=f['Header'].attrs[u'NumFilesPerSnapshot'], NumPart_ThisFile=f['Header'].attrs[u'NumPart_ThisFile'], Omega0=f['Header'].attrs[u'Omega0'], OmegaLambda=f['Header'].attrs[u'OmegaLambda'], HubbleParam=f['Header'].attrs[u'HubbleParam'], MassTable=f['Header'].attrs[u'MassTable'], TotNumPart=f['Header'].attrs[u'NumPart_Total']) if comm.rank == 0: print('attrs:') print(attrs) f.close() # read particles positions if comm.rank == 0: print('Reading particle data from %s' % fname) cat = HDFCatalog( fname, exclude=['PartType1/ParticleIDs', 'PartType1/Velocities'], attrs=attrs) if comm.rank == 0: print('columns: ', cat.columns) # convert kpc/h to Mpc/h cat['Position'] = cat['PartType1/Coordinates'] / 1e3 if comm.rank == 0: print('paint...') catmesh = cat.to_mesh(Nmesh=Nmesh, window='cic', compensated=False, interlaced=False) if False: # crashes when using too many cores (empty arrays) rfield = catmesh.compute() stats = get_cstats_string(rfield) if comm.rank == 0: print('density stats: %s' % stats) # save to bigfile out_fname = '%s_PtcleDensity_z%d_Ng%d' % ( fname, int(catmesh.attrs['Redshift']), Nmesh) if comm.rank == 0: print('Writing to %s' % out_fname) catmesh.save(out_fname) if comm.rank == 0: print('Wrote %s' % out_fname)
def paint_cat_to_gridk(PaintGrid_config, gridx=None, gridk=None, column=None, drop_gridx_column=True, rescalefac=1.0, f_log_growth=None, skip_fft=False, Ngrid=None, boxsize=None, grid_ptcle2grid_deconvolution=None, to_mesh_kwargs=None, kmax=None, cache_path=None): """ TODO: Should reimplement this at some point so it does not use the old PaintGrid_config dict any more but only kwargs. """ from nbodykit import CurrentMPIComm comm = CurrentMPIComm.get() logger = logging.getLogger('paint_utils') logger.info("Rank %d: paint_cat_to_gridk" % comm.rank) # return gridx and gridk if they were not passed as args return_gridx = (gridx is None) return_gridk = (gridk is None) # check if catalog path exists if PaintGrid_config.has_key('DataSource'): if PaintGrid_config['DataSource'].has_key('path'): tmp_path = PaintGrid_config['DataSource']['path'] if not os.path.exists(tmp_path): raise Exception('File does not exist: %s' % str(tmp_path)) # nbkit code is in /Users/msl2/anaconda/anaconda/envs/nbodykit-0.3.7-env/lib/python2.7/site-packages/nbodykit/ if True: # check nbkit version is new enough import nbodykit if comm.rank == 0: logger.info("nbkit version: %s" % str(nbodykit.__version__)) if not nbodykit.__version__.startswith('0.3.'): raise Exception( "Please use nbodykit 0.3.7 or higher. Maybe have to run with python, not pythonw" ) # check PaintGrid_config is supported here if PaintGrid_config['Painter']['plugin'] != 'DefaultPainter': raise Exception("nbkit 0.3 wrapper not implemented for plugin %s" % str(PaintGrid_config['plugin'])) implemented_painter_keys = [ 'plugin', 'normalize', 'setMean', 'paint_mode', 'velocity_column', 'fill_empty_cells', 'randseed_for_fill_empty_cells', 'raise_exception_if_too_many_empty_cells' ] for k in PaintGrid_config['Painter'].keys(): if k not in implemented_painter_keys: raise Exception("config key %s not implemented in wrapper" % str(k)) # TODO: implement keys 'weight' or weight_ptcles_by (different particles contribute with different weight or mass) ## Do the painting # Paint number density (e.g. galaxy overdensity) or divergence of momentum density paint_mode = PaintGrid_config['Painter'].get('paint_mode', None) if PaintGrid_config['DataSource']['plugin'] == 'Subsample': ## Read hdf5 catalog file and paint # TODO: get 'same value error' when using mass weighted columns b/c nbk03 # thinks two columns have same name. probably string issue. cat_nbk = read_ms_hdf5_catalog(PaintGrid_config['DataSource']['path']) # Paint options if grid_ptcle2grid_deconvolution is not None: raise Exception("grid_ptcle2grid_deconvolution not implemented any more") if to_mesh_kwargs is None: raise Exception('must specify to_mesh_kwargs') # to_mesh_kwargs = { # 'window': 'cic', # 'compensated': False, # 'interlaced': False, # 'BoxSize': np.array([boxsize, boxsize, boxsize]), # 'dtype': 'f8' # } if paint_mode in [None, 'overdensity']: # Paint overdensity. # Init CatalogMesh object (not painted yet) #mesh = cat_nbk.to_mesh(Nmesh=Ngrid, weight='log10M', **to_mesh_kwargs) mesh = cat_nbk.to_mesh(Nmesh=Ngrid, **to_mesh_kwargs) if comm.rank == 0: logger.info("mesh type: %s" % str(type(mesh))) logger.info("mesh attrs: %s" % str(mesh.attrs)) # Paint. If normalize=True, outfield = 1+delta; if normalize=False: outfield=rho normalize = PaintGrid_config['Painter'].get('normalize', True) outfield = mesh.to_real_field(normalize=normalize) elif paint_mode == 'momentum_divergence': # Paint momentum divergence div((1+delta)v/(faH)). # See https://nbodykit.readthedocs.io/en/latest/cookbook/painting.html#Painting-the-Line-of-sight-Momentum-Field. assert PaintGrid_config['Painter']['velocity_column'] is not None assert f_log_growth is not None theta_k = None for idir in [0, 1, 2]: vi_label = 'tmp_V_%d' % idir # this is the velocity / (f*a*H). units are Mpc/h cat_nbk[vi_label] = (cat_nbk[PaintGrid_config['Painter'] ['velocity_column']][:, idir] / f_log_growth) to_mesh_kwargs.update(dict(position='Position', value=vi_label)) mesh = cat_nbk.to_mesh(Nmesh=Ngrid, **to_mesh_kwargs) if comm.rank == 0: logger.info("mesh type: %s" % str(type(mesh))) logger.info("mesh attrs: %s" % str(mesh.attrs)) # this is (1+delta)v_i/(faH) (if normalize were False would get rho*v_i/(faH)) outfield = FieldMesh(mesh.to_real_field(normalize=True)) # get nabla_i[(1+delta)v_i/(faH)] def grad_i_fcn(k3vec, val, myidir=idir): return -1.0j * k3vec[myidir] * val outfield = outfield.apply(grad_i_fcn, mode='complex', kind='wavenumber') # sum up to get theta(k) = sum_i nabla_i[(1+delta)v_i/(faH)] if theta_k is None: theta_k = FieldMesh(outfield.compute('complex')) else: theta_k = FieldMesh( theta_k.compute(mode='complex') + outfield.compute(mode='complex')) # save theta(x) in outfield (check data types again) outfield = FieldMesh(theta_k.compute(mode='real')).to_real_field() elif paint_mode == 'velocity_divergence': # paint div v, and fill empty cells according to some rule assert PaintGrid_config['Painter']['velocity_column'] is not None assert PaintGrid_config['Painter']['fill_empty_cells'] is not None assert PaintGrid_config['Painter'][ 'randseed_for_fill_empty_cells'] is not None assert f_log_growth is not None vi_labels = [] for idir in [0, 1, 2]: vi_label = 'tmp_V_%d' % idir vi_labels.append(vi_label) # this is the velocity / (f*a*H). units are Mpc/h cat_nbk[vi_label] = (cat_nbk[PaintGrid_config['Painter'] ['velocity_column']][:, idir] / f_log_growth) # call extra function to do the painting, saving result in gridx.G[chi_cols] gridx = paint_chicat_to_gridx( chi_cols=vi_labels, cat=cat_nbk, weight_ptcles_by=PaintGrid_config['Painter'].get( 'weight_ptcles_by', None), fill_empty_chi_cells=PaintGrid_config['Painter'] ['fill_empty_cells'], RandNeighbSeed=PaintGrid_config['Painter'] ['randseed_for_fill_empty_cells'], raise_exception_if_too_many_empty_cells=PaintGrid_config[ 'Painter'].get('raise_exception_if_too_many_empty_cells', True), gridx=gridx, gridk=gridk, cache_path=cache_path, do_plot=False, Ngrid=Ngrid, boxsize=boxsize, kmax=kmax) # delete catalog columns b/c not needed any more for vi in vi_labels: del cat_nbk[vi] # get divergence theta_k = None for idir, vi in enumerate(vi_labels): # this is v_i/(faH) outfield = FieldMesh(gridx.G[vi].compute(mode='real')) # get nabla_i[(1+delta)v_i/(faH)] def grad_i_fcn(k3vec, val, myidir=idir): return -1.0j * k3vec[myidir] * val outfield = outfield.apply(grad_i_fcn, mode='complex', kind='wavenumber') # sum up to get theta(k) = sum_i nabla_i[(1+delta)v_i/(faH)] if theta_k is None: theta_k = FieldMesh(outfield.compute('complex')) else: theta_k = FieldMesh( theta_k.compute(mode='complex') + outfield.compute(mode='complex')) # save theta(x) in outfield (check data types again) outfield = FieldMesh(theta_k.compute(mode='real')).to_real_field() else: raise Exception('Invalid paint_mode %s' % paint_mode) elif PaintGrid_config['DataSource']['plugin'] == 'BigFileGrid': ## Read bigfile grid (mesh) directly, no painting required, e.g. for linear density. # TODO: generalize this so we can read bigfile catalog, not only bigfile mesh file. if paint_mode not in [None, 'overdensity']: raise Exception( 'Can only paint overdensity when reading BigFileGrid') if comm.rank == 0: logger.info("Try reading %s" % PaintGrid_config['DataSource']['path']) mesh = BigFileMesh(PaintGrid_config['DataSource']['path'], dataset=PaintGrid_config['DataSource']['dataset']) if PaintGrid_config['Painter'].get('value', None) is not None: raise Exception( 'value kwarg not allowed in Painter when reading BigFileGrid') # Paint. # If normalize=True, divide by the mean. In particular: # - If paint_mode=None, overdensity: If normalize=True, outfield = 1+delta; if normalize=False, outfield=rho # - If paint_mode=momentum_divergence: Reading from BigFileGrid not implemented outfield = mesh.to_real_field() normalize = PaintGrid_config['Painter'].get('normalize', True) if normalize: cmean = outfield.cmean() if np.abs(cmean) < 1e-6: #raise Exception( # 'Found cmean=%g. Are you sure you want to normalize?' % # cmean) print('WARNING: Found cmean=%g, not normalizing by mean' % cmean) else: outfield.apply(lambda x, v: v / cmean, kind="relative", out=outfield) #raise Exception('todo: normalize manually') else: raise Exception("Unsupported DataSource plugin %s" % PaintGrid_config['DataSource']['plugin']) # print paint info if comm.rank == 0: if paint_mode in [None, 'overdensity']: if normalize: logger.info('painted 1+delta') else: logger.info('painted rho (normalize=False)') elif paint_mode == 'momentum_divergence': logger.info('painted div[(1+delta)v/(faH)]') else: logger.info('painted with paint_mode %s' % paint_mode) if hasattr(outfield, 'attrs'): logger.info("outfield.attrs: %s" % str(outfield.attrs)) # set the mean if 'setMean' in PaintGrid_config['Painter']: cmean = outfield.cmean() if comm.rank == 0: logger.info("mean: %g" % cmean) outfield.apply(lambda x, v: v - cmean + PaintGrid_config['Painter'][ 'setMean'], kind='relative', out=outfield) new_mean = outfield.cmean() if comm.rank == 0: logger.info("setting mean=%s" % str(PaintGrid_config['Painter']['setMean'])) logger.info("new mean: %s" % str(new_mean)) if comm.rank == 0: if hasattr(outfield, 'attrs'): logger.info("outfield.attrs: %s" % str(outfield.attrs)) # Save attrs. Should simplify, keep for backwards compatibility if hasattr(outfield, 'attrs'): field_attrs = convert_np_arrays_to_lists(outfield.attrs) else: field_attrs = {} Ntot = None if 'mesh' in vars(): Ntot = mesh.attrs.get('Ntot', np.nan) elif 'cat_nbk' in vars(): Ntot = cat_nbk.attrs.get('Ntot', np.nan) infodict = { 'MS_infodict': { 'Nbkit_config': PaintGrid_config }, 'field_attrs': field_attrs, 'Ntot': Ntot } if 'mesh' in vars(): infodict['MS_infodict']['cat_attrs'] = convert_np_arrays_to_lists( mesh.attrs) elif 'cat_nbk' in vars(): infodict['MS_infodict']['cat_attrs'] = convert_np_arrays_to_lists( cat_nbk.attrs) column_info = {'Nbkit_infodict': infodict} if comm.rank == 0: logger.info("Rescale factor: %g" % rescalefac) if rescalefac != 1.0: outfield.apply(lambda x, v: v * rescalefac, kind="relative", out=outfield) column_info['rescalefac'] = rescalefac #print("rescaled delta(x) rms min mean max:", np.mean(delta**2)**0.5, np.min(delta), # np.mean(delta), np.max(delta)) # Store density in RealGrid instance if gridx is None: gridx = RealGrid(meshsource=outfield, column=column, Ngrid=Ngrid, column_info=column_info, boxsize=boxsize) else: gridx.append_column(column, outfield, column_info=column_info) del outfield if False: # TEST FUNCTIONS in new Grid class gridx.append_column('bla', gridx.G[column], column_info=column_info) gridx.save_to_bigfile('test.bigfile', new_dataset_for_each_column=True) tmp_gridx = RealGrid(fname='test.bigfile', read_columns=[column]) tmp_gridx.append_columns_from_bigfile('test.bigfile', ['bla']) tmp_gridk = ComplexGrid(meshsource=tmp_gridx.fft_x2k('bla', drop_column=False), column='bla', Ngrid=tmp_gridx.Ngrid, boxsize=tmp_gridx.boxsize) tmp_gridk.store_smoothed_gridx(col='bla', path='./', fname='test_smoothed_gridx.bigfile', R=50., replace_nan=True, helper_gridx=tmp_gridx) #tmp_gridx.apply_smoothing('bla', 'Gaussian', R=100.) tmp_gridx.convert_to_weighted_uniform_catalog( col='bla', uni_cat_Nptcles_per_dim=gridx.Ngrid, fill_value_negative_mass=0.) if comm.rank == 0: logger.info("column_info: %s" % str(gridx.column_infos[column])) if not skip_fft: # Compute FFT of density and store in ComplexGrid if gridk is None: gridk = ComplexGrid(meshsource=gridx.fft_x2k(column, drop_column=True), column=column, Ngrid=gridx.Ngrid, boxsize=gridx.boxsize, column_info=column_info) else: gridk.append_column(column, gridx.fft_x2k(column, drop_column=drop_gridx_column), column_info=column_info) # Deconvolve CIC from grid if grid_ptcle2grid_deconvolution is not None: raise Exception("not implemented; use compensated Painter") gridk.deconvolve_ptcle2grid_from_grid( column, grid_ptcle2grid_deconvolution=grid_ptcle2grid_deconvolution) # Zero-pad high k if kmax is not None: gridk.apply_smoothing(column, mode='Gaussian', R=0.0, kmax=kmax) # Apply smoothing #if smoothing_kwargs is not None: # raise Exception("not implemented yet") if return_gridx and return_gridk: return gridx, gridk elif return_gridx: return gridx elif return_gridk: return gridk
def calc_and_save_model_errors_at_cat_pos( sim_opts=None, grid_opts=None, power_opts=None, trf_fcn_opts=None, ext_grids_to_load=None, cat_specs=None, trf_specs=None, keep_pickle=False, pickle_file_format='dill', pickle_path='$SCRATCH/perr/pickle/', Pkmeas_helper_columns=None, Pkmeas_helper_columns_calc_crosses=False, cache_base_path=None, code_version_for_pickles=None, shifted_fields_Np=None, shifted_fields_Nmesh=None): """ Calculate the model error for all models specified by trf_specs. Do this by reading out the model at the positions of objects in the catalog. """ # store opts in dict so we can save in pickle later opts = dict( sim_opts=sim_opts, grid_opts=grid_opts, power_opts=power_opts, trf_fcn_opts=trf_fcn_opts, ext_grids_to_load=ext_grids_to_load, #xgrids_in_memory=xgrids_in_memory, #kgrids_in_memory=kgrids_in_memory, cat_specs=cat_specs, trf_specs=trf_specs, keep_pickle=keep_pickle, pickle_file_format=pickle_file_format, pickle_path=pickle_path, Pkmeas_helper_columns=Pkmeas_helper_columns, cache_base_path=cache_base_path, code_version_for_pickles=code_version_for_pickles, shifted_fields_Np=shifted_fields_Np, shifted_fields_Nmesh=shifted_fields_Nmesh) ##################################### # Initialize ##################################### # make sure we keep the pickle if it is a big run and do not plot if grid_opts.Ngrid > 256: keep_pickle = True # load defaults if not set if ext_grids_to_load is None: ext_grids_to_load = sim_opts.get_default_ext_grids_to_load( Ngrid=grid_opts.Ngrid) if cat_specs is None: cat_specs = {} ### derived options (do not move above b/c command line args might ### overwrite some options!) opts['in_path'] = path_utils.get_in_path(opts) # for output densities opts['out_rho_path'] = os.path.join(opts['in_path'], 'out_rho_Ng%d' % grid_opts.Ngrid) # expand environment names in paths paths = {} for key in [ 'in_path', 'in_fname', 'in_fname_PTsim_psi_calibration', 'in_fname_halos_to_displace_by_mchi', 'pickle_path', 'cache_base_path', 'grids4plots_base_path', 'out_rho_path' ]: if opts.has_key(key): if opts[key] is None: paths[key] = None else: paths[key] = os.path.expandvars(opts[key]) setup_logging() comm = CurrentMPIComm.get() logger = logging.getLogger('PerrCalc') # make sure there are no duplicate save_bestfit_field entries model_spec.check_trf_specs_consistency(trf_specs) # Init Pickler instance to save pickle later (this will init pickle fname) pickler = None if comm.rank == 0: pickler = Pickler(path=paths['pickle_path'], base_fname='main_calc_vel_at_halopos_Perr', file_format=pickle_file_format, rand_sleep=(grid_opts.Ngrid > 128)) print("Pickler: ", pickler.full_fname) pickler = comm.bcast(pickler, root=0) paths['cache_path'] = utils.make_cache_path(paths['cache_base_path'], comm) # Get list of all densities actually needed for trf fcns. #densities_needed_for_trf_fcns = utils.get_densities_needed_for_trf_fcns( # trf_specs) # ########################################################################## # Run program. # ########################################################################## # calculate D and f cosmo = CosmoModel(**sim_opts.cosmo_params) calc_Da = generate_calc_Da(cosmo=cosmo) f_log_growth = calc_f_log_growth_rate(a=sim_opts.sim_scale_factor, calc_Da=calc_Da, cosmo=cosmo, do_test=True) # save in opts so we can easily access it throughout code (although strictly # speaking it is not a free option but derived from cosmo_params) opts['f_log_growth'] = f_log_growth # Compute model at catalog positions, and residual to target. pickle_dict = calc_model_errors_at_cat_pos( trf_specs=trf_specs, paths=paths, cat_specs=cat_specs, ext_grids_to_load=ext_grids_to_load, trf_fcn_opts=trf_fcn_opts, grid_opts=grid_opts, sim_opts=sim_opts, power_opts=power_opts, Pkmeas_helper_columns=Pkmeas_helper_columns, Pkmeas_helper_columns_calc_crosses=Pkmeas_helper_columns_calc_crosses, f_log_growth=opts['f_log_growth']) # copy over opts so they are saved assert not pickle_dict.has_key('opts') pickle_dict['opts'] = opts.copy() # save all resutls to pickle if comm.rank == 0: pickler.write_pickle(pickle_dict) # print save_bestfit_fields save_bestfit_fields = [t.save_bestfit_field for t in opts['trf_specs']] print('\nsave_bestfit_fields:\n' + '\n'.join(save_bestfit_fields)) # delete pickle if not wanted any more if comm.rank == 0: if keep_pickle: print("Pickle: %s" % pickler.full_fname) else: pickler.delete_pickle_file() # delete cache dir from shutil import rmtree rmtree(paths['cache_path']) return pickle_dict
def weigh_and_shift_uni_cats( Nptcles_per_dim, out_Ngrid, densities_to_shift, displacement_source=None, RSD_displacement_source=None, PsiOrder=1, RSD=False, RSD_method=None, RSD_line_of_sight=None, basepath=None, out_scale_factor=None, internal_scale_factor_for_weights=None, weighted_CIC_mode=None, boxsize=None, Nmesh_orig=None, cosmo_params=None, save_result=False, verbose=False, plot_slices=False, ): """ Do the following: - create uniform catalog - load deltalin or other density delta on grid (given by densities_to_shift) - weigh ptcles in uniform catalog by 1+delta - compute smoothed Psi_lin on grid (=evaluated at ptcle positions) - shift (displace) particles by Psi - compute density of shifted catalog using weighted CIC - save shifted catalog Parameters ---------- Nptcles_per_dim : int Number of particles per dimension to use for uniform catalog out_Ngrid : int Number of cells per dimension for the output grid to which particles are CIC interpolated. densities_to_shift : dict Dicts specifying which fields to load from disk. Specify file name etc. Nmesh_orig : int Only used for output file name. Number of cells per dimension of the original field that is shifted and of the displacement field. displacement_source : dict Dict specifying how to compute displacement field. """ # init MPI comm = CurrentMPIComm.get() print("%d: Greetings from rank %d" % (comm.rank, comm.rank)) logger = logging.getLogger("shift") outmesh_dict = dict() outfiles_dict = dict() if RSD: # calculate f cosmo = CosmoModel(**cosmo_params) calc_Da = generate_calc_Da(cosmo=cosmo) f_log_growth = calc_f_log_growth_rate(a=out_scale_factor, calc_Da=calc_Da, cosmo=cosmo, do_test=True) else: f_log_growth = None # loop over all densities to be shifted (load from disk) for specs_of_density_to_shift in densities_to_shift: # ###################################################################### # Load density to be shifted and weigh particles by delta # ###################################################################### # todo: could create linearmesh on the fly rather than reading from disk. # but also fine to load from disk for now. # get field from bigmesh file; rescale to delta at internal_scale_factor_for_weights. # note that we get delta, not 1+delta. rfield_density_to_shift = nbkit03_utils.get_rfield_from_bigfilemesh_file( specs_of_density_to_shift['in_fname'], file_scale_factor=specs_of_density_to_shift['file_scale_factor'], desired_scale_factor=internal_scale_factor_for_weights, cosmo_params=cosmo_params) nbkit03_utils.rfield_print_info(rfield_density_to_shift, comm, 'rfield_density_to_shift: ') # compute quadratic field if desired if specs_of_density_to_shift.get('calc_quadratic_field', None) is not None: # get quadratic field rfield_density_to_shift = nbkit03_utils.calc_quadratic_field( base_field_mesh=FieldMesh(rfield_density_to_shift), quadfield=specs_of_density_to_shift['calc_quadratic_field'], smoothing_of_base_field=specs_of_density_to_shift.get( 'smoothing_quadratic_source', None), verbose=verbose).compute(mode='real') # print info nbkit03_utils.rfield_print_info( rfield_density_to_shift, comm, '%s: ' % specs_of_density_to_shift['calc_quadratic_field']) if specs_of_density_to_shift.get('calc_trf_of_field', None) is not None: # calculate some transformation of the field if specs_of_density_to_shift['calc_trf_of_field'] == '1': # replace field by 1 rfield_density_to_shift = 0 * rfield_density_to_shift + 1.0 else: raise Exception( "Invalid calc_trf_of_field %s" % str(specs_of_density_to_shift['calc_trf_of_field'])) if specs_of_density_to_shift.get('external_smoothing', None) is not None: rfield_density_to_shift = nbkit03_utils.apply_smoothing( mesh_source=FieldMesh(rfield_density_to_shift), **specs_of_density_to_shift['external_smoothing']).compute( mode='real') # print("%d: rfield_density_to_shift: type=%s, shape=%s, rms=%g"% ( # comm.rank, str(type(rfield_density_to_shift)), str(rfield_density_to_shift.shape), # np.sqrt(np.mean(rfield_density_to_shift**2)))) # ###################################################################### # Compute displacement Psi on grid, interpolate to particle positions, # and shift them. # ###################################################################### # get displacement source field from bigmesh file; scale to out_scale_factor # to make sure we use full displacement up to low z of output. rfield_displacement_source = nbkit03_utils.get_rfield_from_bigfilemesh_file( displacement_source['in_fname'], file_scale_factor=displacement_source['file_scale_factor'], desired_scale_factor=out_scale_factor, cosmo_params=cosmo_params) RSD_method = specs_of_density_to_shift['RSD_method'] if (RSD == False) or (RSD == True and RSD_method == 'LPT'): Psi_rfields = [None, None, None] for direction in range(3): # Get Psi_i = k_i delta_smoothed/k^2 following http://rainwoodman.github.io/pmesh/intro.html if comm.rank == 0: print("%d: Get Psi_%d: " % (comm.rank, direction)) # Compute displacement Psi_rfields[ direction] = nbkit03_utils.get_displacement_from_density_rfield( rfield_displacement_source, component=direction, Psi_type=displacement_source['Psi_type'], smoothing=displacement_source['smoothing'], smoothing_Psi3LPT=displacement_source.get( 'smoothing_Psi3LPT', None), RSD=RSD, RSD_line_of_sight=RSD_line_of_sight, RSD_f_log_growth=f_log_growth) # ###################################################################### # Shift uniform catalog weighted by delta_for_weights and get shifted # density. # ###################################################################### # get delta_shifted (get 1+delta) delta_shifted, attrs = weigh_and_shift_uni_cat( delta_for_weights=rfield_density_to_shift, displacements=Psi_rfields, Nptcles_per_dim=Nptcles_per_dim, out_Ngrid=out_Ngrid, boxsize=boxsize, internal_scale_factor_for_weights= internal_scale_factor_for_weights, out_scale_factor=out_scale_factor, cosmo_params=cosmo_params, weighted_CIC_mode=weighted_CIC_mode, plot_slices=plot_slices, verbose=verbose) elif RSD_method == 'PotentialOfRSDDisplSource_xPtcles': ## Step 1: Move operators without any RSD displacement. ## Step 2: Apply RSD displacement in EUlerian space by displacing by # k/k^2*delta_RSD_displ_source, where delta_RSD_displ_source=delta_ZA. assert RSD_displacement_source is not None ## Step 1: Move operators without any RSD displacement. Psi_rfields = [None, None, None] for direction in range(3): # Get Psi_i = k_i delta_smoothed/k^2 following http://rainwoodman.github.io/pmesh/intro.html if comm.rank == 0: print("%d: Get Psi_%d: " % (comm.rank, direction)) # Compute displacement used to displace operators from q to x Psi_rfields[ direction] = nbkit03_utils.get_displacement_from_density_rfield( rfield_displacement_source, component=direction, Psi_type=displacement_source['Psi_type'], smoothing=displacement_source['smoothing'], smoothing_Psi3LPT=displacement_source.get( 'smoothing_Psi3LPT', None), RSD=False, # no RSD displacement in step 1 RSD_line_of_sight=RSD_line_of_sight, RSD_f_log_growth=f_log_growth) del rfield_displacement_source # Shift uniform catalog weighted by rfield_density_to_shift. # get delta_shifted (get 1+delta) delta_shifted_noRSD, attrs = weigh_and_shift_uni_cat( delta_for_weights=rfield_density_to_shift, displacements=Psi_rfields, Nptcles_per_dim=Nptcles_per_dim, out_Ngrid=out_Ngrid, boxsize=boxsize, internal_scale_factor_for_weights= internal_scale_factor_for_weights, out_scale_factor=out_scale_factor, cosmo_params=cosmo_params, weighted_CIC_mode=weighted_CIC_mode, plot_slices=plot_slices, verbose=verbose) ## Step 2: Apply RSD displacement in Eulerian x space by displacing # by f*k/k^2*deltaZA. # Generate new uniform catalog in x space, with particle weights # given by delta_shifted(x). Then shift by fk/k^2 delta_ZA(x). # Note k/k^2 delta_ZA != k/k^2 delta_lin = Psi_ZA. # In general, use delta_RSD_displacement_source. del Psi_rfields # get RSD displacement source field from bigmesh file; scale to out_scale_factor rfield_RSD_displacement_source = nbkit03_utils.get_rfield_from_bigfilemesh_file( RSD_displacement_source['in_fname'], file_scale_factor=RSD_displacement_source['file_scale_factor'], desired_scale_factor=out_scale_factor, cosmo_params=cosmo_params) Psi_RSD_rfields = [None, None, None] if RSD_line_of_sight == [0, 0, 1]: # get k_z / k^2 delta_ZA direction = 2 Psi_RSD_rfields[ direction] = nbkit03_utils.get_displacement_from_density_rfield( in_density_rfield=rfield_RSD_displacement_source, component=direction, Psi_type=RSD_displacement_source[ 'Psi_type'], # to get k/k^2*in_density_rfield smoothing=RSD_displacement_source['smoothing']) # 23/03/2020: flipping sign of displacement and using - here # gives worse Perr. Psi_RSD_rfields[direction] *= f_log_growth else: raise Exception('RSD_line_of_sight not implemented: ', RSD_line_of_sight) delta_shifted, attrs = weigh_and_shift_uni_cat( delta_for_weights=delta_shifted_noRSD, displacements=Psi_RSD_rfields, Nptcles_per_dim=Nptcles_per_dim, out_Ngrid=out_Ngrid, boxsize=boxsize, internal_scale_factor_for_weights= internal_scale_factor_for_weights, out_scale_factor=out_scale_factor, cosmo_params=cosmo_params, weighted_CIC_mode=weighted_CIC_mode, plot_slices=plot_slices, verbose=verbose) else: raise Exception('Invalid RSD_method %s' % RSD_method) # print cmean shifted_cmean = delta_shifted.cmean() if comm.rank == 0: print("%d: delta_shifted cmean:" % comm.rank, shifted_cmean) # ###################################################################### # save to bigfile # ###################################################################### if save_result: if not RSD: RSDstring = '' else: if RSD_line_of_sight in [[0, 0, 1], [0, 1, 0], [1, 0, 0]]: RSDstring = '_RSD%d%d%d' % (RSD_line_of_sight[0], RSD_line_of_sight[1], RSD_line_of_sight[2]) else: # actually not implemented above RSDstring = '_RSD_%.2f_%.2f_%.2f' % (RSD_line_of_sight[0], RSD_line_of_sight[1], RSD_line_of_sight[2]) if RSD_method == 'LPT': pass elif RSD_method == 'PotentialOfRSDDisplSource_xPtcles': RSDstring += "_%s" % RSD_displacement_source[ 'id_for_out_fname'] else: raise Exception('Invalid RSD_method %s' % RSD_method) out_fname = os.path.join( basepath, '%s_int%s_ext%s_SHIFTEDBY_%s%s_a%.4f_Np%d_Nm%d_Ng%d_CIC%s%s' % (specs_of_density_to_shift['id_for_out_fname'], smoothing_str( specs_of_density_to_shift.get( 'smoothing_quadratic_source', None)), smoothing_str( specs_of_density_to_shift['external_smoothing']), displacement_source['id_for_out_fname'], smoothing_str(displacement_source['smoothing']), out_scale_factor, Nptcles_per_dim, Nmesh_orig, out_Ngrid, weighted_CIC_mode, RSDstring)) if comm.rank == 0: print("Writing to %s" % out_fname) # delta_shifted contains 1+delta_shifted # Wrongly used FieldMesh(1+delta_shifted) until 31/3/2020. outmesh = FieldMesh(delta_shifted) # copy MeshSource attrs for k, v in attrs.items(): outmesh.attrs['MeshSource_%s' % k] = v if comm.rank == 0: print("outmesh.attrs:\n", outmesh.attrs) outmesh.save(out_fname, mode='real') if comm.rank == 0: print("Wrote %s" % out_fname) else: if comm.rank == 0: print("Not writing result to disk") # return mesh if specs_of_density_to_shift.get('return_mesh', False): mykey = specs_of_density_to_shift['id_for_out_fname'] if mykey in outmesh_dict: raise Exception('Do not overwrite outmesh_dict key') else: outmesh_dict[mykey] = outmesh outfiles_dict[mykey] = out_fname # optionally plot slice if plot_slices: plt.imshow(outmesh.preview(Nmesh=32, axes=(0, 1))) if comm.rank == 0: plt_fname = 'outmesh_Np%d_Nm%d_Ng%d.pdf' % ( Nptcles_per_dim, Nmesh_orig, out_Ngrid) plt.savefig(plt_fname) print("Made %s" % plt_fname) return outmesh_dict, outfiles_dict
# Define command-line arguments parser = argparse.ArgumentParser() parser.add_argument('in_file', help='input bigfile') parser.add_argument('out_file', help='output bigfile') parser.add_argument('n_mesh',type=int, help='resolution of downsampled mesh') # Parse arguments args = parser.parse_args() in_file = args.in_file out_file = args.out_file n_mesh = args.n_mesh comm = CurrentMPIComm.get() # Set start time, and print first status message start_time = time.time() print_status(comm,start_time,'Starting script') # Import mesh from file orig_mesh = nbk.BigFileMesh(in_file,'Field') print_status(comm,start_time,'Imported mesh from %s with Nmesh %d' % (in_file,orig_mesh.attrs['Nmesh'][0])) # Define new FieldMesh from downsampled version of original mesh new_mesh = nbk.FieldMesh(orig_mesh.compute(mode='real',Nmesh=n_mesh)) print_status(comm,start_time,'Created new mesh with Nmesh %d' % new_mesh.attrs['Nmesh'][0]) # Save mesh to disk, as bigfile (actual computation of downsampling happens in this step) new_mesh.save(out_file,dataset='Field',mode='real')
def calc_quadratic_field( base_field_mesh=None, second_base_field_mesh=None, quadfield=None, smoothing_of_base_field=None, smoothing_of_second_base_field=None, #return_in_k_space=False, verbose=False): """ Calculate quadratic field, essentially by squaring base_field_mesh with filters applied before squaring. Parameters ---------- base_field_mesh : MeshSource object, typically a FieldMesh object Input field that will be squared. second_base_field_mesh : MeshSource object, typically a FieldMesh object Use this to multiply two fields, e.g. delta1*delta2 or G2[delta1,delta2]. Only implemented for tidal_G2 at the moment. quadfield : string Represents quadratic field to be calculated. Can be - 'tidal_s2': Get s^2 = 3/2*s_ij*s_ij = 3/2*[d_ij d_ij - 1/3 delta^2] = 3/2*d_ijd_ij - delta^2/2, where s_ij = (k_ik_j/k^2-delta_ij^K/3)basefield(\vk) and d_ij = k_ik_j/k^2*basefield(\vk). - 'tidal_G2': Get G2[delta] = d_ij d_ij - delta^2. This is orthogonal to delta^2 at low k which can be useful; also see Assassi et al (2014). - 'shift': Get shift=\vPsi\cdot\vnabla\basefield(\vx), where vPsi=-ik/k^2*basefield(k). - 'PsiNablaDelta': Same as 'shift' - 'growth': Get delta^2(\vx) - 'F2': Get F2[delta] = 17/21 delta^2 + shift + 4/21 tidal_s2 = delta^2 + shift + 2/7 tidal_G2 Returns ------- Return the calculated (Ngrid,Ngrid,Ngrid) field as a FieldMesh object. """ comm = CurrentMPIComm.get() if second_base_field_mesh is not None: if not quadfield.endswith('two_meshs'): raise Exception( 'second_base_field_mesh not implemented for quadfield %s' % quadfield) # apply smoothing if smoothing_of_base_field is not None: base_field_mesh = apply_smoothing(mesh_source=base_field_mesh, **smoothing_of_base_field) if second_base_field_mesh is not None: if smoothing_of_second_base_field is not None: second_base_field_mesh = apply_smoothing( mesh_source=second_base_field_mesh, **smoothing_of_second_base_field) # compute quadratic (or cubic) field if quadfield == 'growth': out_rfield = base_field_mesh.compute(mode='real')**2 elif quadfield == 'growth_two_meshs': out_rfield = (base_field_mesh.compute(mode='real') * second_base_field_mesh.compute(mode='real')) elif quadfield == 'growth-mean': out_rfield = base_field_mesh.compute(mode='real')**2 mymean = out_rfield.cmean() out_rfield -= mymean elif quadfield == 'cube-mean': out_rfield = base_field_mesh.compute(mode='real')**3 mymean = out_rfield.cmean() out_rfield -= mymean elif quadfield == 'tidal_G2': # Get G2[delta] = d_ijd_ij - delta^2 if second_base_field_mesh is None: # Compute -delta^2(\vx) out_rfield = -base_field_mesh.compute(mode='real')**2 # Compute d_ij(x). It's symmetric in i<->j so only compute j>=i. # d_ij = k_ik_j/k^2*basefield(\vk). for idir in range(3): for jdir in range(idir, 3): def my_transfer_function(k3vec, val, idir=idir, jdir=jdir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * k3vec[jdir] / kk * val dij_k = base_field_mesh.apply(my_transfer_function, mode='complex', kind='wavenumber') del my_transfer_function # do fft and convert field_mesh to RealField object dij_x = dij_k.compute(mode='real') if verbose: rfield_print_info(dij_x, comm, 'd_%d%d: ' % (idir, jdir)) # Add \sum_{i,j=0..2} d_ij(\vx)d_ij(\vx) # = [d_00^2+d_11^2+d_22^2 + 2*(d_01^2+d_02^2+d_12^2)] if jdir == idir: fac = 1.0 else: fac = 2.0 out_rfield += fac * dij_x**2 del dij_x, dij_k else: raise Exception('use tidal_G2_two_meshs') elif quadfield == 'tidal_G2_two_meshs': # use second_base_field_mesh # Compute -delta1*delta2(\vx) out_rfield = -(base_field_mesh.compute(mode='real') * second_base_field_mesh.compute(mode='real')) # Compute d_ij(x). It's not symmetric in i<->j so compute all i,j # d_ij = k_ik_j/k^2*basefield(\vk). for idir in range(3): for jdir in range(3): def my_transfer_function(k3vec, val, idir=idir, jdir=jdir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * k3vec[jdir] / kk * val dij_k = base_field_mesh.apply(my_transfer_function, mode='complex', kind='wavenumber') second_dij_k = second_base_field_mesh.apply( my_transfer_function, mode='complex', kind='wavenumber') del my_transfer_function # do fft and convert field_mesh to RealField object dij_x = dij_k.compute(mode='real') second_dij_x = second_dij_k.compute(mode='real') if verbose: rfield_print_info(dij_x, comm, 'd_%d%d: ' % (idir, jdir)) rfield_print_info(second_dij_x, comm, 'd_%d%d: ' % (idir, jdir)) # Add \sum_{i,j=0..2} d_ij(\vx)second_d_ij(\vx) # 29 Jul 2020: Had bug before, assumign symmetry. out_rfield += dij_x * second_dij_x del dij_x, dij_k, second_dij_x, second_dij_k elif quadfield == 'tidal_s2': # Get s^2 = 3/2*d_ijd_ij - delta^2/2 # Compute -delta^2(\vx)/2 out_rfield = -base_field_mesh.compute(mode='real')**2 / 2.0 # Compute d_ij(x). It's symmetric in i<->j so only compute j>=i. # d_ij = k_ik_j/k^2*basefield(\vk). for idir in range(3): for jdir in range(idir, 3): def my_transfer_function(k3vec, val, idir=idir, jdir=jdir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * k3vec[jdir] * val / kk dij_k = base_field_mesh.apply(my_transfer_function, mode='complex', kind='wavenumber') del my_transfer_function dij_x = dij_k.compute(mode='real') if verbose: rfield_print_info(dij_x, comm, 'd_%d%d: ' % (idir, jdir)) # Add \sum_{i,j=0..2} d_ij(\vx)d_ij(\vx) # = [d_00^2+d_11^2+d_22^2 + 2*(d_01^2+d_02^2+d_12^2)] if jdir == idir: fac = 1.0 else: fac = 2.0 out_rfield += fac * 1.5 * dij_x**2 del dij_x, dij_k elif quadfield in ['shift', 'PsiNablaDelta']: # Get shift = \vPsi\cdot\nabla\delta for idir in range(3): # compute Psi_i def Psi_i_fcn(k3vec, val, idir=idir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return -1.0j * k3vec[idir] * val / kk Psi_i_x = base_field_mesh.apply( Psi_i_fcn, mode='complex', kind='wavenumber').compute(mode='real') # compute nabla_i delta def grad_i_fcn(k3vec, val, idir=idir): return -1.0j * k3vec[idir] * val nabla_i_delta_x = base_field_mesh.apply( grad_i_fcn, mode='complex', kind='wavenumber').compute(mode='real') # multiply and add up in x space if idir == 0: out_rfield = Psi_i_x * nabla_i_delta_x else: out_rfield += Psi_i_x * nabla_i_delta_x elif quadfield in ['shift_two_meshs', 'PsiNablaDelta_two_meshs']: # Get shift = 0.5(\vPsi_1\cdot\nabla\delta_2 + \vPsi_2\cdot\nabla\delta_1) out_rfield = None for idir in range(3): def Psi_i_fcn(k3vec, val, idir=idir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return -1.0j * k3vec[idir] * val / kk def grad_i_fcn(k3vec, val, idir=idir): return -1.0j * k3vec[idir] * val # compute Psi_i of mesh1 Psi_i_x = base_field_mesh.apply( Psi_i_fcn, mode='complex', kind='wavenumber').compute(mode='real') # compute nabla_i delta of mesh2 nabla_i_delta_x = second_base_field_mesh.apply( grad_i_fcn, mode='complex', kind='wavenumber').compute(mode='real') # multiply and add up in x space if out_rfield is None: out_rfield = 0.5 * Psi_i_x * nabla_i_delta_x else: out_rfield += 0.5 * Psi_i_x * nabla_i_delta_x del Psi_i_x, nabla_i_delta_x # also swap mesh1 and mesh2 # compute Psi_i of mesh2 Psi_i_x_prime = second_base_field_mesh.apply( Psi_i_fcn, mode='complex', kind='wavenumber').compute(mode='real') # compute nabla_i delta of mesh1 nabla_i_delta_x_prime = base_field_mesh.apply( grad_i_fcn, mode='complex', kind='wavenumber').compute(mode='real') out_rfield += 0.5 * Psi_i_x_prime * nabla_i_delta_x_prime # elif quadfield in ['tidal_G2_par_LOS100', 'tidal_G2_par_LOS010', 'tidal_G2_par_LOS001']: # if quadfield.endswith('100'): # LOS_dir = 0 # elif quadfield.endswith('010'): # LOS_dir = 1 # elif quadfield.endswith('001'): # LOS_dir = 2 # else: # raise Exception('invalid LOS') # # compute nabla_parallel^2/nabla^2 G2 # G2_mesh = calc_quadratic_field( # quadfield='tidal_G2', # base_field_mesh=base_field_mesh, # smoothing_of_base_field=smoothing_of_base_field, # verbose=verbose) # def my_parallel_transfer_function(k3vec, val, idir=LOS_dir, jdir=LOS_dir): # kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh # kk[kk == 0] = 1 # # MS 25 Jul 2020: Include minus sign because (ik_2)^2=-k_2^2. # return -k3vec[idir] * k3vec[jdir] * val / kk # G2_parallel = G2_mesh.apply(my_parallel_transfer_function, # mode='complex', # kind='wavenumber') # del my_parallel_transfer_function # out_rfield = G2_parallel.compute(mode='real') # del G2_parallel, G2_mesh elif quadfield == 'F2': # F2 = delta^2... out_rfield = calc_quadratic_field( quadfield='growth', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') # ... - shift out_rfield -= calc_quadratic_field( quadfield='shift', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') # ... + 2/7 tidal_G2 out_rfield += 2. / 7. * calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') elif quadfield == 'F2_two_meshs': # F2 = delta^2... out_rfield = calc_quadratic_field( quadfield='growth_two_meshs', base_field_mesh=base_field_mesh, second_base_field_mesh=second_base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') # ... - shift out_rfield -= calc_quadratic_field( quadfield='shift_two_meshs', base_field_mesh=base_field_mesh, second_base_field_mesh=second_base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') # ... + 2/7 tidal_G2 out_rfield += 2. / 7. * calc_quadratic_field( quadfield='tidal_G2_two_meshs', base_field_mesh=base_field_mesh, second_base_field_mesh=second_base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') elif quadfield == 'velocity_G2': # G2 = delta^2... out_rfield = calc_quadratic_field( quadfield='growth', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') # ... - shift out_rfield -= calc_quadratic_field( quadfield='shift', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') # ... + 4/7 tidal_G2 out_rfield += 4. / 7. * calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') elif quadfield == 'velocity_G2_two_meshs': # G2 = delta^2... out_rfield = calc_quadratic_field( quadfield='growth_two_meshs', base_field_mesh=base_field_mesh, second_base_field_mesh=second_base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') # ... - shift out_rfield -= calc_quadratic_field( quadfield='shift_two_meshs', base_field_mesh=base_field_mesh, second_base_field_mesh=second_base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') # ... + 4/7 tidal_G2 out_rfield += 4. / 7. * calc_quadratic_field( quadfield='tidal_G2_two_meshs', base_field_mesh=base_field_mesh, second_base_field_mesh=second_base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') elif quadfield in ['delta_G2', 'G2_delta']: # Get G2[delta] * delta out_rfield = (base_field_mesh.compute(mode='real') * calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real')) # take out the mean (already close to 0 but still subtract) mymean = out_rfield.cmean() if comm.rank == 0: print('Subtract mean of G2*delta: %g' % mymean) out_rfield -= mymean elif quadfield in [ 'delta_G2_par_LOS100', 'delta_G2_par_LOS010', 'delta_G2_par_LOS001' ]: # Get delta * G2_parallel[delta] out_rfield = (base_field_mesh.compute(mode='real') * calc_quadratic_field( quadfield='tidal_G2_par_LOS%s' % (quadfield[-3:]), base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real')) # take out the mean (already close to 0 but still subtract) mymean = out_rfield.cmean() if comm.rank == 0: print('Subtract mean of G2par*delta: %g' % mymean) out_rfield -= mymean elif quadfield in ['tidal_G3', 'tidal_G3_flipped_dij_sign']: # Get G3[delta] # Have 3/2 G2 delta = 3/2 (p1.p2)^2/(p1^2 p2^2) - 3/2 # so # G3 = 3/2 G2 delta + delta^3 - (p1.p2)(p2.p3)(p1.p3)/(p1^2 p2^2 p3^2) # Compute 1 * delta^3(\vx) out_rfield = base_field_mesh.compute(mode='real')**3 # Add 3/2 delta * G2[delta] out_rfield += (3. / 2. * base_field_mesh.compute(mode='real') * calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real')) # Compute ppp = (p1.p2)(p2.p3)(p1.p3)/(p1^2 p2^2 p3^2) # = k.q k.p q.p / (k^2 q^2 p^2) # = k_i q_i k_j p_j q_l p_l / (k^2 q^2 p^2) # = sum_ijl d_ij(k) d_il(q) d_jl(p) # where we denoted p1=k, p2=q, p3=p. # Compute d_ij(x). It's symmetric in i<->j so only compute j>=i. # d_ij = k_ik_j/k^2*basefield(\vk). dij_x_dict = {} for idir in range(3): for jdir in range(idir, 3): def my_transfer_function(k3vec, val, idir=idir, jdir=jdir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * k3vec[jdir] * val / kk dij_k = base_field_mesh.apply(my_transfer_function, mode='complex', kind='wavenumber') del my_transfer_function # do fft and convert field_mesh to RealField object dij_x = dij_k.compute(mode='real') del dij_k if verbose: rfield_print_info(dij_x, comm, 'd_%d%d: ' % (idir, jdir)) if quadfield == 'tidal_G3_flipped_dij_sign': # flip sign of dij dij_x *= -1.0 dij_x_dict[(idir, jdir)] = dij_x del dij_x # get j<i by symmetry def get_dij_x(idir, jdir): if jdir >= idir: return dij_x_dict[(idir, jdir)] else: return dij_x_dict[(jdir, idir)] # Compute - sum_ijl d_ij(k) d_il(q) d_jl(p) for idir in range(3): for jdir in range(3): for ldir in range(3): out_rfield -= (get_dij_x(idir, jdir) * get_dij_x(idir, ldir) * get_dij_x(jdir, ldir)) # take out the mean (already close to 0 but still subtract) mymean = out_rfield.cmean() if comm.rank == 0: print('Subtract mean of G3: %g' % mymean) out_rfield -= mymean if verbose: rfield_print_info(out_rfield, comm, 'G3: ') elif quadfield == 'Gamma3': # Get Gamma3[delta] = -4/7 * G2[G2[delta,delta],delta] tmp_G2 = calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) out_rfield = -4. / 7. * calc_quadratic_field( quadfield='tidal_G2_two_meshs', base_field_mesh=base_field_mesh, second_base_field_mesh=tmp_G2, smoothing_of_base_field=smoothing_of_base_field, smoothing_of_second_base_field=None, verbose=verbose).compute(mode='real') elif quadfield in [ 'delta1_par_LOS100', 'delta1_par_LOS010', 'delta1_par_LOS001', 'tidal_G2_par_LOS100', 'tidal_G2_par_LOS010', 'tidal_G2_par_LOS001', 'Gamma3_par_LOS100', 'Gamma3_par_LOS010', 'Gamma3_par_LOS001', 'tidal_G3_par_LOS100', 'tidal_G3_par_LOS010', 'tidal_G3_par_LOS001', 'tidal_G3_flipped_dij_sign_par_LOS100', 'tidal_G3_flipped_dij_sign_par_LOS010', 'tidal_G3_flipped_dij_sign_par_LOS001' ]: if quadfield.endswith('LOS100'): LOS_dir = 0 elif quadfield.endswith('LOS010'): LOS_dir = 1 elif quadfield.endswith('LOS001'): LOS_dir = 2 else: raise Exception('invalid LOS string in %s' % str(quadfield)) if quadfield.startswith('delta1_par'): tmp_base_quadfield_str = 'delta1' elif quadfield.startswith('tidal_G2_par'): tmp_base_quadfield_str = 'tidal_G2' elif quadfield.startswith('Gamma3_par'): tmp_base_quadfield_str = 'Gamma3' elif quadfield.startswith('tidal_G3_par'): tmp_base_quadfield_str = 'tidal_G3' elif quadfield.startswith('tidal_G3_flipped_dij_sign_par'): tmp_base_quadfield_str = 'tidal_G3_flipped_dij_sign' else: raise Exception('Invalid quadfield %s' % str(quadfield)) # compute nabla_parallel^2/nabla^2 tmp_base_quadfield if tmp_base_quadfield_str == 'delta1': tmp_base_quadfield_mesh = FieldMesh( base_field_mesh.compute(mode='real')) else: tmp_base_quadfield_mesh = calc_quadratic_field( quadfield=tmp_base_quadfield_str, base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) def my_parallel_transfer_function(k3vec, val, idir=LOS_dir, jdir=LOS_dir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 # MS 25 Jul 2020: Include minus sign because (ik_2)^2=-k_2^2. # MS 29 Jul 2020: Drop minus sign because 1/nabla^2 gives another minus sign. return k3vec[idir] * k3vec[jdir] * val / kk tmp_base_quadfield_parallel = tmp_base_quadfield_mesh.apply( my_parallel_transfer_function, mode='complex', kind='wavenumber') del my_parallel_transfer_function out_rfield = tmp_base_quadfield_parallel.compute(mode='real') del tmp_base_quadfield_parallel, tmp_base_quadfield_mesh elif quadfield in ['velocity_G2_par_LOS001_two_meshs']: LOS_dir = 2 tmp_base_quadfield_str = 'velocity_G2_two_meshs' # compute nabla_parallel^2/nabla^2 tmp_base_quadfield tmp_base_quadfield_mesh = calc_quadratic_field( quadfield=tmp_base_quadfield_str, base_field_mesh=base_field_mesh, second_base_field_mesh=second_base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) def my_parallel_transfer_function(k3vec, val, idir=LOS_dir, jdir=LOS_dir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 # MS 29 Jul 2020: No minus sign because -(ik_2)^2/k^2=k_2^2/k^2. return k3vec[idir] * k3vec[jdir] * val / kk tmp_base_quadfield_parallel = tmp_base_quadfield_mesh.apply( my_parallel_transfer_function, mode='complex', kind='wavenumber') del my_parallel_transfer_function out_rfield = tmp_base_quadfield_parallel.compute(mode='real') del tmp_base_quadfield_parallel, tmp_base_quadfield_mesh elif quadfield == 'delta1_par_G2_par_LOS001': out_rfield = (calc_quadratic_field( quadfield='delta1_par_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real') * calc_quadratic_field( quadfield='tidal_G2_par_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose).compute(mode='real')) elif quadfield in ['RSDLPT_Q2_LOS001']: # Q2 = sum_i d_i2 d_i2 # d_ij = -k_ik_j/k^2*basefield(\vk). out_rfield = None jdir = 2 for idir in range(3): def my_transfer_function_dij(k3vec, val, idir=idir, jdir=jdir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * k3vec[jdir] * val / kk dij_k = base_field_mesh.apply(my_transfer_function_dij, mode='complex', kind='wavenumber') del my_transfer_function_dij # do fft and convert field_mesh to RealField object dij_x = dij_k.compute(mode='real') del dij_k if verbose: rfield_print_info(dij_x, comm, 'd_%d%d: ' % (idir, jdir)) if out_rfield is None: out_rfield = dij_x**2 else: out_rfield += dij_x**2 elif quadfield in ['RSDLPT_K3_LOS001']: # K3 = sum_i d_i2 dG2_i2 # where d_ij = k_ik_j/k^2*basefield(\vk) # and dG2_ij = k_ik_j/k^2*G2(\vk). tmp_G2_mesh = calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) out_rfield = None jdir = 2 for idir in range(3): def my_transfer_function_dij(k3vec, val, idir=idir, jdir=jdir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * k3vec[jdir] * val / kk dij_x = base_field_mesh.apply( my_transfer_function_dij, mode='complex', kind='wavenumber').compute(mode='real') dijG2_x = tmp_G2_mesh.apply(my_transfer_function_dij, mode='complex', kind='wavenumber').compute(mode='real') del my_transfer_function_dij if out_rfield is None: out_rfield = dij_x * dijG2_x else: out_rfield += dij_x * dijG2_x elif quadfield in ['RSDLPT_S3Ia', 'RSDLPT_S3IIa_LOS001']: # S3Ia = sum_i d_i where d_i = (k_i/k^2 G2) (k_i delta_1) # (or delta_1 -> delta_1 parallel for S3IIa) tmp_G2_mesh = calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) if quadfield == 'RSDLPT_S3Ia': tmp_delta1_mesh = FieldMesh(base_field_mesh.compute(mode='real')) elif quadfield == 'RSDLPT_S3IIa_LOS001': tmp_delta1_mesh = calc_quadratic_field( quadfield='delta1_par_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) out_rfield = None for idir in range(3): def my_trf_for_G2(k3vec, val, idir=idir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * val / kk def my_trf_for_delta1(k3vec, val, idir=idir): return k3vec[idir] * val G2_filtered = tmp_G2_mesh.apply( my_trf_for_G2, mode='complex', kind='wavenumber').compute(mode='real') delta1_filtered = tmp_delta1_mesh.apply( my_trf_for_delta1, mode='complex', kind='wavenumber').compute(mode='real') del my_trf_for_G2, my_trf_for_delta1 if out_rfield is None: out_rfield = G2_filtered * delta1_filtered else: out_rfield += G2_filtered * delta1_filtered elif quadfield in ['RSDLPT_S3Ib_LOS001', 'RSDLPT_S3IIb_LOS001']: # S3Ia = (k_2/k^2 G2) (k_2 delta_1) tmp_G2_mesh = calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) if quadfield == 'RSDLPT_S3Ib_LOS001': tmp_delta1_mesh = FieldMesh(base_field_mesh.compute(mode='real')) elif quadfield == 'RSDLPT_S3IIb_LOS001': tmp_delta1_mesh = calc_quadratic_field( quadfield='delta1_par_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) idir = 2 def my_trf_for_G2(k3vec, val, idir=idir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * val / kk def my_trf_for_delta1(k3vec, val, idir=idir): return k3vec[idir] * val G2_filtered = tmp_G2_mesh.apply(my_trf_for_G2, mode='complex', kind='wavenumber').compute(mode='real') delta1_filtered = tmp_delta1_mesh.apply( my_trf_for_delta1, mode='complex', kind='wavenumber').compute(mode='real') del my_trf_for_G2, my_trf_for_delta1 out_rfield = G2_filtered * delta1_filtered elif quadfield in ['RSDLPT_calQ2_LOS001']: # calQ2 = Q2 - delta1*delta1_parallel out_rfield = calc_quadratic_field( quadfield='RSDLPT_Q2_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) out_rfield -= (base_field_mesh.compute(mode='real') * calc_quadratic_field( quadfield='delta1_par_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose)) elif quadfield in ['RSDLPT_calQ3_LOS001']: # calQ3 = Q3 - delta1 Q2 - 1/2 delta1_par G2 out_rfield = calc_quadratic_field( quadfield='RSDLPT_Q3_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) out_rfield -= (base_field_mesh.compute(mode='real') * calc_quadratic_field( quadfield='RSDLPT_Q2_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose)) out_rfield -= 0.5 * (calc_quadratic_field( quadfield='delta1_par_LOS001', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose) * calc_quadratic_field( quadfield='tidal_G2', base_field_mesh=base_field_mesh, smoothing_of_base_field=smoothing_of_base_field, verbose=verbose)) elif quadfield in ['RSDLPT_Q3_LOS001']: # Q3 = sum_m,n d_m2 d_mn d_n2 # d_ij = k_ik_j/k^2*basefield(\vk). # Compute d_ij(x). It's symmetric in i<->j so only compute j>=i. dij_x_dict = {} for idir in range(3): for jdir in range(idir, 3): def my_dij_transfer_function(k3vec, val, idir=idir, jdir=jdir): kk = sum(ki**2 for ki in k3vec) # k^2 on the mesh kk[kk == 0] = 1 return k3vec[idir] * k3vec[jdir] * val / kk dij_k = base_field_mesh.apply(my_dij_transfer_function, mode='complex', kind='wavenumber') del my_dij_transfer_function # do fft and convert field_mesh to RealField object dij_x = dij_k.compute(mode='real') del dij_k if verbose: rfield_print_info(dij_x, comm, 'd_%d%d: ' % (idir, jdir)) dij_x_dict[(idir, jdir)] = dij_x del dij_x # get j<i by symmetry def get_dij_x(idir, jdir): if jdir >= idir: return dij_x_dict[(idir, jdir)] else: return dij_x_dict[(jdir, idir)] # Compute sum_nm d_2m(k) d_mn(q) d_n2(p), assume LOS=z direction. out_rfield = 0.0 * base_field_mesh.compute(mode='real') for mdir in range(3): for ndir in range(3): out_rfield += (get_dij_x(2, mdir) * get_dij_x(mdir, ndir) * get_dij_x(ndir, 2)) # take out the mean (already close to 0 but still subtract) mymean = out_rfield.cmean() if comm.rank == 0: print('Subtract mean of %s: %g' % (quadfield, mymean)) out_rfield -= mymean elif quadfield.startswith('PsiDot1_'): # get PsiDot \equiv \sum_n n*Psi^{(n)} up to 1st order assert quadfield in ['PsiDot1_0', 'PsiDot1_1', 'PsiDot1_2'] component = int(quadfield[-1]) out_rfield = get_displacement_from_density_rfield( base_field_mesh.compute(mode='real'), component=component, Psi_type='Zeldovich', smoothing=smoothing_of_base_field, RSD=False) elif quadfield.startswith('PsiDot2_'): # get PsiDot \equiv \sum_n n*Psi^{(n)} up to 2nd order assert quadfield in ['PsiDot2_0', 'PsiDot2_1', 'PsiDot2_2'] component = int(quadfield[-1]) # PsiDot \equiv \sum_n n*Psi^{(n)} out_rfield = get_displacement_from_density_rfield( base_field_mesh.compute(mode='real'), component=component, Psi_type='2LPT', prefac_Psi_1storder=1.0, prefac_Psi_2ndorder=2.0, # include n=2 factor smoothing=smoothing_of_base_field, RSD=False) else: raise Exception("quadfield %s not implemented" % str(quadfield)) return FieldMesh(out_rfield)