From 4baac7619ef69794510f5199a14191a9a91264ea Mon Sep 17 00:00:00 2001 From: Deaglan Bartlett Date: Wed, 16 Apr 2025 13:20:13 +0200 Subject: [PATCH] Example 0 --- example0.py | 295 +++++++++++++++++++++++++++++++++++++++++++++++++++ ini_file.ini | 53 +++++++++ outdir | 1 + run.sh | 27 +++++ 4 files changed, 376 insertions(+) create mode 100644 example0.py create mode 100644 ini_file.ini create mode 120000 outdir create mode 100755 run.sh diff --git a/example0.py b/example0.py new file mode 100644 index 0000000..d01c1e4 --- /dev/null +++ b/example0.py @@ -0,0 +1,295 @@ +import aquila_borg as borg +import numpy as np +import numbers +import jaxlib +import jax.numpy as jnp +import jax + + +# Output stream management +cons = borg.console() +def myprint(x): + if isinstance(x, str): + cons.print_std(x) + else: + cons.print_std(repr(x)) + +class MyLikelihood(borg.likelihood.BaseLikelihood): + """ + HADES likelihood class + """ + + def __init__(self, fwd: borg.forward.BaseForwardModel): + + self.fwd = fwd + self.N = [32, 32, 32] # Number of grid points per side + self.L = [500, 500, 500] # Box size lenght Mpc/h + + self.sigma_dens = 0.1 # 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 = borg.cosmo.CosmologicalParameters() + cpar.default() + self.updateCosmology(cpar) + + # Gradient of the likelihood + self.grad_like = jax.grad(self.dens2like) + + def updateCosmology(self, cosmo: borg.cosmo.CosmologicalParameters) -> None: + self.fwd.setCosmoParams(cosmo) + + 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. + + """ + cpar = state['cosmology'] + 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') + 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) + + + 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 + + +@borg.registerGravityBuilder +def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel) -> 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") + + which_model = 'cola' + ai = 0.05 # Initial scale factor + af = 1.0 # Final scale factor + supersampling = 2 + forcesampling = 2 + nsteps = 20 # 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 + + # 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 + )) + + mod.accumulateAdjoint(True) + chain @= mod + + # Cosmological parameters + cpar = borg.cosmo.CosmologicalParameters() + cpar.default() + chain.setCosmoParams(cpar) + + return chain + +@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. + """ + + # Here you can add cosmology sampling, model parameter sampling, etc. + # For now, we'll just sample the initial conditions so don't need to do anything + + 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) + return likelihood diff --git a/ini_file.ini b/ini_file.ini new file mode 100644 index 0000000..7668197 --- /dev/null +++ b/ini_file.ini @@ -0,0 +1,53 @@ +[system] +console_output = borg_log +VERBOSE_LEVEL = 2 +N0 = 32 +N1 = 32 +N2 = 32 +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 +bias_sampler_blocked= true +ares_heat = 1.0 + +[mcmc] +number_to_generate = 20 +random_ic = false +init_random_scaling = 0.1 + +[hades] +algorithm = HMC +max_epsilon = 0.01 +max_timesteps = 50 +mixing = 1 + +[python] +likelihood_path = /home/bartlett/borg_examples/example0.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 \ No newline at end of file diff --git a/outdir b/outdir new file mode 120000 index 0000000..995db9e --- /dev/null +++ b/outdir @@ -0,0 +1 @@ +/data101/bartlett/borg_examples/ \ No newline at end of file diff --git a/run.sh b/run.sh new file mode 100755 index 0000000..53aa91d --- /dev/null +++ b/run.sh @@ -0,0 +1,27 @@ +#!/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/example0 + +mkdir -p $RUN_DIR +cd $RUN_DIR + +INI_FILE=/home/bartlett/borg_examples/ini_file.ini + +cp $INI_FILE ini_file.ini +$BORG INIT ini_file.ini \ No newline at end of file