fix want_AB bug

abacusorg · Oct 23, 2023 · 4909178 · 4909178
1 parent 6f050ba
commit 4909178
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 7 deletions.
diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py
@@ -123,6 +123,7 @@ def __init__(self, sim_params, HOD_params, clustering_params = None, chunk=-1, n
         self.want_ranks = HOD_params.get('want_ranks', False)
         self.want_AB = HOD_params.get('want_AB', False)
         self.want_shear = HOD_params.get('want_shear', False)
+        self.want_expvel = HOD_params.get('want_expvel', False)
         self.want_rsd = HOD_params['want_rsd']
 
         if clustering_params is not None:
@@ -308,7 +309,10 @@ def staging(self):
             halo_ids = maskedhalos["id"].astype(int) # halo IDs
             halo_pos = maskedhalos["x_L2com"] # halo positions, Mpc / h
             halo_vels = maskedhalos['v_L2com'] # halo velocities, km/s
-            halo_vel_dev = maskedhalos["randoms_gaus_vrms"] # halo velocity dispersions, km/s
+            if self.want_expvel:
+                halo_vel_dev = maskedhalos["randoms_exp"] # halo velocity dispersions, km/s
+            else:
+                halo_vel_dev = maskedhalos["randoms_gaus_vrms"] # halo velocity dispersions, km/s
             halo_sigma3d = maskedhalos["sigmav3d_L2com"] # 3d velocity dispersion
             halo_c = maskedhalos['r98_L2com']/maskedhalos['r25_L2com'] # concentration
             halo_rvir = maskedhalos['r98_L2com'] # rvir but using r98
@@ -445,11 +449,11 @@ def staging(self):
                          "pweights": pweights,
                          "prandoms": prandoms,
                          "pinds": pinds}
-        if self.want_AB:
-            halo_data["hdeltac"] = hdeltac
-            halo_data["hfenv"] = hfenv
-            particle_data["pdeltac"] = pdeltac
-            particle_data["pfenv"] = pfenv
+        # if self.want_AB:
+        halo_data["hdeltac"] = hdeltac
+        halo_data["hfenv"] = hfenv
+        particle_data["pdeltac"] = pdeltac
+        particle_data["pfenv"] = pfenv
         if self.want_shear:
             halo_data["hshear"] = hshear
             particle_data["pshear"] = pshear
@@ -531,7 +535,13 @@ def run_hod(self, tracers = None, want_rsd = True, want_nfw = False, NFW_draw =
             # np.random.seed(reseed)
             mtg = MTGenerator(np.random.PCG64(reseed))
             r1 = mtg.random(size=len(self.halo_data['hrandoms']), nthread=Nthread, dtype=np.float32)
-            r2 = mtg.standard_normal(size=len(self.halo_data['hveldev']), nthread=Nthread, dtype=np.float32)
+            if self.want_expvel:
+                rt = mtg.random(size=len(self.halo_data['hrandoms']), nthread=Nthread, dtype=np.float32)
+                r2 = np.zeros(len(rt), dtype=np.float32)
+                r2[rt >= 0.5] = -np.log(2*(1-rt[rt >= 0.5]))
+                r2[rt < 0.5] = np.log(2*rt[rt < 0.5])            
+            else:
+                r2 = mtg.standard_normal(size=len(self.halo_data['hveldev']), nthread=Nthread, dtype=np.float32)
             r3 = mtg.random(size=len(self.particle_data['prandoms']), nthread=Nthread, dtype=np.float32)
             self.halo_data['hrandoms'] = r1
             self.halo_data['hveldev'] = r2*self.halo_data['hsigma3d']/np.sqrt(3)

diff --git a/abacusnbody/hod/prepare_sim.py b/abacusnbody/hod/prepare_sim.py
@@ -624,6 +624,8 @@ def prepare_slab(i, savedir, simdir, simname, z_mock, z_type, tracer_flags, MT,
     halos['npstartA'] = halos_pstart_new
     halos['npoutA'] = halos_pnum_new
     halos['randoms'] = np.random.random(len(halos)) # attaching random numbers
+    halos['randoms_exp'] = (np.random.randint(0, 2, size = len(halos))*2-1)\
+    *np.random.exponential(scale = halos["sigmav3d_L2com"]/np.sqrt(3), size = len(halos)) # attaching random numbers
     halos['randoms_gaus_vrms'] = np.random.normal(loc = 0,
         scale = halos["sigmav3d_L2com"]/np.sqrt(3), size = len(halos)) # attaching random numbers