borg_velocity/tests/sn_inference.py

import aquila_borg as borg
import pandas as pd
import linecache
import numpy as np
from astropy.coordinates import SkyCoord
import astropy.units as apu
import jax.numpy as jnp
import jax

import corner
import matplotlib.pyplot as plt

import borg_velocity.poisson_process as poisson_process
import borg_velocity.projection as projection
import borg_velocity.utils as utils

import numpyro
import numpyro.distributions as dist
from jax import lax, random

from tfr_inference import get_fields, generateMBData

# Output stream management
cons = borg.console()
myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x))


def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha,
                a_tripp, b_tripp, M_SN, sigma_SN, sigma_m, sigma_stretch, sigma_c, 
                hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_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
        - 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^#
        
    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_mobs (np.ndarrary): Observed apparent magnitudes of the tracers (shape = (Nt,))
        
        - all_xtrue (np.ndarrary): True comoving coordinates of the tracers (Mpc/h) (shape = (3, Nt))
        - vbulk (np.ndarray): The bulk velocity of the box (km/s)
    
    """
    
    # Initialize lists to store valid positions and corresponding sig_mu values
    all_xtrue = np.empty((3, Nt))
    all_mtrue = np.empty(Nt)
    all_stretchtrue = np.empty(Nt)
    all_ctrue = np.empty(Nt)
    all_mobs = np.empty(Nt)
    all_stretchobs = np.empty(Nt)
    all_cobs = np.empty(Nt)
    all_RA = np.empty(Nt)
    all_Dec = np.empty(Nt)

    # Counter for accepted positions
    accepted_count = 0
     
    # Bias model
    phi = (1. + dens + bias_epsilon) ** alpha
    
    # Only use centre of box
    x = np.linspace(xmin, xmin + L, dens.shape[0]+1)
    i0 = np.argmin(np.abs(x + Rmax))
    i1 = np.argmin(np.abs(x - Rmax))
    L_small = x[i1] - x[i0]
    xmin_small = x[i0]
    phi_small = phi[i0:i1, i0:i1, i0:i1]
    
    # Loop until we have Nt valid positions
    while accepted_count < Nt:
    
        # Generate positions (comoving)
        xtrue = poisson_process.sample_3d(phi_small, Nt, L_small, (xmin_small, xmin_small, xmin_small))
        
        # Convert to RA, Dec, Distance (comoving)
        rtrue = np.sqrt(np.sum(xtrue** 2, axis=0))  # Mpc/h
        c = SkyCoord(x=xtrue[0], y=xtrue[1], z=xtrue[2], representation_type='cartesian')
        RA = c.spherical.lon.degree
        Dec = c.spherical.lat.degree
        r_hat = np.array(SkyCoord(ra=RA*apu.deg, dec=Dec*apu.deg).cartesian.xyz)
        
        # Compute cosmological redshift
        zcosmo = utils.z_cos(rtrue, cpar.omega_m)
        
        # Compute luminosity distance
        # DO I NEED TO DO /h???
        dL = (1 + zcosmo) * rtrue / cpar.h  # Mpc
        
        # Compute true distance modulus
        mutrue = 5 * np.log10(dL) + 25
        
        # Sample true stretch and colour (c) from its prior
        stretchtrue = hyper_stretch_mu + hyper_stretch_sigma * np.random.randn(Nt)
        ctrue = hyper_c_mu + hyper_c_sigma * np.random.randn(Nt)
        
        # Obtain muSN from mutrue using the intrinsic scatter
        muSN = mutrue + sigma_SN * np.random.randn(Nt)
        
        # Obtain apparent magnitude from the TFR
        mtrue = muSN - (a_tripp * stretchtrue - b_tripp * ctrue) + M_SN
        
        # Scatter true observed apparent magnitudes and linewidths
        mobs = mtrue + sigma_m * np.random.randn(Nt)
        stretchobs = stretchtrue + sigma_stretch * np.random.randn(Nt)
        cobs = ctrue + sigma_c * np.random.randn(Nt)
        
        # Apply apparement magnitude cut
        m = np.ones(mobs.shape, dtype=bool)
        mtrue = mtrue[m]
        stretchtrue = stretchtrue[m]
        ctrue = ctrue[m]
        mobs = mobs[m]
        stretchobs = stretchobs[m]
        cobs = cobs[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
        imin = accepted_count
        imax = accepted_count + selected_count
        all_xtrue[:,imin:imax] = xtrue[:,:selected_count]
        all_mtrue[imin:imax] = mtrue[:selected_count]
        all_stretchtrue[imin:imax] = stretchtrue[:selected_count]
        all_ctrue[imin:imax] = ctrue[:selected_count]
        all_mobs[imin:imax] = mobs[:selected_count]
        all_stretchobs[imin:imax] = stretchobs[:selected_count]
        all_cobs[imin:imax] = cobs[:selected_count]
        all_RA[imin:imax] = RA[:selected_count]
        all_Dec[imin:imax] = Dec[:selected_count]
        
        # Update the accepted count
        accepted_count += selected_count
        
        myprint(f'\tMade {accepted_count} of {Nt}')
        
    # Obtain a bulk velocity
    vhat = np.random.randn(3)
    vhat = vhat / np.linalg.norm(vhat)
    vbulk = np.random.randn() * utils.get_sigma_bulk(L, cpar)
    vbulk = vhat * vbulk
        
    # Get the radial component of the peculiar velocity at the positions of the objects
    myprint('Obtaining peculiar velocities')
    tracer_vel = projection.interp_field(
                            vel,
                            np.expand_dims(all_xtrue, axis=2),
                            L,
                            np.array([xmin, xmin, xmin]),
                            interp_order
                            ) # km/s
    myprint('Adding bulk velocity')
    tracer_vel = tracer_vel + vbulk[:,None,None]
    myprint('Radial projection')
    vr_true = np.squeeze(projection.project_radial(
                            tracer_vel,
                            np.expand_dims(all_xtrue, axis=2),
                            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)
    czCMB = ((1 + zcosmo) * (1 + vr_noised / utils.speed_of_light) - 1) * utils.speed_of_light
    
    return all_RA, all_Dec, czCMB, all_mtrue, all_stretchtrue, all_ctrue, all_mobs, all_stretchobs, all_cobs, all_xtrue, vbulk


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'
    
    # Get header
    columns = ['SN'] + linecache.getline(fname, 6).strip().split()[1:]
    df = pd.read_csv(fname, sep="\s+", skipinitialspace=True, skiprows=7, names=columns)

    zCMB = df['zCMB']
    m = df['mB']
    m_err = df['mBERR']

    x1 = df['x1']
    hyper_stretch_mu = np.median(x1)
    hyper_stretch_sigma = (np.percentile(x1, 84) - np.percentile(x1, 16)) / 2

    c = df['c']
    hyper_c_mu = np.median(c)
    hyper_c_sigma = (np.percentile(c, 84) - np.percentile(c, 16)) / 2

    sigma_m = np.median(df['mBERR'])
    sigma_stretch = np.median(df['x1ERR'])
    sigma_c = np.median(df['cERR'])
    
    return sigma_m, sigma_stretch, sigma_c, hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_sigma



def likelihood_vel(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, MB_pos):
    """
    Evaluate the terms in the likelihood from the velocity and malmquist bias
    
    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))
        - 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,))
        - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h).
            The shape is (3, Nt, Nsig)
        
    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 
    mutripp = m_true + a_tripp * stretch_true - b_tripp * c_true - M_SN
    d2 = ((mutrue - mutripp[:,None]) / sigma_SN) ** 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_vel = tracer_vel + jnp.squeeze(vbulk)[...,None,None]
    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 = lkl_ind.sum()
    
    return loglike


