Begin testing TFR likelihood

.gitignore

@@ -169,3 +169,5 @@ tests/*.h5
 *fft_wisdom*
 *allocation_stats*
 scripts/out_files
+tests/*.npy
+tests/vel_data

borg_velocity/likelihood.py

@@ -275,7 +275,6 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         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')
@@ -400,13 +399,19 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         # 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()

conf/velmass_ini.ini

@@ -29,8 +29,8 @@ ares_heat = 1.0
 
 [mcmc]
 number_to_generate = 15000
-warmup_model = 500
-warmup_cosmo = 0
+warmup_model = 200
+warmup_cosmo = 800
 random_ic = false
 init_random_scaling = 0.1
 bignum = 1e20
@@ -56,13 +56,14 @@ mvlam_bulk_flow = 20
 mvlam_sig_v = 10
 
 [model]
-gravity = lpt
+gravity = cola
+forcesampling = 4
 supersampling = 2
 velocity = CICModel
 af = 1.0
 ai = 0.05
 nsteps = 20
-smooth_R = 4
+smooth_R = 4.
 bias_epsilon = 1e-7
 interp_order = 1
 rsmooth = 8.

notebooks/Analyse_chain.ipynb

@@ -1,11 +1,9 @@
 {
  "cells": [
   {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "40ad2d23-227d-4101-a195-f3b70fe0e33b",
+   "cell_type": "raw",
+   "id": "274fbd8b-2272-443c-b71d-4fcbc8946167",
    "metadata": {},
-   "outputs": [],
    "source": [
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",

notebooks/analysis.py

@@ -327,11 +327,12 @@ def crop_velmass_to_borg(ini_file, field_type):
     return true_field
 
 
-def get_spectra(ini_name, dirname, nframe, iter_max, iter_min, which_field='delta', cut_field=True, mock_type='borg'):
+def get_spectra(ini_name, dirname, nframe, iter_max, iter_min, which_field='delta', cut_field=True, mock_type='borg', MAS=''):
     
     # Which steps to use
     all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min)
     
+    if MAS == '':
         if which_field == 'delta':
             MAS = "CIC"
         elif which_field == 'ics':
@@ -350,6 +351,11 @@ def get_spectra(ini_name, dirname, nframe, iter_max, iter_min, which_field='delt
             boxsize = get_borg_Lbox(ini_name)
     else:
         raise NotImplementedError
+        
+    # Turn to overdensity
+    if which_field == 'delta':
+        delta1 = (1. + delta1) / (1. + delta1).mean() - 1.
+        
     print("BOXSIZE", boxsize)
     Pk = PKL.Pk(delta1.astype(np.float32), boxsize, axis=0, MAS=MAS, threads=1, verbose=True)
     k      = Pk.k3D

scripts/build_borg.sh

@@ -23,11 +23,11 @@ BUILD_DIR=/data101/bartlett/build_borg/
 
 cd $BORG_DIR
 # git pull
-# rm -rf $BUILD_DIR
+rm -rf $BUILD_DIR
 #bash build.sh --c-compiler $(which x86_64-conda_cos6-linux-gnu-gcc) --cxx-compiler $(which x86_64-conda_cos6-linux-gnu-g++) --python --hades-python --build-dir $BUILD_DIR
 bash build.sh --c-compiler $(which x86_64-conda_cos6-linux-gnu-gcc) --cxx-compiler $(which x86_64-conda_cos6-linux-gnu-g++) --python=$(which python3) --install-system-python --hades-python --use-system-hdf5 --build-dir /data101/bartlett/build_borg/
 cd $BUILD_DIR
-#make -j 32
+make -j 14
 
 
 exit 0

scripts/submit_borg.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-#SBATCH --job-name=velmass_ics_model
+#SBATCH --job-name=velmass_ics_model_cola
 #SBATCH --nodes=1
 #SBATCH --exclusive
 #SBATCH --ntasks=40
@@ -21,14 +21,15 @@ module load cuda/12.6
 source /home/bartlett/.bashrc
 source /home/bartlett/anaconda3/etc/profile.d/conda.sh
 conda deactivate
-conda activate borg_env
+# conda activate borg_env 
+conda activate borg_new
 
 # Kill job if there are any errors
 set -e
 
 # Path variables
 BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python
-RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model
+RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model_cola_v3
 
 mkdir -p $RUN_DIR
 cd $RUN_DIR
@@ -42,9 +43,9 @@ set -x
 
 # Run BORG
 INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/velmass_ini.ini
-cp $INI_FILE ini_file.ini
-$BORG INIT ini_file.ini
-# $BORG RESUME ini_file.ini
+# cp $INI_FILE ini_file.ini
+# $BORG INIT ini_file.ini
+$BORG RESUME ini_file.ini
 
 conda deactivate
 

tests/gravity_models.py

@@ -0,0 +1,195 @@
+import aquila_borg as borg
+import numpy as np
+import configparser
+import borg_velocity.utils as utils
+from borg_velocity.utils import myprint
+import borg_velocity.forwards as forwards
+import Pk_library as PKL
+import matplotlib.pyplot as plt
+import symbolic_pofk.syrenhalofit as syrenhalofit
+
+def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_file=None, gravity='lpt') -> 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.
+
+    """
+    
+    myprint(f"Building gravity model {gravity}")
+
+    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'])
+    print('Supersampling:', supersampling)
+    
+    if gravity in ['pm', 'cola']:
+        forcesampling = int(config['model']['forcesampling'])
+        forcesampling = 2
+        print('FORCESAMPLING:', 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 gravity == 'linear':
+        raise NotImplementedError(gravity)
+    elif 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 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 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 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(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
+
+ini_file = '../conf/velmass_ini.ini'
+
+# Setup config
+config = configparser.ConfigParser()
+config.read(ini_file)
+
+# Cosmology
+cosmo = utils.get_cosmopar(ini_file)
+
+# Input box
+box_in = borg.forward.BoxModel()
+box_in.L = (float(config['system']['L0']), float(config['system']['L1']), float(config['system']['L2']))
+box_in.N = (int(config['system']['N0']), int(config['system']['N1']), int(config['system']['N2']))
+# box_in.N = (128, 128, 128)
+# box_in.N = (256, 256, 256)
+box_in.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2']))
+
+# Make some initial conditions
+s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
+s_real = np.fft.irfftn(s_hat, norm="ortho")
+
+all_gravity = ['lpt', '2lpt', 'pm', 'cola']
+all_dens = [None] * len(all_gravity)
+all_pk = [None] * len(all_gravity)
+
+fig, ax = plt.subplots(1, 1, figsize=(6,4))
+
+for i, gravity in enumerate(all_gravity):
+    # Run BORG density field
+    output_density = np.zeros(box_in.N)
+    fwd_model = build_gravity_model(None, box_in, ini_file=ini_file, gravity=gravity)
+    fwd_model.forwardModel_v2(s_hat)
+    fwd_model.getDensityFinal(output_density)
+    all_dens[i] = output_density.copy()
+    
+    Pk = PKL.Pk(all_dens[i].astype(np.float32), box_in.L[0], axis=0, MAS='CIC', threads=1, verbose=True) 
+    all_pk[i] = (Pk.k3D.copy(), Pk.Pk[:,0].copy())
+
+    ax.loglog(all_pk[i][0], all_pk[i][1], label=gravity)
+    
+# Load VELMASS field
+dirname = config['mock']['velmass_dirname']
+field = np.load(dirname + '/../density.npz')['d']
+field = field / np.mean(field) - 1
+with open(dirname + '/auto-rockstar.cfg') as f:
+    data = [r for r in f]
+Lbox = [r for r in data if r.startswith('BOX_SIZE')][0].strip()
+Lbox = float(Lbox.split('=')[1])
+print('VELMASS box', Lbox)
+Pk = PKL.Pk(field.astype(np.float32), Lbox, axis=0, MAS='CIC', threads=1, verbose=True)
+m = Pk.k3D > all_pk[-1][0].min()
+ax.loglog(Pk.k3D[m], Pk.Pk[m,0], label='velmass')
+
+# Syren
+k = Pk.k3D[m]
+pk_syrenhalofit = syrenhalofit.run_halofit(k, 
+                       cosmo.sigma8, cosmo.omega_m, cosmo.omega_b, cosmo.h, cosmo.n_s, 1.0,
+                       emulator='fiducial', extrapolate=True, which_params='Bartlett', add_correction=True)
+ax.loglog(k, pk_syrenhalofit, label='syren')
+ymin = max(ax.get_ylim()[0], 1e1)
+ax.set_ylim(ymin, None)    
+ax.legend()
+fig.tight_layout()
+fig.savefig('../figs/pk_gravity_models_test.png', facecolor='white', bbox_inches='tight')
+plt.close(fig)

tests/likelihood_forwards.py

@@ -0,0 +1,249 @@
+import borg_velocity.likelihood as likelihood
+import borg_velocity.forwards as forwards
+import borg_velocity.utils as utils
+from borg_velocity.utils import myprint
+
+import aquila_borg as borg
+import configparser
+import h5py as h5
+import numpy as np
+from tqdm import tqdm
+import glob
+import os
+import matplotlib.pyplot as plt
+
+dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics'
+ini_file = f'{dirname}/ini_file.ini'
+
+wdir = os.getcwd()
+os.chdir(dirname)
+
+def get_mcmc_steps(nframe, iter_max, iter_min=0):
+    """
+    Obtain evenly-spaced sample of MCMC steps to make movie from
+    """
+    all_mcmc = glob.glob(dirname + '/mcmc_*.h5')
+    x = [m[len(dirname + '/mcmc_'):-3] for m in all_mcmc]
+    all_mcmc = np.sort([int(m[len(dirname + '/mcmc_'):-3]) for m in all_mcmc])
+    if iter_max >= 0:
+        all_mcmc = all_mcmc[all_mcmc <= iter_max]
+    all_mcmc = all_mcmc[all_mcmc >= iter_min]
+    if nframe > 0:
+        max_out = max(all_mcmc)
+        min_out = min(all_mcmc)
+        step = max(int((max_out - min_out+1) / nframe), 1)
+        all_mcmc = all_mcmc[::step]
+        if max_out not in all_mcmc:
+            all_mcmc = np.concatenate([all_mcmc, [max_out]])
+    return all_mcmc
+
+
+@borg.registerGravityBuilder
+def build_gravity_model(gravity: str, 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:
+        - gravity (str): Which gravity model to use (in ['pm', 'cola', 'lpt', '2lpt']
+        - 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 gravity in ['pm', 'cola']:
+        if 'forcesampling' in config['model']:
+            forcesampling = int(config['model']['forcesampling'])
+        else:
+            myprint('Could not find forcesampling, so will set to 4')
+            forcesampling = 4
+            
+    # 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 gravity == 'linear':
+        raise NotImplementedError(gravity)
+    elif 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 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 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 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(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
+
+
+def build_likelihood(gravity, box_in, ini_file):
+    
+    global fwd_vel
+    
+    model = build_gravity_model(gravity, None, box_in, ini_file=ini_file)
+    cosmo = utils.get_cosmopar(ini_file)
+    model.setCosmoParams(cosmo)
+    fwd_param = forwards.NullForward(box_in)
+    mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
+    
+    # Load the mock data
+    mylike.loadMockData(None)
+    
+    return mylike
+
+ref_gravity = 'lpt'
+ref_gravity = '2lpt'
+
+
+# Input box
+box_in = borg.forward.BoxModel()
+config = configparser.ConfigParser()
+config.read(ini_file)
+box_in.L = (float(config['system']['L0']), float(config['system']['L1']), float(config['system']['L2']))
+box_in.N = (int(config['system']['N0']), int(config['system']['N1']), int(config['system']['N2']))
+box_in.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2']))
+
+# Setup BORG forward model and likelihood
+cola_like = build_likelihood('cola', box_in, ini_file)
+lpt_like = build_likelihood(ref_gravity, box_in, ini_file)
+
+# Load some MCMC samples and see how they compare
+nframe = 100
+iter_min = 1000
+iter_max = -1
+all_mcmc = get_mcmc_steps(nframe, iter_max, iter_min=iter_min)
+
+all_lpt_like = np.empty(len(all_mcmc))
+all_cola_like = np.empty(len(all_mcmc))
+
+for i, mcmc in enumerate(all_mcmc):
+    myprint(f'\nStarting MCMC {i+1} of {len(all_mcmc)}: {mcmc}')
+
+    # Load an s_hat from file
+    with h5.File(f'{dirname}/mcmc_{mcmc}.h5', 'r') as f:
+        s_hat = f['scalars/s_hat_field'][:]
+        all_lpt_like[i] = lpt_like.logLikelihoodComplex(s_hat, None)
+        all_cola_like[i] = cola_like.logLikelihoodComplex(s_hat, None)
+
+fig, axs = plt.subplots(1, 2, figsize=(15,4))
+cmap = plt.cm.spring
+markersize = 50
+
+# # Remove first point
+# all_lpt_like = all_lpt_like[1:]
+# all_cola_like = all_cola_like[1:]
+# all_mcmc = all_mcmc[1:]
+
+# Normalise by the minimum value
+lmin = all_lpt_like.min()
+all_lpt_like -= lmin
+all_cola_like -= lmin
+
+sc0 = axs[0].scatter(all_lpt_like, all_cola_like, c=all_mcmc, cmap=cmap, marker='.', s=markersize)
+xlim = axs[0].get_xlim()
+axs[0].plot(xlim, xlim, color='k')
+axs[0].set_xlim(xlim)
+cbar = fig.colorbar(sc0, ax=axs[0], location='right', aspect=40)
+cbar.set_label("MCMC Step")
+
+sc1 = axs[1].scatter(all_lpt_like, all_cola_like - all_lpt_like, c=all_mcmc, cmap=cmap, marker='.', s=markersize)
+axs[1].axhline(0, color='k')
+cbar = fig.colorbar(sc1, ax=axs[1], location='right', aspect=40)
+cbar.set_label("MCMC Step")
+
+for ax in axs:
+    ax.set_xlabel(f'{ref_gravity} Log-Likelihood')
+axs[0].set_ylabel('COLA Log-Likelihood')
+axs[1].set_ylabel(f'COLA Log-Likelihood - {ref_gravity} Log-Likelihood')
+
+# fig.align_xlabels()
+# fig.align_ylabels()
+# fig.tight_layout()
+fig.savefig(f'{wdir}/likelihood_forwards.png', facecolor='white', bbox_inches='tight')

tests/tfr_inference.py

@@ -0,0 +1,393 @@
+import aquila_borg as borg
+import numpy as np
+from astropy.coordinates import SkyCoord
+import astropy.units as apu
+import astropy.constants
+from astropy.cosmology import LambdaCDM
+import pandas as pd
+from scipy.interpolate import interp1d
+import jax.numpy as jnp
+import jax.scipy.special
+
+import numpyro
+import numpyro.distributions as dist
+from jax import lax
+
+import borg_velocity.poisson_process as poisson_process
+import borg_velocity.projection as projection
+
+# Output stream management
+cons = borg.console()
+myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x))
+
+def build_gravity_model(box, cpar, ai=0.05, af=1.0, nsteps=20, forcesampling=4, supersampling=2, rsmooth=4.0, gravity='lpt', velmodel_name='CICModel'):
+    """
+    Builds the gravity model and returns the forward model chain.
+
+    Args:
+        - box (borg.forward.BoxModel): Box within which to run simulation
+        - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use
+        - ai (float, default=0.05): Scale factor to begin simulation
+        - af (float, default=1.0): Scale factor to end simulation
+        - nsteps (int, default=20): Number of steps to use in the simulation
+        - forcesampling (int, default=4): Sampling factor for force evaluations
+        - supersampling (int, default=2): Supersampling factor of particles
+        - rsmooth (float, default=4.0): Smoothing scale for velocity field (Mpc/h)
+        - gravity (str, default='lpt'): Which gravity model to use
+        - velmodel_name (str, default='CICModel'): Which velocity estimator to use
+    
+    Returns:
+        - chain (borg.forward.BaseForwardModel): The forward model for density
+        - fwd_vel (borg.forward.VelocityBase): The forward model for velocity
+
+    """
+    
+    myprint(f"Building gravity model {gravity}")
+
+    # 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 gravity == 'linear':
+        raise NotImplementedError(gravity)
+    elif 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 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 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=nsteps,
+                      tcola=False))
+    elif 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=nsteps,
+                      tcola=True))
+    else:
+        raise NotImplementedError(gravity)
+        
+    mod.accumulateAdjoint(True)
+    chain @= mod
+        
+    # Cosmological parameters
+    chain.setCosmoParams(cpar)
+    
+    # This is the forward model for velocity
+    velmodel = getattr(borg.forward.velocity, velmodel_name)
+    if velmodel_name == 'LinearModel':
+        fwd_vel = velmodel(box, mod, af)
+    elif velmodel_name == 'CICModel':
+        fwd_vel = velmodel(box, mod, rsmooth)
+    else:
+        fwd_vel = velmodel(box, mod)
+    
+    return chain, fwd_vel
+
+
+def get_fields(L, N, xmin, gravity='lpt', velmodel_name='CICModel'):
+    """
+    Obtain a density and velocity field to use for mock
+
+    Args:
+        - L (float): Box length (Mpc/h)
+        - N (int): Number of grid cells per side
+        - xmin (float): Coordinate of corner of the box (Mpc/h)
+        - gravity (str, default='lpt'): Which gravity model to use
+        - velmodel_name (str, default='CICModel'): Which velocity estimator to use
+    
+    Returns:
+        - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use
+        - output_density (np.ndarray): Over-density field
+        - output_velocity (np.ndarray): Velocity field
+
+    """
+    
+    # Setup box and cosmology
+    cpar = borg.cosmo.CosmologicalParameters()
+    cpar.default()
+    box = borg.forward.BoxModel()
+    box.L = (L, L, L)
+    box.N = (N, N, N)
+    box.xmin = (xmin, xmin, xmin)
+    
+    # Get forward models
+    fwd_model, fwd_vel = build_gravity_model(box, cpar, gravity=gravity, velmodel_name=velmodel_name)
+    
+    # Make some initial conditions
+    s_hat = np.fft.rfftn(np.random.randn(*box.N)) / box.Ntot ** (0.5)
+    
+    # Obtain density and velocity fields
+    output_density = np.zeros(box.N)
+    fwd_model.forwardModel_v2(s_hat)
+    fwd_model.getDensityFinal(output_density)
+    output_velocity = fwd_vel.getVelocityField()
+    
+    return cpar, output_density, output_velocity
+
+
+def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh, 
+                a_TFR, b_TFR, sigma_TFR, sigma_m, sigma_eta, 
+                hyper_eta_mu, hyper_eta_sigma, sigma_v, interp_order=1, bias_epsilon=1e-7):
+    
+    # Initialize lists to store valid positions and corresponding sig_mu values
+    all_xtrue = np.empty((3, Nt))
+    all_mobs = np.empty(Nt)
+    all_etaobs = np.empty(Nt)
+
+    # Counter for accepted positions
+    accepted_count = 0
+
+    # Cosmology object needed for z <-> r
+    cosmo = LambdaCDM(H0 = cpar.h * 100, Om0 = cpar.omega_m, Ode0 = cpar.omega_q)
+    
+    # Precompute redshift-distance mapping
+    z_grid = np.logspace(-4, -1, 10000)  # Covers z ~ 0 to 0.1
+    dL_grid = cosmo.luminosity_distance(z_grid).value  # Luminosity distances in Mpc
+
+    # Create an interpolation function: distance -> redshift
+    dL_to_z = interp1d(dL_grid, z_grid, kind="cubic", fill_value="extrapolate")
+    
+    # Bias model
+    phi = (1. + dens + bias_epsilon) ** alpha
+    
+    # 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 = z_at_value(cosmo.comoving_distance, rtrue * apu.Mpc / cpar.h).value
+        zcosmo = dL_to_z(rtrue / cpar.h)
+        
+        # Compute luminosity distance
+        # DO I NEED TO DO /h???
+        dL = (1 + zcosmo) * rtrue / cpar.h  # Mpc/h
+        
+        # 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
+        mobs = mobs[m]
+        etaobs = etaobs[m]
+        xtrue = xtrue[:,m]
+        
+        # Calculate how many valid positions we need to reach Nt
+        remaining_needed = Nt - accepted_count
+        selected_count = min(xtrue.shape[1], remaining_needed)
+        
+        # Append only the needed number of valid positions
+        all_xtrue[:,accepted_count:accepted_count+selected_count] = xtrue[:,:selected_count]
+        all_mobs = mobs[:selected_count]
+        all_etaobs = etaobs[:selected_count]
+        
+        # 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
+    
+    # Obtain total redshift
+    vr_noised = vr_true + sigma_v * np.random.randn(Nt)
+    zCMB = (1 + zcosmo) * (1 + vr_noised / astropy.constants.c.to('km/s').value) - 1
+    
+    return zCMB, all_mobs, all_etaobs, all_xtrue
+
+
+def estimate_data_parameters():
+    
+    """
+    ID   2MASS XSC ID name (HHMMSSss+DDMMSSs)
+    RAdeg   Right ascension (J2000)
+    DEdeg   Declination (J2000)
+    cz2mrs  Heliocentric redshift from the 2MRS (km/s)
+    Kmag    NIR magnitudes in the K band from the 2MRS (mag)
+    Hmag    NIR magnitudes in the H band from the 2MRS (mag)
+    Jmag    NIR magnitudes in the J band from the 2MRS (mag)
+    e_Kmag    Error of the NIR magnitudes in K band from the (mag)
+    e_Hmag    Error of the NIR magnitudes in H band from the (mag)
+    e_Jmag    Error of the NIR magnitudes in J band from the (mag)
+    WHIc    Corrected HI width (km/s)
+    e_WHIc    Error of corrected HI width (km/s)
+    """
+    
+    columns = ['ID', 'RAdeg', 'DEdeg', 'cz2mrs', 'Kmag', 'Hmag', 'Jmag', 'e_Kmag', 'e_Hmah', 'e_Jmag', 'WHIc', 'e_WHIc']
+    fname = '/data101/bartlett/fsigma8/PV_data/2MASS/table1.dat'
+    df = pd.read_csv(fname, sep='\s+', names=columns)
+    
+    sigma_m = np.median(df['e_Kmag'])
+
+    eta = np.log10(df['WHIc']) - 2.5
+    sigma_eta = np.median(df['e_WHIc'] / df['WHIc'] / np.log(10))
+    
+    hyper_eta_mu = np.median(eta)
+    hyper_eta_sigma = (np.percentile(eta, 84) - np.percentile(eta, 16)) / 2
+    
+    return sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma
+
+
+
+def likelihood(a_TFR, b_TFR, sigma_TFR, eta_true, m_true):
+    
+    loglike = 0
+    
+    # Apparent magnitude terms
+    norm = 0.5 * (1 + jax.scipy.special.erf((mthresh - m_true) / (jnp.sqrt(2) * sigma_m)))
+    loglike += 
+        0.5 * jnp.sum((mobs - m_true) ** 2 / sigma_m ** 2) 
+        + jnp.sum(jnp.log(norm)) 
+        + Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2)
+        
+    # Linewidth terms
+    loglike += 
+        0.5 * jnp.sum((etaobs - eta_true) ** 2 / sigma_eta ** 2) 
+        + Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_eta ** 2)
+        
+    # Peculiar velocity terms
+    
+    return
+
+
+def likelihood_model():
+    
+    # TO DO: Sort out these priors
+    a_TFR =  numpyro.sample("a_TFR", dist.Uniform(*alpha_prior))
+    b_TFR =  numpyro.sample("b_TFR", dist.Uniform(*alpha_prior))
+    sigma_TFR = numpyro.sample("sigma_TFR", dist.Uniform(*alpha_prior))
+    
+    eta_true = numpyro.sample("eta_true", dist.Normal(mean, sigma), sample_shape=(Nt,))
+    m_true = numpyro.sample("m_true", dist.Normal(mean, sigma), sample_shape=(Nt,))
+    
+    numpyro.sample("obs", TFRLikelihood(a_TFR, b_TFR, sigma_TFR, eta_true, m_true))
+    
+    return
+
+
+class TFRLikelihood(dist.Distribution):
+    support = dist.constraints.real
+    
+    def __init__(self, a_TFR, b_TFR, sigma_TFR, eta_true, m_true):
+        self.a_TFR, self.b_TFR, self.sigma_TFR, self.eta_true, self.m_true = dist.util.promote_shapes(a_TFR, b_TFR, sigma_TFR, eta_true, m_true)
+        batch_shape = lax.broadcast_shapes(
+                jnp.shape(a_TFR),
+                jnp.shape(b_TFR),
+                jnp.shape(sigma_TFR),
+                jnp.shape(eta_true),
+                jnp.shape(m_true),
+                )
+        super(TFRLikelihood, self).__init__(batch_shape = batch_shape)
+        
+    def sample(self, key, sample_shape=()):
+        raise NotImplementedError
+    
+    def log_prov(self, value)
+        return likelihood(self.a_TFR, self.b_TFR, self.sigma_TFR, self.eta_true, self.m_true)
+
+
+def main():
+    
+    # Get some parameters from the data
+    sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma = estimate_data_parameters()
+    
+    # Other parameters to use
+    L = 500.0
+    N = 64
+    xmin = -L/2
+    Rmax = 100
+    Nt = 2000
+    alpha = 1.4
+    mthresh = 11.25
+    a_TFR = -23
+    b_TFR = -8.2
+    sigma_TFR = 0.3
+    sigma_v = 150
+    
+    cpar, dens, vel = get_fields(L, N, xmin)
+    
+    zCMB, mobs, etaobs, xtrue = create_mock(
+            Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,
+            a_TFR, b_TFR, sigma_TFR, sigma_m, sigma_eta, 
+            hyper_eta_mu, hyper_eta_sigma, sigma_v)
+    
+
+if __name__ == "__main__":
+    main()

tests/velocity_bug.py

@@ -0,0 +1,144 @@
+import aquila_borg as borg
+import numpy as np
+
+# Output stream management
+cons = borg.console()
+myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x))
+
+def build_gravity_model(box: borg.forward.BoxModel, gravity='lpt', velmodel_name='LinearModel') -> borg.forward.BaseForwardModel:
+    """
+    Builds the gravity model and returns the forward model chain.
+
+    Args:
+        - box (borg.forward.BoxModel): The input box model.
+        - gravity (str, default='lpt'): The gravity model to use.
+        - velmodel_name (str, default='LinearModel'): The velocity model to use.
+
+    Returns:
+        - chain (borg.forward.BaseForwardModel): The forward model.
+        - fwd_vel (borg.forward.BaseForwardModel): The forward model for velocities.
+
+    """
+    
+    myprint(f"Building gravity model {gravity}")
+
+    ai = 0.05
+    af = 1.0
+    supersampling = 2
+    forcesampling = 4
+    rsmooth = 4.  # Mpc/h
+    nsteps = 20
+
+    # 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 gravity == 'linear':
+        raise NotImplementedError(gravity)
+    elif 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 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 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=nsteps,
+                      tcola=False))
+    elif 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=nsteps,
+                      tcola=True))
+    else:
+        raise NotImplementedError(gravity)
+        
+    mod.accumulateAdjoint(True)
+    chain @= mod
+        
+    # Cosmological parameters
+    cpar = borg.cosmo.CosmologicalParameters()
+    cpar.default()
+    chain.setCosmoParams(cpar)
+    
+    # This is the forward model for velocity
+    velmodel = getattr(borg.forward.velocity, velmodel_name)
+    if velmodel_name == 'LinearModel':
+        fwd_vel = velmodel(box, mod, af)
+    elif velmodel_name == 'CICModel':
+        fwd_vel = velmodel(box, mod, rsmooth)
+    else:
+        fwd_vel = velmodel(box, mod)
+    
+    return chain, fwd_vel
+
+
+def get_velocity(fwd: borg.forward.BaseForwardModel, fwd_vel: borg.forward.BaseForwardModel, s_hat: np.ndarray):
+    
+    N = s_hat.shape[0]
+    output_density = np.zeros((N,N,N))
+    fwd.forwardModel_v2(s_hat)
+    myprint("getting density")
+    fwd.getDensityFinal(output_density)
+    myprint("getting velocity")
+    output_velocity = fwd_vel.getVelocityField()
+    
+    return output_velocity
+
+
+# Input box
+box_in = borg.forward.BoxModel()
+box_in.L = (500.0, 500.0, 500.0)
+box_in.N = (64, 64, 64)
+box_in.xmin = (-250.0, -250.0, -250.0)
+
+# Make some initial conditions
+s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
+s_real = np.fft.irfftn(s_hat, norm="ortho")
+
+for gravity in ['lpt', '2lpt', 'pm', 'cola']:
+    
+    print(f'\nGravity: {gravity}')
+    
+    fwd, fwd_vel = build_gravity_model(box_in, gravity=gravity, velmodel_name='LinearModel')
+    v = get_velocity(fwd, fwd_vel, s_hat)
+    print('Linear', v.mean(), v.std(), v.shape)
+    np.save(f'vel_linear_{gravity}.npy', v)
+    
+    fwd, fwd_vel = build_gravity_model(box_in, gravity=gravity, velmodel_name='CICModel')
+    v = get_velocity(fwd, fwd_vel, s_hat)
+    print('CIC', v.mean(), v.std(), v.shape)
+    np.save(f'vel_cic_{gravity}.npy', v)