WIP PR resolving #993

halotools/empirical_models/composite_models/hod_models/tests/test_github_993_regression.py

@@ -0,0 +1,25 @@
+"""
+"""
+from .....sim_manager import FakeSim
+from ....factories import PrebuiltHodModelFactory
+import numpy as np
+
+
+def test_mass_definition_flexibility2():
+    """Regression test for Issue #993."""
+    halocat = FakeSim()
+    halocat.halo_table["halo_mass_custom"] = np.copy(halocat.halo_table["halo_mvir"])
+    halocat.halo_table["halo_radius_custom"] = np.copy(halocat.halo_table["halo_rvir"])
+    halocat.halo_table.remove_column("halo_upid")
+    halocat.halo_table.remove_column("halo_mvir")
+    halocat.halo_table.remove_column("halo_rvir")
+    model = PrebuiltHodModelFactory(
+        "zheng07",
+        mdef="custom",
+        halo_mass_column_key="halo_mass_custom",
+        prim_haloprop_key="halo_mass_custom",
+        halo_boundary_key="halo_radius_custom",
+    )
+    assert model.halo_boundary_key == "halo_mass_custom"
+    assert model.prim_haloprop_key == "halo_mass_custom"
+    model.populate_mock(halocat)

halotools/empirical_models/composite_models/hod_models/tests/test_zheng07.py

@@ -14,8 +14,7 @@
 
 
 def test_zheng07_composite1():
-    """ Ensure that the ``zheng07`` pre-built model does not accept non-Zheng07Cens.
-    """
+    """Ensure that the ``zheng07`` pre-built model does not accept non-Zheng07Cens."""
     model1 = PrebuiltHodModelFactory("zheng07")
     model2 = PrebuiltHodModelFactory("zheng07", modulate_with_cenocc=True)
 
@@ -30,7 +29,7 @@ def test_zheng07_composite1():
 
 
 def test_zheng07_composite2():
-    """ This test ensures that the source code provided in the
+    """This test ensures that the source code provided in the
     ``Advanced usage of the ``zheng07`` model`` tutorial behaves as expected.
     """
 
@@ -76,7 +75,7 @@ def mean_occupation(self, **kwargs):
 
 
 def test_modulate_with_cenocc1():
-    """ Regression test for Issue #646. Verify that the ``modulate_with_cenocc``
+    """Regression test for Issue #646. Verify that the ``modulate_with_cenocc``
     keyword argument is actually passed to the satellite occupation component.
     """
     model1 = PrebuiltHodModelFactory("zheng07", modulate_with_cenocc=True)
@@ -86,7 +85,7 @@ def test_modulate_with_cenocc1():
 
 
 def test_modulate_with_cenocc2():
-    """ Regression test for Issue #646. Verify that the ``modulate_with_cenocc``
+    """Regression test for Issue #646. Verify that the ``modulate_with_cenocc``
     keyword argument results in behavior that is properly modified by the centrals.
     """
     model = PrebuiltHodModelFactory("zheng07", modulate_with_cenocc=True)
@@ -107,8 +106,7 @@ def test_host_centric_distance():
 
 
 def test_mass_definition_flexibility():
-    """ Regression test for Issue #993.
-    """
+    """Regression test for Issue #993."""
     model = PrebuiltHodModelFactory(
         "zheng07", mdef="200m", halo_boundary_key="halo_radius_arbitrary"
     )

halotools/empirical_models/factories/hod_mock_factory.py

@@ -39,7 +39,7 @@
 
 
 class HodMockFactory(MockFactory):
-    """ Class responsible for populating a simulation with a
+    """Class responsible for populating a simulation with a
     population of mock galaxies based on an HOD-style model
     built by the `~halotools.empirical_models.HodModelFactory` class.
 