def likelihood_stretch(stretch_true, stretch_obs, sigma_stretch):
    """
    Evaluate the terms in the likelihood from stretch
    
    Args:
        - stretch_true (np.ndarray): True stretch of the tracers (shape = (Nt,))
        - stretch_obs (np.ndarray): Observed stretch of the tracers (shape = (Nt,))
        - sigma_stretch (float or np.ndarray): Uncertainty on the stretch measurements
        
    Returns:
        - loglike (float): The log-likelihood of the data
    """
    
    Nt = stretch_obs.shape[0]
    norm = jnp.where(
        jnp.ndim(sigma_stretch) == 0,
        Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_stretch ** 2),
        jnp.sum(0.5 * jnp.log(2 * jnp.pi * sigma_stretch ** 2))
    )
    loglike = - (
        0.5 * jnp.sum((stretch_obs - stretch_true) ** 2 / sigma_stretch ** 2) 
        + norm
    )
    
    return loglike


def likelihood_c(c_true, c_obs, sigma_c):
    """
    Evaluate the terms in the likelihood from colour
    
    Args:
        - c_true (np.ndarray): True colours of the tracers (shape = (Nt,))
        - c_obs (np.ndarray): Observed colours of the tracers (shape = (Nt,))
        - sigma_c (float or np.ndarray): Uncertainty on the colours measurements
        
    Returns:
        - loglike (float): The log-likelihood of the data
    """
    
    Nt = c_obs.shape[0]
    norm = jnp.where(
        jnp.ndim(sigma_c) == 0,
        Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_c ** 2),
        jnp.sum(0.5 * jnp.log(2 * jnp.pi * sigma_c ** 2))
    )
    loglike = - (
        0.5 * jnp.sum((c_obs - c_true) ** 2 / sigma_c ** 2) 
        + norm
    )
    
    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 or np.ndarray): Uncertainty on the apparent magnitude measurements
        
    Returns:
        - loglike (float): The log-likelihood of the data
    """
    
    Nt = m_obs.shape[0]
    norm = jnp.where(
        jnp.ndim(sigma_m) == 0,
        Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2),
        jnp.sum(0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2))
    )
    loglike = - (
        0.5 * jnp.sum((m_obs - m_true) ** 2 / sigma_m ** 2) 
        + norm
    )
    
    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):
    """
    Evaluate the likelihood for SN sample
    
    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))
        - 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,))
        - 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 or np.ndarray): Uncertainty on the apparent magnitude measurements
        - sigma_stretch (float or np.ndarray): Uncertainty on the stretch measurements
        - sigma_c (float or np.ndarray): 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)
        
    Returns:
        - loglike (float): The log-likelihood of the data
    """

    
    loglike_vel = likelihood_vel(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, 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_m
    
    return loglike


def test_likelihood_scan(prior, 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,
                         czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos):
    """
    Plot likelihood as we scan through the paramaters [alpha, a_tripp, b_tripp, M_SN, sigma_SN, 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_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
        - 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^#
        - 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 or np.ndarray): Uncertainty on the apparent magnitude measurements
        - sigma_stretch (float or np.ndarray): Uncertainty on the stretch measurements
        - sigma_c (float or np.ndarray): 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)
        
    """
    
    
    pars = [alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v]
    par_names = ['alpha', 'a_tripp', 'b_tripp', 'M_SN', 'sigma_SN', 'sigma_v']
    
    orig_ll = - likelihood(*pars, m_true, stretch_true, c_true, vbulk,
                         dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
                         czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos)
    
    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, stretch_true, c_true, vbulk,
                         dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
                         czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos)
        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'sn_likelihood_scan_{name}.png')
        fig = plt.gcf()
        plt.clf()
        plt.close(fig)
    
    return


def 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,):
    """
    Run MCMC over the model parameters
    
    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 (dict): Upper and lower bounds for a uniform prior for the parameters
        - 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
        - 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^#
        - 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 or np.ndarray): Uncertainty on the apparent magnitude measurements
        - sigma_stretch (float or np.ndarray): Uncertainty on the stretch measurements
        - sigma_c (float or np.ndarray): 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)
        
    Returns:
        - mcmc (numpyro.infer.MCMC): MCMC object which has been run
        
    """
    
    Nt = stretch_obs.shape[0]
    omega_m = cpar.omega_m
    h = cpar.h
    sigma_bulk = utils.get_sigma_bulk(L, cpar)
    
    
    def sn_model():

        alpha = numpyro.sample("alpha", dist.Uniform(*prior['alpha']))
        a_tripp =  numpyro.sample("a_tripp", dist.Uniform(*prior['a_tripp']))
        b_tripp =  numpyro.sample("b_tripp", dist.Uniform(*prior['b_tripp']))
        M_SN = numpyro.sample("M_SN", dist.Uniform(*prior['M_SN']))
        sigma_SN = numpyro.sample("sigma_SN", dist.HalfCauchy(1.0))
        sigma_v =  numpyro.sample("sigma_v", dist.Uniform(*prior['sigma_v']))
    
        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_stretch = numpyro.sample("hyper_mean_stretch", dist.Uniform(*prior['hyper_mean_stretch']))
        hyper_sigma_stretch = numpyro.sample("hyper_sigma_stretch", dist.HalfCauchy(1.0)) # Equivalent to 1/sigma prior
        hyper_mean_c = numpyro.sample("hyper_mean_c", dist.Uniform(*prior['hyper_mean_c']))
        hyper_sigma_c = numpyro.sample("hyper_sigma_c", dist.HalfCauchy(1.0)) # Equivalent to 1/sigma prior
        
        # Sample correlation matrix using LKJ prior
        L_corr = numpyro.sample("L_corr", dist.LKJCholesky(3, 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_stretch, hyper_mean_c])
        hyper_sigma = jnp.array([hyper_sigma_m, hyper_sigma_stretch, hyper_sigma_c])
        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])
        stretch_true = numpyro.deterministic("stretch_true", x[:, 1])
        c_true = numpyro.deterministic("c_true", x[:, 2])
        
        # Sample bulk velocity
        vbulk_x = numpyro.sample("vbulk_x", dist.Normal(0, sigma_bulk / jnp.sqrt(3)))
        vbulk_y = numpyro.sample("vbulk_y", dist.Normal(0, sigma_bulk / jnp.sqrt(3)))
        vbulk_z = numpyro.sample("vbulk_z", dist.Normal(0, sigma_bulk / jnp.sqrt(3)))
        
        # Evaluate the likelihood
        numpyro.sample("obs", SNLikelihood(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk_x, vbulk_y, vbulk_z), obs=jnp.array([m_obs, stretch_obs, c_obs]))
        
        
    class SNLikelihood(dist.Distribution):
        support = dist.constraints.real

        def __init__(self, alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk_x, vbulk_y, vbulk_z):
            self.alpha, self.a_tripp, self.b_tripp, self.M_SN, self.sigma_SN, self.sigma_v, self.m_true, self.stretch_true, self.c_true, self.vbulk_x, self.vbulk_y, self.vbulk_z = dist.util.promote_shapes(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk_x, vbulk_y, vbulk_z)
            batch_shape = lax.broadcast_shapes(
                    jnp.shape(alpha),
                    jnp.shape(a_tripp),
                    jnp.shape(b_tripp),
                    jnp.shape(M_SN),
                    jnp.shape(sigma_SN),
                    jnp.shape(sigma_v),
                    jnp.shape(m_true),
                    jnp.shape(stretch_true),
                    jnp.shape(c_true),
                    jnp.shape(vbulk_x),
                    jnp.shape(vbulk_y),
                    jnp.shape(vbulk_z),
                    )
            super(SNLikelihood, self).__init__(batch_shape = batch_shape)

        def sample(self, key, sample_shape=()):
            raise NotImplementedError

        def log_prob(self, value):
            vbulk = jnp.array([self.vbulk_x, self.vbulk_y, self.vbulk_z])
            loglike =  likelihood(self.alpha, self.a_tripp, self.b_tripp, self.M_SN, self.sigma_SN, self.sigma_v, 
                             self.m_true, self.stretch_true, self.c_true, vbulk,
                             dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon,
                             czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos)
            return loglike
        
    rng_key = random.PRNGKey(6)
    rng_key, rng_key_ = random.split(rng_key)
    values = initial
    values['true_vars'] = jnp.array([m_obs, stretch_obs, c_obs]).T
    values['L_corr'] = jnp.identity(3)
    values['vbulk_x'] = 0.
    values['vbulk_y'] = 0.
    values['vbulk_z'] = 0.
    myprint('Preparing MCMC kernel')
    kernel = numpyro.infer.NUTS(sn_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, true_vars):
    """
    Make summary plots from the MCMC and save these to file
    
    Args:
        - mcmc (numpyro.infer.MCMC): MCMC object which has been run
        - 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):
        if p.startswith('hyper_corr'):
            L_corr = samples['L_corr']
            corr_matrix = jnp.matmul(L_corr, jnp.transpose(L_corr, (0, 2, 1)))
            if p == 'hyper_corr_mx':
                samps = samps.at[:,i].set(corr_matrix[:,0,1])
            elif p == 'hyper_corr_mc':
                samps = samps.at[:,i].set(corr_matrix[:,0,2])
            elif p == 'hyper_corr_xc':
                samps = samps.at[:,i].set(corr_matrix[:,1,2])
            else:
                raise NotImplementedError
        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)
    axs1 = np.atleast_1d(axs1)
    for i in range(samps.shape[1]):
        axs1[i].plot(samps[:,i])
        axs1[i].set_ylabel(param_labels[i])
        if truths[i] is not None:
            axs1[i].axhline(truths[i], color='k')
    axs1[-1].set_xlabel('Step Number')
    fig1.tight_layout()
    fig1.savefig('sn_trace.png')
    
    # Corner plot
    fig2, axs2 = plt.subplots(samps.shape[1], samps.shape[1], figsize=(25,25))
    corner.corner(
        np.array(samps),
        labels=param_labels,
        fig=fig2,
        truths=truths
    )
    fig2.savefig('sn_corner.png')
    
    # True vs predicted
    for var in ['stretch', 'c', 'm']:
        vname = var + '_true'
        if vname in samples.keys():
            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)
            
            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('Predicted')
            axs3[1].set_ylabel('Predicted - True')
            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'sn_true_predicted_{var}.png')
        
    return

    

