Add tests for TFR integral

tests/tfr_integral.py

@@ -0,0 +1,82 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+from tfr_inference import estimate_data_parameters, get_fields, create_mock, myprint, generateMBData
+import borg_velocity.utils as utils
+
+def main():
+    
+    myprint('Beginning')
+    
+    # Get some parameters from the data
+    sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma, hyper_m_sigma = estimate_data_parameters()
+    print('sigma_m', sigma_m)
+    print('sigma_eta', sigma_eta)
+    print('hyper_eta_mu', hyper_eta_mu)
+    print('hyper_eta_sigma', hyper_eta_sigma)
+    
+    # Other parameters to use
+    L = 500.0
+    N = 64
+    xmin = -L/2
+    R_lim = L / 2
+    Rmax = 100
+    Nt = 5_000
+    alpha = 1.4
+    mthresh = 11.25
+    a_TFR = -23
+    b_TFR = -8.2
+    sigma_TFR = 0.3
+    sigma_v = 150
+    Nint_points = 201 
+    Nsig = 10
+    frac_sigma_r = 0.07  # WANT A BETTER WAY OF DOING THIS - ESTIMATE THROUGH SIGMAS FROM TFR
+    interp_order = 1
+    bias_epsilon = 1.e-7
+    
+    # Make mock
+    np.random.seed(123)
+    cpar, dens, vel = get_fields(L, N, xmin)
+    RA, Dec, czCMB, m_true, eta_true, m_obs, eta_obs, xtrue, vbulk = create_mock(
+            Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
+            a_TFR, b_TFR, sigma_TFR, sigma_m, sigma_eta, 
+            hyper_eta_mu, hyper_eta_sigma, sigma_v, 
+            interp_order=interp_order, bias_epsilon=bias_epsilon)
+    MB_pos = generateMBData(RA, Dec, czCMB, L, N, R_lim, Nsig, Nint_points, sigma_v, frac_sigma_r)
+    
+    r = np.sqrt(np.sum(MB_pos**2, axis=0))    
+    zcosmo = utils.z_cos(r, cpar.omega_m)
+    mu_true = 5 * np.log10((1 + zcosmo) * r / cpar.h) + 25 
+
+    gamma_inv2 = 1 + sigma_m ** 2 * (hyper_eta_sigma**2 + sigma_eta**2) / (b_TFR**2 * hyper_eta_sigma**2 * sigma_eta**2 + sigma_TFR**2 * (hyper_eta_sigma**2 + sigma_eta**2))
+    print(gamma_inv2)
+    
+    delta0 = (mthresh - m_obs) / sigma_m
+    den = (b_TFR**2 * hyper_eta_sigma**2 * sigma_eta**2 + (sigma_TFR**2 + sigma_m**2) * (hyper_eta_sigma**2 + sigma_eta**2))
+    delta1 = ((m_obs - a_TFR - b_TFR * eta_obs)[:,None] - mu_true) * sigma_m * hyper_eta_sigma**2 / den
+    delta2 = ((m_obs - a_TFR - b_TFR * hyper_eta_sigma)[:,None] - mu_true) * sigma_m * hyper_eta_sigma**2 / den
+    
+    delta = delta0[:,None] + delta1 + delta2
+    
+    fig, axs = plt.subplots(1, 3, figsize=(15,4))
+    hist_kwargs = dict(density=True, histtype='step', bins=30)
+    x = delta.flatten()
+    print('delta', x.min(), x.max())
+    axs[0].hist(x, **hist_kwargs)
+    x = gamma_inv2
+    print('gamma_inv2', x.min(), x.max())
+    axs[1].hist(x, **hist_kwargs)
+    x = (delta * np.sqrt(gamma_inv2)).flatten()
+    print('delta/gamma', x.min(), x.max())
+    axs[2].hist(x, **hist_kwargs)
+    axs[0].set_title(r'$\delta$')
+    axs[1].set_title(r'$\gamma^{-2}$')
+    axs[2].set_title(r'$\delta / \gamma$')
+    fig.tight_layout()
+    fig.savefig('tfr_integral_params.png')
+    
+    
+    return
+
+if __name__ == "__main__":
+    main()