Compute As and example0 now works

analysis.py

@@ -120,7 +120,7 @@ def get_spectra(ini_name, dirname, mcmc_steps, which_field='BORG_final_density',
         idx = mcmc_steps[i]
 
         with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file:
-            delta2 = np.array(mcmc_file['scalars/{which_field}'][...],dtype=np.float64)
+            delta2 = np.array(mcmc_file[f'scalars/{which_field}'][...],dtype=np.float64)
 
         with suppress_stdout():
             Pk = PKL.XPk([delta1.astype(np.float32),delta2.astype(np.float32)], boxsize, axis=0, MAS=[MAS, MAS], threads=1)

example0.py

@@ -41,6 +41,42 @@ def get_cosmopar(ini_file):
     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
 
 
@@ -73,7 +109,8 @@ class MyLikelihood(borg.likelihood.BaseLikelihood):
         self.grad_like = jax.grad(self.dens2like)
 
     def updateCosmology(self, cosmo: borg.cosmo.CosmologicalParameters) -> None:
-        self.fwd.setCosmoParams(cosmo)
+        cpar = compute_As(cosmo)
+        self.fwd.setCosmoParams(cpar)
     
     def updateMetaParameters(self, state: borg.likelihood.MarkovState) -> None:
         """
@@ -83,7 +120,8 @@ class MyLikelihood(borg.likelihood.BaseLikelihood):
             - state (borg.likelihood.MarkovState): The state object to be used in the likelihood.
 
         """
-        cpar = state['cosmology']
+        cosmo = state['cosmology']
+        cpar = compute_As(cosmo)
         self.fwd.setCosmoParams(cpar)
 
     def initializeLikelihood(self, state: borg.likelihood.MarkovState) -> None:
@@ -115,6 +153,7 @@ class MyLikelihood(borg.likelihood.BaseLikelihood):
         # 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
@@ -128,6 +167,8 @@ class MyLikelihood(borg.likelihood.BaseLikelihood):
         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):
@@ -149,7 +190,7 @@ class MyLikelihood(borg.likelihood.BaseLikelihood):
 
     def logLikelihoodComplex(self, s_hat:np.ndarray, gradientIsNext:bool):
 
-        myprint('Getting density field now')
+        # 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)
@@ -193,7 +234,7 @@ class MyLikelihood(borg.likelihood.BaseLikelihood):
         """
         self.updateCosmology(self.fwd.getCosmoParams())
         self.dens2like(self.delta)
-        state["BORG_final_density"][:] = self.delta
+        state["BORG_final_density"][:] = self.delta.copy()
 
 
 @borg.registerGravityBuilder
@@ -215,7 +256,7 @@ def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.Bo
     myprint("Building gravity model")
 
     # READ FROM INI FILE
-    which_model = 'cola'
+    which_model = 'lpt'
     ai = 0.05 # Initial scale factor
     af = 1.0 # Final scale factor
     supersampling = 2
@@ -228,11 +269,15 @@ def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.Bo
     chain.addModel(borg.forward.models.HermiticEnforcer(box))
 
     # CLASS transfer function
-    # chain @= borg.forward.model_lib.M_PRIMORDIAL_AS(box)
+    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(
@@ -280,12 +325,15 @@ def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.Bo
                       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

ini_file.ini

@@ -53,7 +53,7 @@ beta = 1.5
 z0 = 0
 
 [mock]
-sigma_dens = 0.1
+sigma_dens = 1
 
 [gravity]
 which_model = lpt