Working TFR MCMC without magnitude cuts

tests/tfr_inference.py

@@ -165,6 +165,42 @@ def get_fields(L, N, xmin, gravity='lpt', velmodel_name='CICModel'):
 def 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=1, bias_epsilon=1e-7):
+    """
+    Create mock TFR catalogue from a density and velocity field
+    
+    Args:
+        - Nt (int): Number of tracers to produce
+        - L (float): Box length (Mpc/h)
+        - xmin (float): Coordinate of corner of the box (Mpc/h)
+        - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use
+        - dens (np.ndarray): Over-density field (shape = (N, N, N))
+        - 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
+        - mthresh (float): Threshold absolute magnitude in selection
+        - a_TFR (float): TFR relation intercept
+        - b_TFR (float): TFR relation slope
+        - sigma_TFR (float): Intrinsic scatter in the TFR
+        - sigma_v (float): Uncertainty on the velocity field (km/s)
+        - sigma_m (float): Uncertainty on the apparent magnitude measurements
+        - sigma_eta (float): Uncertainty on the linewidth measurements
+        - hyper_eta_mu (float): Mean of the Gaussian hyper-prior for the true eta values
+        - hyper_eta_sigma (float): Std deviation of the Gaussian hyper-prior for the true eta 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^#
+        
+    Returns:
+        - all_RA (np.ndarrary): Right Ascension (degrees) of the tracers (shape = (Nt,))
+        - all_Dec (np.ndarrary): Dec (np.ndarray): Delination (degrees) of the tracers (shape = (Nt,))
+        - czCMB (np.ndarrary): Observed redshifts (km/s) of the tracers (shape = (Nt,))
+        - all_mtrue (np.ndarrary): True apparent magnitudes of the tracers (shape = (Nt,))
+        - all_etatrue (np.ndarrary): True linewidths of the tracers (shape = (Nt,))
+        - all_mobs (np.ndarrary): Observed apparent magnitudes of the tracers (shape = (Nt,))
+        - all_etaobs (np.ndarrary): Observed linewidths of the tracers (shape = (Nt,))
+        - all_xtrue (np.ndarrary): True comoving coordinates of the tracers (Mpc/h) (shape = (3, Nt))
+    
+    """
     
     # Initialize lists to store valid positions and corresponding sig_mu values
     all_xtrue = np.empty((3, Nt))
@@ -226,7 +262,8 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
         etaobs = etatrue + sigma_eta * np.random.randn(Nt)
         
         # Apply apparement magnitude cut
-        m = mobs <= mthresh
+        # m = mobs <= mthresh
+        m = np.ones(mobs.shape, dtype=bool)
         mtrue = mtrue[m]
         etatrue = etatrue[m]
         mobs = mobs[m]
@@ -283,8 +320,10 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
 
 
 def estimate_data_parameters():
-    
     """
+    Using the 2MASS catalogue, estimate some parameters to use in mock generation.
+    The file contains the following columns:
+    
     ID   2MASS XSC ID name (HHMMSSss+DDMMSSs)
     RAdeg   Right ascension (J2000)
     DEdeg   Declination (J2000)
@@ -297,6 +336,13 @@ def estimate_data_parameters():
     e_Jmag    Error of the NIR magnitudes in J band from the (mag)
     WHIc    Corrected HI width (km/s)
     e_WHIc    Error of corrected HI width (km/s)
+    
+    Returns:
+        - sigma_m (float): Estimate of the uncertainty on the apparent magnitude measurements
+        - sigma_eta (float): Estimate of the uncertainty on the linewidth measurements
+        - hyper_eta_mu (float): Estimate of the mean of the Gaussian hyper-prior for the true eta values
+        - hyper_eta_sigma (float): Estimate of the std deviation of the Gaussian hyper-prior for the true eta values
+        - hyper_m_sigma (float): Estimate of the std deviation of the Gaussian hyper-prior for the true m values
     """
     
     columns = ['ID', 'RAdeg', 'DEdeg', 'cz2mrs', 'Kmag', 'Hmag', 'Jmag', 'e_Kmag', 'e_Hmah', 'e_Jmag', 'WHIc', 'e_WHIc']
