Add radial selection

borg_velocity/likelihood.py

@@ -9,6 +9,7 @@ import jaxlib
 from functools import partial
 import ast
 import numbers
+import h5py
 
 import borg_velocity.utils as utils
 from borg_velocity.utils import myprint
@@ -49,6 +50,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         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)]
         
@@ -57,6 +59,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         
         # 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'])
@@ -91,6 +94,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         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
         }
@@ -224,12 +228,22 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         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)]
         self.generateMBData()
         myprint('From mock')
         self.saved_s_hat = s_hat.copy()
         self.logLikelihoodComplex(s_hat, False)
-    
+
     def dens2like(self, output_density: np.ndarray, output_velocity: np.ndarray):
         """
         Given stored distance tracer data, computes the negative log-likelihood of the data
@@ -256,6 +270,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         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, 
@@ -268,10 +283,12 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
                 omega_m, 
                 muA, 
                 alpha, 
+                lam,
                 self.L[0], 
                 jnp.array(self.fwd.getOutputBoxModel().xmin), 
                 self.interp_order, 
-                self.bias_epsilon
+                self.bias_epsilon,
+                self.R_max
             )
             
         if not jnp.isfinite(lkl):
@@ -377,7 +394,8 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
         
         
 # @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, L_BOX, X_MIN, 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
     """
@@ -416,12 +434,13 @@ 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)
     
+    
     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
+    p_r = r ** 2 * jnp.exp(-0.5 * d2) * los_density * jnp.exp(- lam * r / R_max)
     p_r_norm = jnp.expand_dims(jnp.trapz(p_r, r, axis=1), axis=1)
     
     # Integrate to get likelihood
