Add MCMC to TFR testing

tests/tfr_inference.py

@@ -3,18 +3,20 @@ import numpy as np
 from astropy.coordinates import SkyCoord
 import astropy.units as apu
 import astropy.constants
-from astropy.cosmology import LambdaCDM
 import pandas as pd
-from scipy.interpolate import interp1d
 import jax.numpy as jnp
 import jax.scipy.special
 
+import matplotlib.pyplot as plt
+import corner
+
 import numpyro
 import numpyro.distributions as dist
-from jax import lax
+from jax import lax, random
 
 import borg_velocity.poisson_process as poisson_process
 import borg_velocity.projection as projection
+import borg_velocity.utils as utils
 
 # Output stream management
 cons = borg.console()
@@ -166,22 +168,16 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
     
     # Initialize lists to store valid positions and corresponding sig_mu values
     all_xtrue = np.empty((3, Nt))
+    all_mtrue = np.empty(Nt)
+    all_etatrue = np.empty(Nt)
     all_mobs = np.empty(Nt)
     all_etaobs = np.empty(Nt)
+    all_RA = np.empty(Nt)
+    all_Dec = np.empty(Nt)
 
     # Counter for accepted positions
     accepted_count = 0
-
-    # Cosmology object needed for z <-> r
-    cosmo = LambdaCDM(H0 = cpar.h * 100, Om0 = cpar.omega_m, Ode0 = cpar.omega_q)
-    
-    # Precompute redshift-distance mapping
-    z_grid = np.logspace(-4, -1, 10000)  # Covers z ~ 0 to 0.1
-    dL_grid = cosmo.luminosity_distance(z_grid).value  # Luminosity distances in Mpc
-
-    # Create an interpolation function: distance -> redshift
-    dL_to_z = interp1d(dL_grid, z_grid, kind="cubic", fill_value="extrapolate")
-    
+     
     # Bias model
     phi = (1. + dens + bias_epsilon) ** alpha
     
@@ -207,12 +203,11 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
         r_hat = np.array(SkyCoord(ra=RA*apu.deg, dec=Dec*apu.deg).cartesian.xyz)
         
         # Compute cosmological redshift
-        # zcosmo = z_at_value(cosmo.comoving_distance, rtrue * apu.Mpc / cpar.h).value
-        zcosmo = dL_to_z(rtrue / cpar.h)
+        zcosmo = utils.z_cos(rtrue, cpar.omega_m)
         
         # Compute luminosity distance
         # DO I NEED TO DO /h???
-        dL = (1 + zcosmo) * rtrue / cpar.h  # Mpc/h
+        dL = (1 + zcosmo) * rtrue / cpar.h  # Mpc
         
         # Compute true distance modulus
         mutrue = 5 * np.log10(dL) + 25
@@ -232,18 +227,28 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
         
         # Apply apparement magnitude cut
         m = mobs <= mthresh
+        mtrue = mtrue[m]
+        etatrue = etatrue[m]
         mobs = mobs[m]
         etaobs = etaobs[m]
         xtrue = xtrue[:,m]
+        RA = RA[m]
+        Dec = Dec[m]
         
         # Calculate how many valid positions we need to reach Nt
         remaining_needed = Nt - accepted_count
         selected_count = min(xtrue.shape[1], remaining_needed)
         
         # Append only the needed number of valid positions
-        all_xtrue[:,accepted_count:accepted_count+selected_count] = xtrue[:,:selected_count]
-        all_mobs = mobs[:selected_count]
-        all_etaobs = etaobs[:selected_count]
+        imin = accepted_count
+        imax = accepted_count + selected_count
+        all_xtrue[:,imin:imax] = xtrue[:,:selected_count]
+        all_mtrue[imin:imax] = mtrue[:selected_count]
+        all_etatrue[imin:imax] = etatrue[:selected_count]
+        all_mobs[imin:imax] = mobs[:selected_count]
+        all_etaobs[imin:imax] = etaobs[:selected_count]
+        all_RA[imin:imax] = RA[:selected_count]
+        all_Dec[imin:imax] = Dec[:selected_count]
         
         # Update the accepted count
         accepted_count += selected_count
@@ -266,11 +271,15 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
                             np.zeros(3,)
                             )) # km/s
     
+    # Recompute cosmological redshift
+    rtrue = jnp.sqrt(jnp.sum(all_xtrue ** 2, axis=0))
+    zcosmo = utils.z_cos(rtrue, cpar.omega_m)
+    
     # Obtain total redshift
     vr_noised = vr_true + sigma_v * np.random.randn(Nt)
-    zCMB = (1 + zcosmo) * (1 + vr_noised / astropy.constants.c.to('km/s').value) - 1
+    czCMB = ((1 + zcosmo) * (1 + vr_noised / utils.speed_of_light) - 1) * utils.speed_of_light
     