@@ -311,7 +357,9 @@ def estimate_data_parameters():
     hyper_eta_mu = np.median(eta)
     hyper_eta_sigma = (np.percentile(eta, 84) - np.percentile(eta, 16)) / 2
     
-    return sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma
+    hyper_m_sigma = np.amax(df['Kmag']) - np.percentile(df['Kmag'], 16)
+    
+    return sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma, hyper_m_sigma
 
 
 def generateMBData(RA, Dec, cz_obs, L, N, R_lim, Nsig, Nint_points, sigma_v, frac_sigma_r):
@@ -363,11 +411,11 @@ def generateMBData(RA, Dec, cz_obs, L, N, R_lim, Nsig, Nint_points, sigma_v, fra
     return MB_pos
 
 
-def likelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
+def likelihood_vel(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
                dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
-               cz_obs, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh):
+               cz_obs, MB_pos, mthresh):
     """
-    Evaluate the likelihood for TFR sample
+    Evaluate the terms in the likelihood from the velocity and malmquist bias
     
     Args:
         - alpha (float): Exponent for bias model
@@ -386,10 +434,6 @@ def likelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
         - 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,))
-        - 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
-        - sigma_eta (float): Uncertainty on the apparent linewidth measurements
         - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h).
             The shape is (3, Nt, Nsig)
         - mthresh (float): Threshold absolute magnitude in selection
@@ -397,7 +441,7 @@ def likelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
     Returns:
         - loglike (float): The log-likelihood of the data
     """
-
+    
     # Comoving radii of integration points (Mpc/h)
     r = jnp.sqrt(jnp.sum(MB_pos ** 2, axis=0))
     
@@ -444,26 +488,102 @@ def likelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
     # Integrate to get likelihood
     p_cz = jnp.trapezoid(jnp.exp(-0.5 * d2) * p_r / p_r_norm, r, axis=1)
     lkl_ind = jnp.log(p_cz) - scale / 2 - 0.5 * jnp.log(2 * np.pi * sigma_v**2)
-    loglike_vel = - lkl_ind.sum()
+    loglike = lkl_ind.sum()
+    
+    return loglike
+
+
+def likelihood_m(m_true, m_obs, sigma_m, mthresh):
+    """
+    Evaluate the terms in the likelihood from apparent magnitude
+    
+    Args:
+        - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,))
+        - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,))
+        - sigma_m (float): Uncertainty on the apparent magnitude measurements
+        - mthresh (float): Threshold absolute magnitude in selection
+        
+    Returns:
+        - loglike (float): The log-likelihood of the data
+    """
     
     Nt = m_obs.shape[0]
-    
+    # norm = 2 / (1 + jax.scipy.special.erf((mthresh - m_true) / (jnp.sqrt(2) * sigma_m))) / jnp.sqrt(2 * jnp.pi * sigma_m ** 2)
-    # Apparent magnitude terms
+    norm = jnp.sqrt(2 * jnp.pi * sigma_m ** 2) * jnp.ones(Nt)
-    norm = 0.5 * (1 + jax.scipy.special.erf((mthresh - m_true) / (jnp.sqrt(2) * sigma_m)))
+    loglike = - (
-    loglike_m = (
         0.5 * jnp.sum((m_obs - m_true) ** 2 / sigma_m ** 2) 
         + jnp.sum(jnp.log(norm)) 
         + Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2)
     )