@@ -101,7 +101,7 @@ def __init__(
         self.preprocess_halo_catalog(halocat)
 
     def preprocess_halo_catalog(self, halocat):
-        """ Method to pre-process a halo catalog upon instantiation of
+        """Method to pre-process a halo catalog upon instantiation of
         the mock object. This pre-processing includes identifying the
         catalog columns that will be used by the model to create the mock,
         building lookup tables associated with the halo profile,
@@ -122,10 +122,8 @@ def preprocess_halo_catalog(self, halocat):
             Default is set in `~halotools.empirical_models.model_defaults`.
 
         """
-        try:
-            assert "halo_upid" in list(halocat.halo_table.keys())
-        except AssertionError:
-            raise HalotoolsError(missing_halo_upid_msg)
+        if "halo_upid" not in list(halocat.halo_table.keys()):
+            halocat.halo_table["halo_upid"] = -1
 
         # Make cuts on halo catalog #
         # Select host halos only, since this is an HOD-style model
@@ -401,7 +399,7 @@ def populate(self, seed=None, **kwargs):
             self.galaxy_table = self.galaxy_table[mask]
 
     def allocate_memory(self, seed=None):
-        """ Method allocates the memory for all the numpy arrays
+        """Method allocates the memory for all the numpy arrays
         that will store the information about the mock.
         These arrays are bound directly to the mock object.
 
@@ -509,7 +507,7 @@ def allocate_memory(self, seed=None):
             self.galaxy_table[key] = np.zeros(self.Ngals, dtype=dt[key].type)
 
     def estimate_ngals(self, seed=None):
-        """ Method to estimate the number of galaxies produced by the
+        """Method to estimate the number of galaxies produced by the
         mock.populate() method. It runs one realization of all
         mc_occupation methods and reports the total number of galaxies
         produced. However, no extra memory is allocated for the

halotools/empirical_models/factories/subhalo_mock_factory.py

@@ -19,15 +19,17 @@
 from ...custom_exceptions import HalotoolsError
 
 
-__all__ = ['SubhaloMockFactory']
-__author__ = ['Andrew Hearin']
+__all__ = ["SubhaloMockFactory"]
+__author__ = ["Andrew Hearin"]
 
-unavailable_haloprop_msg = ("Your model requires that the ``%s`` key appear in the halo catalog,\n"
-    "but this column is not available in the catalog you attempted to populate.\n")
+unavailable_haloprop_msg = (
+    "Your model requires that the ``%s`` key appear in the halo catalog,\n"
+    "but this column is not available in the catalog you attempted to populate.\n"
+)
 
 
 class SubhaloMockFactory(MockFactory):
-    """ Class responsible for populating a simulation with a
+    """Class responsible for populating a simulation with a
     population of mock galaxies based on models generated by
     `~halotools.empirical_models.SubhaloModelFactory`.
 
@@ -61,23 +63,27 @@ def __init__(self, **kwargs):
         """
 
         MockFactory.__init__(self, **kwargs)
-        halocat = kwargs['halocat']
+        halocat = kwargs["halocat"]
 
         # Pre-compute any additional halo properties required by the model
         self.preprocess_halo_catalog(halocat)
         self.precompute_galprops()
 
     def preprocess_halo_catalog(self, halocat):
-        """ Method to pre-process a halo catalog upon instantiation of the mock object.
-        """
+        """Method to pre-process a halo catalog upon instantiation of the mock object."""
         halo_table = halocat.halo_table
 
-        if (('halo_hostid' not in self.additional_haloprops) & ('halo_hostid' in list(halo_table.keys()))):
-            self.additional_haloprops.append('halo_hostid')
+        self.additional_haloprops.append("halo_upid")
+
+        if ("halo_hostid" not in self.additional_haloprops) & (
+            "halo_hostid" in list(halo_table.keys())
+        ):
+            self.additional_haloprops.append("halo_hostid")
 
-        if (('halo_mvir_host_halo' not in self.additional_haloprops) &
-                ('halo_mvir_host_halo' in list(halo_table.keys()))):
-            self.additional_haloprops.append('halo_mvir_host_halo')
+        if ("halo_mvir_host_halo" not in self.additional_haloprops) & (
+            "halo_mvir_host_halo" in list(halo_table.keys())
+        ):
+            self.additional_haloprops.append("halo_mvir_host_halo")
 
         self.halo_table = Table()
         for key in self.additional_haloprops:
@@ -87,7 +93,7 @@ def preprocess_halo_catalog(self, halocat):
                 raise HalotoolsError(unavailable_haloprop_msg % key)
 
     def precompute_galprops(self):
-        """ Method pre-processes the input subhalo catalog, and pre-computes
+        """Method pre-processes the input subhalo catalog, and pre-computes
         all halo properties that will be inherited by the ``galaxy_table``.
 
         For example, in subhalo-based models, the phase space coordinates of the
@@ -107,13 +113,15 @@ def precompute_galprops(self):
             self.galaxy_table[key] = self.halo_table[key]
             self._precomputed_galprop_list.append(key)
 
-        phase_space_keys = ['x', 'y', 'z', 'vx', 'vy', 'vz']
+        phase_space_keys = ["x", "y", "z", "vx", "vy", "vz"]
         for newkey in phase_space_keys:
-            self.galaxy_table[newkey] = self.galaxy_table[model_defaults.host_haloprop_prefix+newkey]
+            self.galaxy_table[newkey] = self.galaxy_table[
+                model_defaults.host_haloprop_prefix + newkey
+            ]
             self._precomputed_galprop_list.append(newkey)
 
-        self.galaxy_table['galid'] = np.arange(len(self.galaxy_table))
-        self._precomputed_galprop_list.append('galid')
+        self.galaxy_table["galid"] = np.arange(len(self.galaxy_table))
+        self._precomputed_galprop_list.append("galid")
 
         # Component models may explicitly distinguish between certain types of halos,
         # e.g., subhalos vs. host halos. Since this assignment is not dynamic,
@@ -122,21 +130,24 @@ def precompute_galprops(self):
 
             try:
                 f = component_model.gal_type_func
-                newkey = feature + '_gal_type'
+                newkey = feature + "_gal_type"
                 self.galaxy_table[newkey] = f(table=self.galaxy_table)
                 self._precomputed_galprop_list.append(newkey)
             except AttributeError:
                 pass
             except:
                 clname = component_model.__class__.__name__
-                msg = ("\nThe `gal_type_func` attribute of the " + clname +
-                    "\nraises an unexpected exception when passed a halo table as a "
+                msg = (
+                    "\nThe `gal_type_func` attribute of the "
+                    + clname
+                    + "\nraises an unexpected exception when passed a halo table as a "
                     "table keyword argument. \n"
                     "If the features in your component model have explicit dependence "
                     "on galaxy type, \nthen you must implement the `gal_type_func` mechanism "
                     "in such a way that\nthis function accepts a "
                     "length-N halo table as a ``table`` keyword argument, \n"
-                    "and returns a length-N array of strings.\n")
+                    "and returns a length-N array of strings.\n"
+                )
                 raise HalotoolsError(msg)
 
     def populate(self, seed=None):
@@ -237,17 +248,19 @@ def populate(self, seed=None):
                 seed += 1
             func(table=self.galaxy_table, seed=seed)
 
-        if hasattr(self.model, 'galaxy_selection_func'):
+        if hasattr(self.model, "galaxy_selection_func"):
             mask = self.model.galaxy_selection_func(self.galaxy_table)
             self.galaxy_table = self.galaxy_table[mask]
 
     def _allocate_memory(self, seed=None):
-        """
-        """
+        """"""
         Ngals = len(self.galaxy_table)
 
-        new_column_generator = (key for key in self.model._galprop_dtypes_to_allocate.names
-            if key not in self._precomputed_galprop_list)
+        new_column_generator = (
+            key
+            for key in self.model._galprop_dtypes_to_allocate.names
+            if key not in self._precomputed_galprop_list
+        )
 
         for key in new_column_generator:
             dt = self.model._galprop_dtypes_to_allocate[key]