Working SN test inference

tests/sn_inference.py

@@ -41,9 +41,17 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha,
         - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N))
         - Rmax (float): Maximum allowed comoving radius of a tracer (Mpc/h)
         - alpha (float): Exponent for bias model
-
+        - a_tripp (float): Coefficient of stretch in the Tripp relation
+        - b_tripp (float): Coefficient of colour in the Tripp relation
+        - M_SN (float): Absolute magnitude of supernovae
+        - sigma_SN (float): Intrinsic scatter in the Tripp relation
         - sigma_m (float): Uncertainty on the apparent magnitude measurements
-        
+        - sigma_stretch (float): Uncertainty on the stretch measurements
+        - sigma_c (float): Uncertainty on the colour measurements
+        - hyper_stretch_mu (float): Mean of Gaussian hyper prior for the true stretch values
+        - hyper_stretch_sigma (float): Std of Gaussian hyper prior for the true stretch values
+        - hyper_c_mu (float): Mean of hyper Gaussian prior for the true colour values
+        - hyper_c_sigma (float): Std of Gaussian hyper prior for the true colour values
         - sigma_v (float): Uncertainty on the velocity field (km/s)
         - interp_order (int, default=1): Order of interpolation from grid points to the line of sight
         - bias_epsilon (float, default=1e-7): Small number to add to 1 + delta to prevent 0^#
@@ -197,6 +205,14 @@ def estimate_data_parameters():
     """
     Using Foundation DR1, estimate some parameters to use in mock generation.
     
+    Returns:
+        - sigma_m (float): Uncertainty on the apparent magnitude measurements
+        - sigma_stretch (float): Uncertainty on the stretch measurements
+        - sigma_c (float): Uncertainty on the colour measurements
+        - hyper_stretch_mu (float): Estimate of the mean of Gaussian hyper prior for the true stretch values
+        - hyper_stretch_sigma (float): Estimate of the std of Gaussian hyper prior for the true stretch values
+        - hyper_c_mu (float): Estimate of the mean of hyper Gaussian prior for the true colour values
+        - hyper_c_sigma (float): Estimate of the std of Gaussian hyper prior for the true colour values
     """
 
     fname = '/data101/bartlett/fsigma8/PV_data/Foundation_DR1/Foundation_DR1.FITRES.TEXT'
@@ -233,10 +249,14 @@ def likelihood_vel(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, str
     
     Args:
         - alpha (float): Exponent for bias model
-        
+        - a_tripp (float): Coefficient of stretch in the Tripp relation
+        - b_tripp (float): Coefficient of colour in the Tripp relation
+        - M_SN (float): Absolute magnitude of supernovae
+        - sigma_SN (float): Intrinsic scatter in the Tripp relation
         - sigma_v (float): Uncertainty on the velocity field (km/s)
         - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,))
-
+        - stretch_true (np.ndarray): True stretch values of the tracers (shape = (Nt,))
+        - c_true (np.ndarray): True colour values of the tracers (shape = (Nt,))
         - vbulk (np.ndarray): Bulk velocity of the box (km/s) (shape=(3,))
         - dens (np.ndarray): Over-density field (shape = (N, N, N))
         - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N))
@@ -350,6 +370,28 @@ def likelihood_c(c_true, c_obs, sigma_c):
     return loglike
 
 
+def likelihood_m(m_true, m_obs, sigma_m):
+    """
+    Evaluate the terms in the likelihood from apparent magntiude
+    
+    Args:
+        - m_true (np.ndarray): True apparent magnitude of the tracers (shape = (Nt,))
+        - m_obs (np.ndarray): Observed apparent magnitude of the tracers (shape = (Nt,))
+        - sigma_m (float): Uncertainty on the apparent magnitude measurements
+        
+    Returns:
+        - loglike (float): The log-likelihood of the data
+    """
+    
+    Nt = m_obs.shape[0]
+    loglike = - (
+        0.5 * jnp.sum((m_obs - m_true) ** 2 / sigma_m ** 2) 
+        + Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2)
+    )
+    
+    return loglike
+
+
 def likelihood(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk,
                dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
                cz_obs, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos):
@@ -358,10 +400,14 @@ def likelihood(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch
     
     Args:
         - alpha (float): Exponent for bias model
-        
+        - a_tripp (float): Coefficient of stretch in the Tripp relation
+        - b_tripp (float): Coefficient of colour in the Tripp relation
+        - M_SN (float): Absolute magnitude of supernovae
+        - sigma_SN (float): Intrinsic scatter in the Tripp relation
         - sigma_v (float): Uncertainty on the velocity field (km/s)
         - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,))