+    
+    return loglike
+
+
+def likelihood_eta(eta_true, eta_obs, sigma_eta):
+    """
+    Evaluate the terms in the likelihood from linewidth
+    
+    Args:
+        - eta_true (np.ndarray): True linewidths of the tracers (shape = (Nt,))
+        - eta_obs (np.ndarray): Observed linewidths of the tracers (shape = (Nt,))
+        - sigma_eta (float): Uncertainty on the linewidth measurements
         
-    # Linewidth terms
+    Returns:
-    loglike_eta = (
+        - loglike (float): The log-likelihood of the data
+    """
+    
+    Nt = eta_obs.shape[0]
+    loglike = - (
         0.5 * jnp.sum((eta_obs - eta_true) ** 2 / sigma_eta ** 2) 
         + Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_eta ** 2)
     )
     
-    # loglike = - (loglike_vel + loglike_m + loglike_eta)
+    return loglike
-    loglike = - (loglike_eta + loglike_m)
+
+
+def likelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
+               dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
+               cz_obs, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh):
+    """
+    Evaluate the likelihood for TFR sample
+    
+    Args:
+        - alpha (float): Exponent for bias model
+        - a_TFR (float): TFR relation intercept
+        - b_TFR (float): TFR relation slope
+        - sigma_TFR (float): Intrinsic scatter in the TFR
+        - sigma_v (float): Uncertainty on the velocity field (km/s)
+        - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,))
+        - eta_true (np.ndarray): True linewidths of the tracers (shape = (Nt,))
+        - 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
+        - h (float): Hubble constant H0 = 100 h km/s/Mpc
+        - L (float): Comoving box size (Mpc/h)
+        - 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,))
+        - 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
+        - sigma_eta (float): Uncertainty on the linewidth measurements
+        - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h).
+            The shape is (3, Nt, Nsig)
+        - mthresh (float): Threshold absolute magnitude in selection
+        
+    Returns:
+        - loglike (float): The log-likelihood of the data
+    """
+
+    
+    loglike_vel = likelihood_vel(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
+                   dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
+                   cz_obs, MB_pos, mthresh)
+    loglike_m = likelihood_m(m_true, m_obs, sigma_m, mthresh)
+    loglike_eta = likelihood_eta(eta_true, eta_obs, sigma_eta)
+    
+    loglike = (loglike_vel + loglike_m + loglike_eta)
     
     return loglike
 
@@ -507,7 +627,7 @@ def test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true,
     pars = [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v]
     par_names = ['alpha', 'a_TFR', 'b_TFR', 'sigma_TFR', 'sigma_v']
     
-    orig_ll = likelihood(*pars, m_true, eta_true,
+    orig_ll = - likelihood(*pars, m_true, eta_true,
                          dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
                          czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh)
     
@@ -526,7 +646,7 @@ def test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true,
         orig_x = pars[i]
         for j, xx in enumerate(x):
             pars[i] = xx
-            all_ll[j] = likelihood(*pars, m_true, eta_true,
+            all_ll[j] = - likelihood(*pars, m_true, eta_true,
                          dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
                          czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh)
         pars[i] = orig_x
@@ -546,8 +666,7 @@ def test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true,
 
 
 def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
-             czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh,
+             czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh,):
-            m_true):
     """
     Run MCMC over the model parameters
     
@@ -573,6 +692,9 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, omega_m, h, L,
             The shape is (3, Nt, Nsig)
         - mthresh (float): Threshold absolute magnitude in selection
         
+    Returns:
+        - mcmc
+        
     """
     
     Nt = eta_obs.shape[0]
@@ -583,50 +705,46 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, omega_m, h, L,
         a_TFR =  numpyro.sample("a_TFR", dist.Uniform(*prior['a_TFR']))
         b_TFR =  numpyro.sample("b_TFR", dist.Uniform(*prior['b_TFR']))
         sigma_TFR = numpyro.sample("sigma_TFR", dist.HalfCauchy(1.0))
