diff --git a/Plotting.ipynb b/Plotting.ipynb index 4f4a110..5d99b30 100644 --- a/Plotting.ipynb +++ b/Plotting.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 7, + "execution_count": 3, "id": "faed859b-c6c1-448f-be71-28d6458e279b", "metadata": {}, "outputs": [], @@ -10,9 +10,8 @@ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import Pk_library as PKL\n", - "import os\n", - "import analysis\n", - "import h5py as h5" + "\n", + "import analysis" ] }, { @@ -190,38 +189,10 @@ "plt.legend()" ] }, - { - "cell_type": "code", - "execution_count": 11, - "id": "df7fc61f-05d5-498b-a8de-0ca00afaddf5", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "mock_data.h5 0.521446084349996\n", - "mcmc_0.h5 0.05371104376346312\n", - "mcmc_20.h5 0.05765913633435571\n", - "mcmc_79.h5 0.06103314972961067\n" - ] - } - ], - "source": [ - "dirname = 'outdir/example1'\n", - "\n", - "all_fname = ['mock_data.h5', 'mcmc_0.h5', 'mcmc_20.h5', 'mcmc_79.h5']\n", - "\n", - "for fname in all_fname:\n", - " with h5.File(os.path.join(dirname, fname), 'r') as f:\n", - " sfield = f['scalars/BORG_final_density'][:]\n", - " print(fname, sfield.std())" - ] - }, { "cell_type": "code", "execution_count": null, - "id": "ed3bc75f-b608-4aab-bc44-9798816fada4", + "id": "df7fc61f-05d5-498b-a8de-0ca00afaddf5", "metadata": {}, "outputs": [], "source": [] diff --git a/example2.py b/example2.py deleted file mode 100644 index ddac7d9..0000000 --- a/example2.py +++ /dev/null @@ -1,430 +0,0 @@ -import aquila_borg as borg -import numpy as np -import numbers -import jaxlib -import jax.numpy as jnp -import jax -import configparser -import ast -import warnings -from functools import partial - -# Output stream management -cons = borg.console() -def myprint(x): - if isinstance(x, str): - cons.print_std(x) - else: - cons.print_std(repr(x)) - -def get_cosmopar(ini_file): - """ - Extract cosmological parameters from an ini file - - Args: - :ini_file (str): Path to the ini file - - Returns: - :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters - """ - - config = configparser.ConfigParser() - config.read(ini_file) - - cpar = borg.cosmo.CosmologicalParameters() - cpar.default() - cpar.fnl = float(config['cosmology']['fnl']) - cpar.omega_k = float(config['cosmology']['omega_k']) - cpar.omega_m = float(config['cosmology']['omega_m']) - cpar.omega_b = float(config['cosmology']['omega_b']) - cpar.omega_q = float(config['cosmology']['omega_q']) - cpar.h = float(config['cosmology']['h100']) - cpar.sigma8 = float(config['cosmology']['sigma8']) - cpar.n_s = float(config['cosmology']['n_s']) - cpar.w = float(config['cosmology']['w']) - cpar.wprime = float(config['cosmology']['wprime']) - - cpar = compute_As(cpar) - - return cpar - - -def compute_As(cpar): - """ - Compute As given values of sigma8 - - Args: - :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with wrong As - - Returns: - :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with updated As - - """ - - # requires BORG-CLASS - if not hasattr(borg.cosmo, 'ClassCosmo'): - raise ImportError( - "BORG-CLASS is required to compute As, but is not installed.") - - sigma8_true = jnp.copy(cpar.sigma8) - cpar.sigma8 = 0 - cpar.A_s = 2.3e-9 - k_max, k_per_decade = 10, 100 - extra_class = {} - extra_class['YHe'] = '0.24' - cosmo = borg.cosmo.ClassCosmo(cpar, k_per_decade, k_max, extra=extra_class) - cosmo.computeSigma8() - cos = cosmo.getCosmology() - cpar.A_s = (sigma8_true/cos['sigma_8'])**2*cpar.A_s - cpar.sigma8 = sigma8_true - - print('Updated cosmology:', cpar) - - return cpar - - -class MyLikelihood(borg.likelihood.BaseLikelihood): - """ - HADES likelihood class - """ - - def __init__(self, fwd: borg.forward.BaseForwardModel, ini_fname: str): - - self.fwd = fwd - - # Read the ini file - self.ini_fname = ini_fname - self.config = configparser.ConfigParser() - self.config.read(ini_fname) - self.N = [int(self.config['system'][f'N{i}']) for i in range(3)] # Number of grid points per side - self.L = [float(self.config['system'][f'L{i}']) for i in range(3)] # Box size lenght Mpc/h - - self.sigma_dens = float(self.config['mock']['sigma_dens']) # Density scatter - - myprint(f"Likelihood initialized with {self.N} grid points and box size {self.L} Mpc/h") - super().__init__(fwd, self.N, self.L) - - # Set up cosmoligical parameters - cpar = get_cosmopar(ini_fname) - self.updateCosmology(cpar) - - # Gradient of the likelihood - self.grad_like = jax.grad(self.dens2like) - - def updateCosmology(self, cosmo: borg.cosmo.CosmologicalParameters) -> None: - cpar = compute_As(cosmo) - self.fwd.setCosmoParams(cpar) - - def updateMetaParameters(self, state: borg.likelihood.MarkovState) -> None: - """ - Update the meta parameters of the sampler (not sampled) from the MarkovState. - - Args: - - state (borg.likelihood.MarkovState): The state object to be used in the likelihood. - - """ - cosmo = state['cosmology'] - cpar = compute_As(cosmo) - self.fwd.setCosmoParams(cpar) - - def initializeLikelihood(self, state: borg.likelihood.MarkovState) -> None: - """ - Initialize the likelihood function. - - Args: - - state (borg.likelihood.MarkovState): The state object to be used in the likelihood. - - """ - myprint("Init likelihood") - state.newArray3d("BORG_final_density", *self.fwd.getOutputBoxModel().N, True) - - # Could load real data - # We'll generate mock data which has its own function - - def generateMockData(self, s_hat:np.ndarray, state: borg.likelihood.MarkovState) -> None: - """ - Generates mock data by simulating the forward model with the given white noise - - Args: - - s_hat (np.ndarray): The input (initial) white noise field. - - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood. - """ - myprint('Making mock from BORG') - - # Get density field from the initial conditions - # Could replace with any (better) simulation here - # This version is self-consistnet - dens = np.zeros(self.fwd.getOutputBoxModel().N) - myprint('Running forward model') - myprint(self.fwd.getCosmoParams()) - self.fwd.forwardModel_v2(s_hat) - self.fwd.getDensityFinal(dens) - state["BORG_final_density"][:] = dens - self.true_dens = dens.copy() - - # Add some scatter - myprint('Adding scatter') - self.obs_dens = self.true_dens + np.random.randn(*self.true_dens.shape) * self.sigma_dens - - # Compute the likelihood and print it - myprint('From mock') - self.saved_s_hat = s_hat.copy() - self.logLikelihoodComplex(s_hat, False) - self.commitAuxiliaryFields(state) - myprint('Done') - - - def dens2like(self, output_density: np.ndarray): - """ - Compute the likelihood from the density field - Args: - - output_density (np.ndarray): The density field to be used in the likelihood. - Returns: - - float: The likelihood value. - """ - # Compute the likelihood from the density field - # This is a simple Gaussian likelihood - # Could be replaced with any other likelihood - diff = output_density - self.obs_dens - like = 0.5 * jnp.sum(diff**2) / (self.sigma_dens**2) - - return like - - - def logLikelihoodComplex(self, s_hat:np.ndarray, gradientIsNext:bool): - - # myprint('Getting density field now') - # Get the density field from the forward model - dens = np.zeros(self.fwd.getOutputBoxModel().N) - self.fwd.forwardModel_v2(s_hat) - self.fwd.getDensityFinal(dens) - - L = self.dens2like(dens) - - if isinstance(L, numbers.Number) or isinstance(L, jaxlib.xla_extension.ArrayImpl): - myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}") - - self.delta = dens.copy() - - return L - - - def gradientLikelihoodComplex(self, s_hat:np.ndarray): - - # Run BORG density field - output_density = np.zeros(self.N) - self.fwd.forwardModel_v2(s_hat) - self.fwd.getDensityFinal(output_density) - - # Compute the gradient of the likelihood - # d logL / d dens - mygradient = self.grad_like(output_density) - - # Now get d logL / d s_hat - mygradient = np.array(mygradient, dtype=np.float64) - self.fwd.adjointModel_v2(mygradient) - mygrad_hat = np.zeros(s_hat.shape, dtype=np.complex128) - self.fwd.getAdjointModel(mygrad_hat) - self.fwd.clearAdjointGradient() - - return mygrad_hat - - def commitAuxiliaryFields(self, state: borg.likelihood.MarkovState) -> None: - """ - Commits the final density field to the Markov state. - Args: - - state (borg.state.State): The state object containing the final density field. - """ - self.updateCosmology(self.fwd.getCosmoParams()) - self.dens2like(self.delta) - state["BORG_final_density"][:] = self.delta.copy() - - -@borg.registerGravityBuilder -def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_fname=None) -> 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. - - """ - - global chain - myprint("Building gravity model") - - if ini_fname is None: - ini_fname=borg.getIniConfigurationFilename() - config = configparser.ConfigParser() - config.read(ini_fname) - # READ FROM INI FILE - which_model = config['gravity']['which_model'] - ai = float(config['gravity']['ai']) # Initial scale factor - af = float(config['gravity']['af']) # Final scale factor - supersampling = int(config['gravity']['supersampling']) - forcesampling = int(config['gravity']['forcesampling']) - nsteps = int(config['gravity']['nsteps']) # Number of steps in the PM solver - - chain = borg.forward.ChainForwardModel(box) - - # Make sure that the initial conditions are real in position space - 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'}}) # helps deals with errors with primordial physics in CLASS for weird cosmologies - chain @= transfer_class - - # This one works - # chain @= borg.forward.models.Primordial(box, ai) - # chain @= borg.forward.models.EisensteinHu(box) - - # Gravity model - if which_model == '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 which_model == '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 which_model == 'pm': - mod = borg.forward.model_lib.M_PM_CIC( - box, - opts=dict(a_initial=af, - a_final=af, - pm_start_z=1/ai - 1, - do_rsd=False, - supersampling=supersampling, - forcesampling=forcesampling, - lightcone=False, - part_factor=1.01, - pm_nsteps=nsteps, # Number of steps in the PM solver - tcola=False - )) - elif which_model == 'cola': - mod = borg.forward.model_lib.M_PM_CIC( - box, - opts=dict(a_initial=af, - a_final=af, - pm_start_z=1/ai - 1, - do_rsd=False, - supersampling=supersampling, - forcesampling=forcesampling, - lightcone=False, - part_factor=1.01, - pm_nsteps=nsteps, # Number of steps in the PM solver - tcola=True - )) - else: - raise ValueError(f"Unknown model {which_model}") - - mod.accumulateAdjoint(True) - chain @= mod - - # Cosmological parameters - cpar = get_cosmopar(borg.getIniConfigurationFilename()) - print('Setting cosmo params', cpar) - chain.setCosmoParams(cpar) - - return chain - - -def check_cosmo_sampling(loop): - return loop.getStepID() > begin_cosmo - - -@borg.registerSamplerBuilder -def build_sampler(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo, loop: borg.samplers.MainLoop): - """ - Builds the sampler and returns the main loop. - - Args: - - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood. - - info (borg.likelihood.LikelihoodInfo): The likelihood info object to be used in the likelihood. - - loop (borg.samplers.MainLoop): The main loop object to be used in the likelihood. - - Returns: - borg.samplers.MainLoop: The main loop. - """ - - global begin_cosmo - - myprint("Building sampler") - - # Read from config file what to sample - config = configparser.ConfigParser() - config.read(borg.getIniConfigurationFilename()) - end = '_sampler_blocked' - to_sample = [k[:-len(end)] for (k, v) in config['block_loop'].items() if k.endswith(end) and v.lower() == 'false'] - myprint(f"Sampling {to_sample}") - - all_sampler = [] - - # Add cosmology samplers - params = [] - initial_values = {} - prior = {} - for p in ['omega_m', 'sigma8']: - if p not in to_sample: - continue - if p in config['prior'].keys() and p in config['cosmology'].keys(): - myprint(f'Adding {p} sampler') - params.append(f"cosmology.{p}") - initial_values[f"cosmology.{p}"] = float(config['cosmology'][p]) - prior[f"cosmology.{p}"] = np.array(ast.literal_eval(config['prior'][p])) - else: - s = f'Could not find {p} prior and/or default, so will not sample' - warnings.warn(s, stacklevel=2) - # Remove for later to prevent duplication - to_sample.remove(p) - - begin_cosmo = int(config['mcmc']['warmup_cosmo']) - - if len(params) > 0: - myprint(f"Sampling cosmology parameters {params}") - cosmo_sampler = borg.samplers.ModelParamsSampler("", params, likelihood, chain, initial_values, prior) - cosmo_sampler.setName("cosmo_sampler") - all_sampler.append(cosmo_sampler) - loop.push(cosmo_sampler) - loop.addToConditionGroup("warmup_cosmo", "cosmo_sampler") - loop.addConditionToConditionGroup("warmup_cosmo", partial(check_cosmo_sampling, loop)) - - # Here you can add model parameter sampling, etc. - - return [] - - -@borg.registerLikelihoodBuilder -def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo) -> borg.likelihood.BaseLikelihood: - """ - Builds the likelihood and returns the likelihood object. - - Args: - - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood. - - info (borg.likelihood.LikelihoodInfo): The likelihood info object to be used in the likelihood. - - Returns: - borg.likelihood.BaseLikelihood: The likelihood object. - """ - - myprint("Building likelihood") - global likelihood - likelihood = MyLikelihood(chain, borg.getIniConfigurationFilename()) - return likelihood diff --git a/example3.py b/example3.py deleted file mode 100644 index 4b2ac93..0000000 --- a/example3.py +++ /dev/null @@ -1,516 +0,0 @@ -import aquila_borg as borg -import numpy as np -import numbers -import jaxlib -import jax.numpy as jnp -import jax -import configparser -import ast -import warnings -from functools import partial - -# Output stream management -cons = borg.console() -def myprint(x): - if isinstance(x, str): - cons.print_std(x) - else: - cons.print_std(repr(x)) - -def get_cosmopar(ini_file): - """ - Extract cosmological parameters from an ini file - - Args: - :ini_file (str): Path to the ini file - - Returns: - :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters - """ - - config = configparser.ConfigParser() - config.read(ini_file) - - cpar = borg.cosmo.CosmologicalParameters() - cpar.default() - cpar.fnl = float(config['cosmology']['fnl']) - cpar.omega_k = float(config['cosmology']['omega_k']) - cpar.omega_m = float(config['cosmology']['omega_m']) - cpar.omega_b = float(config['cosmology']['omega_b']) - cpar.omega_q = float(config['cosmology']['omega_q']) - cpar.h = float(config['cosmology']['h100']) - cpar.sigma8 = float(config['cosmology']['sigma8']) - cpar.n_s = float(config['cosmology']['n_s']) - cpar.w = float(config['cosmology']['w']) - cpar.wprime = float(config['cosmology']['wprime']) - - cpar = compute_As(cpar) - - return cpar - - -def compute_As(cpar): - """ - Compute As given values of sigma8 - - Args: - :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with wrong As - - Returns: - :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with updated As - - """ - - # requires BORG-CLASS - if not hasattr(borg.cosmo, 'ClassCosmo'): - raise ImportError( - "BORG-CLASS is required to compute As, but is not installed.") - - sigma8_true = jnp.copy(cpar.sigma8) - cpar.sigma8 = 0 - cpar.A_s = 2.3e-9 - k_max, k_per_decade = 10, 100 - extra_class = {} - extra_class['YHe'] = '0.24' - cosmo = borg.cosmo.ClassCosmo(cpar, k_per_decade, k_max, extra=extra_class) - cosmo.computeSigma8() - cos = cosmo.getCosmology() - cpar.A_s = (sigma8_true/cos['sigma_8'])**2*cpar.A_s - cpar.sigma8 = sigma8_true - - print('Updated cosmology:', cpar) - - return cpar - -class NullForward(borg.forward.BaseForwardModel): - """ - BORG forward model which does nothing but stores - the values of parameters to be used by the likelihood - """ - - def __init__(self, box: borg.forward.BoxModel) -> None: - """ - Initialise the NullForward class - Args: - box (borg.forward.BoxModel): The input box model. - """ - super().__init__(box, box) - self.setName("nullforward") - - self.params = {} - cpar = get_cosmopar(borg.getIniConfigurationFilename()) - self.setCosmoParams(cpar) - - def setModelParams(self, params: dict) -> None: - """ - Change the values of the model parameters to those given by params - - Args: - params (dict): Dictionary of updated model parameters. - """ - for k, v in params.items(): - self.params[k] = v - myprint(f'Updated model parameters: {self.params}') - - def getModelParam(self, model, keyname: str): - """ - This queries the current state of the parameters keyname in model model. - - Args: - model: The model - keyname (str): The name of the parameter of interest - """ - return self.params[keyname] - - -class MyLikelihood(borg.likelihood.BaseLikelihood): - """ - HADES likelihood class - """ - - def __init__(self, fwd: borg.forward.BaseForwardModel, fwd_param: NullForward, ini_fname: str): - - self.fwd = fwd - self.fwd_param = fwd_param - - # Read the ini file - self.ini_fname = ini_fname - self.config = configparser.ConfigParser() - self.config.read(ini_fname) - self.N = [int(self.config['system'][f'N{i}']) for i in range(3)] # Number of grid points per side - self.L = [float(self.config['system'][f'L{i}']) for i in range(3)] # Box size lenght Mpc/h - - self.sigma_dens = float(self.config['mock']['sigma_dens']) # Density scatter - - myprint(f"Likelihood initialized with {self.N} grid points and box size {self.L} Mpc/h") - super().__init__(fwd, self.N, self.L) - - # Set up cosmoligical parameters - cpar = get_cosmopar(ini_fname) - self.updateCosmology(cpar) - - # Set up the forward model parameters - print('Setting up forward model parameters') - self.mu_true = float(self.config['model']['mu']) - model_params = {'mu':self.mu_true} - self.sigma_mu = float(self.config['model']['sigma_mu']) - self.fwd_param.setModelParams(model_params) - - # Gradient of the likelihood - self.grad_like = jax.grad(self.dens2like) - - def updateCosmology(self, cosmo: borg.cosmo.CosmologicalParameters) -> None: - cpar = compute_As(cosmo) - self.fwd.setCosmoParams(cpar) - self.fwd_param.setCosmoParams(cpar) - - def updateMetaParameters(self, state: borg.likelihood.MarkovState) -> None: - """ - Update the meta parameters of the sampler (not sampled) from the MarkovState. - - Args: - - state (borg.likelihood.MarkovState): The state object to be used in the likelihood. - - """ - cosmo = state['cosmology'] - cpar = compute_As(cosmo) - self.fwd.setCosmoParams(cpar) - self.fwd_param.setCosmoParams(cpar) - - def initializeLikelihood(self, state: borg.likelihood.MarkovState) -> None: - """ - Initialize the likelihood function. - - Args: - - state (borg.likelihood.MarkovState): The state object to be used in the likelihood. - - """ - myprint("Init likelihood") - state.newArray3d("BORG_final_density", *self.fwd.getOutputBoxModel().N, True) - - # Could load real data - # We'll generate mock data which has its own function - - def generateMockData(self, s_hat:np.ndarray, state: borg.likelihood.MarkovState) -> None: - """ - Generates mock data by simulating the forward model with the given white noise - - Args: - - s_hat (np.ndarray): The input (initial) white noise field. - - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood. - """ - myprint('Making mock from BORG') - - # Get density field from the initial conditions - # Could replace with any (better) simulation here - # This version is self-consistnet - dens = np.zeros(self.fwd.getOutputBoxModel().N) - myprint('Running forward model') - myprint(self.fwd.getCosmoParams()) - self.fwd.forwardModel_v2(s_hat) - self.fwd.getDensityFinal(dens) - state["BORG_final_density"][:] = dens - self.true_dens = dens.copy() - - # Add some scatter - myprint('Adding scatter') - self.obs_dens = self.true_dens + np.random.randn(*self.true_dens.shape) * self.sigma_dens - - # Get observed mu - self.obs_mu = self.mu_true + np.random.randn() * self.sigma_mu - - # Compute the likelihood and print it - myprint('From mock') - self.saved_s_hat = s_hat.copy() - self.logLikelihoodComplex(s_hat, False) - self.commitAuxiliaryFields(state) - myprint('Done') - - - def dens2like(self, output_density: np.ndarray): - """ - Compute the likelihood from the density field - Args: - - output_density (np.ndarray): The density field to be used in the likelihood. - Returns: - - float: The likelihood value. - """ - # Compute the likelihood from the density field - # This is a simple Gaussian likelihood - # Could be replaced with any other likelihood - diff = output_density - self.obs_dens - like = 0.5 * jnp.sum(diff**2) / (self.sigma_dens**2) - - mu = self.fwd_param.getModelParam('nullforward', 'mu') - like += 0.5 * (mu - self.obs_mu)**2 / (self.sigma_mu**2) - - return like - - - def logLikelihoodComplex(self, s_hat:np.ndarray, gradientIsNext:bool): - - # myprint('Getting density field now') - # Get the density field from the forward model - dens = np.zeros(self.fwd.getOutputBoxModel().N) - self.fwd.forwardModel_v2(s_hat) - self.fwd.getDensityFinal(dens) - - L = self.dens2like(dens) - - if isinstance(L, numbers.Number) or isinstance(L, jaxlib.xla_extension.ArrayImpl): - myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}") - - self.delta = dens.copy() - - return L - - - def gradientLikelihoodComplex(self, s_hat:np.ndarray): - - # Run BORG density field - output_density = np.zeros(self.N) - self.fwd.forwardModel_v2(s_hat) - self.fwd.getDensityFinal(output_density) - - # Compute the gradient of the likelihood - # d logL / d dens - mygradient = self.grad_like(output_density) - - # Now get d logL / d s_hat - mygradient = np.array(mygradient, dtype=np.float64) - self.fwd.adjointModel_v2(mygradient) - mygrad_hat = np.zeros(s_hat.shape, dtype=np.complex128) - self.fwd.getAdjointModel(mygrad_hat) - self.fwd.clearAdjointGradient() - - return mygrad_hat - - def commitAuxiliaryFields(self, state: borg.likelihood.MarkovState) -> None: - """ - Commits the final density field to the Markov state. - Args: - - state (borg.state.State): The state object containing the final density field. - """ - self.updateCosmology(self.fwd.getCosmoParams()) - self.dens2like(self.delta) - state["BORG_final_density"][:] = self.delta.copy() - - -@borg.registerGravityBuilder -def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_fname=None) -> 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. - - """ - - global chain, fwd_param - myprint("Building gravity model") - - if ini_fname is None: - ini_fname=borg.getIniConfigurationFilename() - config = configparser.ConfigParser() - config.read(ini_fname) - # READ FROM INI FILE - which_model = config['gravity']['which_model'] - ai = float(config['gravity']['ai']) # Initial scale factor - af = float(config['gravity']['af']) # Final scale factor - supersampling = int(config['gravity']['supersampling']) - forcesampling = int(config['gravity']['forcesampling']) - nsteps = int(config['gravity']['nsteps']) # Number of steps in the PM solver - - chain = borg.forward.ChainForwardModel(box) - - # Make sure that the initial conditions are real in position space - 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'}}) # helps deals with errors with primordial physics in CLASS for weird cosmologies - chain @= transfer_class - - # This one works - # chain @= borg.forward.models.Primordial(box, ai) - # chain @= borg.forward.models.EisensteinHu(box) - - # Gravity model - if which_model == '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 which_model == '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 which_model == 'pm': - mod = borg.forward.model_lib.M_PM_CIC( - box, - opts=dict(a_initial=af, - a_final=af, - pm_start_z=1/ai - 1, - do_rsd=False, - supersampling=supersampling, - forcesampling=forcesampling, - lightcone=False, - part_factor=1.01, - pm_nsteps=nsteps, # Number of steps in the PM solver - tcola=False - )) - elif which_model == 'cola': - mod = borg.forward.model_lib.M_PM_CIC( - box, - opts=dict(a_initial=af, - a_final=af, - pm_start_z=1/ai - 1, - do_rsd=False, - supersampling=supersampling, - forcesampling=forcesampling, - lightcone=False, - part_factor=1.01, - pm_nsteps=nsteps, # Number of steps in the PM solver - tcola=True - )) - else: - raise ValueError(f"Unknown model {which_model}") - - mod.accumulateAdjoint(True) - chain @= mod - - # Cosmological parameters - cpar = get_cosmopar(borg.getIniConfigurationFilename()) - print('Setting cosmo params', cpar) - chain.setCosmoParams(cpar) - - # Model parameters - fwd_param = borg.forward.ChainForwardModel(box) - fwd_param @= NullForward(box) - fwd_param.setCosmoParams(cpar) - - return chain - - -def check_cosmo_sampling(loop): - return loop.getStepID() > begin_cosmo - -def check_model_sampling(loop): - return loop.getStepID() > begin_model - - -@borg.registerSamplerBuilder -def build_sampler(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo, loop: borg.samplers.MainLoop): - """ - Builds the sampler and returns the main loop. - - Args: - - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood. - - info (borg.likelihood.LikelihoodInfo): The likelihood info object to be used in the likelihood. - - loop (borg.samplers.MainLoop): The main loop object to be used in the likelihood. - - Returns: - borg.samplers.MainLoop: The main loop. - """ - - global begin_cosmo, begin_model - - myprint("Building sampler") - - # Read from config file what to sample - config = configparser.ConfigParser() - config.read(borg.getIniConfigurationFilename()) - end = '_sampler_blocked' - to_sample = [k[:-len(end)] for (k, v) in config['block_loop'].items() if k.endswith(end) and v.lower() == 'false'] - myprint(f"Sampling {to_sample}") - - all_sampler = [] - - # Add cosmology samplers - params = [] - initial_values = {} - prior = {} - for p in ['omega_m', 'sigma8']: - if p not in to_sample: - continue - if p in config['prior'].keys() and p in config['cosmology'].keys(): - myprint(f'Adding {p} sampler') - params.append(f"cosmology.{p}") - initial_values[f"cosmology.{p}"] = float(config['cosmology'][p]) - prior[f"cosmology.{p}"] = np.array(ast.literal_eval(config['prior'][p])) - else: - s = f'Could not find {p} prior and/or default, so will not sample' - warnings.warn(s, stacklevel=2) - # Remove for later to prevent duplication - to_sample.remove(p) - - begin_cosmo = int(config['mcmc']['warmup_cosmo']) - begin_model = int(config['mcmc']['warmup_model']) - - if len(params) > 0: - myprint(f"Sampling cosmology parameters {params}") - cosmo_sampler = borg.samplers.ModelParamsSampler("", params, likelihood, chain, initial_values, prior) - cosmo_sampler.setName("cosmo_sampler") - all_sampler.append(cosmo_sampler) - loop.push(cosmo_sampler) - loop.addToConditionGroup("warmup_cosmo", "cosmo_sampler") - loop.addConditionToConditionGroup("warmup_cosmo", partial(check_cosmo_sampling, loop)) - - # Add model samplers - params = [] - initial_values = {} - prior = {} - for p in to_sample: - if p in config['prior'].keys(): - myprint(f'Adding {p} sampler') - params.append(p) - initial_values[p] = float(config['model'][p]) - prior[p] = np.array(ast.literal_eval(config['prior'][p])) - else: - s = f'Could not find {p} prior and/or default, so will not sample' - warnings.warn(s, stacklevel=2) - - if len(params) > 0: - myprint('Adding model parameter sampler') - model_sampler = borg.samplers.ModelParamsSampler("", params, likelihood, fwd_param, initial_values, prior) - model_sampler.setName("model_sampler") - all_sampler.append(model_sampler) - loop.push(model_sampler) - loop.addToConditionGroup("warmup_model", "model_sampler") - loop.addConditionToConditionGroup("warmup_model", partial(check_model_sampling, loop)) - - return [] - - -@borg.registerLikelihoodBuilder -def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo) -> borg.likelihood.BaseLikelihood: - """ - Builds the likelihood and returns the likelihood object. - - Args: - - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood. - - info (borg.likelihood.LikelihoodInfo): The likelihood info object to be used in the likelihood. - - Returns: - borg.likelihood.BaseLikelihood: The likelihood object. - """ - - myprint("Building likelihood") - global likelihood - likelihood = MyLikelihood(chain, fwd_param, borg.getIniConfigurationFilename()) - return likelihood diff --git a/ini_example0.ini b/ini_example0.ini index 394f7b1..1adb910 100644 --- a/ini_example0.ini +++ b/ini_example0.ini @@ -20,7 +20,7 @@ bias_sampler_blocked= true ares_heat = 1.0 [mcmc] -number_to_generate = 5 +number_to_generate = 50 random_ic = false init_random_scaling = 0.1 diff --git a/ini_example1.ini b/ini_example1.ini index ffb713c..05ae4f3 100644 --- a/ini_example1.ini +++ b/ini_example1.ini @@ -1,9 +1,9 @@ [system] console_output = borg_log -VERBOSE_LEVEL = 1 -N0 = 32 -N1 = 32 -N2 = 32 +VERBOSE_LEVEL = 2 +N0 = 64 +N1 = 64 +N2 = 64 L0 = 500.0 L1 = 500.0 L2 = 500.0 @@ -20,8 +20,8 @@ bias_sampler_blocked= true ares_heat = 1.0 [mcmc] -number_to_generate = 10 -random_ic = true +number_to_generate = 1 +random_ic = false init_random_scaling = 0.1 [hades] @@ -54,7 +54,7 @@ z0 = 0 [mock] sigma_dens = 1. -sigma_vel = 2000 +sigma_vel = 100 [gravity] which_model = lpt @@ -65,5 +65,5 @@ forcesampling = 2 nsteps = 20 [velocity] -which_model = linear +which_model = cic rsmooth = 8. diff --git a/ini_example2.ini b/ini_example2.ini deleted file mode 100644 index 259e42b..0000000 --- a/ini_example2.ini +++ /dev/null @@ -1,71 +0,0 @@ -[system] -console_output = borg_log -VERBOSE_LEVEL = 2 -N0 = 64 -N1 = 64 -N2 = 64 -L0 = 500.0 -L1 = 500.0 -L2 = 500.0 -corner0 = -250.0 -corner1 = -250.0 -corner2 = -250.0 -NUM_MODES = 100 -test_mode = true -seed_cpower = true - -[block_loop] -hades_sampler_blocked = false -sigma8_sampler_blocked = false -omega_m_sampler_blocked = false -bias_sampler_blocked= true -ares_heat = 1.0 - -[prior] -sigma8 = [0.5, 1.2] -omega_m = [0.1, 0.5] - -[mcmc] -number_to_generate = 5 -random_ic = false -init_random_scaling = 0.1 -warmup_cosmo = 2 - -[hades] -algorithm = HMC -max_epsilon = 0.01 -max_timesteps = 50 -mixing = 1 - -[python] -likelihood_path = /home/bartlett/borg_examples/example2.py - -[run] -run_type = mock -NCAT = 0 - -[cosmology] -omega_r = 0 -fnl = 0 -omega_k = 0 -omega_m = 0.315 -omega_b = 0.049 -omega_q = 0.685 -h100 = 0.68 -sigma8 = 0.81 -n_s = 0.97 -w = -1 -wprime = 0 -beta = 1.5 -z0 = 0 - -[mock] -sigma_dens = 1. - -[gravity] -which_model = lpt -ai = 0.05 -af = 1.0 -supersampling = 2 -forcesampling = 2 -nsteps = 20 diff --git a/ini_example3.ini b/ini_example3.ini deleted file mode 100644 index ffc16f1..0000000 --- a/ini_example3.ini +++ /dev/null @@ -1,78 +0,0 @@ -[system] -console_output = borg_log -VERBOSE_LEVEL = 2 -N0 = 64 -N1 = 64 -N2 = 64 -L0 = 500.0 -L1 = 500.0 -L2 = 500.0 -corner0 = -250.0 -corner1 = -250.0 -corner2 = -250.0 -NUM_MODES = 100 -test_mode = true -seed_cpower = true - -[block_loop] -hades_sampler_blocked = false -sigma8_sampler_blocked = true -omega_m_sampler_blocked = true -mu_sampler_blocked = false -bias_sampler_blocked= true -ares_heat = 1.0 - -[prior] -sigma8 = [0.5, 1.2] -omega_m = [0.1, 0.5] -mu = [0.5, 1.5] - -[mcmc] -number_to_generate = 5 -random_ic = false -init_random_scaling = 0.1 -warmup_cosmo = 2 -warmup_model = 2 - -[hades] -algorithm = HMC -max_epsilon = 0.01 -max_timesteps = 50 -mixing = 1 - -[python] -likelihood_path = /home/bartlett/borg_examples/example3.py - -[run] -run_type = mock -NCAT = 0 - -[cosmology] -omega_r = 0 -fnl = 0 -omega_k = 0 -omega_m = 0.315 -omega_b = 0.049 -omega_q = 0.685 -h100 = 0.68 -sigma8 = 0.81 -n_s = 0.97 -w = -1 -wprime = 0 -beta = 1.5 -z0 = 0 - -[mock] -sigma_dens = 1. - -[gravity] -which_model = lpt -ai = 0.05 -af = 1.0 -supersampling = 2 -forcesampling = 2 -nsteps = 20 - -[model] -mu = 1.0 -sigma_mu = 0.1 diff --git a/run_example2.sh b/run_example2.sh deleted file mode 100755 index f852f4b..0000000 --- a/run_example2.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/sh - -# Modules -module purge -module restore myborg -module load cuda/12.6 - -# Environment -source /home/bartlett/.bashrc -source /home/bartlett/anaconda3/etc/profile.d/conda.sh -conda deactivate -conda activate borg_new - -# Kill job if there are any errors -set -e - -# Path variables -BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python -RUN_DIR=/data101/bartlett/borg_examples/example2 - -mkdir -p $RUN_DIR -cd $RUN_DIR - -INI_FILE=/home/bartlett/borg_examples/ini_example2.ini - -cp $INI_FILE ini_file.ini -$BORG INIT ini_file.ini \ No newline at end of file diff --git a/run_example3.sh b/run_example3.sh deleted file mode 100755 index 0785bb1..0000000 --- a/run_example3.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/sh - -# Modules -module purge -module restore myborg -module load cuda/12.6 - -# Environment -source /home/bartlett/.bashrc -source /home/bartlett/anaconda3/etc/profile.d/conda.sh -conda deactivate -conda activate borg_new - -# Kill job if there are any errors -set -e - -# Path variables -BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python -RUN_DIR=/data101/bartlett/borg_examples/example3 - -mkdir -p $RUN_DIR -cd $RUN_DIR - -INI_FILE=/home/bartlett/borg_examples/ini_example3.ini - -cp $INI_FILE ini_file.ini -$BORG INIT ini_file.ini \ No newline at end of file