Sample model parameters (example 3)

Plotting.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 7,
    "id": "faed859b-c6c1-448f-be71-28d6458e279b",
    "metadata": {},
    "outputs": [],
@@ -10,8 +10,9 @@
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
     "import Pk_library as PKL\n",
-    "\n",
-    "import analysis"
+    "import os\n",
+    "import analysis\n",
+    "import h5py as h5"
    ]
   },
   {
@@ -191,9 +192,37 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "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",
+   "metadata": {},
    "outputs": [],
    "source": []
   }

example2.py

@@ -0,0 +1,430 @@
+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

example3.py

@@ -0,0 +1,516 @@
+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

ini_example0.ini

@@ -20,7 +20,7 @@ bias_sampler_blocked= true
 ares_heat = 1.0
 
 [mcmc]
-number_to_generate = 50
+number_to_generate = 5
 random_ic = false
 init_random_scaling = 0.1
 

ini_example1.ini

@@ -1,9 +1,9 @@
 [system]
 console_output = borg_log
-VERBOSE_LEVEL = 2
-N0 = 64
-N1 = 64
-N2 = 64
+VERBOSE_LEVEL = 1
+N0 = 32
+N1 = 32
+N2 = 32
 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 = 1
-random_ic = false
+number_to_generate = 10
+random_ic = true
 init_random_scaling = 0.1
 
 [hades]
@@ -54,7 +54,7 @@ z0 = 0
 
 [mock]
 sigma_dens = 1.
-sigma_vel = 100
+sigma_vel = 2000
 
 [gravity]
 which_model = lpt
@@ -65,5 +65,5 @@ forcesampling = 2
 nsteps = 20
 
 [velocity]
-which_model = cic
+which_model = linear
 rsmooth = 8.

ini_example2.ini

@@ -0,0 +1,71 @@
+[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

ini_example3.ini

@@ -0,0 +1,78 @@
+[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

run_example2.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/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

run_example3.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/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