-        sigma_v = numpyro.sample("sigma_v", dist.HalfCauchy(1.0))
+        # sigma_v = numpyro.sample("sigma_v", dist.HalfCauchy(1.0))
-
+        sigma_v =  numpyro.sample("sigma_v", dist.Uniform(*prior['sigma_v']))
-#         # Sample the means with a uniform prior
-#         hyper_mean_m = numpyro.sample("hyper_mean_m", dist.Uniform(*prior['hyper_mean_m']))
-#         hyper_mean_eta = numpyro.sample("hyper_mean_eta", dist.Uniform(*prior['hyper_mean_eta']))
-#         hyper_mean = jnp.array([hyper_mean_m, hyper_mean_eta])
-
-#         # Sample standard deviations with a 1/sigma prior (Jeffreys prior approximation)
-#         hyper_sigma_m = numpyro.sample("hyper_sigma_m", dist.HalfCauchy(1.0)) # Equivalent to 1/sigma prior
-#         hyper_sigma_eta = numpyro.sample("hyper_sigma_eta", dist.HalfCauchy(1.0))  
-#         hyper_sigma = jnp.array([hyper_sigma_m, hyper_sigma_eta])
-
-#         # Sample correlation matrix using LKJ prior
-#         L_corr = numpyro.sample("L_corr", dist.LKJCholesky(2, concentration=1.0))  # Cholesky factor of correlation matrix
-#         corr_matrix = L_corr @ L_corr.T  # Convert to full correlation matrix
-        
-#         # Construct full covariance matrix: Σ = D * Corr * D
-#         hyper_cov = jnp.diag(hyper_sigma) @ corr_matrix @ jnp.diag(hyper_sigma)
-
-#         # Sample the true eta and m
-#         x = numpyro.sample("x", dist.MultivariateNormal(hyper_mean, hyper_cov), sample_shape=(Nt,))
-#         m_true = numpyro.deterministic("m_true", x[:, 0])
-#         eta_true = numpyro.deterministic("eta_true", x[:, 1])
     
+        hyper_mean_m = numpyro.sample("hyper_mean_m", dist.Uniform(*prior['hyper_mean_m']))
+        hyper_sigma_m = numpyro.sample("hyper_sigma_m", dist.HalfCauchy(1.0)) # Equivalent to 1/sigma prior
         hyper_mean_eta = numpyro.sample("hyper_mean_eta", dist.Uniform(*prior['hyper_mean_eta']))
         hyper_sigma_eta = numpyro.sample("hyper_sigma_eta", dist.HalfCauchy(1.0)) # Equivalent to 1/sigma prior
-        eta_true =  numpyro.sample("eta_true", dist.Normal(hyper_mean_eta, hyper_sigma_eta), sample_shape=(Nt,))
+        
-
+        # Sample correlation matrix using LKJ prior
+        L_corr = numpyro.sample("L_corr", dist.LKJCholesky(2, concentration=1.0))  # Cholesky factor of correlation matrix
+        corr_matrix = L_corr @ L_corr.T  # Convert to full correlation matrix
+        
+        # Construct full covariance matrix: Σ = D * Corr * D
+        hyper_mean = jnp.array([hyper_mean_m, hyper_mean_eta])
+        hyper_sigma = jnp.array([hyper_sigma_m, hyper_sigma_eta])
+        hyper_cov = jnp.diag(hyper_sigma) @ corr_matrix @ jnp.diag(hyper_sigma)
+        
+        # Sample m_true and eta_true
+        x = numpyro.sample("true_vars", dist.MultivariateNormal(hyper_mean, hyper_cov), sample_shape=(Nt,))
+        m_true = numpyro.deterministic("m_true", x[:, 0])
+        eta_true = numpyro.deterministic("eta_true", x[:, 1])
+        
         # Evaluate the likelihood
-        numpyro.sample("obs", TFRLikelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, eta_true), obs=jnp.array([m_obs, eta_obs]))
+        numpyro.sample("obs", TFRLikelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, hyper_mean_eta, hyper_sigma_eta, m_true, eta_true), obs=jnp.array([m_obs, eta_obs]))
 
 
     class TFRLikelihood(dist.Distribution):
         support = dist.constraints.real
 
