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)