-    return zCMB, all_mobs, all_etaobs, all_xtrue
+    return all_RA, all_Dec, czCMB, all_mtrue, all_etatrue, all_mobs, all_etaobs, all_xtrue
 
 
 def estimate_data_parameters():
@@ -305,66 +314,414 @@ def estimate_data_parameters():
     return sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma
 
 
-
-def likelihood(a_TFR, b_TFR, sigma_TFR, eta_true, m_true):
+def generateMBData(RA, Dec, cz_obs, L, N, R_lim, Nsig, Nint_points, sigma_v, frac_sigma_r):
+    """
+    Generate points along the line of sight of each tracer to enable marginalisation
+    over distance uncertainty. The central distance is given such that the observed
+    redshift equals the cosmological redshift at this distance. The range is then
+    +/- Nsig * sig, where 
+    sig^2 = (sig_v/100)^2 + sig_r^2
+    and sig_v is the velocity uncertainty in km/s
     
-    loglike = 0
+    Args:
+        - RA (np.ndarray): Right Ascension (degrees) of the tracers (shape = (Nt,))
+        - Dec (np.ndarray): Delination (degrees) of the tracers (shape = (Nt,))
+        - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,))
+        - L (float): Box length (Mpc/h)
+        - N (int): Number of grid cells per side
+        - R_lim (float): Maximum allowed (true) comoving distance of a tracer (Mpc/h)
+        - Nsig (float): ???
+        - Nint_points (int): Number of radii over which to integrate the likelihood
+        - sigma_v (float): Uncertainty on the velocity field (km/s)
+        - frac_sigma_r (float): An estimate of the fractional uncertainty on the positions of tracers
+        
+    Returns:
+        - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h).
+            The shape is (3, Nt, Nsig)
+    
+    """
+    
+    myprint(f"Making MB data")
+    
+    # Convert RA, DEC to radial vector
+    r_hat = np.array(SkyCoord(ra=RA*apu.deg, dec=Dec*apu.deg).cartesian.xyz).T
+
+    # Get min and max distance to integrate over
+    # cz = 100 h r, so sigma_v corresponds to a sigma_r of ~ sigma_v / 100
+    robs = cz_obs / 100 
+    sigr = np.sqrt((sigma_v / 100) ** 2 + (frac_sigma_r * robs)**2)
+    rmin = robs - Nsig * sigr
+    rmin = rmin.at[rmin <= 0].set(L / N / 100.)
+    rmax = robs + Nsig * sigr
+    rmax = rmax.at[rmax > R_lim].set(R_lim)
+
+    # Compute coordinates of integration points
+    r_integration = np.linspace(rmin, rmax, Nint_points)
+    MB_pos = np.expand_dims(r_integration, axis=2) * r_hat[None,...]
+    MB_pos = jnp.transpose(MB_pos, (2, 1, 0))
+    
+    return MB_pos
+
+
+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 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
+        
+    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))
+    
+    # p_r = r^2 n(r) N(mutrue; muTFR, sigmaTFR)
+    # Multiply by arbitrary number for numerical stability (cancels in p_r / p_r_norm)
+    number_density = projection.interp_field(
+                                dens,
+                                MB_pos,
+                                L,
+                                jnp.array([xmin, xmin, xmin]),
+                                interp_order,
+                                use_jitted=True,
+                            )
+    number_density = jax.nn.relu(1. + number_density)
+    number_density = jnp.power(number_density + bias_epsilon, alpha)
+    zcosmo = utils.z_cos(r, omega_m)
+    mutrue = 5 * jnp.log10((1 + zcosmo) * r / h) + 25 
+    muTFR = m_true - (a_TFR + b_TFR * eta_true)
+    d2 = ((mutrue - muTFR[:,None]) / sigma_TFR) ** 2
+    best = jnp.amin(jnp.abs(d2), axis=1)
+    d2 = d2 - jnp.expand_dims(jnp.nanmin(d2, axis=1), axis=1)
+    p_r = r ** 2 * jnp.exp(-0.5 * d2) * number_density
+    p_r_norm = jnp.expand_dims(jnp.trapezoid(p_r, r, axis=1), axis=1)    
+    
+    # Peculiar velocity term
+    tracer_vel = projection.interp_field(
+                                vel,
+                                MB_pos,
+                                L,
+                                jnp.array([xmin, xmin, xmin]),
+                                interp_order,
+                                use_jitted=True,
+                                )
+    tracer_vr = projection.project_radial(
+                    tracer_vel,
+                    MB_pos,
+                    jnp.zeros(3,)
+                    )
+    cz_pred = ((1 + zcosmo) * (1 + tracer_vr / utils.speed_of_light) - 1) * utils.speed_of_light
+    d2 = ((cz_pred - jnp.expand_dims(cz_obs, axis=1)) / sigma_v)**2
+    scale = jnp.nanmin(d2, axis=1)
+    d2 = d2 - jnp.expand_dims(scale, axis=1)
+    
+    # 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()
+    
+    Nt = m_obs.shape[0]
     
     # Apparent magnitude terms
     norm = 0.5 * (1 + jax.scipy.special.erf((mthresh - m_true) / (jnp.sqrt(2) * sigma_m)))
