Add some likelihood tests

tests/test_likelihood.py

@@ -0,0 +1,157 @@
+import aquila_borg as borg
+import configparser
+import numpy as np
+import matplotlib.pyplot as plt
+
+import borg_velocity.likelihood as likelihood
+import borg_velocity.forwards as forwards
+import borg_velocity.utils as utils
+
+ini_file = '../conf/basic_ini.ini'
+test_scaling = False
+test_sigma8 = False
+test_omegam = False
+test_alpha = False
+test_muA = True
+
+# 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)
+mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, ini_file)
+
+# Create mock data
+state = borg.likelihood.MarkovState()
+mylike.initializeLikelihood(state)
+mylike.updateCosmology(cosmo)
+s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
+mylike.generateMockData(s_hat, state)
+
+if test_scaling:
+    all_scale = np.linspace(0.5, 1.5, 100)
+    all_lkl = np.empty(all_scale.shape)
+    for i, scale in enumerate(all_scale):
+        all_lkl[i] = mylike.logLikelihoodComplex(scale * s_hat, None)
+    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_scale, all_lkl)
+    ax.axhline(y=0, color='k')
+    ax.axvline(x=1, color='k')
+    ax.set_xlabel(r'$\hat{s}$ scaling')
+    ax.set_ylabel(r'$\mathcal{L}$')
+    fig.tight_layout()
+    fig.savefig('../figs/scaling_test.png')
+    fig.clf()
+    plt.close(fig)
+    
+# Test sigma8
+if test_sigma8:
+    all_sigma8 = np.linspace(0.5, 1.2, 40)
+    all_lkl = np.empty(all_sigma8.shape)
+    cosmo_true = mylike.fwd.getCosmoParams()
+    cosmo = mylike.fwd.getCosmoParams()
+    for i, sigma8 in enumerate(all_sigma8):
+        cosmo.sigma8 = sigma8
+        mylike.updateCosmology(cosmo)
+        all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
+    mylike.updateCosmology(cosmo_true)
+    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_sigma8, all_lkl)
+    ax.axhline(y=0, color='k')
+    ax.axvline(x=cosmo_true.sigma8, color='k')
+    ax.set_xlabel(r'$\sigma_8$')
+    ax.set_ylabel(r'$\mathcal{L}$')
+    fig.tight_layout()
+    fig.savefig('../figs/sigma8_test.png')
+    fig.clf()
+    plt.close(fig)
+    
+    
+# Test sigma8
+if test_omegam:
+    all_omegam = np.linspace(0.1, 0.6, 40)
+    all_lkl = np.empty(all_omegam.shape)
+    cosmo_true = mylike.fwd.getCosmoParams()
+    cosmo = mylike.fwd.getCosmoParams()
+    for i, omegam in enumerate(all_omegam):
+        cosmo.omega_m = omegam
+        mylike.updateCosmology(cosmo)
+        all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
+    mylike.updateCosmology(cosmo_true)
+    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_omegam, all_lkl)
+    ax.axhline(y=0, color='k')
+    ax.axvline(x=cosmo_true.omega_m, color='k')
+    ax.set_xlabel(r'$\Omega_{\rm m}$')
+    ax.set_ylabel(r'$\mathcal{L}$')
+    fig.tight_layout()
+    fig.savefig('../figs/omegam_test.png')
+    fig.clf()
+    plt.close(fig)
+    
+# Test bias model 
+if test_alpha:
+    all_alpha = np.linspace(-1.0, 5.0, 50)
+    all_lkl = np.empty(all_alpha.shape)
+    for i, alpha in enumerate(all_alpha):
+        mylike.fwd_param.setModelParams({'alpha0':alpha})
+        all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
+    mylike.fwd_param.setModelParams({'alpha0':mylike.alpha[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_alpha, all_lkl)
+    ax.axhline(y=0, color='k')
+    ax.axvline(x=mylike.alpha[0], color='k')
+    ax.set_xlabel(r'$\alpha_0$')
+    ax.set_ylabel(r'$\mathcal{L}$')
+    fig.tight_layout()
+    fig.savefig('../figs/alpha_test.png')
+    fig.clf()
+    plt.close(fig)
+    
+    
+# Test bias model 
+if test_muA:
+    all_muA = np.linspace(0.5, 1.5, 50)
+    all_lkl = np.empty(all_muA.shape)
+    for i, muA in enumerate(all_muA):
+        mylike.fwd_param.setModelParams({'muA0':muA})
+        all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
+    mylike.fwd_param.setModelParams({'muA0':mylike.muA[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_muA, all_lkl)
+    ax.axhline(y=0, color='k')
+    ax.axvline(x=mylike.muA[0], color='k')
+    ax.set_xlabel(r'$\mu_0$')
+    ax.set_ylabel(r'$\mathcal{L}$')
+    fig.tight_layout()
+    fig.savefig('../figs/muA_test.png')
+    fig.clf()
+    plt.close(fig)