Add TFR and SN mock generation to BORG scripts

borg_velocity/borg_mock.py

@@ -0,0 +1,426 @@
+import aquila_borg as borg
+import numpy as np
+from astropy.coordinates import SkyCoord
+import astropy.units as apu
+import jax.numpy as jnp
+import jax
+import configparser
+import ast
+
+import borg_velocity.poisson_process as poisson_process
+import borg_velocity.projection as projection
+import borg_velocity.utils as utils
+from borg_velocity.utils import myprint
+
+
+def tfr_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_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))
+    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
+
+    # 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 linewidth (eta) from its prior
+        etatrue = hyper_eta_mu + hyper_eta_sigma * np.random.randn(Nt)
+        
+        # Obtain muTFR from mutrue using the intrinsic scatter
+        muTFR = mutrue + sigma_TFR * np.random.randn(Nt)
+        
+        # Obtain apparent magnitude from the TFR
+        mtrue = muTFR + (a_TFR + b_TFR * etatrue)
+        
+        # Scatter true observed apparent magnitudes and linewidths
+        mobs = mtrue + sigma_m * np.random.randn(Nt)
+        etaobs = etatrue + sigma_eta * np.random.randn(Nt)
+        
+        # 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
+        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
+        
+        myprint(f'\tMade {accepted_count} of {Nt}')
+        
+    # 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('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_etatrue, all_mobs, all_etaobs, all_xtrue
+
+
+
+def sn_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))
+    
+    """
+    
+    # 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)
+        
+        # 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}')
+
+    # 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('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
+
+
+
+def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
+    """
+    WRITE DOCSTRING
+    """
+    
+    myprint('Making mock from BORG')
+    
+    config = configparser.ConfigParser()
+    config.read(ini_file)
+    
+    # Ensure cosmology is correct
+    cosmo = utils.get_cosmopar(ini_file)
+    fwd_model.setCosmoParams(cosmo)
+    
+    # Run BORG density field
+    dens = np.zeros(fwd_model.getOutputBoxModel().N)
+    fwd_model.forwardModel_v2(s_hat)
+    fwd_model.getDensityFinal(dens)
+    state["BORG_final_density"][:] = dens
+    
+    # Get velocity field
+    cosmo = utils.get_cosmopar(ini_file)
+    vel = fwd_vel.getVelocityField()
+    
+    # Obtain a bulk velocity
+    vbulk = np.array(ast.literal_eval(config['model']['bulk_flow'])).reshape((3, 1, 1, 1))
+    
+    if seed is not None:
+        np.random.seed(seed)
+        
+    L = float(config['system']['L0'])
+    xmin = float(config['system']['corner0'])
+    nsamp = int(config['run']['nsamp'])
+    Rmax = float(config['mock']['R_max'])
+    sigma_v = float(config['model']['sig_v'])
+    interp_order = int(config['model']['interp_order'])
+    bias_epsilon = float(config['model']['bias_epsilon'])
+    
+    tracer_type = [None] * nsamp
+    RA = [None] * nsamp
+    Dec = [None] * nsamp
+    czCMB = [None] * nsamp
+    m_true = [None] * nsamp
+    m_obs = [None] * nsamp
+    eta_true = [None] * nsamp
+    eta_obs = [None] * nsamp
+    stretch_true = [None] * nsamp
+    stretch_obs = [None] * nsamp
+    c_true = [None] * nsamp
+    c_obs = [None] * nsamp
+    xtrue = [None] * nsamp
+        
+    for i in range(nsamp):
+        
+        myprint(f'Making mock for sample {i}')
+       
+        Nt = int(config[f'sample_{i}']['Nt'])
+        alpha = float(config[f'sample_{i}']['alpha'])
+        tracer_type[i] = config[f'sample_{i}']['tracer_type']
+        
+        myprint(f'Tracer type: {tracer_type[i]}')
+         
+        if tracer_type[i] == 'tfr':
+            mthresh = float(config[f'sample_{i}']['mthresh'])
+            a_TFR = float(config[f'sample_{i}']['a_tfr'])
+            b_TFR = float(config[f'sample_{i}']['b_tfr'])
+            sigma_TFR = float(config[f'sample_{i}']['sigma_tfr'])
+            sigma_m = float(config[f'sample_{i}']['sigma_m'])
+            sigma_eta = float(config[f'sample_{i}']['sigma_eta'])
+            hyper_eta_mu = float(config[f'sample_{i}']['hyper_eta_mu'])
+            hyper_eta_sigma = float(config[f'sample_{i}']['hyper_eta_sigma'])
+            mthresh = float(config[f'sample_{i}']['mthresh'])
+            RA[i], Dec[i], czCMB[i], m_true[i], eta_true[i], m_obs[i], eta_obs[i], xtrue[i] = tfr_create_mock(
+                    Nt, L, xmin, cosmo, dens, vel + vbulk, 
+                    Rmax, alpha, mthresh, 
+                    a_TFR, b_TFR, sigma_TFR, sigma_m, sigma_eta, 
+                    hyper_eta_mu, hyper_eta_sigma, sigma_v, 
+                    interp_order=interp_order, bias_epsilon=bias_epsilon)
+        elif tracer_type[i] == 'sn':
+            a_tripp = float(config[f'sample_{i}']['a_tripp'])
+            b_tripp = float(config[f'sample_{i}']['b_tripp'])
+            M_SN = float(config[f'sample_{i}']['m_sn'])
+            sigma_SN = float(config[f'sample_{i}']['sigma_sn'])
+            sigma_m = float(config[f'sample_{i}']['sigma_m'])
+            sigma_stretch = float(config[f'sample_{i}']['sigma_stretch'])
+            sigma_c = float(config[f'sample_{i}']['sigma_c'])
+            hyper_stretch_mu = float(config[f'sample_{i}']['hyper_stretch_mu'])
+            hyper_stretch_sigma = float(config[f'sample_{i}']['hyper_stretch_sigma'])
+            hyper_c_mu = float(config[f'sample_{i}']['hyper_c_mu'])
+            hyper_c_sigma = float(config[f'sample_{i}']['hyper_c_sigma'])
+            RA[i], Dec[i], czCMB[i], m_true[i], stretch_true[i], c_true[i], m_obs[i], stretch_obs[i], c_obs[i], xtrue[i] = sn_create_mock(
+                Nt, L, xmin, cosmo, dens, vel + vbulk, 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)
+        else:
+            raise NotImplementedError
+    
+    return tracer_type, RA, Dec, czCMB, m_true, m_obs, eta_true, eta_obs, stretch_true, stretch_obs, c_true, c_obs, xtrue

borg_velocity/likelihood.py

@@ -1,27 +1,19 @@
-import numpy as np
 import aquila_borg as borg
+import jax
+import numpy as np
 import jax.numpy as jnp
 import configparser
-import warnings
-import aquila_borg as borg
-#import symbolic_pofk.linear
-import jax
-from jax import lax
-import jaxlib
-from functools import partial
-import ast
-import numbers
 import h5py
+import warnings
+import ast
 import re
-import os
+from functools import partial
 
 import borg_velocity.utils as utils
-from borg_velocity.utils import myprint, compute_As, get_sigma_bulk
+from borg_velocity.utils import myprint, generateMBData, compute_As, get_sigma_bulk
+from borg_velocity.borg_mock import borg_mock
 import borg_velocity.forwards as forwards
-import borg_velocity.mock_maker as mock_maker
-import borg_velocity.projection as projection
 from borg_velocity.samplers import HMCBiasSampler, derive_prior, MVSliceBiasSampler, BlackJaxBiasSampler, TransformedBlackJaxBiasSampler
-import borg_velocity.samplers as samplers
 
 class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
     """