-    loglike += 
-        0.5 * jnp.sum((mobs - m_true) ** 2 / sigma_m ** 2) 
+    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)
+    )
         
     # Linewidth terms
-    loglike += 
-        0.5 * jnp.sum((etaobs - eta_true) ** 2 / sigma_eta ** 2) 
+    loglike_eta = (
+        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)
+    loglike = - (loglike_eta + loglike_m)
+    
+    return loglike
+
+
+def test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, 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):
+    """
+    Plot likelihood as we scan through the paramaters [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v]
+    to verify that the likelihood shape looks reasonable
+    
+    Args:
+        - prior (dict): Upper and lower bounds for a uniform prior for the parameters
+        - 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 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
         
-    # Peculiar velocity terms
+    """
+    
+    
+    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,
+                         dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
+                         czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh)
+    
+    for i, name in enumerate(par_names):
+        
+        myprint(f'Scanning {name}')
+        
+        if name in prior:
+            x = np.linspace(*prior[name], 20)
+        else:
+            pmin = pars[i] * 0.2
+            pmax = pars[i] * 2.0
+            x = np.linspace(pmin, pmax, 20)
+            
+        all_ll = np.empty(x.shape)
+        orig_x = pars[i]
+        for j, xx in enumerate(x):
+            pars[i] = xx
+            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
+            
+        plt.figure()
+        plt.plot(x, all_ll, '.')
+        plt.axvline(orig_x, ls='--', color='k')
+        plt.axhline(orig_ll, ls='--', color='k')
+        plt.xlabel(name)
+        plt.ylabel('Negative log-likelihood')
+        plt.savefig(f'likelihood_scan_{name}.png')
+        fig = plt.gcf()
+        plt.clf()
+        plt.close(fig)
     
     return
 
 
-def likelihood_model():
+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,
+            m_true):
+    """
+    Run MCMC over the model parameters
     
-    # TO DO: Sort out these priors
-    a_TFR =  numpyro.sample("a_TFR", dist.Uniform(*alpha_prior))
-    b_TFR =  numpyro.sample("b_TFR", dist.Uniform(*alpha_prior))
-    sigma_TFR = numpyro.sample("sigma_TFR", dist.Uniform(*alpha_prior))
-    
-    eta_true = numpyro.sample("eta_true", dist.Normal(mean, sigma), sample_shape=(Nt,))
-    m_true = numpyro.sample("m_true", dist.Normal(mean, sigma), sample_shape=(Nt,))
-    
-    numpyro.sample("obs", TFRLikelihood(a_TFR, b_TFR, sigma_TFR, eta_true, m_true))
-    
-    return
-
-
-class TFRLikelihood(dist.Distribution):
-    support = dist.constraints.real
-    
-    def __init__(self, a_TFR, b_TFR, sigma_TFR, eta_true, m_true):
-        self.a_TFR, self.b_TFR, self.sigma_TFR, self.eta_true, self.m_true = dist.util.promote_shapes(a_TFR, b_TFR, sigma_TFR, eta_true, m_true)
-        batch_shape = lax.broadcast_shapes(
-                jnp.shape(a_TFR),
-                jnp.shape(b_TFR),
-                jnp.shape(sigma_TFR),
-                jnp.shape(eta_true),
-                jnp.shape(m_true),
-                )
-        super(TFRLikelihood, self).__init__(batch_shape = batch_shape)
+    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
+        - initial
+        - 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 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
         
-    def sample(self, key, sample_shape=()):
-        raise NotImplementedError
+    """
     
