Add VELMASS mock maker

README.md

@@ -17,7 +17,7 @@ pip install -U "jax[cuda12]=0.4.31”
 pip install blackjax
 
 cd /home/bartlett/borg
-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/
+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 /data101/bartlett/build_borg/
 make -j 32
 make python-install

borg_velocity/likelihood.py

@@ -245,7 +245,8 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         if self.run_type == 'data':
             raise NotImplementedError
         elif self.run_type == 'velmass':
-            raise NotImplementedError
+            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)
@@ -357,7 +358,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
             # 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) + self.bulk_flow ** 2 / 2. / self.sigma_bulk ** 2)
+        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(

borg_velocity/mock_maker.py

@@ -6,7 +6,7 @@ import configparser
 import ast
 
 import borg_velocity.utils as utils
-from borg_velocity.utils import myprint
+from borg_velocity.utils import myprint, parse_file_to_dict
 import borg_velocity.forwards as forwards
 import borg_velocity.poisson_process as poisson_process
 import borg_velocity.projection as projection
@@ -47,6 +47,8 @@ def radially_scatter(xtrue, frac_sig_x):
 
 def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
     
+    myprint('Making mock from BORG')
+    
     config = configparser.ConfigParser()
     config.read(ini_file)
     
@@ -89,7 +91,7 @@ def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
         
         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'])
+        lam = float(config[f'sample_{i}']['lam'])
         Nt = int(config[f'sample_{i}']['Nt'])
         phi = (1. + output_density + bias_epsilon) ** alpha
         
@@ -186,3 +188,153 @@ def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
         coord_meas[i] = coord_meas[i] * muA
         
     return coord_true, coord_meas, sig_mu, vr_true, cz_obs
+
+
+def velmass_mock(ini_file, seed=None):
+    
+    myprint('Making mock from VELMASS')
+    
+    if seed is not None:
+        np.random.seed(seed)
+    
+    config = configparser.ConfigParser()
+    config.read(ini_file)
+    
+    dirname = config['mock']['velmass_dirname']
+    
+    with open(dirname + "/halos_0.0.ascii", "rb") as f:
+        for _ in range(30):
+            r = str(f.readline())
+            if "Units" in r:
+                myprint(r)
+
+    # Get box size
+    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])
+    origin = np.array([Lbox/2, Lbox/2, Lbox/2])    
+    omega_m = [r for r in data if r.startswith('Om')][0].strip()
+    omega_m = float(omega_m.split('=')[1])
+    
+    halo_file = dirname + '/out.npy'
+    halos = np.load(halo_file)
+
+    # Cut by mass
+    Mmin = float(config['mock']['velmass_Mmin'])
+    halos = halos[halos['Mvir'] > Mmin]
+
+    # Load positions
+    xtrue = halos['X']  - origin[0] # Mpc / h
+    ytrue = halos['Y']  - origin[1] # Mpc / h
+    ztrue = halos['Z']  - origin[2] # Mpc / h
+    rtrue = np.sqrt(xtrue**2 + ytrue**2 + ztrue**2) # Mpc / h
+    xtrue = np.array([xtrue, ytrue, ztrue]) # shape = (3, nhalos)
+    
+    # Load velocities
+    vtrue = np.array([halos['VX'], halos['VY'], halos['VZ']])
+    
+    # Get relative probability of each object according to selection function
+    R_max = float(config['mock']['R_max'])
+    
+    # Define probability array. This is needed to prevent an object being in two samples
+    prob = np.ones(len(rtrue))
+    
+    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'])
+    
+    # Sample 
+    nsamp = int(config['run']['nsamp'])
+    coord_true = [None] * nsamp
+    coord_meas = [None] * nsamp
+    vr_true = [None] * nsamp
+    sig_mu = [None] * nsamp
+    for i in range(nsamp):
+        
+        myprint(f'Making mock for sample {i}')
+        lam = float(config[f'sample_{i}']['lam'])
+        frac_sig_x = float(config[f'sample_{i}']['frac_sig_rhMpc'])
+        Nt = int(config[f'sample_{i}']['Nt'])
+        
+        # Radial selection function
+        m = prob != 0  # objects already found
+        prob[m] = np.exp(- lam * rtrue[m] / R_max)
+        prob /= np.sum(prob)
+        
+        # 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))
+        tracer_vel = np.empty((3, Nt))
+
+        # Counter for accepted positions
+        accepted_count = 0
+
+        # Loop until we have Nt valid positions
+        while accepted_count < Nt:
+
+            ids = np.random.choice(len(rtrue), size=Nt, p=prob, replace=False)
+
+            # Apply radial scattering
+            xmeas, sig_mu[i] = radially_scatter(xtrue[:,ids], 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
+            ids = ids[valid_indices]
+            xmeas = xmeas[:,valid_indices]
+            radial_distances = radial_distances[valid_indices]
+
+            # Calculate how many valid positions we need to reach Nt
+            remaining_needed = Nt - accepted_count
+            selected_count = min(len(ids), remaining_needed)
+            ids = ids[:selected_count]
+            
+            # Append only the needed number of valid halos
+            coord_meas_valid[:,accepted_count:accepted_count+selected_count] = xmeas[:,:selected_count]
+            coord_true_valid[:,accepted_count:accepted_count+selected_count] = xtrue[:,ids]
+            tracer_vel[:,accepted_count:accepted_count+selected_count] = vtrue[:,ids]
+
+            # Update the accepted count
+            accepted_count += selected_count
+
+            myprint(f'\tMade {accepted_count} of {Nt}')
+
+            # Set up for next iteration
+            prob[ids] = 0
+            prob /= np.sum(prob)
+
+        coord_true[i] = coord_true_valid.copy()
+        coord_meas[i] = coord_meas_valid.copy()
+        
+        vr_true[i] = np.squeeze(projection.project_radial(
+                                np.expand_dims(tracer_vel, axis=2),
+                                np.expand_dims(coord_true[i], axis=2),
+                                np.zeros(3,)
+                                ))
+        
+    # Get the correct omega_m
+    config_dict = parse_file_to_dict(f'{dirname}/rockstar.cfg')
+    omega_m = config_dict['Om']
+    
+    # Compute observed redshifts (including noise - ACTUALLY NOT YET)
+    sig_v = float(config['model']['sig_v'])
+    cz_obs = [None] * nsamp
+    for i in range(nsamp):
+        rtrue = np.sqrt(np.sum(coord_true[i] ** 2, axis=0))
+        zco = utils.z_cos(rtrue, omega_m)
+        cz_obs[i] = utils.speed_of_light * zco + (1 + zco) * vr_true[i]
+        # cz_obs[i] += np.random.normal(size=cz_obs[i].shape) * sig_v  # CHECK THIS LINE!!!!!
+        
+    # Add observational systematic due to incorrect distance estimate 
+    # \mu -> \mu + 5 log10(A), or equivalently d -> A d
+    for i in range(nsamp):
+        muA = float(config[f'sample_{i}']['muA'])
+        coord_meas[i] = coord_meas[i] * muA
+        
+    return coord_true, coord_meas, sig_mu, vr_true, cz_obs
+ 

borg_velocity/samplers.py

@@ -246,8 +246,8 @@ class MVSliceBiasSampler(borg.samplers.PyBaseSampler):
             self._update_attrs(x)
             return -self.likelihood.logLikelihoodComplex(x_hat, False)
 
-        self.y[:] = borg.samplers.mv_slice_sampler(
-            state, _callback, self.y, self.mvs_state, self.lam
+        self.y[:] = borg.samplers.mpi_mv_slice_sampler(
+            state, _callback, self.y, self.mvs_state, self.lam, mpi=None
         )
         self._update_attrs(self.y)
 

borg_velocity/utils.py

@@ -169,6 +169,31 @@ def get_sigma_bulk(R, cpar):
     return sigma
 
 
+def parse_file_to_dict(file_path):
+    config_dict = {}
+    with open(file_path, 'r') as file:
+        for line in file:
+            line = line.strip()
+            if not line or line.startswith("#"):  # Skip empty lines and comments
+                continue
+            key, value = line.split("=", 1)
+            key = key.strip()
+            value = value.strip()
+
+            # Convert the value to the appropriate type (int, float, or leave as string)
+            if value.isdigit():
+                value = int(value)
+            else:
+                try:
+                    value = float(value)
+                except ValueError:
+                    pass  # Leave it as a string if it cannot be converted
+
+            config_dict[key] = value
+
+    return config_dict
+
+
 if __name__ == "__main__":
     
     cpar = get_cosmopar('../conf/basic_ini.ini')

conf/basic_ini.ini

@@ -75,7 +75,7 @@ bulk_flow = [-200.0, 200.0]
 omega_r = 0
 fnl = 0
 omega_k = 0
-omega_m = 0.335
+omega_m = 0.315
 omega_b = 0.049
 omega_q = 0.685
 h100 = 0.68

conf/supranta_ini.ini

@@ -15,24 +15,24 @@ test_mode = true
 seed_cpower = true
 
 [block_loop]
-hades_sampler_blocked = false
+hades_sampler_blocked = true
 bias_sampler_blocked= true
 nmean_sampler_blocked= true
 sigma8_sampler_blocked = true
 omega_m_sampler_blocked = true
-muA_sampler_blocked = false
-alpha_sampler_blocked = false
-lam_sampler_blocked = false
+muA_sampler_blocked = true
+alpha_sampler_blocked = true
+lam_sampler_blocked = true
 sig_v_sampler_blocked = false
-bulk_flow_sampler_blocked = false
+bulk_flow_sampler_blocked = true
 ares_heat = 1.0
 
 [mcmc]
 number_to_generate = 15000
-warmup_model = 500
+warmup_model = 0
 warmup_cosmo = 0
 random_ic = false
-init_random_scaling = 0.1
+init_random_scaling = 1.0
 bignum = 1e20
 
 [hades]
@@ -42,7 +42,7 @@ max_timesteps = 50
 mixing = 1
 
 [sampling]
-algorithm = TransformedBlackJax
+algorithm = mvslice
 epsilon = 0.005
 Nsteps = 20
 refresh = 0.1
@@ -51,7 +51,7 @@ warmup_nsteps = 400
 warmup_target_acceptance_rate = 0.7
 mvlam_mua = 0.01
 mvlam_alpha = 0.5
-mvlam_lam = 1.0
+mvlam_lam = 2.0
 mvlam_bulk_flow = 20
 mvlam_sig_v = 10
 
@@ -85,7 +85,7 @@ bulk_flow = [-200.0, 200.0]
 omega_r = 0
 fnl = 0
 omega_k = 0
-omega_m = 0.335
+omega_m = 0.315
 omega_b = 0.049
 omega_q = 0.685
 h100 = 0.68
@@ -104,6 +104,8 @@ NSAMP = 4
 [mock]
 seed = 12345
 R_max = 100
+velmass_dirname = /data54/lavaux/VELMASS/halo_central/halos_new/
+velmass_Mmin = 3e12
 
 [python]
 likelihood_path = /home/bartlett/fsigma8/borg_velocity/borg_velocity/likelihood.py

conf/velmass_ini.ini

@@ -0,0 +1,139 @@
+[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
+muA_sampler_blocked = true
+alpha_sampler_blocked = true
+lam_sampler_blocked = true
+sig_v_sampler_blocked = true
+bulk_flow_sampler_blocked = true
+ares_heat = 1.0
+
+[mcmc]
+number_to_generate = 15000
+warmup_model = 0
+warmup_cosmo = 0
+random_ic = false
+init_random_scaling = 0.1
+bignum = 1e20
+
+[hades]
+algorithm = HMC
+max_epsilon = 0.01
+max_timesteps = 50
+mixing = 1
+
+[sampling]
+algorithm = transformedblackjax
+epsilon = 0.005
+Nsteps = 20
+refresh = 0.1
+rng_seed = 1
+warmup_nsteps = 400
+warmup_target_acceptance_rate = 0.7
+mvlam_mua = 0.01
+mvlam_alpha = 0.5
+mvlam_lam = 2.0
+mvlam_bulk_flow = 20
+mvlam_sig_v = 10
+
+[model]
+gravity = lpt
+supersampling = 2
+velocity = CICModel
+af = 1.0
+ai = 0.05
+nsteps = 20
+smooth_R = 4
+bias_epsilon = 1e-7
+interp_order = 1
+rsmooth = 8.
+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]
+muA = [0.5, 1.5]
+alpha = [0.0, 10.0]
+lam = [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 = velmass
+NCAT = 0
+NSAMP = 4
+
+[mock]
+seed = 12345
+R_max = 100
+velmass_dirname = /data54/lavaux/VELMASS/halo_central/halos_new/
+velmass_Mmin = 3e12
+
+[python]
+likelihood_path = /home/bartlett/fsigma8/borg_velocity/borg_velocity/likelihood.py
+
+[sample_0]
+Nt = 345
+muA = 1.0
+alpha = 1.4
+lam = 5
+frac_sig_rhMpc = 0.07
+
+[sample_1]
+Nt = 1682
+muA = 1.0
+alpha = 1.4
+lam = 5
+frac_sig_rhMpc = 0.07
+
+[sample_2]
+Nt = 556
+muA = 1.0
+alpha = 1.4
+lam = 5
+frac_sig_rhMpc = 0.07
+
+[sample_3]
+Nt = 1225
+muA = 1.0
+alpha = 1.4
+lam = 5
+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/blackjax_model_ic_bf_prior
+RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/velmass_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/supranta_ini.ini
-# cp $INI_FILE ini_file.ini
-# $BORG INIT ini_file.ini
-$BORG RESUME ini_file.ini
+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

scripts/submit_borg.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-#SBATCH --job-name=blackjax_model_ic_bf_prior
+#SBATCH --job-name=velmass_ics
 #SBATCH --nodes=1
 #SBATCH --exclusive
 #SBATCH --ntasks=40
@@ -28,7 +28,7 @@ set -e
 
 # Path variables
 BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python
-RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_bf_prior
+RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/velmass_ics
 
 mkdir -p $RUN_DIR
 cd $RUN_DIR
@@ -41,10 +41,10 @@ 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
-$BORG RESUME ini_file.ini
+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
 
 conda deactivate