Merge pull request #41 from abacusorg/fix-origin

Fix origin

Author: lgarrison

abacusnbody/data/compaso_halo_catalog.py

@@ -421,7 +421,7 @@ def __init__(self, path, cleaned=True, subsamples=False, convert_units=True, unp
         if halo_lc == None:
             halo_lc = self.is_path_halo_lc(path)
             if verbose and halo_lc:
-                print('Detected halo light cone catalog')
+                print('Detected halo light cone catalog.')
         self.halo_lc = halo_lc
 
         # If loading halo light cones, turn off cleaning and bit unpacking because done already
@@ -820,6 +820,8 @@ def _read_halo_info(self, halo_fns, fields, cleaned_fns=None, cleaned_fields=Non
             print(f'{len(fields)} halo catalog fields ({len(cleaned_fields)} cleaned) requested. '
                 f'Reading {len(raw_dependencies)} fields from disk. '
                 f'Computing {len(extra_fields)} intermediate fields.')
+            if self.halo_lc:
+                print('\nFor more information on the halo light cone catalog fields, see https://abacussummit.readthedocs.io/en/latest/data-products.html#halo-light-cone-catalogs')
 
         self.halos = Table(cols, copy=False)
         self.halos.meta.update(self.header)
@@ -969,9 +971,26 @@ def _sigmav_loader(m,raw,halos):
         pat = re.compile(r'SO(?:_L2max)?(?:_central_density)')
         self.halo_field_loaders[pat] = lambda m,raw,halos: raw[m[0]]
 
-        # Halo light cone catalog specific fields
-        pat = re.compile(r'index_halo|origin|pos_avg|pos_interp|vel_avg|vel_interp|redshift_interp|N_interp')
+        # loader for halo light cone catalog specific fields
+        pat = re.compile(r'index_halo|pos_avg|vel_avg|redshift_interp|N_interp')
         self.halo_field_loaders[pat] = lambda m,raw,halos: raw[m[0]]
+
+        # loader for halo light cone catalog field `origin`
+        pat = re.compile(r'origin')
+        self.halo_field_loaders[pat] = lambda m,raw,halos: raw[m[0]]%3
+
+        # loader for halo light cone catalog fields: interpolated position and velocity        
+        pat = re.compile(r'(?P<pv>pos|vel)_interp')
+        def lc_interp_loader(m, raw, halos):
+            columns = {}
+            interped = (raw['origin'] // 3).astype(bool)
+            if m[0] == 'pos_interp' or 'pos_interp' in halos.colnames:
+                columns['pos_interp'] = np.where(interped[:, None], raw['pos_avg'], raw['pos_interp'])
+            if m[0] == 'vel_interp' or 'vel_interp' in halos.colnames:
+                columns['vel_interp'] = np.where(interped[:, None], raw['vel_avg'], raw['vel_interp'])
+            return columns
+
+        self.halo_field_loaders[pat] = lc_interp_loader
 
         # eigvecs loader
         pat = re.compile(r'(?P<rnv>sigma(?:r|n|v)_eigenvecs)(?P<which>Min|Mid|Maj)(?P<com>_(?:L2)?com)')
@@ -1722,3 +1741,4 @@ def unpack_euler16(bin_this):
                     ('sigmavtan_L2com', np.float32),
                     ('rvcirc_max_L2com', np.float32),
 ], align=True)
+ 

abacusnbody/hod/prepare_sim.py

@@ -26,7 +26,8 @@
 import multiprocessing
 from multiprocessing import Pool
 
-from sklearn.neighbors import KDTree
+from scipy.spatial import cKDTree
+
 
 DEFAULTS = {}
 DEFAULTS['path2config'] = 'config/abacus_hod.yaml'
@@ -298,7 +299,7 @@ def prepare_slab(i, savedir, simdir, simname, z_mock, tracer_flags, MT, want_ran
                 if randpos.shape[0] > 0:
                     # random points on the edges
                     rand_N = randpos.shape[0]
-                    randpos_tree = KDTree(randpos) # TODO: needs to be periodic, fix bug
+                    randpos_tree = KDTree(randpos)
                     randinds_inner = randpos_tree.query_radius(allpos[index_bounds], r = halos['r98_L2com'][index_bounds])
                     randinds_outer = randpos_tree.query_radius(allpos[index_bounds], r = rad_outer)
                     rand_norm = np.zeros(len(index_bounds))
@@ -308,9 +309,17 @@ def prepare_slab(i, savedir, simdir, simname, z_mock, tracer_flags, MT, want_ran
                 else:
                     rand_norm = np.ones(len(index_bounds))
 
-        allpos_tree = KDTree(allpos)
-        allinds_inner = allpos_tree.query_radius(allpos, r = halos['r98_L2com'])
-        allinds_outer = allpos_tree.query_radius(allpos, r = rad_outer)
+        if halo_lc:
+            # periodicity not needed for halo light cones
+            allpos_tree = cKDTree(allpos)
+            allinds_inner = allpos_tree.query_ball_point(allpos, r = halos['r98_L2com'])
+            allinds_outer = allpos_tree.query_ball_point(allpos, r = rad_outer)
+        else:
+            # note that periodicity exists only in y and z directions
+            tmp = allpos+Lbox/2.
+            allpos_tree = cKDTree(tmp, boxsize=Lbox) # needs to be within 0 and Lbox for periodicity
+            allinds_inner = allpos_tree.query_ball_point(tmp, r = halos['r98_L2com'])
+            allinds_outer = allpos_tree.query_ball_point(tmp, r = rad_outer)
         print("computing m stacks")
         Menv = np.array([np.sum(allmasses[allinds_outer[ind]]) - np.sum(allmasses[allinds_inner[ind]]) \
             for ind in np.arange(len(halos))])

tests/data_mocks_summit_new/Mini_N64_L32/z0.000/galaxies_rsd/ELGs.dat

@@ -1,4 +1,4 @@
-# %ECSV 1.0
+# %ECSV 0.9
 # ---
 # datatype:
 # - {name: x, datatype: float64}

tests/data_mocks_summit_new/Mini_N64_L32/z0.000/galaxies_rsd/LRGs.dat

@@ -1,4 +1,4 @@
-# %ECSV 1.0
+# %ECSV 0.9
 # ---
 # datatype:
 # - {name: x, datatype: float64}