Add radial selection

tests/evaluate_likelihood.py

@@ -0,0 +1,84 @@
+import borg_velocity.likelihood as likelihood
+import borg_velocity.forwards as forwards
+import borg_velocity.utils as utils
+
+import aquila_borg as borg
+import configparser
+import h5py as h5
+import numpy as np
+from tqdm import tqdm
+import glob
+
+dirname = '/data101/bartlett/fsigma8/borg_velocity/basic_run_ic'
+ini_file = f'{dirname}/basic_ini.ini'
+
+def get_mcmc_steps(nframe, iter_max, iter_min=0):
+    """
+    Obtain evenly-spaced sample of MCMC steps to make movie from
+    """
+    all_mcmc = glob.glob(dirname + '/mcmc_*.h5')
+    x = [m[len(dirname + '/mcmc_'):-3] for m in all_mcmc]
+    all_mcmc = np.sort([int(m[len(dirname + '/mcmc_'):-3]) for m in all_mcmc])
+    if iter_max >= 0:
+        all_mcmc = all_mcmc[all_mcmc <= iter_max]
+    all_mcmc = all_mcmc[all_mcmc >= iter_min]
+    if nframe > 0:
+        max_out = max(all_mcmc)
+        min_out = min(all_mcmc)
+        step = max(int((max_out - min_out+1) / nframe), 1)
+        all_mcmc = all_mcmc[::step]
+        if max_out not in all_mcmc:
+            all_mcmc = np.concatenate([all_mcmc, [max_out]])
+    return all_mcmc
+
+# Input box
+box_in = borg.forward.BoxModel()
+config = configparser.ConfigParser()
+config.read(ini_file)
+box_in.L = (float(config['system']['L0']), float(config['system']['L1']), float(config['system']['L2']))
+box_in.N = (int(config['system']['N0']), int(config['system']['N1']), int(config['system']['N2']))
+box_in.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2']))
+
+# Setup BORG forward model and likelihood
+model = likelihood.build_gravity_model(None, box_in, ini_file=ini_file)
+cosmo = utils.get_cosmopar(ini_file)
+model.setCosmoParams(cosmo)
+fwd_param = forwards.NullForward(box_in)
+fwd_vel = likelihood.fwd_vel
+mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
+
+# Load s_hat from file
+with h5.File(f'{dirname}/mock_data.h5', 'r') as f:
+    s_hat = f['scalars/s_hat_field'][:]
+    print(f['scalars/hmc_Elh'][:], f['scalars/hmc_Eprior'][:])
+
+# Make mock data
+state = borg.likelihood.MarkovState()
+mylike.initializeLikelihood(state)
+mylike.updateCosmology(cosmo)
+mylike.generateMockData(s_hat, state)
+
+print(dir(mylike))
+quit()
+
+# Evaluate likelihood
+initial_logL = mylike.logLikelihoodComplex(s_hat, None)
+print(initial_logL)
+print(np.var(s_hat))
+# print('var(s_hat): 1.0002235630494603, Call to logLike: 2945.19580078125')
+
+# Load some MCMC samples and see how they compare
+nframe = 10
+iter_min = 0
+iter_max = 1748
+all_mcmc = get_mcmc_steps(nframe, iter_max, iter_min=0)
+all_logL = np.zeros(len(all_mcmc))
+for i in tqdm(range(len(all_mcmc))):
+    with h5.File(f'{dirname}/mcmc_{all_mcmc[i]}.h5', 'r') as f:
+        s_hat = f['scalars/s_hat_field'][:]
+        # print(f['scalars'].keys())
+        print(f['scalars/hmc_Elh'][:], f['scalars/hmc_Eprior'][:])
+    # all_logL[i] = mylike.logLikelihoodComplex(s_hat, None)
+
+print(all_logL)
+    

tests/test_likelihood.py

@@ -7,11 +7,12 @@ import borg_velocity.likelihood as likelihood
 import borg_velocity.forwards as forwards
 import borg_velocity.utils as utils
 
-ini_file = '../conf/basic_ini.ini'
+ini_file = '../conf/supranta_ini.ini'
 test_scaling = True
 test_sigma8 = True
 test_omegam = True
 test_alpha = True
+test_lam = True
 test_muA = True
 
 # Input box
@@ -34,6 +35,7 @@ mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
 state = borg.likelihood.MarkovState()
 mylike.initializeLikelihood(state)
 mylike.updateCosmology(cosmo)
+np.random.seed(2)
 s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
 mylike.generateMockData(s_hat, state)
 
@@ -84,8 +86,7 @@ if test_sigma8:
     fig.clf()
     plt.close(fig)
     
-    
-# Test sigma8
+# Test omegam
 if test_omegam:
     all_omegam = np.linspace(0.1, 0.6, 100)
     all_lkl = np.empty(all_omegam.shape)
@@ -134,6 +135,29 @@ if test_alpha:
     fig.clf()
     plt.close(fig)
     
+# Test radial sampling
+if test_lam:
+    all_lam = np.linspace(0.0, 10.0, 100)
+    all_lkl = np.empty(all_lam.shape)
+    for i, lam in enumerate(all_lam):
+        mylike.fwd_param.setModelParams({'lam0':lam})
+        all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
+    mylike.fwd_param.setModelParams({'lam0':mylike.lam[0]})
+    fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
+    all_lkl -= fid_lkl
+    all_lkl = np.exp(-all_lkl)
+
+    fig, ax = plt.subplots(1, 1, figsize=(5,5))
+    ax.plot(all_lam, all_lkl)
+    ax.axhline(y=0, color='k')
+    ax.axvline(x=mylike.lam[0], color='k')
+    ax.set_xlabel(r'$\lambda_0$')
+    ax.set_ylabel(r'$\mathcal{L}$')
+    fig.tight_layout()
+    fig.savefig('../figs/lam_test.png')
+    fig.clf()
+    plt.close(fig)
+    
     
 # Test bias model 
 if test_muA:
@@ -157,3 +181,4 @@ if test_muA:
     fig.savefig('../figs/muA_test.png')
     fig.clf()
     plt.close(fig)
+