-        
+        - stretch_true (np.ndarray): True stretch values of the tracers (shape = (Nt,))
+        - c_true (np.ndarray): True colour values of the tracers (shape = (Nt,))
         - vbulk (np.ndarray): Bulk velocity of the box (km/s) (shape=(3,))
         - dens (np.ndarray): Over-density field (shape = (N, N, N))
         - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N))
@@ -373,9 +419,11 @@ def likelihood(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch
         - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^#
         - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
         - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,))
-
+        - stretch_obs (np.ndarray): Observed stretch values of the tracers (shape = (Nt,))
+        - c_obs (np.ndarray): Observed colour values of the tracers (shape = (Nt,))
         - sigma_m (float): Uncertainty on the apparent magnitude measurements
-        - sigma_eta (float): Uncertainty on the linewidth measurements
+        - sigma_stretch (float): Uncertainty on the stretch measurements
+        - sigma_c (float): Uncertainty on the colour measurements
         - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h).
             The shape is (3, Nt, Nsig)
         
@@ -389,8 +437,9 @@ def likelihood(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch
                    cz_obs, MB_pos)
     loglike_stretch = likelihood_stretch(stretch_true, stretch_obs, sigma_stretch)
     loglike_c = likelihood_c(c_true, c_obs, sigma_c)
+    loglike_m = likelihood_m(m_true, m_obs, sigma_m)
     
-    loglike = (loglike_vel + loglike_stretch + loglike_c)
+    loglike = loglike_vel + loglike_stretch + loglike_c + loglike_m
     
     return loglike
 
@@ -405,10 +454,15 @@ def test_likelihood_scan(prior, alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v
     Args:
         - prior (dict): Upper and lower bounds for a uniform prior for the parameters
         - alpha (float): Exponent for bias model
-    
+        - a_tripp (float): Coefficient of stretch in the Tripp relation
+        - b_tripp (float): Coefficient of colour in the Tripp relation
+        - M_SN (float): Absolute magnitude of supernovae
+        - sigma_SN (float): Intrinsic scatter in the Tripp relation
         - sigma_v (float): Uncertainty on the velocity field (km/s)
         - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,))
-
+        - stretch_true (np.ndarray): True stretch values of the tracers (shape = (Nt,))
+        - c_true (np.ndarray): True colour values of the tracers (shape = (Nt,))
+        - vbulk (np.ndarray): Bulk velocity of the box (km/s) (shape=(3,))
         - dens (np.ndarray): Over-density field (shape = (N, N, N))
         - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N))
         - omega_m (float): Matter density parameter Om
@@ -417,11 +471,13 @@ def test_likelihood_scan(prior, alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v
         - xmin (float): Coordinate of corner of the box (Mpc/h)
         - interp_order (int): Order of interpolation from grid points to the line of sight
         - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^#
-        - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
+        - czCMB (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
         - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,))
-
+        - stretch_obs (np.ndarray): Observed stretch values of the tracers (shape = (Nt,))
+        - c_obs (np.ndarray): Observed colour values of the tracers (shape = (Nt,))
         - sigma_m (float): Uncertainty on the apparent magnitude measurements
-        - sigma_eta (float): Uncertainty on the apparent linewidth measurements
+        - sigma_stretch (float): Uncertainty on the stretch measurements
+        - sigma_c (float): Uncertainty on the colour measurements
         - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h).
             The shape is (3, Nt, Nsig)
         
@@ -477,8 +533,8 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, cpar, L, xmin,
     Args:
         - num_warmup (int): Number of warmup steps to take in the MCMC
         - num_samples (int): Number of samples to take in the MCMC
-        - prior
+        - prior (dict): Upper and lower bounds for a uniform prior for the parameters
-        - initial
+        - initial (dict): Initial values for the MCMC
         - dens (np.ndarray): Over-density field (shape = (N, N, N))
         - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N))
         - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use
@@ -486,11 +542,13 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, cpar, L, xmin,
         - xmin (float): Coordinate of corner of the box (Mpc/h)
         - interp_order (int): Order of interpolation from grid points to the line of sight
         - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^#
-        - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
+        - czCMB (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
         - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,))
