Merge pull request #126 from abacusorg/sandydev

bug fix keys()
abacusorg · Jan 9, 2024 · c1e1467 · c1e1467
2 parents 126f8a2 + ccea5f8
commit c1e1467
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py
@@ -153,16 +153,25 @@ def __init__(self, sim_params, HOD_params, clustering_params = None, chunk=-1, n
             self.shearbins = np.linspace(-0.5, 0.5, 101)
 
             self.halo_mass_func, edges = np.histogramdd(
-                np.vstack((np.log10(self.halo_data['hmass']), self.halo_data['hdeltac'], self.halo_data['hfenv'])).T,
+                np.vstack((np.log10(self.halo_data['hmass']), 
+                           self.halo_data.get('hdeltac', np.zeros(len(self.halo_data['hmass']))), 
+                           self.halo_data.get('hfenv', np.zeros(len(self.halo_data['hmass']))))).T,
                 bins = [self.logMbins, self.deltacbins, self.fenvbins],
                 weights = self.halo_data['hmultis'])
         else:
             from abacusnbody.metadata import get_meta
             meta = get_meta(self.sim_name, redshift=0.1)
             self.lbox = meta['BoxSize']
 
+        if self.want_AB:
+            assert 'hfenv' in self.halo_data.keys()
+            assert 'hdeltac' in self.halo_data.keys()
+        if self.want_shear:
+            assert 'hshear' in self.halo_data.keys()
+
         self.halo_mass_func_wshear, edges = np.histogramdd(
-            np.vstack((np.log10(self.halo_data['hmass']), self.halo_data.get('hdeltac', np.zeros(len(self.halo_data['hmass']))),
+            np.vstack((np.log10(self.halo_data['hmass']), 
+                       self.halo_data.get('hdeltac', np.zeros(len(self.halo_data['hmass']))),
                        self.halo_data.get('hfenv', np.zeros(len(self.halo_data['hmass']))),
                        self.halo_data.get('hshear', np.zeros(len(self.halo_data['hmass']))))).T,
             bins = [self.logMbins, self.deltacbins, self.fenvbins, self.shearbins],