@@ -39,50 +31,16 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
             ADD FW_VEL HERE
         """
         
+        # Load configuration file
         self.ini_file = ini_file
-        
         myprint("Reading from configuration file: " + ini_file)
         config = configparser.ConfigParser()
         config.read(ini_file)
         
-        # Grid parameters
+        # Initialise grid parameters
         self.N = [int(config['system'][f'N{i}']) for i in range(3)]
         self.L = [float(config['system'][f'L{i}']) for i in range(3)]
         
-        # Catalogues
-        self.nsamp = int(config['run']['nsamp'])
-        assert self.nsamp > 0, "Need at least one sample to run"
-        self.Nt = [int(config[f'sample_{i}']['Nt']) for i in range(self.nsamp)]
-        self.alpha = [float(config[f'sample_{i}']['alpha']) for i in range(self.nsamp)]
-        self.lam = [float(config[f'sample_{i}']['lam']) for i in range(self.nsamp)]
-        self.muA = [float(config[f'sample_{i}']['muA']) for i in range(self.nsamp)]
-        self.frac_sig_rhMpc = [float(config[f'sample_{i}']['frac_sig_rhMpc']) for i in range(self.nsamp)]
-        
-        # What type of run we're doing
-        self.run_type = config['run']['run_type']
-        
-        # Seed if creating mocks
-        self.mock_seed = int(config['mock']['seed'])
-        self.R_max = float(config['mock']['r_max'])
-        
-        # Model parameters
-        self.ai = float(config['model']['ai'])
-        self.af = float(config['model']['af'])
-        self.sig_v = float(config['model']['sig_v'])
-        if config['model']['R_lim'] == 'none':
-            self.R_lim = self.L[0]/2
-        else:
-            self.R_lim = float(config['model']['R_lim'])
-        self.Nint_points = int(config['model']['nint_points'])
-        self.Nsig = float(config['model']['Nsig'])
-        self.sig_v = float(config['model']['sig_v'])
-        self.smooth_R = float(config['model']['smooth_R'])
-        self.bias_epsilon = float(config['model']['bias_epsilon'])
-        self.interp_order = int(config['model']['interp_order'])
-        
-        # Load bulk flow
-        self.bulk_flow = np.array(ast.literal_eval(config['model']['bulk_flow']))
-        
         # For log-likelihood values
         self.bignum = float(config['mcmc']['bignum'])
         
@@ -94,16 +52,6 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         self.fwd_param = param_model
         self.fwd_vel = fwd_vel
         
-        # Initialise model parameters
-        self.model_params = {
-            **{f'mua{i}':self.muA[i] for i in range(self.nsamp)},
-            **{f'alpha{i}':self.alpha[i] for i in range(self.nsamp)},
-            **{f'lam{i}':self.lam[i] for i in range(self.nsamp)},
-            **{f'bulk_flow_{d}':self.bulk_flow[i] for i, d in enumerate(['x', 'y', 'z'])},
-            'sig_v':self.sig_v
-        }
-        self.fwd_param.setModelParams(self.model_params)
-        
         # Initialise cosmological parameters
         cpar = utils.get_cosmopar(self.ini_file)
         self.fwd.setCosmoParams(cpar)
@@ -111,6 +59,10 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         self.updateCosmology(cpar)
         myprint(f"Original cosmological parameters: {self.fwd.getCosmoParams()}")
         
+        # Initialise model parameters
+        self.sig_v = float(config['model']['sig_v'])
+        self.bulk_flow = np.array(ast.literal_eval(config['model']['bulk_flow']))
+        
         # Initialise derivative
         self.grad_like = jax.grad(self.dens2like, argnums=(0,1))
         
@@ -119,6 +71,25 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
             self.action = utils.get_action() 
         else:
             self.action = None
+            
+        # What type of run we're doing (mock/velmass/data)
+        self.run_type = config['run']['run_type']
+        
+        # Seed if creating mocks
+        self.mock_seed = int(config['mock']['seed'])
+        
+        # Number of samples to distance tracers
+        self.nsamp = int(config['run']['nsamp'])
+        
+        # Initialise integration variables
+        if config['model']['R_lim'] == 'none':
+            self.R_lim = self.L[0]/2
+        else:
+            self.R_lim = float(config['model']['R_lim'])
+        self.Nint_points = int(config['model']['nint_points'])
+        self.Nsig = float(config['model']['Nsig'])
+        self.MB_pos = [None] * self.nsamp
+        self.frac_sigma_r = [float(config[f'sample_{i}']['frac_sig_rhMpc']) for i in range(self.nsamp)]
         
         
     def initializeLikelihood(self, state: borg.likelihood.MarkovState) -> None:
@@ -130,21 +101,10 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         """
         
         myprint("Init likelihood")
-
-        # for i in range(self.nsamp):
-        #     state.newArray1d(f"tracer_vr_{i}", self.Nt[i], True)
-        #     if self.run_type != 'data':
-        #         state.newArray1d(f"data_x_{i}", self.Nt[i], True)
-        #         state.newArray1d(f"data_y_{i}", self.Nt[i], True)
-        #         state.newArray1d(f"data_z_{i}", self.Nt[i], True)
-        #         state.newArray1d(f"true_x_{i}", self.Nt[i], True)
-        #         state.newArray1d(f"true_y_{i}", self.Nt[i], True)
-        #         state.newArray1d(f"true_z_{i}", self.Nt[i], True)
-        # self.data = [state[f"tracer_vr_{i}"] for i in range(self.nsamp)]
         state.newArray3d("BORG_final_density", *self.fwd.getOutputBoxModel().N, True)
         
         if self.run_type == 'data':
-            self.loadObservedData(make_plot=False)
+            self.loadObservedData()
         elif borg.EMBEDDED:
             if self.action == 'INIT':
                 pass  # Data will be loaded later
@@ -181,8 +141,6 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         
         # Convert sigma8 to As
         cpar = compute_As(cpar)
-        # cpar.A_s = 1.e-9 * symbolic_pofk.linear.sigma8_to_As(
-            # cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s)
         myprint(f"Updating cosmology Om = {cosmo.omega_m}, sig8 = {cosmo.sigma8}, As = {cosmo.A_s}")
         
         cpar.omega_q = 1. - cpar.omega_m - cpar.omega_k
@@ -190,44 +148,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         self.fwd_param.setCosmoParams(cpar)
         
         self.sigma_bulk = get_sigma_bulk(self.L[0], self.fwd.getCosmoParams())
-        
+                
-        
-    def generateMBData(self) -> None:
-        """
-        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
-        """
-        
-        self.MB_pos = [None] * self.nsamp
-        self.r_integration = [None] * self.nsamp
-        self.RA = [None] * self.nsamp
-        self.DEC = [None] * self.nsamp
-        
-        for i in range(self.nsamp):
-            myprint(f"Making MB data for sample {i}")
-            
-            # Get angular coordinates of all points
-            r_hat = projection.get_radial_vectors(self.coord_meas[i])
-            
-            # Get min and max distance to integrate over
-            robs = self.cz_obs[i] / 100
-            sigr = np.sqrt((self.sig_v / 100) ** 2 + (self.frac_sig_rhMpc[i] * robs)**2)
-            rmin = robs - self.Nsig * sigr
-            rmin = rmin.at[rmin <= 0].set(self.L[0] / self.N[0] / 100.)
-            rmax = robs + self.Nsig * sigr
-            rmax = rmax.at[rmax > self.R_lim].set(self.R_lim)
-            
-            # Compute coordinates of integration points
-            self.r_integration[i] = np.linspace(rmin, rmax, self.Nint_points)
-            cartesian_pos_MB = np.expand_dims(self.r_integration[i], axis=2) * r_hat
-            self.r_integration[i] = self.r_integration[i].T
-            self.MB_pos[i] = cartesian_pos_MB
-            self.MB_pos[i] = jnp.transpose(self.MB_pos[i], (2, 1, 0))
-        
         
     def generateMockData(self, s_hat: np.ndarray, state: borg.likelihood.MarkovState,) -> None:
         """
@@ -245,63 +166,39 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         if self.run_type == 'data':
             raise NotImplementedError
         elif self.run_type == 'velmass':
-            self.coord_true, self.coord_meas, self.sig_mu, self.vr_true, self.cz_obs = \
+            raise NotImplementedError
-                mock_maker.velmass_mock(self.ini_file, seed=self.mock_seed)
         elif self.run_type == 'mock':
-            self.coord_true, self.coord_meas, self.sig_mu, self.vr_true, self.cz_obs = \
+            self.tracer_type, self.RA, self.Dec, self.cz_obs, self.m_true, self.m_obs, self.eta_true, self.eta_obs, self.stretch_true, self.stretch_obs, self.c_true, self.c_obs, self.xtrue = borg_mock(s_hat, state, self.fwd, self.fwd_vel, self.ini_file, seed=self.mock_seed)
-                mock_maker.borg_mock(s_hat, state, self.fwd, self.fwd_vel, self.ini_file, seed=self.mock_seed)
         else:
             raise NotImplementedError
             
-        # Save mock to file
+        for i in range(self.nsamp):
-        with h5py.File(f'tracer_data_{self.run_type}.h5', 'w') as h5file:
+            self.MB_pos[i] = utils.generateMBData(
-            for i in range(self.nsamp):
+                self.RA[i], self.Dec[i], self.cz_obs[i], self.L[0], self.N[0], 
-                sample_group = h5file.create_group(f'sample_{i}')
+                self.R_lim, self.Nsig, self.Nint_points, self.sig_v, self.frac_sigma_r[i])
-                sample_group.create_dataset('coord_true', data=self.coord_true[i])
-                sample_group.create_dataset('coord_meas', data=self.coord_meas[i])
-                sample_group.create_dataset('sig_mu', data=self.sig_mu[i])
-                sample_group.create_dataset('vr_true', data=self.vr_true[i])
-                sample_group.create_dataset('cz_obs', data=self.cz_obs[i])
             
-        self.r_hMpc = [np.sqrt(np.sum(self.coord_meas[i] ** 2, axis=0)) for i in range(self.nsamp)]
+        # Save mock to file
+        # with h5py.File(f'tracer_data_{self.run_type}.h5', 'w') as h5file:
+        #     for i in range(self.nsamp):
+        #         sample_group = h5file.create_group(f'sample_{i}')
+        #         sample_group.create_dataset('coord_true', data=self.coord_true[i])
         
         # Check that the measured coordinates all lie within the box
-        for i in range(self.nsamp):
+        # for i in range(self.nsamp):
-            if np.amax(self.r_hMpc[i]) > self.R_lim:
+        #     if np.amax(self.r_hMpc[i]) > self.R_lim:
-                raise ValueError('All tracers must have measured coordinates within R_lim')
+        #         raise ValueError('All tracers must have measured coordinates within R_lim')
         
-        self.generateMBData()
         myprint('From mock')
         self.saved_s_hat = s_hat.copy()
         self.logLikelihoodComplex(s_hat, False)
         
+        quit()
+        
     def loadMockData(self, state: borg.likelihood.MarkovState) -> None:
         
         myprint('Loading previously made mock data')
         
-        self.coord_true = [None] * self.nsamp
+        raise NotImplementedError
-        self.coord_meas = [None] * self.nsamp
-        self.sig_mu = [None] * self.nsamp
-        self.vr_true = [None] * self.nsamp
-        self.cz_obs = [None] * self.nsamp
-        self.r_hMpc = [None] * self.nsamp
-
-        with h5py.File(f'tracer_data_{self.run_type}.h5', 'r') as f:
-            for i in range(self.nsamp):
-                self.coord_true[i] = jnp.array(f[f'sample_{i}/coord_true'][:])
-                self.coord_meas[i] = jnp.array(f[f'sample_{i}/coord_meas'][:])
-                self.sig_mu[i] = f[f'sample_{i}/sig_mu'][()]
-                self.vr_true[i] = jnp.array(f[f'sample_{i}/vr_true'][:])
-                self.cz_obs[i] = jnp.array(f[f'sample_{i}/cz_obs'][:])
-                self.r_hMpc[i] = np.sqrt(np.sum(self.coord_meas[i] ** 2, axis=0))
-                
-        # Check that the measured coordinates all lie within the box
-        for i in range(self.nsamp):
-            if np.amax(self.r_hMpc[i]) > self.R_lim:
-                raise ValueError('All tracers must have measured coordinates within R_lim')
-
-        self.generateMBData()
-
 
     def dens2like(self, output_density: np.ndarray, output_velocity: np.ndarray):
         """
