def make_lagfields(nc,
seed,
bs=1536,
T=40,
B=2,
simpath=sc_simpath,
outpath=sc_outpath,
Rsm=0):
fname = get_filename(nc, seed, T=T, B=B)
spath = simpath + fname
opath = outpath + fname
try:
os.makedirs(opath)
except:
pass
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
rank = pm.comm.rank
lin = BigFileMesh(spath + '/linear', 'LinearDensityK').paint()
lin -= lin.cmean()
print(rank, 'lin field read')
header = '1,b1,b2,bg,bk'
names = header.split(',')
lag_fields = tools.getlagfields(
pm, lin, R=Rsm) # use the linear field at the desired redshift
print(rank, 'lag field created')
for i, ff in enumerate(lag_fields):
x = FieldMesh(ff)
x.save(opath + 'lag', dataset=names[i], mode='real')
def evaluate1(model, data, norm=True, kmin=None, dk=None):
'''return position order:
px,p1,p2
'''
from nbodykit.lab import FieldMesh, FFTPower, ProjectedFFTPower
if kmin is None: kmin = 0
if dk is None: dk = 2 * numpy.pi / model.BoxSize[0]
if norm:
mod = model / model.cmean()
dat = data / data.cmean()
else:
mod, dat = model, data
px = FFTPower(first=FieldMesh(mod),
second=FieldMesh(dat),
mode='1d',
kmin=kmin,
dk=dk)
p1 = FFTPower(first=FieldMesh(mod), mode='1d', kmin=kmin, dk=dk)
p2 = FFTPower(first=FieldMesh(dat), mode='1d', kmin=kmin, dk=dk)
return px.power, p1.power, p2.power
def evaluate(self, model, data):
from nbodykit.lab import FieldMesh, FFTPower, ProjectedFFTPower
xm = FFTPower(first=FieldMesh(model.map3d),
second=FieldMesh(data.map3d),
mode='1d')
xs = FFTPower(first=FieldMesh(model.s),
second=FieldMesh(data.s),
mode='1d')
pm1 = FFTPower(first=FieldMesh(model.map3d), mode='1d')
ps1 = FFTPower(first=FieldMesh(model.s), mode='1d')
pm2 = FFTPower(first=FieldMesh(data.map3d), mode='1d')
ps2 = FFTPower(first=FieldMesh(data.s), mode='1d')
data_preview = dict(s=[], d=[], map3d=None)
model_preview = dict(s=[], d=[], map3d=None)
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview['d'].append(data.d.preview(axes=axes))
data_preview['s'].append(data.s.preview(axes=axes))
model_preview['d'].append(model.d.preview(axes=axes))
model_preview['s'].append(model.s.preview(axes=axes))
data_preview['map3d'] = data.map3d.preview(axes=(0, 1))
model_preview['map3d'] = model.map3d.preview(axes=(0, 1))
return xm.power, xs.power, pm1.power, pm2.power, ps1.power, ps2.power, data_preview, model_preview
def write_power(d, path, a):
meshsource = FieldMesh(d)
r = FFTPower(meshsource, mode='1d')
if config.pm.comm.rank == 0:
print('Writing matter power spectrum at %s' % path)
# only root rank saves
numpy.savetxt(path,
numpy.array([
r.power['k'], r.power['power'].real, r.power['modes'],
r.power['power'].real / solver.cosmology.scale_independent_growth_factor(1.0 / a - 1) ** 2,
]).T,
comments='# k p N p/D**2')
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
rank = pm.comm.rank
#grid = pm.mesh_coordinates()*bs/nc
lin = BigFileMesh(lpath+ '/linear', 'LinearDensityK').paint()
lin -= lin.cmean()
print(rank, 'lin field read')
header = '1,b1,b2,bg,bk'
names = header.split(',')
lag_fields = tools.getlagfields(pm, lin, R=Rsm) # use the linear field at the desired redshift
print(rank, 'lag field created')
for i, ff in enumerate(lag_fields):
x = FieldMesh(ff)
x.save(lpath + 'lag', dataset=names[i], mode='real')
###
dyn = BigFileCatalog(dpath + '/fastpm_%0.4f/1'%aa)
print(rank, 'files read')
#
fpos = dyn['Position'].compute()
idd = dyn['ID'].compute()
attrs = dyn.attrs
grid = tools.getqfromid(idd, attrs, nc)
def evaluate(model, data, M0=0, stoch=False, kmin=None, dk=None):
'''return position order:
xm.power, xs.power, xd.power,
pm1.power, pm2.power,
ps1.power, ps2.power,
pd1.power, pd2.power,
data_preview, model_preview
'''
from nbodykit.lab import FieldMesh, FFTPower, ProjectedFFTPower
if kmin is None: kmin = 0
if dk is None: dk = 2 * numpy.pi / model.s.BoxSize[0]
model.mapp += M0
data.mapp += M0
xm = FFTPower(first=FieldMesh(model.mapp / model.mapp.cmean()),
second=FieldMesh(data.mapp / data.mapp.cmean()),
mode='1d',
kmin=kmin,
dk=dk)
xd = FFTPower(first=FieldMesh(model.d),
second=FieldMesh(data.d),
mode='1d',
kmin=kmin,
dk=dk)
xs = FFTPower(first=FieldMesh(model.s),
second=FieldMesh(data.s),
mode='1d',
kmin=kmin,
dk=dk)
pm1 = FFTPower(first=FieldMesh(model.mapp / model.mapp.cmean()),
mode='1d',
kmin=kmin,
dk=dk)
pd1 = FFTPower(first=FieldMesh(model.d), mode='1d', kmin=kmin, dk=dk)
ps1 = FFTPower(first=FieldMesh(model.s), mode='1d', kmin=kmin, dk=dk)
pm2 = FFTPower(first=FieldMesh(data.mapp / data.mapp.cmean()),
mode='1d',
kmin=kmin,
dk=dk)
pd2 = FFTPower(first=FieldMesh(data.d), mode='1d', kmin=kmin, dk=dk)
ps2 = FFTPower(first=FieldMesh(data.s), mode='1d', kmin=kmin, dk=dk)
if stoch:
psd = FFTPower(first=FieldMesh(data.s),
second=FieldMesh(model.s),
mode='1d',
kmin=kmin,
dk=dk)
pdd = FFTPower(first=FieldMesh(data.d),
second=FieldMesh(model.d),
mode='1d',
kmin=kmin,
dk=dk)
pmd = FFTPower(first=FieldMesh(data.mapp / data.mapp.cmean()),
second=FieldMesh(model.mapp / model.mapp.cmean()),
mode='1d',
kmin=kmin,
dk=dk)
data_preview = dict(s=[], d=[], mapp=[])
model_preview = dict(s=[], d=[], mapp=[])
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview['d'].append(data.d.preview(axes=axes))
data_preview['s'].append(data.s.preview(axes=axes))
data_preview['mapp'].append(data.mapp.preview(axes=axes))
model_preview['d'].append(model.d.preview(axes=axes))
model_preview['s'].append(model.s.preview(axes=axes))
model_preview['mapp'].append(model.mapp.preview(axes=axes))
#data_preview['mapp'] = data.mapp.preview(axes=(0, 1))
#model_preview['mapp'] = model.mapp.preview(axes=(0, 1))
if stoch:
return xm.power, xs.power, xd.power, pm1.power, pm2.power, ps1.power, ps2.power, pd1.power, pd2.power, \
psd.power, pdd.power, pmd.power, data_preview, model_preview
return xm.power, xs.power, xd.power, pm1.power, pm2.power, ps1.power, ps2.power, pd1.power, pd2.power, data_preview, model_preview
def evaluate2d(model,
data,
M0=0,
kmin=None,
dk=None,
Nmu=5,
retmesh=False,
los=[0, 0, 1]):
'''return position order:
xm.power, xs.power, xd.power,
pm1.power, pm2.power,
ps1.power, ps2.power,
pd1.power, pd2.power,
data_preview, model_preview
'''
from nbodykit.lab import FieldMesh, FFTPower, ProjectedFFTPower
if kmin is None: kmin = 0
if dk is None: dk = 2 * numpy.pi / model.s.BoxSize[0]
model.mapp += M0
data.mapp += M0
xm = FFTPower(first=FieldMesh(model.mapp / model.mapp.cmean()),
second=FieldMesh(data.mapp / data.mapp.cmean()),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
xd = FFTPower(first=FieldMesh(model.d),
second=FieldMesh(data.d),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
xs = FFTPower(first=FieldMesh(model.s),
second=FieldMesh(data.s),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
pm1 = FFTPower(first=FieldMesh(model.mapp / model.mapp.cmean()),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
pd1 = FFTPower(first=FieldMesh(model.d),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
ps1 = FFTPower(first=FieldMesh(model.s),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
pm2 = FFTPower(first=FieldMesh(data.mapp / data.mapp.cmean()),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
pd2 = FFTPower(first=FieldMesh(data.d),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
ps2 = FFTPower(first=FieldMesh(data.s),
mode='2d',
kmin=kmin,
dk=dk,
Nmu=Nmu,
los=[0, 0, 1])
if retmesh:
data_preview = dict(s=[], d=[], mapp=[])
model_preview = dict(s=[], d=[], mapp=[])
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview['d'].append(data.d.preview(axes=axes))
data_preview['s'].append(data.s.preview(axes=axes))
data_preview['mapp'].append(data.mapp.preview(axes=axes))
model_preview['d'].append(model.d.preview(axes=axes))
model_preview['s'].append(model.s.preview(axes=axes))
model_preview['mapp'].append(model.mapp.preview(axes=axes))
return xm.power, xs.power, xd.power, pm1.power, pm2.power, ps1.power, ps2.power, pd1.power, pd2.power,\
data_preview, model_preview
else:
return xm.power, xs.power, xd.power, pm1.power, pm2.power, ps1.power, ps2.power, pd1.power, pd2.power
def load_TNG_map(TNG_basepath, snapNum, field, pm):
assert field in ['dm', 'Mstar', 'ne', 'T', 'nHI', 'neVz', 'MstarVz']
#Try directly loading map. If it does not exist, paint the map and save it for future use.
address = TNG_basepath + '/snapdir_' + str(snapNum).zfill(
3) + '/' + field + 'map_Nmesh' + str(pm.Nmesh[0])
try:
TNGmap = BigFileMesh(address, dataset='Field').to_real_field()
except:
if field == 'dm':
pos = []
for mdi in range(3):
pos.append(
load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType='dm',
field='Coordinates',
mdi=mdi))
pos = np.array(pos).T
mass = 1.0 * pm.Nmesh.prod() / comm.allreduce(len(pos), op=MPI.SUM)
layout = pm.decompose(pos)
pos1 = layout.exchange(pos)
del pos
TNGmap = pm.create(type="real")
TNGmap.paint(pos1, mass=mass, layout=None, hold=False)
del pos1
elif field == 'Mstar':
mass = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType='stars',
field='Masses')
pos = np.zeros((len(mass), 3), dtype=np.float32)
for mdi in range(3):
pos[:, mdi] = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType='stars',
field='Coordinates',
mdi=mdi)
layout = pm.decompose(pos)
pos1 = layout.exchange(pos)
del pos
mass1 = layout.exchange(mass)
del mass
TNGmap = pm.create(type="real")
TNGmap.paint(pos1, mass=mass1, layout=None, hold=False)
del pos1, mass1
elif field == 'ne':
raise NotImplementedError
elif field == 'T':
raise NotImplementedError
elif field == 'nHI':
def nHI(mass, XHI):
h = 0.6774
XH = 0.76
mp = 1.6726219e-27
Msun10 = 1.989e40
BoxSize = 205.
Mpc_cm = 3.085677581e24
return Msun10 / h / (BoxSize / args.Nmesh * Mpc_cm /
h)**3 * mass * XH / mp * XHI
gasmass = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType='gas',
field='Masses')
XHI = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType='gas',
field='NeutralHydrogenAbundance')
mass = nHI(gasmass, XHI)
pos = np.zeros((len(mass), 3), dtype=np.float32)
for mdi in range(3):
pos[:, mdi] = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType='gas',
field='Coordinates',
mdi=mdi)
layout = pm.decompose(pos)
pos1 = layout.exchange(pos)
del pos
mass1 = layout.exchange(mass)
del mass
TNGmap = pm.create(type="real")
TNGmap.paint(pos1, mass=mass1, layout=None, hold=False)
del pos1, mass1
elif field == 'ne_vz':
raise NotImplementedError
elif field == 'Mstar_vz':
raise NotImplementedError
FieldMesh(TNGmap).save(address)
return TNGmap
noise=noise,
title=title,
fsize=15)
if rank == 0:
numpy.savetxt(optfolder + '/error_ps.txt',
numpy.array([kerror, perror]).T,
header='kerror, perror')
if stage2:
ipkmodel = interp1d(kerror,
perror,
bounds_error=False,
fill_value=(perror[0], perror[-1]))
ivarmesh = truth_noise_model.get_ivarmesh(data_p, ipkmodel)
FieldMesh(ivarmesh).save(optfolder + 'ivarmesh',
dataset='ivar',
mode='real')
kerror, perror = eval_bfit(data_n.mapp,
bmod,
optfolder,
noise=noise,
title=title,
fsize=15,
suff='-noise')
if rank == 0:
numpy.savetxt(optfolder + '/error_psn.txt',
numpy.array([kerror, perror]).T,
header='kerror, perror')
else:
ivarmesh = None
def evaluate1(model, data, field, M0=0, stoch=False, kmin=None, dk=None):
'''return position order:
xm.power, xs.power, xd.power,
pm1.power, pm2.power,
ps1.power, ps2.power,
pd1.power, pd2.power,
data_preview, model_preview
'''
from nbodykit.lab import FieldMesh, FFTPower, ProjectedFFTPower
if kmin is None: kmin = 0
if dk is None: dk = 2*numpy.pi/model.s.BoxSize[0]
model.mapp += M0
data.mapp += M0
#print('Means are : ', model.mapp.cmean(), data.mapp.cmean())
if abs(model.mapp.cmean()) > 1e-3: modmappmean = model.mapp.cmean()
else: modmappmean = 1.
if abs(data.mapp.cmean()) > 1e-3: datmappmean = data.mapp.cmean()
else: datmappmean = 1.
modmappmean, datmappmean = 1., 1.
#if abs(model.mapp.cmean()) > 1e-3: modmappmean = model.mapp.cmean()
#else: modmappmean = 1.
#if abs(data.mapp.cmean()) > 1e-3: datmappmean = data.mapp.cmean()
#else: datmappmean = 1.
data_preview, model_preview = [], []
if field == 'mapp':
x = FFTPower(first=FieldMesh(model.mapp/modmappmean), second=FieldMesh(data.mapp/datmappmean), mode='1d', kmin=kmin, dk=dk)
p1 = FFTPower(first=FieldMesh(model.mapp/modmappmean), mode='1d', kmin=kmin, dk=dk)
p2 = FFTPower(first=FieldMesh(data.mapp/datmappmean), mode='1d', kmin=kmin, dk=dk)
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview.append(data.mapp.preview(axes=axes))
model_preview.append(model.mapp.preview(axes=axes))
elif field == 's':
x = FFTPower(first=FieldMesh(model.s), second=FieldMesh(data.s), mode='1d', kmin=kmin, dk=dk)
p1 = FFTPower(first=FieldMesh(model.s), mode='1d', kmin=kmin, dk=dk)
p2 = FFTPower(first=FieldMesh(data.s), mode='1d', kmin=kmin, dk=dk)
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview.append(data.s.preview(axes=axes))
model_preview.append(model.s.preview(axes=axes))
elif field == 'd':
x = FFTPower(first=FieldMesh(model.d), second=FieldMesh(data.d), mode='1d', kmin=kmin, dk=dk)
p1 = FFTPower(first=FieldMesh(model.d), mode='1d', kmin=kmin, dk=dk)
p2 = FFTPower(first=FieldMesh(data.d), mode='1d', kmin=kmin, dk=dk)
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview.append(data.d.preview(axes=axes))
model_preview.append(model.d.preview(axes=axes))
else:
print('Field not recongnized')
return 0
return x.power, p1.power, p2.power, data_preview, model_preview
def evaluate2d1(model, data, field, M0=0, kmin=None, dk=None, Nmu=5, retmesh=False, los=[0, 0, 1]):
'''return position order:
xm.power, xs.power, xd.power,
pm1.power, pm2.power,
ps1.power, ps2.power,
pd1.power, pd2.power,
data_preview, model_preview
'''
from nbodykit.lab import FieldMesh, FFTPower, ProjectedFFTPower
if kmin is None: kmin = 0
if dk is None: dk = 2*numpy.pi/model.s.BoxSize[0]
model.mapp += M0
data.mapp += M0
if abs(model.mapp.cmean()) > 1e-3: modmappmean = model.mapp.cmean()
else: modmappmean = 1.
if abs(data.mapp.cmean()) > 1e-3: datmappmean = data.mapp.cmean()
else: datmappmean = 1.
modmappmean, datmappmean = 1., 1.
if field == 'mapp':
x = FFTPower(first=FieldMesh(model.mapp/modmappmean),
second=FieldMesh(data.mapp/datmappmean), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
p1 = FFTPower(first=FieldMesh(model.mapp/modmappmean), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
p2 = FFTPower(first=FieldMesh(data.mapp/datmappmean), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
elif field == 's':
x = FFTPower(first=FieldMesh(model.s), second=FieldMesh(data.s), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
p1 = FFTPower(first=FieldMesh(model.s), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
p2 = FFTPower(first=FieldMesh(data.s), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
elif field == 'd':
x = FFTPower(first=FieldMesh(model.d), second=FieldMesh(data.d), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
p1 = FFTPower(first=FieldMesh(model.d), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
p2 = FFTPower(first=FieldMesh(data.d), mode='2d', kmin=kmin, dk=dk, Nmu=Nmu, los=[0, 0, 1])
else:
print('Field not recongnized')
return 0
return x.power, p1.power, p2.power
cencat,
satcat,
los=los)
if rank == 0: print('Creating HI mesh in redshift space')
h1mesh = HImodelz.createmesh(bs,
nc,
halocat,
cencat,
satcat,
mode=mode,
position='RSDpos',
weight='HImass',
tofield=True)
FieldMesh(h1mesh).save(scratchcm + sim +
'/fastpm_%0.4f/HImeshz-N%04d/' % (aa, nc),
dataset=modelname,
mode='real')
if rank == 0: print('Creating HI mesh in real space for bias')
h1mesh = HImodelz.createmesh(bs,
nc,
halocat,
cencat,
satcat,
mode=mode,
position='Position',
weight='HImass',
tofield=True)
FieldMesh(h1mesh).save(scratchcm + sim +
'/fastpm_%0.4f/HImesh-N%04d/' % (aa, nc),
dataset=modelname,
def calc_model_errors_at_cat_pos(trf_specs=None,
paths=None,
cat_specs=None,
ext_grids_to_load=None,
trf_fcn_opts=None,
grid_opts=None,
sim_opts=None,
power_opts=None,
Pkmeas_helper_columns=None,
Pkmeas_helper_columns_calc_crosses=False,
f_log_growth=None,
debug=True):
"""
Compute models at catalog positions, and residual to target, for each
trf_spec.
"""
comm = CurrentMPIComm.get()
result_of_trf_spec = OrderedDict()
for trf_spec in trf_specs:
cats = OrderedDict()
if True:
## Load target catalog
cat_id = trf_spec.target_field
cat_opts = cat_specs[cat_id]
# read file
fname = os.path.join(paths['in_path'], cat_opts['in_fname'])
if comm.rank == 0:
print('Read %s' % fname)
# Read rho, not dividing or subtracting mean, to get velocity correctly
cat = BigFileCatalog(fname, dataset='./', header='Header')
# cuts
for cut_column, cut_instruction in cat_opts['cuts'].items():
cut_op, cut_value = cut_instruction
if cut_op == 'min':
cat = cat[cat[cut_column] >= cut_value]
elif cut_op == 'max':
cat = cat[cat[cut_column] < cut_value]
else:
raise Exception('Invalid cut operation %s' % str(cut_op))
# Set Position column
cat['Position'] = cat[cat_opts['position_column']]
# compute the value we are interested in, save in 'val' column
component = cat_opts['val_component']
if component is None:
cat['val'] = cat[cat_opts['val_column']][:]
else:
cat['val'] = cat[cat_opts['val_column']][:, component]
# catalog rescale factor
if cat_opts.get('rescale_factor', None) is not None:
if cat_opts['rescale_factor'] == 'RSDFactor':
cat['val'] *= cat.attrs['RSDFactor'][0]
else:
raise Exception('Invalid rescale_factor %s' %
cat_opts['rescale_factor'])
# additional rescale factors or post processing options for catalog
if hasattr(trf_spec, 'field_opts'):
if cat_id in trf_spec.field_opts:
this_field_opts = trf_spec.field_opts[cat_id]
if this_field_opts.get('additional_rescale_factor',
1.0) != 1.0:
resc_fac = this_field_opts['additional_rescale_factor']
if type(resc_fac) == float:
cat['val'] *= resc_fac
else:
raise Exception("Invalid rescale factor: %s" %
resc_fac)
# keep only Position and val columns, delete all other columns
cat2 = catalog_persist(cat, columns=['Position', 'val'])
del cat
cstats = get_cstats_string(cat2['val'].compute())
if comm.rank == 0:
print('CATALOG %s: %s\n' % (cat_id, cstats))
# Store in cats dict
cats[cat_id] = cat2
del cat2
target_cat = cats[trf_spec.target_field]
# Make sure we have no linear source (trf fcns not implemented here)
if trf_spec.linear_sources is not None:
raise Exception(
'Trf fcns not implemented for model error at cat pos')
## Load remaining fields, assuming they are bigfiles meshs, and get
# residual at target positions.
residual_cat = target_cat.copy()
for mesh_id in trf_spec.fixed_linear_sources:
fname = os.path.join(paths['in_path'], mesh_id)
if comm.rank == 0:
print('Read %s' % fname)
if hasattr(trf_spec, 'field_opts'):
fopts = trf_spec.field_opts[mesh_id]
else:
# default options for meshs
fopts = {'read_mode': 'velocity', 'readout_window': 'cic'}
# read mesh
if fopts['read_mode'] == 'velocity':
# get rho (don't divide by mean)
mesh = read_vel_from_bigfile(fname)
elif fopts['read_mode'] == 'density':
# compute fractional delta (taking out mean)
mesh = read_delta_from_bigfile(fname)
else:
raise Exception('Invalid read_mode: %s' % fopts['read_mode'])
# more post processing of mesh
if fopts.get('additional_rescale_factor', 1.0) != 1.0:
resc_fac = fopts['additional_rescale_factor']
if comm.rank == 0:
print('Apply rescale fac to %s: %s' %
(mesh_id, str(resc_fac)))
if type(resc_fac) == float:
mesh = FieldMesh(mesh.compute(mode='real') * resc_fac)
elif resc_fac == 'f_log_growth':
mesh = FieldMesh(mesh.compute(mode='real') * f_log_growth)
else:
raise Exception("Invalid rescale factor: %s" % resc_fac)
cstats = get_cstats_string(mesh.compute())
if comm.rank == 0:
print('MESH %s: %s\n' % (mesh_id, cstats))
# Read out mesh at Position of target catalog
if comm.rank == 0:
print('Read out mesh at target pos: %s' % mesh_id)
mesh_at_target_pos = readout_mesh_at_cat_pos(
mesh=mesh,
cat=target_cat,
readout_window=fopts['readout_window'])
## Compute residual = target - fixed_linear_sources at target position
residual_cat['val'] -= mesh_at_target_pos
cstats = get_cstats_string(residual_cat['val'].compute())
if comm.rank == 0:
print('RESIDUAL %s - %s: %s\n' %
(trf_spec.target_field, trf_spec.save_bestfit_field, cstats))
# plot histogram
if debug:
import matplotlib.pyplot as plt
plt.hist(residual_cat['val'].compute(),
bins=50,
range=(-10, 10),
alpha=0.5)
plt.hist(target_cat['val'].compute(),
bins=50,
range=(-10, 10),
alpha=0.5)
plt.hist(mesh_at_target_pos, bins=50, range=(-10, 10), alpha=0.5)
plt.show()
## paint to mesh
to_mesh_kwargs = {
#'Nmesh': grid_opts.Ngrid,
'window': 'cic',
'compensated': False,
'interlaced': False,
#'BoxSize': residual_cat.attrs['BoxSize'],
'dtype': 'f8'
}
# rho mesh of target, taking avg value instead of summing
target_mesh = FieldMesh(
mass_avg_weighted_paint_cat_to_rho(target_cat,
weight='val',
Nmesh=grid_opts.Ngrid,
to_mesh_kwargs=to_mesh_kwargs,
rho_of_empty_cells=0.0,
verbose=True)[0])
# rho mesh of residual, taking avg value instead of summing
residual_mesh = FieldMesh(
mass_avg_weighted_paint_cat_to_rho(residual_cat,
weight='val',
Nmesh=grid_opts.Ngrid,
to_mesh_kwargs=to_mesh_kwargs,
rho_of_empty_cells=0.0,
verbose=True)[0])
## Compute power spectra
pow_kwargs = {
'mode': power_opts.Pk_1d_2d_mode,
'k_bin_width': power_opts.k_bin_width
}
## Save results of this trf_spec in dict
trf_res = OrderedDict()
# Power spectra
trf_res['power'] = OrderedDict()
trf_res['power']['Target-Model'] = calc_power(residual_mesh,
**pow_kwargs)
trf_res['power']['Target'] = calc_power(target_mesh, **pow_kwargs)
if debug:
# plot power spectra
import matplotlib.pyplot as plt
print('plot power')
k = trf_res['power']['Target'].power['k']
plt.loglog(k, k**2 * trf_res['power']['Target'].power['power'])
k = trf_res['power']['Target-Model'].power['k']
plt.loglog(k,
k**2 * trf_res['power']['Target-Model'].power['power'])
plt.show()
# target catalog results
trf_res['target_cat_attrs'] = target_cat.attrs
trf_res['target_cat_csize'] = target_cat.csize
# todo: save rms of fields and target
result_of_trf_spec[trf_spec.save_bestfit_field] = trf_res
# save all results in a big pickle_dict
pickle_dict = OrderedDict()
pickle_dict['result_of_trf_spec'] = result_of_trf_spec
print('pickle_dict:', pickle_dict)
return pickle_dict
def evaluate(model, data, map2=False):
from nbodykit.lab import FieldMesh, FFTPower, ProjectedFFTPower
if abs(model.mapp.cmean()) > 1e-3: modmappmean = model.mapp.cmean()
else: modmappmean = 1.
if abs(data.mapp.cmean()) > 1e-3: datmappmean = data.mapp.cmean()
else: datmappmean = 1.
modmappmean, datmappmean = 1., 1.
modmappmean2, datmappmean2 = 1., 1.
# try:
# data.mapp2.cmean()
# map2 = True
# except: map2 = False
#
xm = FFTPower(first=FieldMesh(model.mapp/modmappmean), second=FieldMesh(data.mapp/datmappmean), mode='1d')
xd = FFTPower(first=FieldMesh(model.d), second=FieldMesh(data.d), mode='1d')
xs = FFTPower(first=FieldMesh(model.s), second=FieldMesh(data.s), mode='1d')
pm1 = FFTPower(first=FieldMesh(model.mapp/modmappmean), mode='1d')
pd1 = FFTPower(first=FieldMesh(model.d), mode='1d')
ps1 = FFTPower(first=FieldMesh(model.s), mode='1d')
pm2 = FFTPower(first=FieldMesh(data.mapp/datmappmean), mode='1d')
pd2 = FFTPower(first=FieldMesh(data.d), mode='1d')
ps2 = FFTPower(first=FieldMesh(data.s), mode='1d')
if map2 :
xm2 = FFTPower(first=FieldMesh(model.mapp2/modmappmean2), second=FieldMesh(data.mapp2/datmappmean2), mode='1d')
pm12 = FFTPower(first=FieldMesh(model.mapp2/modmappmean2), mode='1d')
pm22 = FFTPower(first=FieldMesh(data.mapp2/datmappmean2), mode='1d')
data_preview = dict(s=[], d=[], mapp=[])
model_preview = dict(s=[], d=[], mapp=[])
data_preview2 = dict(s=[], d=[], mapp=[])
model_preview2 = dict(s=[], d=[], mapp=[])
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview['d'].append(data.d.preview(axes=axes))
data_preview['s'].append(data.s.preview(axes=axes))
data_preview['mapp'].append(data.mapp.preview(axes=axes))
model_preview['d'].append(model.d.preview(axes=axes))
model_preview['s'].append(model.s.preview(axes=axes))
model_preview['mapp'].append(model.mapp.preview(axes=axes))
if map2:
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview2['d'].append(data.d.preview(axes=axes))
data_preview2['s'].append(data.s.preview(axes=axes))
data_preview2['mapp'].append(data.mapp2.preview(axes=axes))
model_preview2['d'].append(model.d.preview(axes=axes))
model_preview2['s'].append(model.s.preview(axes=axes))
model_preview2['mapp'].append(model.mapp2.preview(axes=axes))
#data_preview['mapp'] = data.mapp.preview(axes=(0, 1))
#model_preview['mapp'] = model.mapp.preview(axes=(0, 1))
if map2:
return [xm.power, xs.power, xd.power, pm1.power, pm2.power, ps1.power, ps2.power, pd1.power, pd2.power, data_preview, model_preview], \
[xm2.power, xs.power, xd.power, pm12.power, pm22.power, ps1.power, ps2.power, pd1.power, pd2.power, data_preview2, model_preview2],
else: return xm.power, xs.power, xd.power, pm1.power, pm2.power, ps1.power, ps2.power, pd1.power, pd2.power, data_preview, model_preview
#pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
#rank = pm.comm.rank
dyn = BigFileCatalog(dpath + '/fastpm_%0.4f/1' % aa)
#
fpos = dyn['Position'].compute()
idd = dyn['ID'].compute()
attrs = dyn.attrs
grid = tools.getqfromid(idd, attrs, nc)
print('grid computed')
del dyn, idd
header = '1,b1,b2,bg,bk'
names = header.split(',')
for i in range(len(names)):
ff = BigFileMesh(lpath + '/lag', names[i]).paint()
pm = ff.pm
rank = pm.comm.rank
glay, play = pm.decompose(grid), pm.decompose(fpos)
wts = ff.readout(grid, layout=glay, resampler='nearest')
print(rank, ' got weights')
x = FieldMesh(pm.paint(fpos, mass=wts, layout=play))
x.save(lpath + 'eul-z%03d-R%d' % (zz * 100, Rsm),
dataset=names[i],
mode='real')
del pm, ff, wts, x
#PM
#whitec = pm.generate_whitenoise(seed, mode='complex', unitary=False)
#lineark = whitec.apply(lambda k, v:Planck15.get_pklin(sum(ki ** 2 for ki in k)**0.5, 0) ** 0.5 * v / v.BoxSize.prod() ** 0.5)
#linear = lineark.c2r()
linear = BigFileMesh(
'/project/projectdirs/astro250/chmodi/cosmo4d/data/z00/L0400_N0128_S0100_05step/mesh/',
's').paint()
lineark = linear.r2c()
state = solver.lpt(lineark, grid, conf['stages'][0], order=2)
solver.nbody(state, leapfrog(conf['stages']))
final = pm.paint(state.X)
if pm.comm.rank == 0:
print('X, V computed')
cat = ArrayCatalog({
'Position': state.X,
'Velocity': state.V
},
BoxSize=pm.BoxSize,
Nmesh=pm.Nmesh)
kdd = KDDensity(cat).density
cat['KDDensity'] = kdd
#Save
path = '/project/projectdirs/astro250/chmodi/cosmo4d/data/'
ofolder = path + 'z%02d/L%04d_N%04d_S%04d_%02dstep_fpm/' % (
z * 10, bs, nc, seed, nsteps)
FieldMesh(linear).save(ofolder + 'mesh', dataset='s', mode='real')
FieldMesh(final).save(ofolder + 'mesh', dataset='d', mode='real')
cat.save(ofolder + 'dynamic/1', ('Position', 'Velocity', 'KDDensity'))
def evaluate(model, data):
from nbodykit.lab import FieldMesh, FFTPower, ProjectedFFTPower
#print('Means are : ', model.mapp.cmean(), data.mapp.cmean())
if abs(model.mapp.cmean()) > 1e-3: modmappmean = model.mapp.cmean()
else: modmappmean = 1.
if abs(data.mapp.cmean()) > 1e-3: datmappmean = data.mapp.cmean()
else: datmappmean = 1.
modmappmean, datmappmean = 1., 1.
xm = FFTPower(first=FieldMesh(model.mapp / modmappmean),
second=FieldMesh(data.mapp / datmappmean),
mode='1d')
xd = FFTPower(first=FieldMesh(model.d),
second=FieldMesh(data.d),
mode='1d')
xs = FFTPower(first=FieldMesh(model.s),
second=FieldMesh(data.s),
mode='1d')
pm1 = FFTPower(first=FieldMesh(model.mapp / modmappmean), mode='1d')
pd1 = FFTPower(first=FieldMesh(model.d), mode='1d')
ps1 = FFTPower(first=FieldMesh(model.s), mode='1d')
pm2 = FFTPower(first=FieldMesh(data.mapp / datmappmean), mode='1d')
pd2 = FFTPower(first=FieldMesh(data.d), mode='1d')
ps2 = FFTPower(first=FieldMesh(data.s), mode='1d')
data_preview = dict(s=[], d=[], mapp=[])
model_preview = dict(s=[], d=[], mapp=[])
for axes in [[1, 2], [0, 2], [0, 1]]:
data_preview['d'].append(data.d.preview(axes=axes))
data_preview['s'].append(data.s.preview(axes=axes))
data_preview['mapp'].append(data.mapp.preview(axes=axes))
model_preview['d'].append(model.d.preview(axes=axes))
model_preview['s'].append(model.s.preview(axes=axes))
model_preview['mapp'].append(model.mapp.preview(axes=axes))
#data_preview['mapp'] = data.mapp.preview(axes=(0, 1))
#model_preview['mapp'] = model.mapp.preview(axes=(0, 1))
return xm.power, xs.power, xd.power, pm1.power, pm2.power, ps1.power, ps2.power, pd1.power, pd2.power, data_preview, model_preview
method='L-BFGS-B',
args=(residue_model, loss_train_model, loss_validate_model),
jac=loss.derivative,
bounds=bounds,
options={
'maxiter': 2000,
'disp': False
})
if comm.rank == 0:
print('Finished optimization.')
print('Best validate loss:', loss.best_value)
print('Best param:', loss.best_x)
param = loss.best_x
else:
param = np.loadtxt(args.restore)
if args.target == 'dm':
assert len(param) == 5 * args.Nstep
baryon = False
else:
assert len(param) == 5 * args.Nstep + 3
baryon = True
#evaluate
model = LDL.build(X=X, pm=pm, Nstep=args.Nstep, baryon=baryon)
LDLmap = model.compute('F', init=dict(param=param))
save = args.save + '/' + args.target + '_snap' + str(args.snapNum).zfill(
3) + '_Nmesh%d_Nstep%d_n%.2f_map' % (args.Nmesh, args.Nstep, args.n)
FieldMesh(LDLmap).save(save)
cencat,
satcat,
los=los)
if rank == 0: print('Creating HI mesh in redshift space')
h1mesh = HImodelz.createmesh(bs,
nc,
halocat,
cencat,
satcat,
mode=mode,
position='RSDpos',
weight='HImass',
tofield=True)
FieldMesh(h1mesh).save(args['outfolder'] % aa + '/HImeshz-N%04d/' %
(nc),
dataset=modelname,
mode='real')
if rank == 0: print('Creating HI mesh in real space for bias')
h1mesh = HImodelz.createmesh(bs,
nc,
halocat,
cencat,
satcat,
mode=mode,
position='Position',
weight='HImass',
tofield=True)
FieldMesh(h1mesh).save(args['outfolder'] % aa + '/HImesh-N%04d/' %
(nc),
dataset=modelname,
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
rank = pm.comm.rank
#grid = pm.mesh_coordinates()*bs/nc
lin = BigFileMesh(lpath + '/linear', 'LinearDensityK').paint()
lin -= lin.cmean()
print(rank, 'lin field read')
header = '1,b1,b2,bg,bk'
names = header.split(',')
lag_fields = tools.getlagfields(
pm, lin, R=Rsm) # use the linear field at the desired redshift
print(rank, 'lag field created')
for i, ff in enumerate(lag_fields):
x = FieldMesh(ff)
x.save(lpath + 'lag', dataset=names[i], mode='real')
## ###
## dyn = BigFileCatalog(dpath + '/fastpm_%0.4f/1'%aa)
##
## print(rank, 'files read')
##
## #
## fpos = dyn['Position'].compute()
## idd = dyn['ID'].compute()
## attrs = dyn.attrs
##
## grid = tools.getqfromid(idd, attrs, nc)
## print(rank, 'grid computed')
##
rhomass = rho.readout(position, layout=layout)
if rank == 0: print(rhomass)
random = truth_pm.paint(position, layout=layout)
randommass = random.readout(position, layout=layout)
if rank == 0: print(randommass)
basemesh = std.gauss(data_p.mapp, Rbao)
dispmesh = std.calc_displist(base=basemesh, b=bias, kmin=kmin, angle=angle)
if rank == 0: print([d[...].std() for d in dispmesh])
hpshift = std.displace(truth_pm,
dispmesh,
position,
rsd=True,
f=ff,
beta=beta,
mass=rhomass)
rshift = std.displace(truth_pm,
dispmesh,
position,
rsd=True,
f=ff,
beta=beta,
mass=randommass)
recon = hpshift - rshift
if upsample: fname = 'stdrecon_up'
else: fname = 'stdrecon'
if usenoise: fname += '-noise'
FieldMesh(recon).save(optfolder + fname, dataset='std', mode='real')
def load_TNG_map(TNG_basepath, snapNum, field, pm):
assert field in ['dm', 'Mstar', 'ne', 'nT', 'nHI', 'neVz', 'MstarVz']
#Try directly loading map. If it does not exist, paint the map and save it for future use.
address = TNG_basepath + '/snapdir_' + str(snapNum).zfill(
3) + '/' + field + 'map_Nmesh' + str(pm.Nmesh[0])
try:
TNGmap = BigFileMesh(address, dataset='Field').to_real_field()
except:
if field == 'dm':
partType = 'dm'
elif field in ['Mstar', 'MstarVz']:
partType = 'stars'
elif field in ['ne', 'nT', 'nHI', 'neVz']:
partType = 'gas'
if field == 'Mstar':
mass = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Masses')
elif field == 'ne':
gasmass = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Masses')
Xe = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='ElectronAbundance')
a = scalefactor(TNG_basepath, snapNum)
mass = n_cm3(gasmass, Xe, a, pm.Nmesh[0])
del gasmass, Xe
elif field == 'nT':
def temperature(Xe, u):
XH = 0.76
kb = 1.38064852e-23
mp = 1.6726219e-27
ufac = 1e6 #u: (km/s)^2 -> (m/s)^2
mu = 4. / (1. + 3. * XH + 4. * XH * Xe) * mp
T = 2. / 3. * ufac * u / kb * mu
return T
Xe = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='ElectronAbundance')
u = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='InternalEnergy')
T = temperature(Xe, u)
del u
gasmass = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Masses')
a = scalefactor(TNG_basepath, snapNum)
ne = n_cm3(gasmass, Xe, a, pm.Nmesh[0])
del gasmass, Xe
mass = ne * T
del ne, T
elif field == 'nHI':
gasmass = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Masses')
XHI = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='NeutralHydrogenAbundance')
a = scalefactor(TNG_basepath, snapNum)
mass = n_cm3(gasmass, XHI, a, pm.Nmesh[0])
del gasmass, XHI
elif field == 'neVz':
gasmass = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Masses')
Xe = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='ElectronAbundance')
a = scalefactor(TNG_basepath, snapNum)
ne = n_cm3(gasmass, Xe, a, pm.Nmesh[0])
vz = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Velocities',
mdi=2) * a**0.5
mass = ne * vz
del ne, vz
elif field == 'MstarVz':
mass = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Masses')
a = scalefactor(TNG_basepath, snapNum)
vz = load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Velocities',
mdi=2) * a**0.5
mass = mass * vz
del vz
pos = []
for mdi in range(3):
pos.append(
load_TNG_data(TNG_basepath=TNG_basepath,
snapNum=snapNum,
partType=partType,
field='Coordinates',
mdi=mdi))
pos = np.array(pos).T
layout = pm.decompose(pos)
pos1 = layout.exchange(pos)
if field == 'dm':
mass1 = 1.0 * pm.Nmesh.prod() / pm.comm.allreduce(len(pos),
op=MPI.SUM)
else:
mass1 = layout.exchange(mass)
del mass
del pos
TNGmap = pm.create(type="real")
TNGmap.paint(pos1, mass=mass1, layout=None, hold=False)
del pos1, mass1
FieldMesh(TNGmap).save(address)
return TNGmap