-    def log_prov(self, value)
-        return likelihood(self.a_TFR, self.b_TFR, self.sigma_TFR, self.eta_true, self.m_true)
+    Nt = eta_obs.shape[0]
+    
+    def tfr_model():
+
+        alpha = numpyro.sample("alpha", dist.Uniform(*prior['alpha']))
+        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))
+
+#         # 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_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,))
+
+        # 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]))
+
+
+    class TFRLikelihood(dist.Distribution):
+        support = dist.constraints.real
+
+        def __init__(self, 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.eta_true = dist.util.promote_shapes(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, 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(eta_true),
+                    )
+            super(TFRLikelihood, self).__init__(batch_shape = batch_shape)
+
+        def sample(self, key, sample_shape=()):
+            raise NotImplementedError
+
+        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,
+                             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
+    
+    rng_key = random.PRNGKey(6)
+    rng_key, rng_key_ = random.split(rng_key)
+    values = initial
+    myprint('Preparing MCMC kernel')
+    kernel = numpyro.infer.NUTS(tfr_model,
+            init_strategy=numpyro.infer.initialization.init_to_value(values=initial)
+    )
+    mcmc = numpyro.infer.MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
+    myprint('Running MCMC')
+    mcmc.run(rng_key_)
+    mcmc.print_summary()
+        
+    return mcmc
+
+
+def process_mcmc_run(mcmc, param_labels, truths, obs):
+    
+    # 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])
+
+    # Trace plot of non-redshift quantities
+    fig1, axs1 = plt.subplots(samps.shape[1], 1, figsize=(6,3*samps.shape[1]), sharex=True)
+    axs1 = np.atleast_1d(axs1)
+    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')
+    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))
+    corner.corner(
+        np.array(samps),
+        labels=param_labels,
+        fig=fig2,
+        truths=truths
+    )
+    fig2.savefig('corner.png')
+    
+    # True vs predicted
+    for var in ['eta', 'm']:
+        vname = var + '_true'
+        if vname in samples.keys():
+            xtrue = obs[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)
+            
+            fig3, axs3 = plt.subplots(2, 1, figsize=(10,8), sharex=True)
+            plot_kwargs = {'fmt':'.', 'markersize':3, 'zorder':10,
+                     'capsize':1, 'elinewidth':1, 'alpha':1}
+            axs3[0].errorbar(xtrue, xpred_median, yerr=[xpred_minus, xpred_plus], **plot_kwargs)
+            axs3[1].errorbar(xtrue, xpred_median - xtrue, yerr=[xpred_minus, xpred_plus], **plot_kwargs)
+            axs3[1].set_xlabel('True')
+            axs3[0].set_ylabel('True')
+            axs3[1].set_ylabel('True - Predicted')
+            xlim = axs3[0].get_xlim()
+            ylim = axs3[0].get_ylim()
+            axs3[0].plot(xlim, xlim, color='k', zorder=0)
+            axs3[0].set_xlim(xlim)
+            axs3[0].set_ylim(ylim)
+            axs3[1].axhline(0, color='k', zorder=0)
+            fig3.suptitle(var)
+            fig3.align_labels()
+            fig3.tight_layout()
+            fig3.savefig(f'true_predicted_{var}.png')
+        
+        
+    return
 
 
 def main():
     
+    myprint('Beginning')
+    
     # Get some parameters from the data
     sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma = estimate_data_parameters()
     
@@ -372,22 +729,83 @@ def main():
     L = 500.0
     N = 64
     xmin = -L/2
+    R_lim = L / 2
     Rmax = 100
-    Nt = 2000
+    Nt = 100
     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
+    num_warmup = 1000
+    num_samples = 1000
     
+    prior = {
+        'alpha': [0.5, 2.5],
+        '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]
+    }
+    initial = {
+        'alpha': alpha,
+        'a_TFR': a_TFR,
+        'b_TFR': b_TFR,
+        'hyper_mean_eta': hyper_eta_mu,
+        'hyper_sigma_eta': hyper_eta_sigma,
+        # 'hyper_mean_m': mthresh,
+        'sigma_TFR': sigma_TFR,
+        'sigma_v': sigma_v,
+    }
+    
+    # Make mock
+    np.random.seed(123)
     cpar, dens, vel = get_fields(L, N, xmin)
-    
-    zCMB, mobs, etaobs, xtrue = create_mock(
+    RA, Dec, czCMB, m_true, eta_true, m_obs, eta_obs, xtrue = 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)
+            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)
     
+    # Test likelihood
+    loglike = likelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
+                         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)
+    myprint(f'loglike {loglike}')
+    
+    # Scan over parameters to make plots verifying behaviour
+    test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true,
+                         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)
+    
+    # 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,
+                   m_true)
+    
+    param_labels = ['alpha', 'a_TFR', 'b_TFR', 'sigma_TFR', 'sigma_v', 'hyper_mean_eta', 'hyper_sigma_eta']
+    truths = [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, hyper_eta_mu, hyper_eta_sigma]
+    param_labels = ['hyper_mean_eta', 'hyper_sigma_eta']
+    truths = [hyper_eta_mu, hyper_eta_sigma]
+    obs = {'m':m_obs, 'eta':eta_obs}
+    process_mcmc_run(mcmc, param_labels, truths, obs)
 
 if __name__ == "__main__":
-    main()
+    main()
+    
+"""
+TO DO
+
+- Fails to initialise currently when loglike includes the BORG term
+- Runs MCMC with this likelihood
+- Add bulk velocity
+- Deal with case where sigma_eta and sigma_m could be floats vs arrays
+
+"""