def main():
    
    myprint('Beginning')
    
    sigma_m, sigma_stretch, sigma_c, hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_sigma =  estimate_data_parameters()

    # Other parameters to use
    L = 500.0
    N = 64
    xmin = -L/2
    R_lim = L / 2
    Rmax = 100
    Nt = 100
    alpha = 1.4
    sigma_v = 150
    interp_order = 1
    bias_epsilon = 1.e-7
    Nint_points = 201 
    Nsig = 10
    frac_sigma_r = 0.07  # WANT A BETTER WAY OF DOING THIS - ESTIMATE THROUGH SIGMAS FROM Tripp formula
    
    # These values are from Table 6 of Boruah et al. 2020
    a_tripp = 0.140
    b_tripp = 2.78
    M_SN =  - 18.558
    sigma_SN = 0.082
    
    num_warmup = 1000
    num_samples = 2000
    
    # Make mock
    np.random.seed(123)
    cpar, dens, vel = get_fields(L, N, xmin)
    RA, Dec, czCMB, m_true, stretch_true, c_true, m_obs, stretch_obs, c_obs, xtrue, vbulk = create_mock(
                Nt, L, xmin, cpar, dens, vel, Rmax, alpha,
                a_tripp, b_tripp, M_SN, sigma_SN, sigma_m, sigma_stretch, sigma_c, 
                hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_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)
    
    initial = {
        'a_tripp': a_tripp,
        'b_tripp': b_tripp,
        'M_SN': M_SN,
        'sigma_SN': sigma_SN,
        'sigma_v': sigma_v,
        'hyper_mean_stretch': hyper_stretch_mu,
        'hyper_sigma_stretch': hyper_stretch_sigma,
        'hyper_mean_c': hyper_c_mu,
        'hyper_sigma_c': hyper_c_sigma,
        'hyper_mean_m': np.median(m_obs),
        'hyper_sigma_m': (np.percentile(m_obs, 84) - np.percentile(m_obs, 16)) / 2,
    }
    prior = {
        'alpha': [0.5, 6],
        'a_tripp': [0.01, 0.2],
        'b_tripp': [2.5, 4.5],
        'M_SN': [-19.5, -17.5],
        'hyper_mean_stretch': [hyper_stretch_mu - hyper_stretch_sigma, hyper_stretch_mu + hyper_stretch_sigma],
        'hyper_mean_c':[hyper_c_mu - hyper_c_sigma, hyper_c_mu + hyper_c_sigma],
        'hyper_mean_m':[initial['hyper_mean_m'] - initial['hyper_sigma_m'], initial['hyper_mean_m'] + initial['hyper_sigma_m']],
        'sigma_v': [10, 3000],
    }
    
    # Test likelihood
    loglike = likelihood(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk,
                         dens, vel, cpar.omega_m, cpar.h, L, xmin, interp_order, bias_epsilon,
                         czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos)
    myprint(f'loglike {loglike}')
    
    # Scan over parameters to make plots verifying behaviour
    test_likelihood_scan(prior, alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk,
                         dens, vel, cpar.omega_m, cpar.h, L, xmin, interp_order, bias_epsilon,
                         czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos)
    
    
    # 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)
    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_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, 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)
    
    return


if __name__ == "__main__":
    main()