Example 0

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

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

outdir

@@ -0,0 +1 @@
+/data101/bartlett/borg_examples/

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