def __exit__(self, exc_type, exc_value, exc_traceback):
"""
Exit gracefully by closing and freeing the MPI-related variables
"""
if exc_value is not None:
trace = ''.join(
traceback.format_exception(exc_type,
exc_value,
exc_traceback,
limit=5))
self.logger.error("an exception has occurred on rank %d:\n%s" %
(self.rank, trace))
# bit of hack that forces mpi4py to exit all ranks
# see https://groups.google.com/forum/embed/#!topic/mpi4py/RovYzJ8qkbc
os._exit(1)
# wait and exit
self.logger.debug("rank %d process finished" % self.rank)
self.basecomm.Barrier()
if self.is_root():
self.logger.debug("master is finished; terminating")
CurrentMPIComm.pop()
if self.comm is not None:
self.comm.Free()
def test_download_failure(comm):
CurrentMPIComm.set(comm)
# initialize with bad redshift
BAD_REDSHIFT = 100.0
with pytest.raises(Exception):
cat = DemoHaloCatalog('bolshoi', 'rockstar', BAD_REDSHIFT)
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)
CurrentMPIComm.push(self.comm)
return self
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_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 test_download(comm):
from halotools.sim_manager import UserSuppliedHaloCatalog
CurrentMPIComm.set(comm)
# download and load the cached catalog
cat = DemoHaloCatalog('bolshoi', 'rockstar', 0.5)
assert all(col in cat for col in ['Position', 'Velocity'])
# convert to halotools catalog
halotools_cat = cat.to_halotools()
assert isinstance(halotools_cat, UserSuppliedHaloCatalog)
# bad simulation name
with pytest.raises(Exception):
cat = DemoHaloCatalog('BAD', 'rockstar', 0.5)
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)
CurrentMPIComm.push(self.comm)
return self
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 __exit__(self, exc_type, exc_value, exc_traceback):
"""
Exit gracefully by closing and freeing the MPI-related variables
"""
if exc_value is not None:
trace = ''.join(traceback.format_exception(exc_type, exc_value, exc_traceback, limit=5))
self.logger.error("an exception has occurred on rank %d:\n%s" %(self.rank, trace))
# bit of hack that forces mpi4py to exit all ranks
# see https://groups.google.com/forum/embed/#!topic/mpi4py/RovYzJ8qkbc
os._exit(1)
# wait and exit
self.logger.debug("rank %d process finished" %self.rank)
self.basecomm.Barrier()
if self.is_root():
self.logger.debug("master is finished; terminating")
CurrentMPIComm.pop()
if self.comm is not None:
self.comm.Free()
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 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 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 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 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 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 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():
"""
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 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
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)
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():
"""
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 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 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
# 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_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