@@ -326,36 +223,8 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         
         omega_m = self.fwd.getCosmoParams().omega_m
         
-#         self.lkl_ind = [None] * self.nsamp
+        raise NotImplementedError
         
-#         global temp_lkl_ind
-        
-        for i in range(self.nsamp):
-            muA = self.model_params[f'mua{i}']
-            alpha = self.model_params[f'alpha{i}']
-            lam = self.model_params[f'lam{i}']
-            lkl += vel2like(
-                self.cz_obs[i], 
-                v, 
-                output_density, 
-                self.MB_pos[i], 
-                self.r_integration[i], 
-                self.r_hMpc[i], 
-                self.sig_mu[i], 
-                sig_v, 
-                omega_m, 
-                muA, 
-                alpha, 
-                lam,
-                self.L[0], 
-                jnp.array(self.fwd.getOutputBoxModel().xmin), 
-                self.interp_order, 
-                self.bias_epsilon,
-                self.R_max
-            )
-            
-            # self.lkl_ind[i] = temp_lkl_ind.copy()
-            
         # Add in bulk flow prior
         lkl += jnp.sum(0.5 * jnp.log(2 * np.pi) + jnp.log(self.sigma_bulk / jnp.sqrt(3)) + self.bulk_flow ** 2 / 2. / (self.sigma_bulk / jnp.sqrt(3)) ** 2)
             
@@ -392,11 +261,6 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
             for k in self.model_params.keys():
                 self.model_params[k] = self.fwd_param.getModelParam('nullforward', k)
             self.updateCosmology(self.fwd.getCosmoParams())  # for sigma8 -> As
-            
-        # fname = f'{os.getcwd()}/s_hat.npy'
-        # myprint(f'Saving s_hat field to {fname}')
-        # np.save(fname, s_hat)
-        # myprint('Done')
 
         if not skip_density:
             
@@ -407,7 +271,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
             self.fwd.forwardModel_v2(s_hat)
             self.fwd.getDensityFinal(output_density)
             
-            myprint('Getting velocity field now')
+            myprint('Getting velocity field')
         
             # Get velocity field
             output_velocity = self.fwd_vel.getVelocityField()
@@ -476,123 +340,6 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         self.dens2like(self.delta, self.vel)
         state["BORG_final_density"][:] = self.delta
         
-        
-# @partial(jax.jit, static_argnames=['L_BOX', 'interp_order', 'bias_epsilon'])
-def vel2like(cz_obs, v, MB_field, MB_pos, r, r_hMpc, sig_mu, sig_v, omega_m, muA, alpha, lam, 
-             L_BOX, X_MIN, interp_order, bias_epsilon, R_max):
-    """
-    Jitted part of dens2like
-    """
-    
-    tracer_vel = projection.interp_field(v,
-                                MB_pos,
-                                L_BOX,
-                                X_MIN,
-                                interp_order,
-                                use_jitted=True,
-                                )
-            
-    tracer_vr = projection.project_radial(
-                    tracer_vel,
-                    MB_pos,
-                    jnp.zeros(3,)
-                    )
-    
-    # Remove the velocity term
-    # tracer_vr = jnp.zeros(tracer_vr.shape)
-    
-    # Convert velocities to redshifts
-    zco = utils.z_cos(r, omega_m)
-    cz_pred = utils.speed_of_light * zco + (1 + zco) * tracer_vr
-    delta_cz_sigv = (cz_pred - jnp.expand_dims(cz_obs, axis=1)) / sig_v
-
-    # Get p_r
-    delta_mu = 5. * jnp.log10(r / jnp.expand_dims(r_hMpc,axis=1) * muA)
-
-    # Get los biased density
-    los_density = projection.interp_field(
-                                MB_field,
-                                MB_pos,
-                                L_BOX,
-                                X_MIN,
-                                interp_order,
-                                use_jitted=True,
-                            )
-    los_density = jax.nn.relu(1. + los_density)
-    los_density = jnp.power(los_density + bias_epsilon, alpha)
-    
-    # Remove bias term
-    # los_density = jnp.ones(los_density.shape)
-    
-    d2 = (delta_mu / sig_mu) ** 2
-    best = jnp.amin(jnp.abs(d2), axis=1)
-
-    # Multiply p_r by arbitrary number for numerical stability (cancels in p_r / p_r_norm)
-    d2 = d2 - jnp.expand_dims(jnp.nanmin(d2, axis=1), axis=1)
-    p_r = r ** 2 * jnp.exp(-0.5 * d2) * los_density * jnp.exp(- lam * r / R_max)
-    p_r_norm = jnp.expand_dims(jnp.trapezoid(p_r, r, axis=1), axis=1)
-    
-    # Integrate to get likelihood
-    d2 = delta_cz_sigv**2
-    scale = jnp.nanmin(d2, axis=1)
-    d2 = d2 - jnp.expand_dims(scale, axis=1)
-    exp_delta_cz = jnp.exp(-0.5 * d2)
-    
-    p_cz = jnp.trapezoid(exp_delta_cz * p_r / p_r_norm, r, axis=1)
-    lkl_ind = jnp.log(p_cz) - scale / 2 - 0.5 * jnp.log(2 * np.pi * sig_v**2)
-    lkl = - lkl_ind.sum()
-    
-    # global temp_lkl_ind
-    # temp_lkl_ind = lkl_ind
-        
-    return lkl
-
-
-@derive_prior
-def transform_mua(x):
-    a, b = model_params_prior['mua']
-    return samplers.transform_uniform(x, a, b)
-    
-def inv_transform_mua(alpha):
-    a, b = model_params_prior['mua']
-    return samplers.inv_transform_uniform(alpha, a, b)
-
-@derive_prior
-def transform_alpha(x):
-    a, b = model_params_prior['alpha']
-    return samplers.transform_uniform(x, a, b)
-    
-def inv_transform_alpha(alpha):
-    a, b = model_params_prior['alpha']
-    return samplers.inv_transform_uniform(alpha, a, b)
-
-@derive_prior
-def transform_lam(x):
-    a, b = model_params_prior['lam']
-    return samplers.transform_uniform(x, a, b)
-    
-def inv_transform_lam(lam):
-    a, b = model_params_prior['lam']
-    return samplers.inv_transform_uniform(lam, a, b)
-
-@derive_prior
-def transform_sig_v(x):
-    a, b = model_params_prior['sig_v']
-    return samplers.transform_uniform(x, a, b)
-    
-def inv_transform_sig_v(alpha):
-    a, b = model_params_prior['sig_v']
-    return samplers.inv_transform_uniform(alpha, a, b)
-
-@derive_prior
-def transform_bulk_flow(x):
-    a, b = model_params_prior['bulk_flow']
-    return samplers.transform_uniform(x, a, b)
-    
-def inv_transform_bulk_flow(alpha):
-    a, b = model_params_prior['bulk_flow']
-    return samplers.inv_transform_uniform(alpha, a, b)
-    
 
 @borg.registerGravityBuilder
 def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_file=None) -> borg.forward.BaseForwardModel:
