195 lines
No EOL
7.1 KiB
Python
195 lines
No EOL
7.1 KiB
Python
import aquila_borg as borg
|
|
import numpy as np
|
|
import configparser
|
|
import borg_velocity.utils as utils
|
|
from borg_velocity.utils import myprint
|
|
import borg_velocity.forwards as forwards
|
|
import Pk_library as PKL
|
|
import matplotlib.pyplot as plt
|
|
import symbolic_pofk.syrenhalofit as syrenhalofit
|
|
|
|
def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_file=None, gravity='lpt') -> borg.forward.BaseForwardModel:
|
|
"""
|
|
Builds the gravity model and returns the forward model chain.
|
|
|
|
Args:
|
|
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
|
- box (borg.forward.BoxModel): The input box model.
|
|
- ini_file (str, default=None): The location of the ini file. If None, use borg.getIniConfigurationFilename()
|
|
|
|
Returns:
|
|
borg.forward.BaseForwardModel: The forward model.
|
|
|
|
"""
|
|
|
|
myprint(f"Building gravity model {gravity}")
|
|
|
|
if ini_file is None:
|
|
myprint("Reading from configuration file: " + borg.getIniConfigurationFilename())
|
|
config = configparser.ConfigParser()
|
|
config.read(borg.getIniConfigurationFilename())
|
|
else:
|
|
myprint("Reading from configuration file: " + ini_file)
|
|
config = configparser.ConfigParser()
|
|
config.read(ini_file)
|
|
ai = float(config['model']['ai'])
|
|
af = float(config['model']['af'])
|
|
supersampling = int(config['model']['supersampling'])
|
|
print('Supersampling:', supersampling)
|
|
|
|
if gravity in ['pm', 'cola']:
|
|
forcesampling = int(config['model']['forcesampling'])
|
|
forcesampling = 2
|
|
print('FORCESAMPLING:', forcesampling)
|
|
|
|
# Setup forward model
|
|
chain = borg.forward.ChainForwardModel(box)
|
|
chain.addModel(borg.forward.models.HermiticEnforcer(box))
|
|
|
|
# CLASS transfer function
|
|
chain @= borg.forward.model_lib.M_PRIMORDIAL_AS(box)
|
|
transfer_class = borg.forward.model_lib.M_TRANSFER_CLASS(box, opts=dict(a_transfer=1.0))
|
|
transfer_class.setModelParams({"extra_class_arguments":{'YHe':'0.24'}})
|
|
chain @= transfer_class
|
|
|
|
if gravity == 'linear':
|
|
raise NotImplementedError(gravity)
|
|
elif gravity == 'lpt':
|
|
mod = borg.forward.model_lib.M_LPT_CIC(
|
|
box,
|
|
opts=dict(a_initial=af,
|
|
a_final=af,
|
|
do_rsd=False,
|
|
supersampling=supersampling,
|
|
lightcone=False,
|
|
part_factor=1.01,))
|
|
elif gravity == '2lpt':
|
|
mod = borg.forward.model_lib.M_2LPT_CIC(
|
|
box,
|
|
opts=dict(a_initial=af,
|
|
a_final=af,
|
|
do_rsd=False,
|
|
supersampling=supersampling,
|
|
lightcone=False,
|
|
part_factor=1.01,))
|
|
elif gravity == 'pm':
|
|
mod = borg.forward.model_lib.M_PM_CIC(
|
|
box,
|
|
opts=dict(a_initial=af,
|
|
a_final=af,
|
|
do_rsd=False,
|
|
supersampling=supersampling,
|
|
part_factor=1.01,
|
|
forcesampling=forcesampling,
|
|
pm_start_z=1/ai - 1,
|
|
pm_nsteps=int(config['model']['nsteps']),
|
|
tcola=False))
|
|
elif gravity == 'cola':
|
|
mod = borg.forward.model_lib.M_PM_CIC(
|
|
box,
|
|
opts=dict(a_initial=af,
|
|
a_final=af,
|
|
do_rsd=False,
|
|
supersampling=supersampling,
|
|
part_factor=1.01,
|
|
forcesampling=forcesampling,
|
|
pm_start_z=1/ai - 1,
|
|
pm_nsteps=int(config['model']['nsteps']),
|
|
tcola=True))
|
|
else:
|
|
raise NotImplementedError(gravity)
|
|
|
|
mod.accumulateAdjoint(True)
|
|
chain @= mod
|
|
|
|
# Cosmological parameters
|
|
if ini_file is None:
|
|
cpar = utils.get_cosmopar(borg.getIniConfigurationFilename())
|
|
else:
|
|
cpar = utils.get_cosmopar(ini_file)
|
|
chain.setCosmoParams(cpar)
|
|
|
|
# This is the forward model for the model parameters
|
|
fwd_param = borg.forward.ChainForwardModel(box)
|
|
mod_null = forwards.NullForward(box)
|
|
fwd_param.addModel(mod_null)
|
|
fwd_param.setCosmoParams(cpar)
|
|
|
|
# This is the forward model for velocity
|
|
velmodel_name = config['model']['velocity']
|
|
velmodel = getattr(borg.forward.velocity, velmodel_name)
|
|
if velmodel_name == 'LinearModel':
|
|
fwd_vel = velmodel(box, mod, af)
|
|
elif velmodel_name == 'CICModel':
|
|
rsmooth = float(config['model']['rsmooth']) # Mpc/h
|
|
fwd_vel = velmodel(box, mod, rsmooth)
|
|
else:
|
|
fwd_vel = velmodel(box, mod)
|
|
|
|
return chain
|
|
|
|
ini_file = '../conf/velmass_ini.ini'
|
|
|
|
# Setup config
|
|
config = configparser.ConfigParser()
|
|
config.read(ini_file)
|
|
|
|
# Cosmology
|
|
cosmo = utils.get_cosmopar(ini_file)
|
|
|
|
# Input box
|
|
box_in = borg.forward.BoxModel()
|
|
box_in.L = (float(config['system']['L0']), float(config['system']['L1']), float(config['system']['L2']))
|
|
box_in.N = (int(config['system']['N0']), int(config['system']['N1']), int(config['system']['N2']))
|
|
# box_in.N = (128, 128, 128)
|
|
# box_in.N = (256, 256, 256)
|
|
box_in.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2']))
|
|
|
|
# Make some initial conditions
|
|
s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
|
|
s_real = np.fft.irfftn(s_hat, norm="ortho")
|
|
|
|
all_gravity = ['lpt', '2lpt', 'pm', 'cola']
|
|
all_dens = [None] * len(all_gravity)
|
|
all_pk = [None] * len(all_gravity)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(6,4))
|
|
|
|
for i, gravity in enumerate(all_gravity):
|
|
# Run BORG density field
|
|
output_density = np.zeros(box_in.N)
|
|
fwd_model = build_gravity_model(None, box_in, ini_file=ini_file, gravity=gravity)
|
|
fwd_model.forwardModel_v2(s_hat)
|
|
fwd_model.getDensityFinal(output_density)
|
|
all_dens[i] = output_density.copy()
|
|
|
|
Pk = PKL.Pk(all_dens[i].astype(np.float32), box_in.L[0], axis=0, MAS='CIC', threads=1, verbose=True)
|
|
all_pk[i] = (Pk.k3D.copy(), Pk.Pk[:,0].copy())
|
|
|
|
ax.loglog(all_pk[i][0], all_pk[i][1], label=gravity)
|
|
|
|
# Load VELMASS field
|
|
dirname = config['mock']['velmass_dirname']
|
|
field = np.load(dirname + '/../density.npz')['d']
|
|
field = field / np.mean(field) - 1
|
|
with open(dirname + '/auto-rockstar.cfg') as f:
|
|
data = [r for r in f]
|
|
Lbox = [r for r in data if r.startswith('BOX_SIZE')][0].strip()
|
|
Lbox = float(Lbox.split('=')[1])
|
|
print('VELMASS box', Lbox)
|
|
Pk = PKL.Pk(field.astype(np.float32), Lbox, axis=0, MAS='CIC', threads=1, verbose=True)
|
|
m = Pk.k3D > all_pk[-1][0].min()
|
|
ax.loglog(Pk.k3D[m], Pk.Pk[m,0], label='velmass')
|
|
|
|
# Syren
|
|
k = Pk.k3D[m]
|
|
pk_syrenhalofit = syrenhalofit.run_halofit(k,
|
|
cosmo.sigma8, cosmo.omega_m, cosmo.omega_b, cosmo.h, cosmo.n_s, 1.0,
|
|
emulator='fiducial', extrapolate=True, which_params='Bartlett', add_correction=True)
|
|
ax.loglog(k, pk_syrenhalofit, label='syren')
|
|
ymin = max(ax.get_ylim()[0], 1e1)
|
|
ax.set_ylim(ymin, None)
|
|
ax.legend()
|
|
fig.tight_layout()
|
|
fig.savefig('../figs/pk_gravity_models_test.png', facecolor='white', bbox_inches='tight')
|
|
plt.close(fig) |