249 lines
No EOL
8.5 KiB
Python
249 lines
No EOL
8.5 KiB
Python
import borg_velocity.likelihood as likelihood
|
||
import borg_velocity.forwards as forwards
|
||
import borg_velocity.utils as utils
|
||
from borg_velocity.utils import myprint
|
||
|
||
import aquila_borg as borg
|
||
import configparser
|
||
import h5py as h5
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
import glob
|
||
import os
|
||
import matplotlib.pyplot as plt
|
||
|
||
dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics'
|
||
ini_file = f'{dirname}/ini_file.ini'
|
||
|
||
wdir = os.getcwd()
|
||
os.chdir(dirname)
|
||
|
||
def get_mcmc_steps(nframe, iter_max, iter_min=0):
|
||
"""
|
||
Obtain evenly-spaced sample of MCMC steps to make movie from
|
||
"""
|
||
all_mcmc = glob.glob(dirname + '/mcmc_*.h5')
|
||
x = [m[len(dirname + '/mcmc_'):-3] for m in all_mcmc]
|
||
all_mcmc = np.sort([int(m[len(dirname + '/mcmc_'):-3]) for m in all_mcmc])
|
||
if iter_max >= 0:
|
||
all_mcmc = all_mcmc[all_mcmc <= iter_max]
|
||
all_mcmc = all_mcmc[all_mcmc >= iter_min]
|
||
if nframe > 0:
|
||
max_out = max(all_mcmc)
|
||
min_out = min(all_mcmc)
|
||
step = max(int((max_out - min_out+1) / nframe), 1)
|
||
all_mcmc = all_mcmc[::step]
|
||
if max_out not in all_mcmc:
|
||
all_mcmc = np.concatenate([all_mcmc, [max_out]])
|
||
return all_mcmc
|
||
|
||
|
||
@borg.registerGravityBuilder
|
||
def build_gravity_model(gravity: str, state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_file=None) -> borg.forward.BaseForwardModel:
|
||
"""
|
||
Builds the gravity model and returns the forward model chain.
|
||
|
||
Args:
|
||
- gravity (str): Which gravity model to use (in ['pm', 'cola', 'lpt', '2lpt']
|
||
- 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.
|
||
|
||
"""
|
||
|
||
global chain, fwd_param, fwd_vel
|
||
myprint("Building gravity model")
|
||
|
||
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'])
|
||
|
||
if gravity in ['pm', 'cola']:
|
||
if 'forcesampling' in config['model']:
|
||
forcesampling = int(config['model']['forcesampling'])
|
||
else:
|
||
myprint('Could not find forcesampling, so will set to 4')
|
||
forcesampling = 4
|
||
|
||
# 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
|
||
|
||
|
||
def build_likelihood(gravity, box_in, ini_file):
|
||
|
||
global fwd_vel
|
||
|
||
model = build_gravity_model(gravity, None, box_in, ini_file=ini_file)
|
||
cosmo = utils.get_cosmopar(ini_file)
|
||
model.setCosmoParams(cosmo)
|
||
fwd_param = forwards.NullForward(box_in)
|
||
mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
|
||
|
||
# Load the mock data
|
||
mylike.loadMockData(None)
|
||
|
||
return mylike
|
||
|
||
ref_gravity = 'lpt'
|
||
ref_gravity = '2lpt'
|
||
|
||
|
||
# Input box
|
||
box_in = borg.forward.BoxModel()
|
||
config = configparser.ConfigParser()
|
||
config.read(ini_file)
|
||
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.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2']))
|
||
|
||
# Setup BORG forward model and likelihood
|
||
cola_like = build_likelihood('cola', box_in, ini_file)
|
||
lpt_like = build_likelihood(ref_gravity, box_in, ini_file)
|
||
|
||
# Load some MCMC samples and see how they compare
|
||
nframe = 100
|
||
iter_min = 1000
|
||
iter_max = -1
|
||
all_mcmc = get_mcmc_steps(nframe, iter_max, iter_min=iter_min)
|
||
|
||
all_lpt_like = np.empty(len(all_mcmc))
|
||
all_cola_like = np.empty(len(all_mcmc))
|
||
|
||
for i, mcmc in enumerate(all_mcmc):
|
||
myprint(f'\nStarting MCMC {i+1} of {len(all_mcmc)}: {mcmc}')
|
||
|
||
# Load an s_hat from file
|
||
with h5.File(f'{dirname}/mcmc_{mcmc}.h5', 'r') as f:
|
||
s_hat = f['scalars/s_hat_field'][:]
|
||
all_lpt_like[i] = lpt_like.logLikelihoodComplex(s_hat, None)
|
||
all_cola_like[i] = cola_like.logLikelihoodComplex(s_hat, None)
|
||
|
||
fig, axs = plt.subplots(1, 2, figsize=(15,4))
|
||
cmap = plt.cm.spring
|
||
markersize = 50
|
||
|
||
# # Remove first point
|
||
# all_lpt_like = all_lpt_like[1:]
|
||
# all_cola_like = all_cola_like[1:]
|
||
# all_mcmc = all_mcmc[1:]
|
||
|
||
# Normalise by the minimum value
|
||
lmin = all_lpt_like.min()
|
||
all_lpt_like -= lmin
|
||
all_cola_like -= lmin
|
||
|
||
sc0 = axs[0].scatter(all_lpt_like, all_cola_like, c=all_mcmc, cmap=cmap, marker='.', s=markersize)
|
||
xlim = axs[0].get_xlim()
|
||
axs[0].plot(xlim, xlim, color='k')
|
||
axs[0].set_xlim(xlim)
|
||
cbar = fig.colorbar(sc0, ax=axs[0], location='right', aspect=40)
|
||
cbar.set_label("MCMC Step")
|
||
|
||
sc1 = axs[1].scatter(all_lpt_like, all_cola_like - all_lpt_like, c=all_mcmc, cmap=cmap, marker='.', s=markersize)
|
||
axs[1].axhline(0, color='k')
|
||
cbar = fig.colorbar(sc1, ax=axs[1], location='right', aspect=40)
|
||
cbar.set_label("MCMC Step")
|
||
|
||
for ax in axs:
|
||
ax.set_xlabel(f'{ref_gravity} Log-Likelihood')
|
||
axs[0].set_ylabel('COLA Log-Likelihood')
|
||
axs[1].set_ylabel(f'COLA Log-Likelihood - {ref_gravity} Log-Likelihood')
|
||
|
||
# fig.align_xlabels()
|
||
# fig.align_ylabels()
|
||
# fig.tight_layout()
|
||
fig.savefig(f'{wdir}/likelihood_forwards.png', facecolor='white', bbox_inches='tight') |