-        def __init__(self, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, eta_true):
+        def __init__(self, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, hyper_mean_eta, hyper_sigma_eta, m_true, eta_true):
-            self.alpha, self.a_TFR, self.b_TFR, self.sigma_TFR, self.sigma_v, self.eta_true = dist.util.promote_shapes(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, eta_true)
+            self.alpha, self.a_TFR, self.b_TFR, self.sigma_TFR, self.sigma_v, self.hyper_mean_eta, self.hyper_sigma_eta, self.m_true, self.eta_true = dist.util.promote_shapes(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, hyper_mean_eta, hyper_sigma_eta, m_true, eta_true)
             batch_shape = lax.broadcast_shapes(
                     jnp.shape(alpha),
                     jnp.shape(a_TFR),
                     jnp.shape(b_TFR),
                     jnp.shape(sigma_TFR),
                     jnp.shape(sigma_v),
-                    # jnp.shape(m_true),
+                    jnp.shape(hyper_mean_eta),
+                    jnp.shape(hyper_sigma_eta),
+                    jnp.shape(m_true),
                     jnp.shape(eta_true),
                     )
             super(TFRLikelihood, self).__init__(batch_shape = batch_shape)
@@ -636,7 +754,7 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, omega_m, h, L,
 
         def log_prob(self, value):
             loglike =  likelihood(self.alpha, self.a_TFR, self.b_TFR, self.sigma_TFR, self.sigma_v, 
-                             m_true, self.eta_true,
+                             self.m_true, self.eta_true,
                              dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
                              czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh)
             return loglike
@@ -644,6 +762,8 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, omega_m, h, L,
     rng_key = random.PRNGKey(6)
     rng_key, rng_key_ = random.split(rng_key)
     values = initial