-
+        - stretch_obs (np.ndarray): Observed stretch values of the tracers (shape = (Nt,))
+        - c_obs (np.ndarray): Observed colour values of the tracers (shape = (Nt,))
         - sigma_m (float): Uncertainty on the apparent magnitude measurements
-  
+        - sigma_stretch (float): Uncertainty on the stretch measurements
+        - sigma_c (float): Uncertainty on the colour measurements
         - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h).
             The shape is (3, Nt, Nsig)
         
@@ -641,7 +699,7 @@ def process_mcmc_run(mcmc, param_labels, truths, true_vars):
     fig1.savefig('sn_trace.png')
     
     # Corner plot
-    fig2, axs2 = plt.subplots(samps.shape[1], samps.shape[1], figsize=(20,20))
+    fig2, axs2 = plt.subplots(samps.shape[1], samps.shape[1], figsize=(25,25))
     corner.corner(
         np.array(samps),
         labels=param_labels,
@@ -736,7 +794,7 @@ def main():
         'hyper_sigma_m': (np.percentile(m_obs, 84) - np.percentile(m_obs, 16)) / 2,
     }
     prior = {
-        'alpha': [0.5, 4.5],
+        'alpha': [0.5, 6],
         'a_tripp': [0.01, 0.2],
         'b_tripp': [2.5, 4.5],
         'M_SN': [-19.5, -17.5],
@@ -760,14 +818,16 @@ def main():
     
     # Run a MCMC
     mcmc = run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, cpar, L, xmin, interp_order, bias_epsilon,
-             czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos,)
+             czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos)
     param_labels = ['alpha', 'a_tripp', 'b_tripp', 'M_SN', 'sigma_SN', 'sigma_v', 
-                    'hyper_mean_m', 'hyper_sigma_m', 'hyper_mean_stretch', 'hyper_sigma_stretch', 
+                    'hyper_mean_m', 'hyper_sigma_m', 
-                    'hyper_mean_c', 'hyper_sigma_c', 'hyper_corr_mx', 'hyper_corr_mc', 'hyper_corr_xc',
+                    'hyper_mean_stretch', 'hyper_sigma_stretch', 'hyper_mean_c', 'hyper_sigma_c',
+                    'hyper_corr_mx', 'hyper_corr_mc', 'hyper_corr_xc',
                     'vbulk_x', 'vbulk_y', 'vbulk_z']
     truths = [alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, 
-              None, None, hyper_stretch_mu, hyper_stretch_sigma, 
+              None, None, 
-              hyper_c_mu, hyper_c_sigma, None, None, None,
+              hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_sigma, 
+              None, None, None,
               vbulk[0], vbulk[1], vbulk[2]]
     true_vars = {'m':m_true, 'stretch':stretch_true, 'c': c_true}
     process_mcmc_run(mcmc, param_labels, truths, true_vars)
@@ -782,8 +842,6 @@ if __name__ == "__main__":
 """
 TO DO
 
-- Fix SN inference - poor sampling and Tripp variables not constrained
 - Deal with case where sigma_eta and sigma_m could be floats vs arrays
-- Add in selection cuts for the supernovae
 
 """

tests/tfr_inference.py

@@ -621,7 +621,7 @@ def test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true,
         - xmin (float): Coordinate of corner of the box (Mpc/h)
         - interp_order (int): Order of interpolation from grid points to the line of sight
         - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^#
-        - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
+        - czCMB (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
         - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,))
         - eta_obs (np.ndarray): Observed linewidths of the tracers (shape = (Nt,))
         - sigma_m (float): Uncertainty on the apparent magnitude measurements
@@ -701,8 +701,8 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, cpar, L, xmin,
     Args:
         - num_warmup (int): Number of warmup steps to take in the MCMC
         - num_samples (int): Number of samples to take in the MCMC
-        - prior
+        - prior (dict): Upper and lower bounds for a uniform prior for the parameters
-        - initial
+        - initial (dict): Initial values for the MCMC
         - dens (np.ndarray): Over-density field (shape = (N, N, N))
         - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N))
         - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use
@@ -710,7 +710,7 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, cpar, L, xmin,
         - xmin (float): Coordinate of corner of the box (Mpc/h)
         - interp_order (int): Order of interpolation from grid points to the line of sight
         - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^#
-        - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
+        - czCMB (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
         - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,))
         - eta_obs (np.ndarray): Observed linewidths of the tracers (shape = (Nt,))
         - sigma_m (float): Uncertainty on the apparent magnitude measurements