Merge pull request 'main' (#2) from main into mosman_changes

Reviewed-on: #2

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",
+    "\n",
-    "import analysis\n",
+    "import analysis"
-    "import h5py as h5"
    ]
   },
   {
@@ -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": []

example2.py

@@ -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

example3.py

@@ -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

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
 

ini_example1.ini

@@ -1,9 +1,9 @@
 [system]
 console_output = borg_log
-VERBOSE_LEVEL = 1
+VERBOSE_LEVEL = 2
-N0 = 32
+N0 = 64
-N1 = 32
+N1 = 64
-N2 = 32
+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
+number_to_generate = 1
-random_ic = true
+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.

ini_example2.ini

@@ -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

ini_example3.ini

@@ -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

run_example2.sh

@@ -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

run_example3.sh

@@ -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