+    values['true_vars'] = jnp.array([m_obs, eta_obs]).T
+    values['L_corr'] = jnp.identity(2)
     myprint('Preparing MCMC kernel')
     kernel = numpyro.infer.NUTS(tfr_model,
             init_strategy=numpyro.infer.initialization.init_to_value(values=initial)
@@ -656,14 +776,28 @@ def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, omega_m, h, L,
     return mcmc
 
 
-def process_mcmc_run(mcmc, param_labels, truths, obs):
+def process_mcmc_run(mcmc, param_labels, truths, true_vars):
+    """
+    Make summary plots from the MCMC and save these to file
+    
+    Args:
+        - mcmc 
+        - param_labels (list[str]): Names of the parameters to plot
+        - truths (list[float]): True values of the parameters to plot. If unknown, then entry is None
+        - true_vars (dict): True values of the observables to compare against inferred ones
+    """
     
     # Convert samples into a single array
     samples = mcmc.get_samples()
     
     samps = jnp.empty((len(samples[param_labels[0]]), len(param_labels)))
     for i, p in enumerate(param_labels):
-        samps = samps.at[:,i].set(samples[p])
+        if p == 'hyper_corr':
+            L_corr = samples['L_corr']
+            corr_matrix = jnp.matmul(L_corr, jnp.transpose(L_corr, (0, 2, 1)))
+            samps = samps.at[:,i].set(corr_matrix[:,0,1])
+        else:
+            samps = samps.at[:,i].set(samples[p])
 
     # Trace plot of non-redshift quantities
     fig1, axs1 = plt.subplots(samps.shape[1], 1, figsize=(6,3*samps.shape[1]), sharex=True)
@@ -671,13 +805,14 @@ def process_mcmc_run(mcmc, param_labels, truths, obs):
     for i in range(samps.shape[1]):
         axs1[i].plot(samps[:,i])
         axs1[i].set_ylabel(param_labels[i])
-        axs1[i].axhline(truths[i], color='k')
+        if truths[i] is not None:
+            axs1[i].axhline(truths[i], color='k')
     axs1[-1].set_xlabel('Step Number')
     fig1.tight_layout()
     fig1.savefig('trace.png')
     
     # Corner plot
-    fig2, axs2 = plt.subplots(samps.shape[1], samps.shape[1], figsize=(10,10))
+    fig2, axs2 = plt.subplots(samps.shape[1], samps.shape[1], figsize=(15,15))
     corner.corner(
         np.array(samps),
         labels=param_labels,
@@ -690,7 +825,7 @@ def process_mcmc_run(mcmc, param_labels, truths, obs):
     for var in ['eta', 'm']:
         vname = var + '_true'
         if vname in samples.keys():
-            xtrue = obs[var]
+            xtrue = true_vars[var]
             xpred_median = np.median(samples[vname], axis=0)
             xpred_plus = np.percentile(samples[vname], 84, axis=0) - xpred_median
             xpred_minus = xpred_median - np.percentile(samples[vname], 16, axis=0)
@@ -723,7 +858,7 @@ def main():
     myprint('Beginning')
     
     # Get some parameters from the data
-    sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma = estimate_data_parameters()
+    sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma, hyper_m_sigma = estimate_data_parameters()
     
     # Other parameters to use
     L = 500.0
@@ -751,7 +886,8 @@ def main():
         'a_TFR': [-25, -20],
         'b_TFR': [-10, -5],
         'hyper_mean_eta': [hyper_eta_mu - 0.5, hyper_eta_mu + 0.5],
-        # 'hyper_mean_m':[mthresh - 5, mthresh + 5]
+        'hyper_mean_m':[mthresh - 5, mthresh + 5],
+        'sigma_v': [50, 300],
     }
     initial = {
         'alpha': alpha,
@@ -759,7 +895,8 @@ def main():
         'b_TFR': b_TFR,
         'hyper_mean_eta': hyper_eta_mu,
         'hyper_sigma_eta': hyper_eta_sigma,
-        # 'hyper_mean_m': mthresh,
+        'hyper_mean_m': mthresh,
+        'hyper_sigma_m': hyper_m_sigma,
         'sigma_TFR': sigma_TFR,
         'sigma_v': sigma_v,
     }
@@ -787,15 +924,12 @@ def main():
     
     # Run a MCMC
     mcmc = run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, cpar.omega_m, cpar.h, L, xmin, interp_order, bias_epsilon,
-             czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh,
+             czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh,)
-                   m_true)
     
-    param_labels = ['alpha', 'a_TFR', 'b_TFR', 'sigma_TFR', 'sigma_v', 'hyper_mean_eta', 'hyper_sigma_eta']
+    param_labels = ['alpha', 'a_TFR', 'b_TFR', 'sigma_TFR', 'sigma_v', 'hyper_mean_m', 'hyper_sigma_m', 'hyper_mean_eta', 'hyper_sigma_eta', 'hyper_corr']
-    truths = [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, hyper_eta_mu, hyper_eta_sigma]
+    truths = [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, None, None, hyper_eta_mu, hyper_eta_sigma, None]
-    param_labels = ['hyper_mean_eta', 'hyper_sigma_eta']
+    true_vars = {'m':m_true, 'eta':eta_true}
-    truths = [hyper_eta_mu, hyper_eta_sigma]
+    process_mcmc_run(mcmc, param_labels, truths, true_vars)
-    obs = {'m':m_obs, 'eta':eta_obs}
-    process_mcmc_run(mcmc, param_labels, truths, obs)
 
 if __name__ == "__main__":
     main()
@@ -803,8 +937,7 @@ if __name__ == "__main__":
 """
 TO DO
 
-- Fails to initialise currently when loglike includes the BORG term
+- Reinsert magnitude cut
-- Runs MCMC with this likelihood
 - Add bulk velocity
 - Deal with case where sigma_eta and sigma_m could be floats vs arrays