Ensure tracers are inside box

borg_velocity/likelihood.py

@@ -13,6 +13,7 @@ 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
@@ -144,7 +145,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         
         if self.run_type == 'data':
             self.loadObservedData(make_plot=False)
-        else:
+        elif borg.EMBEDDED:
             if self.action == 'INIT':
                 pass  # Data will be loaded later
             elif self.action == 'RESUME':
@@ -260,6 +261,12 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
                 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()
@@ -285,6 +292,11 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
                 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()
 
@@ -312,6 +324,10 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         
         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}']
@@ -336,12 +352,13 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
                 self.R_max
             )
             
+            # self.lkl_ind[i] = temp_lkl_ind.copy()
+            
         # Add in bulk flow prior
         # sigma_bulk = get_sigma_bulk(self.L[0], self.fwd.getCosmoParams())
         # lkl += jnp.sum(0.5 * jnp.log(2 * np.pi) + jnp.log(sigma_bulk) + self.bulk_flow ** 2 / 2. / sigma_bulk ** 2)
             
-        # if not jnp.isfinite(lkl):
-        #     lkl = self.bignum
+        # lkl = jnp.clip(lkl, -self.bignum, self.bignum)
         lkl = lax.cond(
             jnp.isfinite(lkl),
             lambda _: lkl,         # If True (finite), return lkl
@@ -374,6 +391,11 @@ 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:
             # Run BORG density field
@@ -469,6 +491,9 @@ def vel2like(cz_obs, v, MB_field, MB_pos, r, r_hMpc, sig_mu, sig_v, omega_m, muA
                     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
@@ -489,6 +514,8 @@ def vel2like(cz_obs, v, MB_field, MB_pos, r, r_hMpc, sig_mu, sig_v, omega_m, muA
     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)
@@ -503,9 +530,13 @@ def vel2like(cz_obs, v, MB_field, MB_pos, r, r_hMpc, sig_mu, sig_v, omega_m, muA
     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
 

borg_velocity/mock_maker.py

@@ -6,6 +6,7 @@ import configparser
 import ast
 
 import borg_velocity.utils as utils
+from borg_velocity.utils import myprint
 import borg_velocity.forwards as forwards
 import borg_velocity.poisson_process as poisson_process
 import borg_velocity.projection as projection
@@ -75,10 +76,17 @@ def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
     # Sample positions according to bias model
     bias_epsilon = float(config['model']['bias_epsilon'])
     R_max = float(config['mock']['R_max'])
+    if config['model']['R_lim'] == 'none':
+        R_lim = fwd_model.getOutputBoxModel().L[0]/2
+    else:
+        R_lim = float(config['model']['R_lim'])
     coord_true = [None] * nsamp
     coord_meas = [None] * nsamp
     sig_mu = [None] * nsamp
     for i in range(nsamp):
+        
+        myprint(f'Making mock for sample {i}')
+        
         frac_sig_x = float(config[f'sample_{i}']['frac_sig_rhMpc'])
         alpha = float(config[f'sample_{i}']['alpha'])
         lam = int(config[f'sample_{i}']['lam'])
@@ -99,11 +107,52 @@ def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
         phi *= np.exp(- lam * r / R_max)
 
         # Sample coordinates
-        coord_true[i] = poisson_process.sample_3d(phi, Nt, 
-              fwd_model.getOutputBoxModel().L[0], fwd_model.getOutputBoxModel().xmin)
-        coord_meas[i], sig_mu[i] = radially_scatter(coord_true[i], frac_sig_x)
-
         
+        # coord_true[i] = poisson_process.sample_3d(phi, Nt, 
+        #       fwd_model.getOutputBoxModel().L[0], fwd_model.getOutputBoxModel().xmin)
+        # coord_meas[i], sig_mu[i] = radially_scatter(coord_true[i], frac_sig_x)
+
+        # Initialize lists to store valid positions and corresponding sig_mu values
+        coord_meas_valid = np.empty((3, Nt))
+        coord_true_valid = np.empty((3, Nt))
+
+        # Counter for accepted positions
+        accepted_count = 0
+
+        # Loop until we have Nt valid positions
+        while accepted_count < Nt:
+            # Generate positions
+            xtrue = poisson_process.sample_3d(phi, Nt, 
+                                                   fwd_model.getOutputBoxModel().L[0], 
+                                                   fwd_model.getOutputBoxModel().xmin)
+
+            # Apply radial scattering
+            xmeas, sig_mu[i] = radially_scatter(xtrue, frac_sig_x)
+
+            # Compute radial distances of each measured coordinate
+            radial_distances = np.sqrt(np.sum(xmeas ** 2, axis=0))
+            
+            # Filter coordinates that meet the R_lim condition
+            valid_indices = radial_distances < R_lim
+            valid_coords = xmeas[:,valid_indices]
+            valid_coords_true = xtrue[:,valid_indices]
+
+            # Calculate how many valid positions we need to reach Nt
+            remaining_needed = Nt - accepted_count
+            selected_count = min(valid_coords.shape[1], remaining_needed)
+
+            # Append only the needed number of valid positions
+            coord_meas_valid[:,accepted_count:accepted_count+selected_count] = valid_coords[:,:selected_count]
+            coord_true_valid[:,accepted_count:accepted_count+selected_count] = valid_coords_true[:,:selected_count]
+            
+            # Update the accepted count
+            accepted_count += selected_count
+            
+            myprint(f'\tMade {accepted_count} of {Nt}')
+        
+        coord_true[i] = coord_true_valid.copy()
+        coord_meas[i] = coord_meas_valid.copy()
+
     # Interpolate velocities to tracers
     interp_order = int(config['model']['interp_order'])
     vr_true = [None] * nsamp

borg_velocity/samplers.py

@@ -365,7 +365,6 @@ class BlackJaxBiasSampler(borg.samplers.PyBaseSampler):
             self.y[:] = bj_state.position
             
             # Save results to the borg state object
-            # state.newArray1d("inverse_mass_matrix", len(self.parameters['inverse_mass_matrix']), True)
             state["inverse_mass_matrix"][:] = self.parameters['inverse_mass_matrix']
             state["step_size"] = float(self.parameters['step_size'])
 
@@ -527,6 +526,11 @@ class TransformedBlackJaxBiasSampler(borg.samplers.PyBaseSampler):
             (state, self.parameters), info = warmup.run(warmup_key, self.y, num_steps=self.warmup_nsteps)
             # x = info.state.position
             self.y[:] = state.position
+            
+            # Save results to the borg state object
+            state["inverse_mass_matrix"][:] = self.parameters['inverse_mass_matrix']
+            state["step_size"] = float(self.parameters['step_size'])
+
 
         nuts = blackjax.nuts(logdensity_fn, **self.parameters)
         bias_state = nuts.init(self.y)

conf/supranta_ini.ini

@@ -15,25 +15,25 @@ test_mode = true
 seed_cpower = true
 
 [block_loop]
-hades_sampler_blocked = true
+hades_sampler_blocked = false
 bias_sampler_blocked= true
 nmean_sampler_blocked= true
 sigma8_sampler_blocked = true
 omega_m_sampler_blocked = true
-muA_sampler_blocked = true
-alpha_sampler_blocked = true
-lam_sampler_blocked = true
-sig_v_sampler_blocked = true
+muA_sampler_blocked = false
+alpha_sampler_blocked = false
+lam_sampler_blocked = false
+sig_v_sampler_blocked = false
 bulk_flow_sampler_blocked = false
 ares_heat = 1.0
 
 [mcmc]
-number_to_generate = 3
-warmup_model = 0
+number_to_generate = 15000
+warmup_model = 1000
 warmup_cosmo = 0
 random_ic = false
-init_random_scaling = 1.0
-bignum = 1e300
+init_random_scaling = 0.1
+bignum = 1e20
 
 [hades]
 algorithm = HMC
@@ -65,9 +65,9 @@ nsteps = 20
 smooth_R = 4
 bias_epsilon = 1e-7
 interp_order = 1
-rsmooth = 4.
+rsmooth = 8.
 sig_v = 150.
-R_lim = none
+R_lim = 220
 Nint_points = 201 
 Nsig = 10
 bulk_flow = [0.0, 0.0, 0.0]

scripts/run_borg.sh

@@ -9,8 +9,6 @@ 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 compiler_test
 conda activate borg_new
 
 # Kill job if there are any errors
@@ -18,7 +16,7 @@ set -e
 
 # Path variables
 BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python
-RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/test_dir_blackjax_restart
+RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v2
 
 mkdir -p $RUN_DIR
 cd $RUN_DIR

scripts/submit_borg.sh

@@ -1,14 +1,21 @@
-#!/bin/sh
-#PBS -S /bin/sh
-#PBS -N supranta_N64_hmc_modelpar
-#PBS -j oe
-#PBS -m ae
-#PBS -l nodes=h10:has1gpu:ppn=40,walltime=5:00:00
+#!/bin/bash
+#SBATCH --job-name=blackjax_model_ic_v2
+#SBATCH --nodes=1
+#SBATCH --exclusive
+#SBATCH --ntasks=40
+#SBATCH --ntasks-per-node=40
+#SBATCH --time=48:00:00
+#SBATCH --partition=pscomp
+#SBATCH --gres=gpu
+#SBATCH --output=/data101/bartlett/fsigma8/borg_velocity/out_files/velocity_inference_%j.out
+#SBATCH --error=/data101/bartlett/fsigma8/borg_velocity/out_files/velocity_inference_%j.err
+#SBATCH --mail-user=deaglan.bartlett@iap.fr
+#SBATCH --mail-type=END,FAIL  
 
 # Modules
 module purge
 module restore myborg
-module load cuda/12.3
+module load cuda/12.6
 
 # Environment
 source /home/bartlett/.bashrc
@@ -21,7 +28,7 @@ set -e
 
 # Path variables
 BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python
-RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/test_dir_hmc_modelpar
+RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v2
 
 mkdir -p $RUN_DIR
 cd $RUN_DIR
@@ -37,6 +44,7 @@ set -x
 INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/supranta_ini.ini
 cp $INI_FILE ini_file.ini
 $BORG INIT ini_file.ini
+# $BORG RESUME ini_file.ini
 
 conda deactivate
 

scripts/torque_submit_borg.sh

@@ -0,0 +1,47 @@
+#!/bin/sh
+#PBS -S /bin/sh
+#PBS -N supranta_N64_hmc_modelpar
+#PBS -j oe
+#PBS -m ae
+#PBS -l nodes=h10:has1gpu:ppn=40,walltime=5:00:00
+
+# Modules
+module purge
+module restore myborg
+module load cuda/12.3
+
+# Environment
+source /home/bartlett/.bashrc
+source /home/bartlett/anaconda3/etc/profile.d/conda.sh
+conda deactivate
+conda activate borg_env
+
+# 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/test_dir_hmc_modelpar
+
+mkdir -p $RUN_DIR
+cd $RUN_DIR
+
+# Create a custom file descriptor (3) for tracing
+exec 3>$RUN_DIR/trace_file.txt
+
+# Redirect trace output to the custom file descriptor
+BASH_XTRACEFD="3"
+set -x
+
+# Run BORG
+INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/supranta_ini.ini
+cp $INI_FILE ini_file.ini
+$BORG INIT ini_file.ini
+
+conda deactivate
+
+# Disable tracing and close the custom file descriptor
+set +x
+exec 3>&-
+
+exit 0

tests/test_likelihood.py

@@ -2,19 +2,26 @@ import aquila_borg as borg
 import configparser
 import numpy as np
 import matplotlib.pyplot as plt
+import h5py
+import os
+import jax.numpy as jnp
 
 import borg_velocity.likelihood as likelihood
 import borg_velocity.forwards as forwards
 import borg_velocity.utils as utils
+import borg_velocity.projection as projection
 
 ini_file = '../conf/supranta_ini.ini'
-test_scaling = True
-test_sigma8 = True
-test_omegam = True
-test_alpha = True
-test_lam = True
-test_muA = True
-test_bulk = True
+mock_dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic/'
+# mock_dirname = None
+test_scaling = False
+test_sigma8 = False
+test_omegam = False
+test_alpha = False
+test_lam = False
+test_muA = False
+test_bulk = False
+test_sample = True
 
 # Input box
 box_in = borg.forward.BoxModel()
@@ -36,12 +43,27 @@ mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
 state = borg.likelihood.MarkovState()
 mylike.initializeLikelihood(state)
 mylike.updateCosmology(cosmo)
-np.random.seed(2)
-s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
-mylike.generateMockData(s_hat, state)
+if mock_dirname is None:
+    print('Creating mock data')
+    np.random.seed(2)
+    s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
+    mylike.generateMockData(s_hat, state)
+else:
+    print('Loading mock data from', mock_dirname)
+    cwd = os.getcwd()
+    os.chdir(mock_dirname)
+    with h5py.File('mock_data.h5') as f:
+        s_hat = f['scalars/s_hat_field'][:]
+    mylike.loadMockData(state)
+    os.chdir(cwd)
+
+mylike.logLikelihoodComplex(s_hat, None)
+mylike.gradientLikelihoodComplex(s_hat) 
+
+quit()
 
 if test_scaling:
-    all_scale = np.linspace(0.5, 1.5, 100)
+    all_scale = np.linspace(0.2, 1.8, 100)
     all_lkl = np.empty(all_scale.shape)
     for i, scale in enumerate(all_scale):
         all_lkl[i] = mylike.logLikelihoodComplex(scale * s_hat, None)
@@ -214,4 +236,75 @@ if test_bulk:
     fig.savefig('../figs/bulk_test.png')
     fig.clf()
     plt.close(fig)
+    
+    
+if test_sample:
+    with h5py.File(f'{mock_dirname}/mcmc_96.h5') as f:
+        sample_s_hat = f['scalars/s_hat_field'][:]
+    mylike.logLikelihoodComplex(s_hat, None)
+    np.savez(f'{mock_dirname}/mock_lkl_ind.npz', *mylike.lkl_ind)
+    mylike.logLikelihoodComplex(sample_s_hat, None)
+    np.savez(f'{mock_dirname}/mcmc_lkl_ind.npz', *mylike.lkl_ind)
+
+    print(s_hat.shape)
+    N = s_hat.shape[0]
+    
+    # Run BORG density field - s_hat
+    print('\nTrue')
+    output_density = np.zeros((N,N,N))
+    mylike.fwd.forwardModel_v2(s_hat)
+    mylike.fwd.getDensityFinal(output_density)
+    output_velocity = mylike.fwd_vel.getVelocityField()
+    np.save(f'{mock_dirname}/mock_vel.npy', output_velocity)
+    tracer_vr = [None] * mylike.nsamp
+    tracer_vel = [None] * mylike.nsamp
+    print('3D', output_velocity.min(), output_velocity.max())
+    for i in range(mylike.nsamp):
+        tracer_vel[i] = projection.interp_field(
+                                output_velocity,
+                                mylike.MB_pos[i],
+                                mylike.L[0], 
+                                jnp.array(mylike.fwd.getOutputBoxModel().xmin),
+                                1,
+                                use_jitted=True,
+                                )
+        tracer_vr[i] = projection.project_radial(
+                    tracer_vel[i],
+                    mylike.MB_pos[i],
+                    jnp.zeros(3,)
+                    )
+        print(i, tracer_vr[i].min(), tracer_vr[i].max(), tracer_vel[i].min(), tracer_vel[i].max())
+    np.savez(f'{mock_dirname}/mock_tracer_vr.npz', *tracer_vr)
+    np.savez(f'{mock_dirname}/mock_tracer_vel.npz', *tracer_vel)
+    
+    # Run BORG density field - sample_s_hat
+    print('\nMCMC Sample')
+    output_density = np.zeros((N,N,N))
+    mylike.fwd.forwardModel_v2(sample_s_hat)
+    mylike.fwd.getDensityFinal(output_density)
+    output_velocity = mylike.fwd_vel.getVelocityField()
+    np.save(f'{mock_dirname}/mcmc_vel.npy', output_velocity)
+    tracer_vr = [None] * mylike.nsamp
+    tracer_vel = [None] * mylike.nsamp
+    print('3D', output_velocity.min(), output_velocity.max())
+    for i in range(mylike.nsamp):
+        tracer_vel[i] = projection.interp_field(
+                                output_velocity,
+                                mylike.MB_pos[i],
+                                mylike.L[0], 
+                                jnp.array(mylike.fwd.getOutputBoxModel().xmin),
+                                1,
+                                use_jitted=True,
+                                )
+        tracer_vr[i] = projection.project_radial(
+                    tracer_vel[i],
+                    mylike.MB_pos[i],
+                    jnp.zeros(3,)
+                    )
+        print(i, tracer_vr[i].min(), tracer_vr[i].max(), tracer_vel[i].min(), tracer_vel[i].max())
+    np.savez(f'{mock_dirname}/mcmc_tracer_vr.npz', *tracer_vr)
+    np.savez(f'{mock_dirname}/mcmc_tracer_vel.npz', *tracer_vel)
+    
+    # Save the MB points
+    np.savez(f'{mock_dirname}/MB_pos.npz', *mylike.MB_pos)