@@ -725,7 +472,6 @@ def check_model_sampling(loop):
 def check_cosmo_sampling(loop):
     return loop.getStepID() > begin_cosmo
             
-
 @borg.registerSamplerBuilder
 def build_sampler(
     state: borg.likelihood.MarkovState,
@@ -947,3 +693,16 @@ def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.L
     likelihood = VelocityBORGLikelihood(chain, fwd_param, fwd_vel, borg.getIniConfigurationFilename())
     return likelihood
 
+
+"""
+To Do:
+- Add in the true colours etc. to the model parameter sampling
+- Call to the correct log-likelihoods
+- Sample the variables from priors
+- Check global variables are used
+- Create minimal ini file
+- Docstrings
+- Don't want all these samplers - only gradient ones
+- Save mock to file
+"""
+

borg_velocity/supranta_likelihood/likelihood.py

@@ -0,0 +1,949 @@
+import numpy as np
+import aquila_borg as borg
+import jax.numpy as jnp
+import configparser
+import warnings
+import aquila_borg as borg
+#import symbolic_pofk.linear
+import jax
+from jax import lax
+import jaxlib
+from functools import partial
+import ast
+import numbers
+import h5py
+import re
+import os
+
+import borg_velocity.utils as utils
+from borg_velocity.utils import myprint, compute_As, get_sigma_bulk
+import borg_velocity.forwards as forwards
+import borg_velocity.mock_maker as mock_maker
+import borg_velocity.projection as projection
+from borg_velocity.samplers import HMCBiasSampler, derive_prior, MVSliceBiasSampler, BlackJaxBiasSampler, TransformedBlackJaxBiasSampler
+import borg_velocity.samplers as samplers
+
+class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
+    """
+    HADES likelihood for distance-tracers
+    """
+
+    def __init__(self, fwd: borg.forward.BaseForwardModel, param_model: forwards.NullForward, fwd_vel: borg.forward.BaseForwardModel, ini_file: str) -> None:
+        """
+        Initialises the VelocityBORGLikelihood class
+
+        Args:
+            - fwd (borg.forward.BaseForwardModel): The forward model to be used in the likelihood.
+            - param_model (forwards.NullForward): An empty forward model for storing model parameters.
+            - ini_file (str): The name of the ini file containing the model and borg parameters. 
+            ADD FW_VEL HERE
+        """
+        
+        self.ini_file = ini_file
+        
+        myprint("Reading from configuration file: " + ini_file)
+        config = configparser.ConfigParser()
+        config.read(ini_file)
+        
+        # Grid parameters
+        self.N = [int(config['system'][f'N{i}']) for i in range(3)]
+        self.L = [float(config['system'][f'L{i}']) for i in range(3)]
+        
+        # Catalogues
+        self.nsamp = int(config['run']['nsamp'])
+        assert self.nsamp > 0, "Need at least one sample to run"
+        self.Nt = [int(config[f'sample_{i}']['Nt']) for i in range(self.nsamp)]
+        self.alpha = [float(config[f'sample_{i}']['alpha']) for i in range(self.nsamp)]
+        self.lam = [float(config[f'sample_{i}']['lam']) for i in range(self.nsamp)]
+        self.muA = [float(config[f'sample_{i}']['muA']) for i in range(self.nsamp)]
+        self.frac_sig_rhMpc = [float(config[f'sample_{i}']['frac_sig_rhMpc']) for i in range(self.nsamp)]
+        
+        # What type of run we're doing
+        self.run_type = config['run']['run_type']
+        
+        # Seed if creating mocks
+        self.mock_seed = int(config['mock']['seed'])
+        self.R_max = float(config['mock']['r_max'])
+        
+        # Model parameters
+        self.ai = float(config['model']['ai'])
+        self.af = float(config['model']['af'])
+        self.sig_v = float(config['model']['sig_v'])
+        if config['model']['R_lim'] == 'none':
+            self.R_lim = self.L[0]/2
+        else:
+            self.R_lim = float(config['model']['R_lim'])
+        self.Nint_points = int(config['model']['nint_points'])
+        self.Nsig = float(config['model']['Nsig'])
+        self.sig_v = float(config['model']['sig_v'])
+        self.smooth_R = float(config['model']['smooth_R'])
+        self.bias_epsilon = float(config['model']['bias_epsilon'])
+        self.interp_order = int(config['model']['interp_order'])
+        
+        # Load bulk flow
+        self.bulk_flow = np.array(ast.literal_eval(config['model']['bulk_flow']))
+        
+        # For log-likelihood values
+        self.bignum = float(config['mcmc']['bignum'])
+        
+        myprint(f" Init {self.N}, {self.L}")
+        super().__init__(fwd, self.N, self.L)
+        
+        # Define the forward models
+        self.fwd = fwd
+        self.fwd_param = param_model
+        self.fwd_vel = fwd_vel
+        
+        # Initialise model parameters
+        self.model_params = {
+            **{f'mua{i}':self.muA[i] for i in range(self.nsamp)},
+            **{f'alpha{i}':self.alpha[i] for i in range(self.nsamp)},
+            **{f'lam{i}':self.lam[i] for i in range(self.nsamp)},
+            **{f'bulk_flow_{d}':self.bulk_flow[i] for i, d in enumerate(['x', 'y', 'z'])},
+            'sig_v':self.sig_v
+        }
+        self.fwd_param.setModelParams(self.model_params)
+        
+        # Initialise cosmological parameters
+        cpar = utils.get_cosmopar(self.ini_file)
+        self.fwd.setCosmoParams(cpar)
+        self.fwd_param.setCosmoParams(cpar)
+        self.updateCosmology(cpar)
+        myprint(f"Original cosmological parameters: {self.fwd.getCosmoParams()}")
+        
+        # Initialise derivative
+        self.grad_like = jax.grad(self.dens2like, argnums=(0,1))
+        
+        # Find out what kind of run this is
+        if borg.EMBEDDED:
+            self.action = utils.get_action() 
+        else:
+            self.action = None
+        
+        
+    def initializeLikelihood(self, state: borg.likelihood.MarkovState) -> None:
+        """
+        Initialise the likelihood internal variables and MarkovState variables.
+
+        Args:
+            - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
+        """
+        
+        myprint("Init likelihood")
+
+        # for i in range(self.nsamp):
+        #     state.newArray1d(f"tracer_vr_{i}", self.Nt[i], True)
+        #     if self.run_type != 'data':
+        #         state.newArray1d(f"data_x_{i}", self.Nt[i], True)
+        #         state.newArray1d(f"data_y_{i}", self.Nt[i], True)
+        #         state.newArray1d(f"data_z_{i}", self.Nt[i], True)
+        #         state.newArray1d(f"true_x_{i}", self.Nt[i], True)
+        #         state.newArray1d(f"true_y_{i}", self.Nt[i], True)
+        #         state.newArray1d(f"true_z_{i}", self.Nt[i], True)
+        # self.data = [state[f"tracer_vr_{i}"] for i in range(self.nsamp)]
+        state.newArray3d("BORG_final_density", *self.fwd.getOutputBoxModel().N, True)
+        
+        if self.run_type == 'data':
+            self.loadObservedData(make_plot=False)
+        elif borg.EMBEDDED:
+            if self.action == 'INIT':
+                pass  # Data will be loaded later
+            elif self.action == 'RESUME':
+                self.loadMockData(state)  # load data from mock_data.h5
+            else:
+                myprint(f'Unknown action: {self.action}') 
+                raise NotImplementedError 
+        
+        
+    def updateMetaParameters(self, state: borg.likelihood.MarkovState) -> None:
+        """
+        Update the meta parameters of the sampler (not sampled) from the MarkovState.
+
+        Args:
+            - state (borg.likelihood.MarkovState): The state object to be used in the likelihood.
+
+        """
+        cpar = state['cosmology']
+        cpar.omega_q = 1. - cpar.omega_m - cpar.omega_k
+        self.fwd.setCosmoParams(cpar)
+        self.fwd_param.setCosmoParams(cpar)
+
+        
+    def updateCosmology(self, cosmo: borg.cosmo.CosmologicalParameters) -> None:
+        """
+        Updates the forward model's cosmological parameters with the given values.
+
+        Args:
+            - cosmo (borg.cosmo.CosmologicalParameters): The cosmological parameters.
+
+        """
+        cpar = cosmo
+        
+        # Convert sigma8 to As
+        cpar = compute_As(cpar)
+        # cpar.A_s = 1.e-9 * symbolic_pofk.linear.sigma8_to_As(
+            # cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s)
+        myprint(f"Updating cosmology Om = {cosmo.omega_m}, sig8 = {cosmo.sigma8}, As = {cosmo.A_s}")
+        
+        cpar.omega_q = 1. - cpar.omega_m - cpar.omega_k
+        self.fwd.setCosmoParams(cpar)
+        self.fwd_param.setCosmoParams(cpar)
+        
+        self.sigma_bulk = get_sigma_bulk(self.L[0], self.fwd.getCosmoParams())
+        
+        
+    def generateMBData(self) -> None:
+        """
+        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
+        """
+        
+        self.MB_pos = [None] * self.nsamp
+        self.r_integration = [None] * self.nsamp
+        self.RA = [None] * self.nsamp
+        self.DEC = [None] * self.nsamp
+        
+        for i in range(self.nsamp):
+            myprint(f"Making MB data for sample {i}")
+            
+            # Get angular coordinates of all points
+            r_hat = projection.get_radial_vectors(self.coord_meas[i])
+            
+            # Get min and max distance to integrate over
+            robs = self.cz_obs[i] / 100
+            sigr = np.sqrt((self.sig_v / 100) ** 2 + (self.frac_sig_rhMpc[i] * robs)**2)
+            rmin = robs - self.Nsig * sigr
+            rmin = rmin.at[rmin <= 0].set(self.L[0] / self.N[0] / 100.)
+            rmax = robs + self.Nsig * sigr
+            rmax = rmax.at[rmax > self.R_lim].set(self.R_lim)
+            
+            # Compute coordinates of integration points
+            self.r_integration[i] = np.linspace(rmin, rmax, self.Nint_points)
+            cartesian_pos_MB = np.expand_dims(self.r_integration[i], axis=2) * r_hat
+            self.r_integration[i] = self.r_integration[i].T
+            self.MB_pos[i] = cartesian_pos_MB
+            self.MB_pos[i] = jnp.transpose(self.MB_pos[i], (2, 1, 0))
+        
+        
+    def generateMockData(self, s_hat: np.ndarray, state: borg.likelihood.MarkovState,) -> None:
+        """
+        Generates mock data by simulating the forward model with the given white noise,
+        drawing distance tracers from the density field, computing their distance
+        moduli and radial velocities, and adding Gaussian noise to the appropriate
+        variables. Also calculates the initial negative log-likelihood of the data.
+
+        Args:
+             - s_hat (np.ndarray): The input (initial) density field.
+             - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
+             - make_plot (bool, default=True): Whether to make diagnostic plots for the mock data generation
+        """
+        
+        if self.run_type == 'data':
+            raise NotImplementedError
+        elif self.run_type == 'velmass':
+            self.coord_true, self.coord_meas, self.sig_mu, self.vr_true, self.cz_obs = \
+                mock_maker.velmass_mock(self.ini_file, seed=self.mock_seed)
+        elif self.run_type == 'mock':
+            self.coord_true, self.coord_meas, self.sig_mu, self.vr_true, self.cz_obs = \
+                mock_maker.borg_mock(s_hat, state, self.fwd, self.fwd_vel, self.ini_file, seed=self.mock_seed)
+        else:
+            raise NotImplementedError
+            
+        # Save mock to file
+        with h5py.File(f'tracer_data_{self.run_type}.h5', 'w') as h5file:
+            for i in range(self.nsamp):
+                sample_group = h5file.create_group(f'sample_{i}')
+                sample_group.create_dataset('coord_true', data=self.coord_true[i])
+                sample_group.create_dataset('coord_meas', data=self.coord_meas[i])
+                sample_group.create_dataset('sig_mu', data=self.sig_mu[i])
+                sample_group.create_dataset('vr_true', data=self.vr_true[i])
+                sample_group.create_dataset('cz_obs', data=self.cz_obs[i])
+            
+        self.r_hMpc = [np.sqrt(np.sum(self.coord_meas[i] ** 2, axis=0)) for i in range(self.nsamp)]
+        
+        # Check that the measured coordinates all lie within the box
+        for i in range(self.nsamp):
+            if np.amax(self.r_hMpc[i]) > self.R_lim:
+                raise ValueError('All tracers must have measured coordinates within R_lim')
+        
+        self.generateMBData()
+        myprint('From mock')
+        self.saved_s_hat = s_hat.copy()
+        self.logLikelihoodComplex(s_hat, False)
+        
+    def loadMockData(self, state: borg.likelihood.MarkovState) -> None:
+        
+        myprint('Loading previously made mock data')
+        
+        self.coord_true = [None] * self.nsamp
+        self.coord_meas = [None] * self.nsamp
+        self.sig_mu = [None] * self.nsamp
+        self.vr_true = [None] * self.nsamp
+        self.cz_obs = [None] * self.nsamp
+        self.r_hMpc = [None] * self.nsamp
+
+        with h5py.File(f'tracer_data_{self.run_type}.h5', 'r') as f:
+            for i in range(self.nsamp):
+                self.coord_true[i] = jnp.array(f[f'sample_{i}/coord_true'][:])
+                self.coord_meas[i] = jnp.array(f[f'sample_{i}/coord_meas'][:])
+                self.sig_mu[i] = f[f'sample_{i}/sig_mu'][()]
+                self.vr_true[i] = jnp.array(f[f'sample_{i}/vr_true'][:])
+                self.cz_obs[i] = jnp.array(f[f'sample_{i}/cz_obs'][:])
+                self.r_hMpc[i] = np.sqrt(np.sum(self.coord_meas[i] ** 2, axis=0))
+                
+        # Check that the measured coordinates all lie within the box
+        for i in range(self.nsamp):
+            if np.amax(self.r_hMpc[i]) > self.R_lim:
+                raise ValueError('All tracers must have measured coordinates within R_lim')
+
+        self.generateMBData()
+
+
+    def dens2like(self, output_density: np.ndarray, output_velocity: np.ndarray):
+        """
+        Given stored distance tracer data, computes the negative log-likelihood of the data
+        for this final density field.
+
+        Args:
+            output_density (np.ndarray): The z=0 density field.
+        Return:
+            lkl (float): The negative log-likelihood of the data.
+        """
+        
+        lkl = 0
+        
+        sig_v = self.model_params['sig_v']
+        
+        # Compute velocity field
+        self.bulk_flow = jnp.array([self.model_params['bulk_flow_x'],
+                          self.model_params['bulk_flow_y'],
+                          self.model_params['bulk_flow_z']])
+        v = output_velocity + self.bulk_flow.reshape((3, 1, 1, 1))
+        
+        omega_m = self.fwd.getCosmoParams().omega_m
+        
+#         self.lkl_ind = [None] * self.nsamp
+        
+#         global temp_lkl_ind
+        
+        for i in range(self.nsamp):
+            muA = self.model_params[f'mua{i}']
+            alpha = self.model_params[f'alpha{i}']
+            lam = self.model_params[f'lam{i}']
+            lkl += vel2like(
+                self.cz_obs[i], 
+                v, 
+                output_density, 
+                self.MB_pos[i], 
+                self.r_integration[i], 
+                self.r_hMpc[i], 
+                self.sig_mu[i], 
+                sig_v, 
+                omega_m, 
+                muA, 
+                alpha, 
+                lam,
+                self.L[0], 
+                jnp.array(self.fwd.getOutputBoxModel().xmin), 
+                self.interp_order, 
+                self.bias_epsilon,
+                self.R_max
+            )
+            
+            # self.lkl_ind[i] = temp_lkl_ind.copy()
+            
+        # Add in bulk flow prior
+        lkl += jnp.sum(0.5 * jnp.log(2 * np.pi) + jnp.log(self.sigma_bulk / jnp.sqrt(3)) + self.bulk_flow ** 2 / 2. / (self.sigma_bulk / jnp.sqrt(3)) ** 2)
+            
+        # lkl = jnp.clip(lkl, -self.bignum, self.bignum)
+        lkl = lax.cond(
+            jnp.isfinite(lkl),
+            lambda _: lkl,         # If True (finite), return lkl
+            lambda _: self.bignum, # If False (not finite), return self.bignum
+            operand=None           # No additional operand needed here
+        )
+            
+        return lkl
+        
+    def logLikelihoodComplex(self, s_hat: np.ndarray, gradientIsNext: bool, skip_density: bool=False, update_from_model: bool=True):
+        """
+        Calculates the negative log-likelihood of the data.
+
+        Args:
+            - s_hat (np.ndarray): The input white noise.
+            - gradientIsNext (bool): If True, prepares the forward model for gradient calculations.
+            - skip_density (bool, default=False): If True, do not repeat the s_hat -> density, velocity computation
+                and use the stored result
+            - update_from_model (bool, default=True): If True, update self.model_params with self.fwd_param.getModelParam
+
+        Returns:
+            The negative log-likelihood value.
+
+        """
+
+        N = self.fwd.getBoxModel().N[0]
+        L = self.fwd.getOutputBoxModel().L[0]
+        
+        if update_from_model:
+            for k in self.model_params.keys():
+                self.model_params[k] = self.fwd_param.getModelParam('nullforward', k)
+            self.updateCosmology(self.fwd.getCosmoParams())  # for sigma8 -> As
+            
+        # fname = f'{os.getcwd()}/s_hat.npy'
+        # myprint(f'Saving s_hat field to {fname}')
+        # np.save(fname, s_hat)
+        # myprint('Done')
+
+        if not skip_density:
+            
+            myprint('Getting density field now')
+            
+            # Run BORG density field
+            output_density = np.zeros((N,N,N))
+            self.fwd.forwardModel_v2(s_hat)
+            self.fwd.getDensityFinal(output_density)
+            
+            myprint('Getting velocity field now')
+        
+            # Get velocity field
+            output_velocity = self.fwd_vel.getVelocityField()
+            
+        else:
+            output_density = self.delta.copy()
+            output_velocity = self.vel.copy()
+
+        L = self.dens2like(output_density, output_velocity)
+        if isinstance(L, numbers.Number) or isinstance(L, jaxlib.xla_extension.ArrayImpl):
+            myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}")
+            myprint(self.model_params)
+            myprint(self.fwd.getCosmoParams())
+                    
+        self.delta = output_density.copy()
+        self.vel = output_velocity.copy()
+
+        return L
+    
+    
+    def gradientLikelihoodComplex(self, s_hat: np.ndarray):
+        """
+        Calculates the adjoint negative log-likelihood of the data.
+
+        Args:
+            - s_hat (np.ndarray): The input density field.
+
+        Returns:
+            The adjoint negative log-likelihood gradient.
+
+        """
+
+        N = self.fwd.getBoxModel().N[0]
+        L = self.fwd.getOutputBoxModel().L[0]
+
+        # Run BORG density field
+        output_density = np.zeros((N,N,N))
+        self.fwd.forwardModel_v2(s_hat)
+        self.fwd.getDensityFinal(output_density)
+        
+        # Get velocity field
+        output_velocity = self.fwd_vel.getVelocityField()
+
+        # getlike(dens, vel)
+        mygradient, velgradient = self.grad_like(output_density, output_velocity)
+        
+        mygradient = np.array(mygradient, dtype=np.float64)
+        vgradient = np.array(velgradient, dtype=np.float64)
+        
+        self.fwd_vel.computeAdjointModel(vgradient)
+        self.fwd.adjointModel_v2(mygradient)
+        mygrad_hat = np.zeros(s_hat.shape, dtype=np.complex128)
+        self.fwd.getAdjointModel(mygrad_hat)
+        self.fwd.clearAdjointGradient()
+        
+        return mygrad_hat
+
+
+    def commitAuxiliaryFields(self, state: borg.likelihood.MarkovState) -> None:
+        """
+        Commits the final density field to the Markov state.
+        Args:
+            - state (borg.state.State): The state object containing the final density field.
+        """
+        self.updateCosmology(self.fwd.getCosmoParams())
+        self.dens2like(self.delta, self.vel)
+        state["BORG_final_density"][:] = self.delta
+        
+        
+# @partial(jax.jit, static_argnames=['L_BOX', 'interp_order', 'bias_epsilon'])
+def vel2like(cz_obs, v, MB_field, MB_pos, r, r_hMpc, sig_mu, sig_v, omega_m, muA, alpha, lam, 
+             L_BOX, X_MIN, interp_order, bias_epsilon, R_max):
+    """
+    Jitted part of dens2like
+    """
+    
+    tracer_vel = projection.interp_field(v,
+                                MB_pos,
+                                L_BOX,
+                                X_MIN,
+                                interp_order,
+                                use_jitted=True,
+                                )
+            
+    tracer_vr = projection.project_radial(
+                    tracer_vel,
+                    MB_pos,
+                    jnp.zeros(3,)
+                    )
+    
+    # Remove the velocity term
+    # tracer_vr = jnp.zeros(tracer_vr.shape)
+    
+    # Convert velocities to redshifts
+    zco = utils.z_cos(r, omega_m)
+    cz_pred = utils.speed_of_light * zco + (1 + zco) * tracer_vr
+    delta_cz_sigv = (cz_pred - jnp.expand_dims(cz_obs, axis=1)) / sig_v
+
+    # Get p_r
+    delta_mu = 5. * jnp.log10(r / jnp.expand_dims(r_hMpc,axis=1) * muA)
+
+    # Get los biased density
+    los_density = projection.interp_field(
+                                MB_field,
+                                MB_pos,
+                                L_BOX,
+                                X_MIN,
+                                interp_order,
+                                use_jitted=True,
+                            )
+    los_density = jax.nn.relu(1. + los_density)
+    los_density = jnp.power(los_density + bias_epsilon, alpha)
+    
+    # Remove bias term
+    # los_density = jnp.ones(los_density.shape)
+    
+    d2 = (delta_mu / sig_mu) ** 2
+    best = jnp.amin(jnp.abs(d2), axis=1)
+
+    # Multiply p_r by arbitrary number for numerical stability (cancels in p_r / p_r_norm)
+    d2 = d2 - jnp.expand_dims(jnp.nanmin(d2, axis=1), axis=1)
+    p_r = r ** 2 * jnp.exp(-0.5 * d2) * los_density * jnp.exp(- lam * r / R_max)
+    p_r_norm = jnp.expand_dims(jnp.trapezoid(p_r, r, axis=1), axis=1)
+    
+    # Integrate to get likelihood
+    d2 = delta_cz_sigv**2
+    scale = jnp.nanmin(d2, axis=1)
+    d2 = d2 - jnp.expand_dims(scale, axis=1)
+    exp_delta_cz = jnp.exp(-0.5 * d2)
+    
+    p_cz = jnp.trapezoid(exp_delta_cz * p_r / p_r_norm, r, axis=1)
+    lkl_ind = jnp.log(p_cz) - scale / 2 - 0.5 * jnp.log(2 * np.pi * sig_v**2)
+    lkl = - lkl_ind.sum()
+    
+    # global temp_lkl_ind
+    # temp_lkl_ind = lkl_ind
+        
+    return lkl
+
+
+@derive_prior
+def transform_mua(x):
+    a, b = model_params_prior['mua']
+    return samplers.transform_uniform(x, a, b)
+    
+def inv_transform_mua(alpha):
+    a, b = model_params_prior['mua']
+    return samplers.inv_transform_uniform(alpha, a, b)
+
+@derive_prior
+def transform_alpha(x):
+    a, b = model_params_prior['alpha']
+    return samplers.transform_uniform(x, a, b)
+    
+def inv_transform_alpha(alpha):
+    a, b = model_params_prior['alpha']
+    return samplers.inv_transform_uniform(alpha, a, b)
+
+@derive_prior
+def transform_lam(x):
+    a, b = model_params_prior['lam']
+    return samplers.transform_uniform(x, a, b)
+    
+def inv_transform_lam(lam):
+    a, b = model_params_prior['lam']
+    return samplers.inv_transform_uniform(lam, a, b)
+
+@derive_prior
+def transform_sig_v(x):
+    a, b = model_params_prior['sig_v']
+    return samplers.transform_uniform(x, a, b)
+    
+def inv_transform_sig_v(alpha):
+    a, b = model_params_prior['sig_v']
+    return samplers.inv_transform_uniform(alpha, a, b)
+
+@derive_prior
+def transform_bulk_flow(x):
+    a, b = model_params_prior['bulk_flow']
+    return samplers.transform_uniform(x, a, b)
+    
+def inv_transform_bulk_flow(alpha):
+    a, b = model_params_prior['bulk_flow']
+    return samplers.inv_transform_uniform(alpha, a, b)
+    
+
+@borg.registerGravityBuilder
+def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_file=None) -> borg.forward.BaseForwardModel:
+    """
+    Builds the gravity model and returns the forward model chain.
+
+    Args:
+        - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
+        - box (borg.forward.BoxModel): The input box model.
+        - ini_file (str, default=None): The location of the ini file. If None, use borg.getIniConfigurationFilename()
+
+    Returns:
+        borg.forward.BaseForwardModel: The forward model.
+
+    """
+
+    global chain, fwd_param, fwd_vel
+    myprint("Building gravity model")
+
+    if ini_file is None:
+        myprint("Reading from configuration file: " + borg.getIniConfigurationFilename())
+        config = configparser.ConfigParser()
+        config.read(borg.getIniConfigurationFilename())
+    else:
+        myprint("Reading from configuration file: " + ini_file)
+        config = configparser.ConfigParser()
+        config.read(ini_file)
+    ai = float(config['model']['ai'])
+    af = float(config['model']['af'])
+    supersampling = int(config['model']['supersampling'])
+    
+    if config['model']['gravity'] in ['pm', 'cola']:
+        forcesampling = int(config['model']['forcesampling'])
+
+    # Setup forward model
+    chain = borg.forward.ChainForwardModel(box)
+    chain.addModel(borg.forward.models.HermiticEnforcer(box))
+  
+    # CLASS transfer function
+    chain @= borg.forward.model_lib.M_PRIMORDIAL_AS(box)
+    transfer_class = borg.forward.model_lib.M_TRANSFER_CLASS(box, opts=dict(a_transfer=1.0))
+    transfer_class.setModelParams({"extra_class_arguments":{'YHe':'0.24'}})
+    chain @= transfer_class
+    
+    if config['model']['gravity'] == 'linear':
+        raise NotImplementedError(config['model']['gravity'])
+    elif config['model']['gravity'] == 'lpt':
+        mod = borg.forward.model_lib.M_LPT_CIC(
+            box, 
+            opts=dict(a_initial=af,
+                      a_final=af,
+                      do_rsd=False,
+                      supersampling=supersampling,
+                      lightcone=False,
+                      part_factor=1.01,))
+    elif config['model']['gravity'] == '2lpt':
+        mod = borg.forward.model_lib.M_2LPT_CIC(
+            box, 
+            opts=dict(a_initial=af,
+                      a_final=af,
+                      do_rsd=False,
+                      supersampling=supersampling,
+                      lightcone=False,
+                      part_factor=1.01,))
+    elif config['model']['gravity'] == 'pm':
+        mod =  borg.forward.model_lib.M_PM_CIC(
+            box, 
+            opts=dict(a_initial=af,
+                      a_final=af,
+                      do_rsd=False,
+                      supersampling=supersampling,
+                      part_factor=1.01,
+                      forcesampling=forcesampling,
+                      pm_start_z=1/ai - 1,
+                      pm_nsteps=int(config['model']['nsteps']),
+                      tcola=False))
+    elif config['model']['gravity'] == 'cola':
+        mod = borg.forward.model_lib.M_PM_CIC(
+            box, 
+            opts=dict(a_initial=af,
+                      a_final=af,
+                      do_rsd=False,
+                      supersampling=supersampling,
+                      part_factor=1.01,
+                      forcesampling=forcesampling,
+                      pm_start_z=1/ai - 1,
+                      pm_nsteps=int(config['model']['nsteps']),
+                      tcola=True))
+    else:
+        raise NotImplementedError(config['model']['gravity'])
+        
+    mod.accumulateAdjoint(True)
+    chain @= mod
+        
+    # Cosmological parameters
+    if ini_file is None:
+        cpar = utils.get_cosmopar(borg.getIniConfigurationFilename())
+    else:
+        cpar = utils.get_cosmopar(ini_file)
+    chain.setCosmoParams(cpar)
+    
+    # This is the forward model for the model parameters
+    fwd_param = borg.forward.ChainForwardModel(box)
+    mod_null = forwards.NullForward(box)
+    fwd_param.addModel(mod_null)
+    fwd_param.setCosmoParams(cpar)
+    
+    # This is the forward model for velocity
+    velmodel_name = config['model']['velocity']
+    velmodel = getattr(borg.forward.velocity, velmodel_name)
+    if velmodel_name == 'LinearModel':
+        fwd_vel = velmodel(box, mod, af)
+    elif velmodel_name == 'CICModel':
+        rsmooth = float(config['model']['rsmooth'])  # Mpc/h
+        fwd_vel = velmodel(box, mod, rsmooth)
+    else:
+        fwd_vel = velmodel(box, mod)
+    
+    return chain
+
+_glob_model = None
+_glob_cosmo = None
+begin_model = None
+begin_cosmo = None
+model_params_prior = {}
+
+def check_model_sampling(loop):
+    return loop.getStepID() > begin_model
+
+def check_cosmo_sampling(loop):
+    return loop.getStepID() > begin_cosmo
+            
+
+@borg.registerSamplerBuilder
+def build_sampler(
+    state: borg.likelihood.MarkovState,
+    info: borg.likelihood.LikelihoodInfo,
+    loop: borg.samplers.MainLoop
+):
+    """
+    Builds the sampler and returns it.
+    Which parameters to sample are given in the ini file.
+    We assume all parameters are NOT meant to be sampled, unless we find "XX_sampler_blocked = false" in the ini file
+
+    Args:
+        - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
+        - info (borg.likelihood.LikelihoodInfo): The likelihood information.
+        ADD LOOP - DEFINE CONDITIONS AND DEPENDENCIES
+
+    Returns:
+        List of samplers to use.
+
+    """
+    global _glob_model, _glob_cosmo, begin_model, begin_cosmo, model_params_prior
+    borg.print_msg(borg.Level.std, "Hello sampler, loop is {l}, step_id={s}", l=loop, s=loop.getStepID())
+    
+    myprint("Building sampler")
+    
+    myprint("Reading from configuration file: " + borg.getIniConfigurationFilename())
+    config = configparser.ConfigParser()
+    config.read(borg.getIniConfigurationFilename())
+    end = '_sampler_blocked'
+    to_sample = [k[:-len(end)] for (k, v) in config['block_loop'].items() if k[-len(end):] == end and v.lower() == 'false']
+    myprint(f'Parameters to sample: {to_sample}')
+    nsamp = int(config['run']['nsamp'])
+
+    all_sampler = []
+
+    # Cosmology sampler arguments
+    prefix = ""
+    params = []
+    initial_values = {}
+    prior = {}
+    for p in ["omega_m", "sigma8"]:
+        if p not in to_sample:
+            continue
+        if p in config['prior'].keys() and p in config['cosmology'].keys():
+            myprint(f'Adding {p} sampler')
+            params.append(f"cosmology.{p}")
+            initial_values[f"cosmology.{p}"] = float(config['cosmology'][p])
+            prior[f"cosmology.{p}"] = np.array(ast.literal_eval(config['prior'][p]))
+        else:
+            s = f'Could not find {p} prior and/or default, so will not sample'
+            warnings.warn(s, stacklevel=2)
+        # Remove for later to prevent duplication 
+        to_sample.remove(p)
+       
+    begin_cosmo = int(config['mcmc']['warmup_cosmo'])
+    
+    if len(params) > 0:
+        myprint('Adding cosmological parameter sampler')
+        cosmo_sampler = borg.samplers.ModelParamsSampler(prefix, params, likelihood, chain, initial_values, prior)
+        cosmo_sampler.setName("cosmo_sampler")
+        _glob_cosmo = cosmo_sampler
+        all_sampler.append(cosmo_sampler)
+        loop.push(cosmo_sampler)
+        loop.addToConditionGroup("warmup_cosmo", "cosmo_sampler")
+        loop.addConditionToConditionGroup("warmup_cosmo", partial(check_cosmo_sampling, loop))
+
+    # Model parameter sampler
+    prefix = ""
+    params = []
+    initial_values = {}
+    prior = {}
+    transform_attributes = []
+    inv_transform_attributes = []
+    
+    for p in to_sample:
+        if p in config['prior'].keys():
+            if p == 'sig_v':
+                myprint(f'Adding {p} sampler')
+                params.append(p)
+                initial_values[p] = float(config['model'][p])
+                if 'inf' in config['prior'][p]:
+                    x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))    
+                    prior[p] = np.array([xx if xx != 'inf' else np.inf for xx in x])
+                else:
+                    prior[p] = np.array(ast.literal_eval(config['prior'][p]))
+                model_params_prior[p] = prior[p]
+                transform_attributes.append(globals().get(f'transform_{p}'))
+                inv_transform_attributes.append(globals().get(f'inv_transform_{p}'))                            
+            elif p == 'bulk_flow':
+                for i, d in enumerate(['_x', '_y', '_z']):
+                    myprint(f'Adding {p}{d} sampler')
+                    params.append(f'{p}{d}')
+                    initial_values[f'{p}{d}'] = np.array(ast.literal_eval(config['model']['bulk_flow']))[i]
+                    if 'inf' in config['prior'][p]:
+                        x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))    
+                        prior[f'{p}{d}'] = np.array([xx if xx != 'inf' else np.inf for xx in x])
+                    else:
+                        prior[f'{p}{d}'] = np.array(ast.literal_eval(config['prior'][p]))
+                    transform_attributes.append(globals().get(f'transform_{p}'))
+                    inv_transform_attributes.append(globals().get(f'inv_transform_{p}')) 
+                model_params_prior[p] = prior[f'{p}_x']
+            else:
+                for i in range(nsamp):
+                    myprint(f'Adding {p}{i} sampler')
+                    params.append(f'{p}{i}')
+                    initial_values[f'{p}{i}'] = float(config[f'sample_{i}'][p])
+                    if 'inf' in config['prior'][p]:
+                        x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))    
+                        prior[f'{p}{i}'] = np.array([xx if xx != 'inf' else np.inf for xx in x])
+                    else:
+                        prior[f'{p}{i}'] = np.array(ast.literal_eval(config['prior'][p]))
+                    transform_attributes.append(globals().get(f'transform_{p}'))
+                    inv_transform_attributes.append(globals().get(f'inv_transform_{p}')) 
+                model_params_prior[p] = prior[f'{p}{0}']
+        else:
+            s = f'Could not find {p} prior, so will not sample'
+            warnings.warn(s, stacklevel=2)
+            
+    begin_model = int(config['mcmc']['warmup_model'])
+   
+    if len(params) > 0:
+        myprint('Adding model parameter sampler')
+        if config['sampling']['algorithm'].lower() == 'slice':
+            model_sampler = borg.samplers.ModelParamsSampler(prefix, params, likelihood, fwd_param, initial_values, prior)
+        elif config['sampling']['algorithm'].lower() == 'hmc':
+            model_sampler = HMCBiasSampler(
+                "model_params",
+                likelihood,
+                params,
+                transform_attributes=transform_attributes,
+                inv_transform_attributes=inv_transform_attributes,
+                prior = [p.prior for p in transform_attributes],
+                Nsteps = int(config['sampling']['nsteps']),
+                epsilon = float(config['sampling']['epsilon']),
+                refresh = float(config['sampling']['refresh']),
+            )
+            with open('model_params.txt', 'w') as file:
+                for item in params:
+                    file.write(f"{item}\n")
+        elif config['sampling']['algorithm'].lower() == 'mvslice':
+            lam = [None] * len(params)
+            for i, p in enumerate(params):
+                if p.startswith('bulk_flow'):
+                    lam[i] = float(config['sampling'][f'mvlam_bulk_flow'])
+                elif p[-1].isdigit():
+                    # Search for a non-digit before end of reversed string
+                    match = re.search(r'\D(?=\d*$)', p[::-1])
+                    lam[i] = float(config['sampling'][f'mvlam_{p[:match.start()]}'])
+                else:
+                    lam[i] = float(config['sampling'][f'mvlam_{p}'])
+            model_sampler =  MVSliceBiasSampler(
+                "model_params",
+                likelihood,
+                params,
+                lam
+            )
+            with open('model_params.txt', 'w') as file:
+                for item in params:
+                    file.write(f"{item}\n")
+        elif config['sampling']['algorithm'].lower() == 'blackjax':
+            # state.newArray1d("inverse_mass_matrix", len(params), True)
+            # state.newScalar('step_size', 0., True)
+            model_sampler = BlackJaxBiasSampler(
+                "model_params",
+                likelihood,
+                params,
+                int(config['sampling']['rng_seed']),
+                int(config['sampling']['warmup_nsteps']),
+                float(config['sampling']['warmup_target_acceptance_rate']),
+            )
+            with open('model_params.txt', 'w') as file:
+                for item in params:
+                    file.write(f"{item}\n")
+        elif config['sampling']['algorithm'].lower() == 'transformedblackjax':
+            model_sampler = TransformedBlackJaxBiasSampler(
+                "model_params",
+                likelihood,
+                params,
+                transform_attributes=transform_attributes,
+                inv_transform_attributes=inv_transform_attributes,
+                prior = [p.prior for p in transform_attributes],
+                rng_seed = int(config['sampling']['rng_seed']),
+                warmup_nsteps = int(config['sampling']['warmup_nsteps']),
+                warmup_target_acceptance_rate = float(config['sampling']['warmup_target_acceptance_rate']),
+            )
+            with open('model_params.txt', 'w') as file:
+                for item in params:
+                    file.write(f"{item}\n")
+        else:
+            raise NotImplementedError
+        model_sampler.setName("model_sampler")
+        _glob_model = model_sampler
+        loop.push(model_sampler)
+        all_sampler.append(model_sampler)
+        loop.addToConditionGroup("warmup_model", "model_sampler")
+        loop.addConditionToConditionGroup("warmup_model", partial(check_model_sampling, loop))
+        myprint('Done')
+        
+    return []
+
+
+@borg.registerLikelihoodBuilder
+def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo) -> borg.likelihood.BaseLikelihood:
+    """
+    Builds the likelihood object and returns it.
+
+    Args:
+        - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
+        - info (borg.likelihood.LikelihoodInfo): The likelihood information.
+
+    Returns:
+        borg.likelihood.BaseLikelihood: The likelihood object.
+
+    """
+    global likelihood, fwd_param
+    myprint("Building likelihood")
+    myprint(chain.getCosmoParams())
+    boxm = chain.getBoxModel()
+    likelihood = VelocityBORGLikelihood(chain, fwd_param, fwd_vel, borg.getIniConfigurationFilename())
+    return likelihood
+