@@ -454,6 +473,15 @@ 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']
@@ -731,7 +759,7 @@ def build_sampler(
             model_sampler = borg.samplers.ModelParamsSampler(prefix, params, likelihood, fwd_param, initial_values, prior)
         elif config['sampling']['algorithm'].lower() == 'hmc':
             model_sampler = HMCBiasSampler(
-                prefix,
+                "model_params",
                 likelihood,
                 params,
                 transform_attributes=transform_attributes,
@@ -741,6 +769,9 @@ def build_sampler(
                 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")
         else:
             raise NotImplementedError
         model_sampler.setName("model_sampler")
@@ -749,6 +780,7 @@ def build_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_velocity/mock_maker.py

@@ -51,6 +51,10 @@ def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
     
     nsamp = int(config['run']['nsamp'])
     
+    # Ensure cosmology is correct
+    cosmo = utils.get_cosmopar(ini_file)
+    fwd_model.setCosmoParams(cosmo)
+    
     # Run BORG density field
     output_density = np.zeros(fwd_model.getOutputBoxModel().N)
     fwd_model.forwardModel_v2(s_hat)
@@ -77,28 +81,28 @@ def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
     for i in range(nsamp):
         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'])
         Nt = int(config[f'sample_{i}']['Nt'])
         phi = (1. + output_density + bias_epsilon) ** alpha
-        coord_true[i] = np.zeros((3,Nt))
-        coord_meas[i] = np.zeros((3,Nt))
-        nmade = 0
-        while (nmade < Nt):
-            ctrue = poisson_process.sample_3d(phi, Nt, 
-                  fwd_model.getOutputBoxModel().L[0], fwd_model.getOutputBoxModel().xmin)
-            cmeas, sig_mu[i] = radially_scatter(ctrue, frac_sig_x)
-            
-            # Only use tracers which lie within R_max
-            r = np.sqrt(np.sum(ctrue**2, axis=0))
-            m = r < R_max
-            
-            nnew = m.sum()
-            if nmade + nnew > Nt:
-                coord_true[i][:,nmade:] = ctrue[:,m][:,:Nt-nmade]
-                coord_meas[i][:,nmade:] = cmeas[:,m][:,:Nt-nmade]
-            else:
-                coord_true[i][:,nmade:nmade+nnew] = ctrue[:,m]
-                coord_meas[i][:,nmade:nmade+nnew] = cmeas[:,m]
-            nmade = min(nmade + m.sum(), Nt)
+        
+        # Weight poisson mean by radial selection P(r) \propto r^2 exp(-lam r/R_max)
+        # Do not include r^2 factor as we are sampling in Cartesian 3D,
+        # so do not need the Jacobian term
+        N = fwd_model.getOutputBoxModel().N
+        L = fwd_model.getOutputBoxModel().L
+        xmin = fwd_model.getOutputBoxModel().xmin
+        x = np.linspace(xmin[0], xmin[0] + L[0], N[0])
+        y = np.linspace(xmin[1], xmin[1] + L[1], N[1])
+        z = np.linspace(xmin[2], xmin[2] + L[2], N[2])
+        X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
+        r = np.sqrt(X**2 + Y**2 + Z**2)
+        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)
+
         
     # Interpolate velocities to tracers
     interp_order = int(config['model']['interp_order'])

borg_velocity/samplers.py

@@ -162,8 +162,8 @@ class HMCBiasSampler(borg.samplers.PyBaseSampler):
     def restore(self, state: borg.likelihood.MarkovState):
         
         # Define attribute names
-        attrname_real = f"{self.prefix}attributes"
-        attrname_var = f"{self.prefix}hmc_var"
+        attrname_real = self.prefix
+        attrname_var = f"{self.prefix}_hmc_var"
 
         # Initialize attributes in the state
         state.newArray1d(attrname_real, len(self.attributes), True)

conf/basic_ini.ini

@@ -16,11 +16,11 @@ seed_cpower = true
 
 [block_loop]
 hades_sampler_blocked = true
-bias_sampler_blocked= false
+bias_sampler_blocked= true
 nmean_sampler_blocked= true
-sigma8_sampler_blocked = true
+sigma8_sampler_blocked = false
 muA_sampler_blocked = false
-omega_m_sampler_blocked = true
+omega_m_sampler_blocked = false
 alpha_sampler_blocked = false
 sig_v_sampler_blocked = false
 bulk_flow_sampler_blocked = false
@@ -108,4 +108,4 @@ frac_sig_rhMpc = 0.07
 Nt = 230
 muA = 1.0
 alpha = 1.4
-frac_sig_rhMpc = 0.07
+frac_sig_rhMpc = 0.07

conf/supranta_ini.ini

@@ -0,0 +1,129 @@
+[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 = true
+bias_sampler_blocked= true
+nmean_sampler_blocked= true
+sigma8_sampler_blocked = false
+muA_sampler_blocked = false
+omega_m_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 = 10000
+warmup_model = 0
+warmup_cosmo = 0
+random_ic = false
+init_random_scaling = 1.0
+bignum = 1e300
+
+[hades]
+algorithm = HMC
+max_epsilon = 0.01
+max_timesteps = 50
+mixing = 1
+
+[sampling]
+algorithm = HMC
+epsilon = 0.001
+Nsteps = 20
+refresh = 0.1
+
+[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 = 4.
+sig_v = 150.
+R_lim = none
+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.335
+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]
+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/basic_run
+RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/supranta_pars_N64
 
 mkdir -p $RUN_DIR
 cd $RUN_DIR
@@ -29,6 +29,6 @@ BASH_XTRACEFD="3"
 set -x
 
 # Just ICs
-INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/basic_ini.ini
-cp $INI_FILE basic_ini.ini
-$BORG INIT basic_ini.ini
+INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/supranta_ini.ini
+cp $INI_FILE ini_file.ini
+$BORG INIT ini_file.ini

scripts/submit_borg.sh

@@ -1,9 +1,9 @@
 #!/bin/sh
 #PBS -S /bin/sh
-#PBS -N model_hmc_test
+#PBS -N supranta_pars_N64
 #PBS -j oe
 #PBS -m ae
-#PBS -l nodes=h13:has1gpu:ppn=40,walltime=24:00:00
+#PBS -l nodes=h11:has1gpu:ppn=40,walltime=72:00:00
 
 # Modules
 module purge
@@ -21,7 +21,7 @@ set -e
 
 # Path variables
 BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python
-RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/basic_run
+RUN_DIR=/data101/bartlett/fsigma8/borg_velocity/supranta_pars_N64
 
 mkdir -p $RUN_DIR
 cd $RUN_DIR
@@ -34,9 +34,9 @@ BASH_XTRACEFD="3"
 set -x
 
 # Run BORG
-INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/basic_ini.ini
-cp $INI_FILE basic_ini.ini
-$BORG INIT basic_ini.ini
+INI_FILE=/home/bartlett/fsigma8/borg_velocity/conf/supranta_ini.ini
+cp $INI_FILE ini_file.ini
+$BORG INIT ini_file.ini
 
 conda deactivate
 
@@ -44,4 +44,4 @@ conda deactivate
 set +x
 exec 3>&-
 
-exit 0
+exit 0

tests/evaluate_likelihood.py

@@ -0,0 +1,84 @@
+import borg_velocity.likelihood as likelihood
+import borg_velocity.forwards as forwards
+import borg_velocity.utils as utils
+
+import aquila_borg as borg
+import configparser
+import h5py as h5
+import numpy as np
+from tqdm import tqdm
+import glob
+
+dirname = '/data101/bartlett/fsigma8/borg_velocity/basic_run_ic'
+ini_file = f'{dirname}/basic_ini.ini'
+
+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
+
+# 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
+model = likelihood.build_gravity_model(None, box_in, ini_file=ini_file)
+cosmo = utils.get_cosmopar(ini_file)
+model.setCosmoParams(cosmo)
+fwd_param = forwards.NullForward(box_in)
+fwd_vel = likelihood.fwd_vel
+mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
+
+# Load s_hat from file
+with h5.File(f'{dirname}/mock_data.h5', 'r') as f:
+    s_hat = f['scalars/s_hat_field'][:]
+    print(f['scalars/hmc_Elh'][:], f['scalars/hmc_Eprior'][:])
+
+# Make mock data
+state = borg.likelihood.MarkovState()
+mylike.initializeLikelihood(state)
+mylike.updateCosmology(cosmo)
+mylike.generateMockData(s_hat, state)
+
+print(dir(mylike))
+quit()
+
+# Evaluate likelihood
+initial_logL = mylike.logLikelihoodComplex(s_hat, None)
+print(initial_logL)
+print(np.var(s_hat))
+# print('var(s_hat): 1.0002235630494603, Call to logLike: 2945.19580078125')
+
+# Load some MCMC samples and see how they compare
+nframe = 10
+iter_min = 0
+iter_max = 1748
+all_mcmc = get_mcmc_steps(nframe, iter_max, iter_min=0)
+all_logL = np.zeros(len(all_mcmc))
+for i in tqdm(range(len(all_mcmc))):
+    with h5.File(f'{dirname}/mcmc_{all_mcmc[i]}.h5', 'r') as f:
+        s_hat = f['scalars/s_hat_field'][:]
+        # print(f['scalars'].keys())
+        print(f['scalars/hmc_Elh'][:], f['scalars/hmc_Eprior'][:])
+    # all_logL[i] = mylike.logLikelihoodComplex(s_hat, None)
+
+print(all_logL)
+    

tests/test_likelihood.py

@@ -7,11 +7,12 @@ import borg_velocity.likelihood as likelihood
 import borg_velocity.forwards as forwards
 import borg_velocity.utils as utils
 
-ini_file = '../conf/basic_ini.ini'
+ini_file = '../conf/supranta_ini.ini'
 test_scaling = True
 test_sigma8 = True
 test_omegam = True
 test_alpha = True
+test_lam = True
 test_muA = True
 
 # Input box
@@ -34,6 +35,7 @@ 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)
 
@@ -84,8 +86,7 @@ if test_sigma8:
     fig.clf()
     plt.close(fig)
     
-    
-# Test sigma8
+# Test omegam
 if test_omegam:
     all_omegam = np.linspace(0.1, 0.6, 100)
     all_lkl = np.empty(all_omegam.shape)
@@ -134,6 +135,29 @@ if test_alpha:
     fig.clf()
     plt.close(fig)
     
+# Test radial sampling
+if test_lam:
+    all_lam = np.linspace(0.0, 10.0, 100)
+    all_lkl = np.empty(all_lam.shape)
+    for i, lam in enumerate(all_lam):
+        mylike.fwd_param.setModelParams({'lam0':lam})
+        all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
+    mylike.fwd_param.setModelParams({'lam0':mylike.lam[0]})
+    fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
+    all_lkl -= fid_lkl
+    all_lkl = np.exp(-all_lkl)
+
+    fig, ax = plt.subplots(1, 1, figsize=(5,5))
+    ax.plot(all_lam, all_lkl)
+    ax.axhline(y=0, color='k')
+    ax.axvline(x=mylike.lam[0], color='k')
+    ax.set_xlabel(r'$\lambda_0$')
+    ax.set_ylabel(r'$\mathcal{L}$')
+    fig.tight_layout()
+    fig.savefig('../figs/lam_test.png')
+    fig.clf()
+    plt.close(fig)
+    
     
 # Test bias model 
 if test_muA:
@@ -157,3 +181,4 @@ if test_muA:
     fig.savefig('../figs/muA_test.png')
     fig.clf()
     plt.close(fig)
+