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

Reviewed-on: #2

analysis.py

@@ -100,6 +100,8 @@ def get_spectra(ini_name, dirname, mcmc_steps, which_field='BORG_final_density',
             MAS = "CIC"
         elif which_field == 'ics':
             MAS = None
+        elif which_field.startwith('BORG_final_velocity'):
+            MAS = "CIC"
         else:
             raise NotImplementedError
     

example1.py

@@ -0,0 +1,418 @@
+import aquila_borg as borg
+import numpy as np
+import numbers
+import jaxlib
+import jax.numpy as jnp
+import jax
+import configparser
+
+# 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, 
+        fwd_vel: borg.forward.BaseForwardModel, 
+        ini_fname: str):
+
+        self.fwd = fwd
+        self.fwd_vel = fwd_vel
+
+        # 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
+        self.sigma_vel = float(self.config['mock']['sigma_vel'])  # Velocity 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, argnums=(0, 1))
+
+    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)
+        state.newArray3d("BORG_final_velocity_x", *self.fwd.getOutputBoxModel().N, True)
+        state.newArray3d("BORG_final_velocity_y", *self.fwd.getOutputBoxModel().N, True)
+        state.newArray3d("BORG_final_velocity_z", *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()
+
+        # Get velocity field
+        vel = self.fwd_vel.getVelocityField()
+        self.true_vel = vel.copy()
+
+        # Add some scatter
+        myprint('Adding scatter')
+        self.obs_dens = self.true_dens + np.random.randn(*self.true_dens.shape) * self.sigma_dens
+        self.obs_vel = self.true_vel + np.random.randn(*self.true_vel.shape) * self.sigma_vel
+
+        # 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, output_velocity: np.ndarray) -> float:
+        """
+        Compute the likelihood from the density field
+        Args:
+            - output_density (np.ndarray): The density field to be used in the likelihood.
+            - output_velocity (np.ndarray): The velocity 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
+        diff_vel = output_velocity - self.obs_vel
+        like = 0.5 * jnp.sum(diff**2) / (self.sigma_dens**2)
+        like += 0.5 * jnp.sum(diff_vel**2) / (self.sigma_vel**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)
+
+        # Get the velocity field from the forward model
+        vel = self.fwd_vel.getVelocityField()
+
+        L = self.dens2like(dens, vel)
+
+        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()
+        self.vel = vel.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)
+
+        # Run BORG velocity field
+        vel = self.fwd_vel.getVelocityField()
+ 
+        # Compute the gradient of the likelihood
+        # d logL / d dens, d logL / d vel
+        dens_gradient, vel_gradient = self.grad_like(output_density, vel)
+
+        # Now get d logL / d s_hat
+        dens_gradient = np.array(dens_gradient, dtype=np.float64)
+        vel_gradient = np.array(vel_gradient, dtype=np.float64)
+        
+        self.fwd_vel.computeAdjointModel(vel_gradient)
+        
+        self.fwd.adjointModel_v2(dens_gradient)
+        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, self.vel)
+        state["BORG_final_density"][:] = self.delta.copy()
+        state["BORG_final_velocity_x"][:] = self.vel[0].copy()
+        state["BORG_final_velocity_y"][:] = self.vel[1].copy()
+        state["BORG_final_velocity_z"][:] = self.vel[2].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_vel
+    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
+
+    # 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)
+
+    # Set the forward model for velocities
+    vel_model = config['velocity']['which_model']
+    if vel_model == 'linear':
+        fwd_vel = borg.forward.velocity.LinearModel(box, mod, af)
+    elif vel_model == 'cic':
+        rsmooth = float(config['velocity']['rsmooth'])
+        fwd_vel = borg.forward.velocity.CICModel(box, mod, rsmooth)
+    elif vel_model == 'sic':
+        fwd_vel = borg.forward.velocity.SICModel(box, mod)
+    else:
+        raise ValueError(f"Unknown model {vel_model}")
+
+    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, fwd_vel, borg.getIniConfigurationFilename())
+    return likelihood

ini_example1.ini

@@ -0,0 +1,69 @@
+[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
+bias_sampler_blocked= true
+ares_heat = 1.0
+
+[mcmc]
+number_to_generate = 1
+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/example1.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.
+sigma_vel = 100
+
+[gravity]
+which_model = lpt
+ai = 0.05
+af = 1.0
+supersampling = 2
+forcesampling = 2
+nsteps = 20
+
+[velocity]
+which_model = cic
+rsmooth = 8.

run_example0.sh

@@ -21,7 +21,7 @@ RUN_DIR=/data101/bartlett/borg_examples/example0
 mkdir -p $RUN_DIR
 cd $RUN_DIR
 
-INI_FILE=/home/bartlett/borg_examples/ini_file.ini
+INI_FILE=/home/bartlett/borg_examples/ini_example0.ini
 
 cp $INI_FILE ini_file.ini
 $BORG INIT ini_file.ini

run_example1.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/example1
+
+mkdir -p $RUN_DIR
+cd $RUN_DIR
+
+INI_FILE=/home/bartlett/borg_examples/ini_example1.ini
+
+cp $INI_FILE ini_file.ini
+$BORG INIT ini_file.ini