borg_velocity/utils.py

@@ -8,6 +8,10 @@ import jax.numpy as jnp
 import scipy.integrate
 import math
 import re
+import numpy as np
+import jax.numpy as jnp
+from astropy.coordinates import SkyCoord
+import astropy.units as apu
 
 # Output stream management
 cons = borg.console()
@@ -194,6 +198,55 @@ def parse_file_to_dict(file_path):
     return config_dict
 
 
+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
+    
+    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): How many standard deviations away from the predicted radius to use
+        - 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
+
+
 if __name__ == "__main__":
     
     cpar = get_cosmopar('../conf/basic_ini.ini')

conf/borg_mock_ini.ini

@@ -0,0 +1,162 @@
+[system]
+console_output = borg_log
+VERBOSE_LEVEL = 2
+N0 = 64
+N1 = 64
+N2 = 64
+L0 = 500.0
+L1 = 500.0
+L2 = 500.0
+corner0 = -250.0
+corner1 = -250.0
+corner2 = -250.0
+NUM_MODES = 100
+test_mode = true
+seed_cpower = true
+
+[block_loop]
+hades_sampler_blocked = false
+bias_sampler_blocked= true
+nmean_sampler_blocked= true
+sigma8_sampler_blocked = true
+omega_m_sampler_blocked = true
+alpha_sampler_blocked = false
+sig_v_sampler_blocked = false
+bulk_flow_sampler_blocked = false
+ares_heat = 1.0
+
+[mcmc]
+number_to_generate = 4
+warmup_model = 200
+warmup_cosmo = 800
+random_ic = false
+init_random_scaling = 0.1
+bignum = 1e20
+
+[hades]
+algorithm = HMC
+max_epsilon = 0.01
+max_timesteps = 50
+mixing = 1
+
+[sampling]
+algorithm = slice
+epsilon = 0.005
+Nsteps = 20
+refresh = 0.1
+rng_seed = 1
+warmup_nsteps = 400
+warmup_target_acceptance_rate = 0.7
+
+[model]
+gravity = lpt
+forcesampling = 4
+supersampling = 2
+velocity = CICModel
+af = 1.0
+ai = 0.05
+nsteps = 20
+Rsmooth = 4.
+bias_epsilon = 1e-7
+interp_order = 1
+sig_v = 150.
+R_lim = 220
+Nint_points = 201 
+Nsig = 10
+bulk_flow = [0.0, 0.0, 0.0]
+
+[prior]
+omega_m = [0.1, 0.8]
+sigma8 = [0.1, 1.5]
+alpha = [0.0, 10.0]
+sig_v = [50.0, 200.0]
+bulk_flow = [-200.0, 200.0]
+
+[cosmology]
+omega_r = 0
+fnl = 0
+omega_k = 0
+omega_m = 0.315
+omega_b = 0.049
+omega_q = 0.685
+h100 = 0.68
+sigma8 = 0.81
+n_s = 0.97
+w = -1
+wprime = 0
+beta = 1.5
+z0 = 0
+
+[run]
+run_type = mock
+NCAT = 0
+NSAMP = 4
+
+[mock]
+seed = 12345
+R_max = 100
+
+[python]
+likelihood_path = /home/bartlett/fsigma8/borg_velocity/borg_velocity/likelihood.py
+
+[sample_0]
+tracer_type = tfr
+Nt = 1247
+alpha = 1.4
+a_TFR = -23
+b_TFR = -8.2
+sigma_TFR = 0.3
+sigma_m = 0.042
+sigma_eta = 0.01
+hyper_eta_mu = 0.08
+hyper_eta_sigma = 0.13
+mthresh = 11.25
+frac_sig_rhMpc = 0.07
+
+[sample_1]
+tracer_type = tfr
+Nt = 2010
+alpha = 1.4
+a_TFR = -23
+b_TFR = -8.2
+sigma_TFR = 0.3
+sigma_m = 0.042
+sigma_eta = 0.01
+hyper_eta_mu = 0.08
+hyper_eta_sigma = 0.13
+mthresh = 11.25
+frac_sig_rhMpc = 0.07
+
+[sample_2]
+tracer_type = sn
+Nt = 125
+alpha = 1.4
+a_tripp = 0.140
+b_tripp = 2.78
+M_SN =  -18.558
+sigma_SN = 0.082
+sigma_m = 0.04
+sigma_stretch = 0.17
+sigma_c = 0.04
+hyper_stretch_mu = -0.04
+hyper_stretch_sigma = 1.23
+hyper_c_mu = -0.02
+hyper_c_sigma = 0.08
+frac_sig_rhMpc = 0.07
+
+[sample_3]
+tracer_type = sn
+Nt = 55
+alpha = 1.4
+a_tripp = 0.140
+b_tripp = 2.78
+M_SN =  -18.558
+sigma_SN = 0.082
+sigma_m = 0.04
+sigma_stretch = 0.17
+sigma_c = 0.04
+hyper_stretch_mu = -0.04
+hyper_stretch_sigma = 1.23
+hyper_c_mu = -0.02
+hyper_c_sigma = 0.08
+frac_sig_rhMpc = 0.07

scripts/run_borg.sh

@@ -16,7 +16,7 @@ set -e
 
 # Path variables
 BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python
-RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/velmass_test
+RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/borg_mock_test
 
 mkdir -p $RUN_DIR
 cd $RUN_DIR
@@ -29,7 +29,7 @@ BASH_XTRACEFD="3"
 set -x
 
 # Just ICs
-INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/velmass_ini.ini
+INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/borg_mock_ini.ini
 cp $INI_FILE ini_file.ini
 $BORG INIT ini_file.ini
 # $BORG RESUME ini_file.ini

tests/sn_inference.py

@@ -760,7 +760,7 @@ 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