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')