Add minimum step before sampling cosmology or model parameters

borg_velocity/likelihood.py

@@ -167,7 +167,6 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         f = cosmology.gplus(self.af)  # dD / da
         f *= self.af / cosmology.d_plus(self.af)  # f = dlnD / dlna
         self.f = f
-
         
     def generateMBData(self) -> None:
         """
@@ -530,6 +529,17 @@ def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.Bo
     
     return chain
 
+_glob_model = None
+_glob_cosmo = None
+begin_model = None
+begin_cosmo = None
+
+def check_model_sampling(loop):
+    return loop.getStepID() > begin_model
+
+def check_cosmo_sampling(loop):
+    return loop.getStepID() > begin_cosmo
+            
 
 @borg.registerSamplerBuilder
 def build_sampler(
@@ -551,6 +561,9 @@ def build_sampler(
         List of samplers to use.
 
     """
+    global _glob_model, _glob_cosmo, begin_model, begin_cosmo
+    borg.print_msg(borg.Level.std, "Hello sampler, loop is {l}, step_id={s}", l=loop, s=loop.getStepID())
+    
     myprint("Building sampler")
     
     myprint("Reading from configuration file: " + borg.getIniConfigurationFilename())
@@ -581,17 +594,19 @@ def build_sampler(
             warnings.warn(s, stacklevel=2)
         # Remove for later to prevent duplication 
         to_sample.remove(p)
-    
+       
-    def myfun(loop):
+    begin_cosmo = int(config['mcmc']['warmup_cosmo'])
-        return loop.getStepId() > 1000 # put this in ini file
+
     
     if len(params) > 0:
         myprint('Adding cosmological parameter sampler')
         cosmo_sampler = borg.samplers.ModelParamsSampler(prefix, params, likelihood, chain, initial_values, prior)
         cosmo_sampler.setName("cosmo_sampler")
+        _glob_cosmo = cosmo_sampler
         all_sampler.append(cosmo_sampler)
+        loop.push(cosmo_sampler)
         loop.addToConditionGroup("warmup_cosmo", "cosmo_sampler")
-        loop.addConditionToConditionGroup("warmup_cosmo", partial(myfun, loop))
+        loop.addConditionToConditionGroup("warmup_cosmo", partial(check_cosmo_sampling, loop))
 
     # Model parameter sampler
     prefix = ""
@@ -632,11 +647,26 @@ def build_sampler(
         else:
             s = f'Could not find {p} prior, so will not sample'
             warnings.warn(s, stacklevel=2)
+            
+    begin_model = int(config['mcmc']['warmup_model'])
+    
+   
     if len(params) > 0:
         myprint('Adding model parameter sampler')
-        all_sampler.append(borg.samplers.ModelParamsSampler(prefix, params, likelihood, fwd_param, initial_values, prior))
+        model_sampler = borg.samplers.ModelParamsSampler(prefix, params, likelihood, fwd_param, initial_values, prior)
+        model_sampler.setName("model_sampler")
+        _glob_model = model_sampler
+        loop.push(model_sampler)
+        all_sampler.append(model_sampler)
+        loop.addToConditionGroup("warmup_model", "model_sampler")
+        loop.addConditionToConditionGroup("warmup_model", partial(check_model_sampling, loop))
+        
+    print('Warmups:', begin_cosmo, begin_model)
+    print(dir(loop))
+    # quit()
 
-    return all_sampler
+    # return all_sampler
+    return []
 
 
 @borg.registerLikelihoodBuilder
@@ -659,8 +689,3 @@ def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.L
     likelihood = VelocityBORGLikelihood(chain, fwd_param, fwd_vel, borg.getIniConfigurationFilename())
     return likelihood
 
-
-"""
-TO DO
-- samplers.MainLoop - add condition that something begins sampling after some number of steps
-"""

conf/basic_ini.ini

@@ -1,4 +1,5 @@
 [system]
+console_output = borg_log
 VERBOSE_LEVEL = 2
 N0 = 32
 N1 = 32
@@ -17,16 +18,18 @@ seed_cpower = true
 hades_sampler_blocked = false
 bias_sampler_blocked= true
 nmean_sampler_blocked= true
-sigma8_sampler_blocked = true
+sigma8_sampler_blocked = false
-muA_sampler_blocked = true
+muA_sampler_blocked = false
-omega_m_sampler_blocked = true
+omega_m_sampler_blocked = false
-alpha_sampler_blocked = true
+alpha_sampler_blocked = false
-sig_v_sampler_blocked = true
+sig_v_sampler_blocked = false
-bulk_flow_sampler_blocked = true
+bulk_flow_sampler_blocked = false
 ares_heat = 1.0
 
 [mcmc]
 number_to_generate = 10
+warmup_model = 3
+warmup_cosmo = 7
 random_ic = false
 init_random_scaling = 0.1
 bignum = 1e300