Merge pull request #66 from dangilman/sidm_evolution

update update GC

pyHalo/PresetModels/cdm.py

@@ -18,7 +18,8 @@ def CDM(z_lens, z_source, sigma_sub=0.025, log_mlow=6., log_mhigh=10., log10_sig
         shmf_log_slope=-1.9, cone_opening_angle_arcsec=6.,
         log_m_host=13.3,  r_tidal=0.25, LOS_normalization=1.0, two_halo_contribution=True,
         delta_power_law_index=0.0, geometry_type='DOUBLE_CONE', kwargs_cosmo=None, host_scaling_factor=0.55,
-        redshift_scaling_factor=0.37, two_halo_Lazar_correction=True, draw_poisson=True, c_host=6.0):
+        redshift_scaling_factor=0.37, two_halo_Lazar_correction=True, draw_poisson=True, c_host=6.0,
+        add_globular_clusters=False, kwargs_globular_clusters=None):
     """
     This class generates realizations of dark matter structure in Cold Dark Matter
 
@@ -62,9 +63,10 @@ class for details
     :param two_halo_Lazar_correction: bool; if True, adds the correction to the two-halo contribution from around the
     main deflector presented by Lazar et al. (2021)
     :param c_host: manually set host halo concentration
+    :param add_globular_clusters: bool; include a population of globular clusters around image positions
+    :param kwargs_globular_clusters: keyword arguments for the GC population; see documentation in RealizationExtensions
     :return: a realization of dark matter halos
     """
-
     # FIRST WE CREATE AN INSTANCE OF PYHALO, WHICH SETS THE COSMOLOGY
     pyhalo = pyHalo(z_lens, z_source, kwargs_cosmo)
     # WE ALSO SPECIFY THE GEOMETRY OF THE RENDERING VOLUME
@@ -73,7 +75,6 @@ class for details
     if infall_redshift_model == 'HYBRID_INFALL':
         kwargs_infall_model['log_m_host'] = log_m_host
     pyhalo.lens_cosmo.setup_infall_model(infall_redshift_model, kwargs_infall_model)
-
     # NOW WE SET THE MASS FUNCTION CLASSES FOR SUBHALOS AND FIELD HALOS
     # NOTE: MASS FUNCTION CLASSES SHOULD NOT BE INSTANTIATED HERE
     mass_function_model_subhalos = CDMPowerLaw
@@ -159,13 +160,17 @@ class for details
                          'truncation_model_field_halos': truncation_model_fieldhalos,
                          'concentration_model_field_halos': concentration_model_fieldhalos,
                          'kwargs_density_profile': {}}
-
-    realization_list = pyhalo.render(population_model_list, mass_function_class_list, kwargs_mass_function_list,
+    realization = pyhalo.render(population_model_list, mass_function_class_list, kwargs_mass_function_list,
                                      spatial_distribution_class_list, kwargs_spatial_distribution_list,
                                      geometry, mdef_subhalos, mdef_field_halos, kwargs_halo_model,
                                      two_halo_Lazar_correction,
-                                     nrealizations=1)
-    return realization_list[0]
+                                     nrealizations=1)[0]
+    if add_globular_clusters:
+        from pyHalo.realization_extensions import RealizationExtensions
+        ext = RealizationExtensions(realization)
+        kwargs_globular_clusters['host_halo_mass'] = 10 ** log_m_host
+        realization = ext.add_globular_clusters(**kwargs_globular_clusters)
+    return realization
 
 def CDMCorrelatedStructure(z_lens, z_source, log_mlow=6., log_mhigh=10.,
         concentration_model='DIEMERJOYCE19', kwargs_concentration_model={},

pyHalo/PresetModels/wdm.py

@@ -19,7 +19,7 @@ def WDM(z_lens, z_source, log_mc, sigma_sub=0.025, log_mlow=6., log_mhigh=10., l
         LOS_normalization=1.0, geometry_type='DOUBLE_CONE', kwargs_cosmo=None,
         mdef_subhalos='TNFW', mdef_field_halos='TNFW', kwargs_density_profile={},
         host_scaling_factor=0.5, redshift_scaling_factor=0.3, two_halo_Lazar_correction=True,
-        draw_poisson=True, c_host=None):
+        draw_poisson=True, c_host=None, add_globular_clusters=False, kwargs_globular_clusters=None):
 
     """
     This class generates realizations of dark matter structure in Warm Dark Matter
@@ -68,6 +68,8 @@ def WDM(z_lens, z_source, log_mc, sigma_sub=0.025, log_mlow=6., log_mhigh=10., l
     :param two_halo_Lazar_correction: bool; if True, adds the correction to the two-halo contribution from around the
     main deflector presented by Lazar et al. (2021)
     :param c_host: manually set host halo concentration
+    :param add_globular_clusters: bool; include a population of globular clusters around image positions
+    :param kwargs_globular_clusters: keyword arguments for the GC population; see documentation in RealizationExtensions
     :return: a realization of dark matter halos
     """
     # FIRST WE CREATE AN INSTANCE OF PYHALO, WHICH SETS THE COSMOLOGY
@@ -167,12 +169,17 @@ def WDM(z_lens, z_source, log_mc, sigma_sub=0.025, log_mlow=6., log_mhigh=10., l
                          'truncation_model_field_halos': truncation_model_fieldhalos,
                          'concentration_model_field_halos': concentration_model_fieldhalos,
                          'kwargs_density_profile': kwargs_density_profile}
-    realization_list = pyhalo.render(population_model_list, mass_function_class_list, kwargs_mass_function_list,
+    realization = pyhalo.render(population_model_list, mass_function_class_list, kwargs_mass_function_list,
                                      spatial_distribution_class_list, kwargs_spatial_distribution_list,
                                      geometry, mdef_subhalos, mdef_field_halos, kwargs_halo_model,
                                      two_halo_Lazar_correction,
-                                     nrealizations=1)
-    return realization_list[0]
+                                     nrealizations=1)[0]
+    if add_globular_clusters:
+        from pyHalo.realization_extensions import RealizationExtensions
+        ext = RealizationExtensions(realization)
+        kwargs_globular_clusters['host_halo_mass'] = 10 ** log_m_host
+        realization = ext.add_globular_clusters(**kwargs_globular_clusters)
+    return realization
 
 
 def WDM_mixed(z_lens, z_source, log_mc, mixed_DM_frac, sigma_sub=0.025, log_mlow=6., log_mhigh=10.,
@@ -257,7 +264,8 @@ def WDMGeneral(z_lens, z_source, log_mc, dlogT_dlogk, sigma_sub=0.025, log_mlow=
         shmf_log_slope=-1.9, cone_opening_angle_arcsec=6., log_m_host=13.3, r_tidal=0.25,
         LOS_normalization=1.0, geometry_type='DOUBLE_CONE', kwargs_cosmo=None,
         mdef_subhalos='TNFW', mdef_field_halos='TNFW', kwargs_density_profile={},
-        host_scaling_factor=0.55, redshift_scaling_factor=0.37, two_halo_Lazar_correction=True, c_host=None):
+        host_scaling_factor=0.55, redshift_scaling_factor=0.37, two_halo_Lazar_correction=True, c_host=None,
+        add_globular_clusters=False, kwargs_globular_clusters=None):
 
     """
     This preset model implements a generalized treatment of warm dark matter, or any theory that produces a cutoff in
@@ -301,6 +309,8 @@ def WDMGeneral(z_lens, z_source, log_mc, dlogT_dlogk, sigma_sub=0.025, log_mlow=
     :param two_halo_Lazar_correction: bool; if True, adds the correction to the two-halo contribution from around the
     main deflector presented by Lazar et al. (2021)
     :param c_host: manually fix the host halo concentration
+    :param add_globular_clusters: bool; include a population of globular clusters around image positions
+    :param kwargs_globular_clusters: keyword arguments for the GC population; see documentation in RealizationExtensions
     :return:
     """
     # FIRST WE CREATE AN INSTANCE OF PYHALO, WHICH SETS THE COSMOLOGY
@@ -385,11 +395,15 @@ def WDMGeneral(z_lens, z_source, log_mc, dlogT_dlogk, sigma_sub=0.025, log_mlow=
                          'truncation_model_field_halos': truncation_model_fieldhalos,
                          'concentration_model_field_halos': concentration_model_fieldhalos,
                          'kwargs_density_profile': kwargs_density_profile}
-
     realization_list = pyhalo.render(population_model_list, mass_function_class_list, kwargs_mass_function_list,
                                      spatial_distribution_class_list, kwargs_spatial_distribution_list,
                                      geometry, mdef_subhalos, mdef_field_halos, kwargs_halo_model,
                                      two_halo_Lazar_correction,
                                      nrealizations=1)
-
-    return realization_list[0]
+    realization = realization_list[0]
+    if add_globular_clusters:
+        from pyHalo.realization_extensions import RealizationExtensions
+        ext = RealizationExtensions(realization)
+        kwargs_globular_clusters['host_halo_mass'] = 10 ** log_m_host
+        realization = ext.add_globular_clusters(**kwargs_globular_clusters)
+    return realization

pyHalo/realization_extensions.py

@@ -52,28 +52,38 @@ def add_globular_clusters(self, log10_mgc_mean, log10_mgc_sigma, rendering_radiu
         :param coordinate_center_y: center of rendering area in arcsec
         :return: an instance of Realization that includes the GC's
         """
-        n = self.number_globular_clusters(log10_mgc_mean, log10_mgc_sigma, rendering_radius_arcsec, galaxy_Re,
-                                          host_halo_mass, f)
-        mfunc = Gaussian(n, log10_mgc_mean, log10_mgc_sigma)
-        m = mfunc.draw()
-        z = self._realization.lens_cosmo.z_lens
-        uniform_spatial_distribution = Uniform(rendering_radius_arcsec, self._realization.geometry)
-        x, y = uniform_spatial_distribution.draw(len(m), z, center_x, center_y)
-        lens_cosmo = self._realization.lens_cosmo
-        GCS = []
-        for (m_gc, x_center_gc, y_center_gc) in zip(m, x, y):
-            unique_tag = np.random.rand()
-            profile = GlobularCluster(m_gc, x_center_gc, y_center_gc, lens_cosmo.z_lens, lens_cosmo,
-                                      gc_profile_args, unique_tag)
-            GCS.append(profile)
-        gcs_realization = Realization.from_halos(GCS, self._realization.lens_cosmo,
-                                                 self._realization.kwargs_halo_model,
-                                                 self._realization.apply_mass_sheet_correction,
-                                                 self._realization.rendering_classes,
-                                                 self._realization._rendering_center_x,
-                                                 self._realization._rendering_center_y,
-                                                 self._realization.geometry)
-        new_realization = self._realization.join(gcs_realization)
+        if isinstance(center_x, int) or isinstance(center_x, float):
+            center_x = [center_x]
+            center_y = [center_y]
+        assert len(center_x) == len(center_y)
+        GC_realization = None
+        for x_center, y_center in zip(center_x, center_y):
+            n = self.number_globular_clusters(log10_mgc_mean, log10_mgc_sigma, rendering_radius_arcsec, galaxy_Re,
+                                              host_halo_mass, f)
+            mfunc = Gaussian(n, log10_mgc_mean, log10_mgc_sigma)
+            m = mfunc.draw()
+            z = self._realization.lens_cosmo.z_lens
+            uniform_spatial_distribution = Uniform(rendering_radius_arcsec, self._realization.geometry)
+            x, y = uniform_spatial_distribution.draw(len(m), z, x_center, y_center)
+            lens_cosmo = self._realization.lens_cosmo
+            GCS = []
+            for (m_gc, x_center_gc, y_center_gc) in zip(m, x, y):
+                unique_tag = np.random.rand()
+                profile = GlobularCluster(m_gc, x_center_gc, y_center_gc, lens_cosmo.z_lens, lens_cosmo,
+                                          gc_profile_args, unique_tag)
+                GCS.append(profile)
+            gcs_realization = Realization.from_halos(GCS, self._realization.lens_cosmo,
+                                                     self._realization.kwargs_halo_model,
+                                                     self._realization.apply_mass_sheet_correction,
+                                                     self._realization.rendering_classes,
+                                                     self._realization._rendering_center_x,
+                                                     self._realization._rendering_center_y,
+                                                     self._realization.geometry)
+            if GC_realization is None:
+                GC_realization = gcs_realization
+            else:
+                GC_realization = GC_realization.join(gcs_realization)
+        new_realization = self._realization.join(GC_realization)
         return new_realization
 
     def number_globular_clusters(self, log10_mgc_mean, log10_mgc_sigma, rendering_radius_arcsec,

tests/test_preset_models.py

@@ -40,6 +40,23 @@ def test_CDM(self):
         self._test_default_infall_model(cdm, 'hybrid')
         self._test_default_infall_model(cdm3, 'direct')
 
+        gc_profile_args = {'gamma': 2.5,
+                           'gc_size_lightyear': 250,
+                           'r_core_fraction': 0.05}
+        kwargs_globular_clusters = {'log10_mgc_mean': 5.0,
+                                    'log10_mgc_sigma': 0.5,
+                                    'rendering_radius_arcsec': 0.1,
+                                    'gc_profile_args': gc_profile_args,
+                                    'galaxy_Re': 6.0,
+                                    'host_halo_mass': 10**13.3,
+                                    'f': 3.4e-5,
+                                    'center_x': [0.5, 1.0],
+                                    'center_y': [-1, 0]}
+        cdm = CDM(0.5, 1.5,
+                  add_globular_clusters=True,
+                  kwargs_globular_clusters=kwargs_globular_clusters)
+        _ = cdm.lensing_quantities()
+
     def test_CDM_correlated_structure_only(self):
 
         cdm = CDMCorrelatedStructure(0.5, 1.5)
@@ -51,6 +68,23 @@ def test_WDM(self):
         _ = wdm.lensing_quantities()
         _ = preset_model_from_name('WDM')
         self._test_default_infall_model(wdm, 'hybrid')
+        gc_profile_args = {'gamma': 2.5,
+                           'gc_size_lightyear': 250,
+                           'r_core_fraction': 0.05}
+        kwargs_globular_clusters = {'log10_mgc_mean': 5.0,
+                                    'log10_mgc_sigma': 0.5,
+                                    'rendering_radius_arcsec': 0.1,
+                                    'gc_profile_args': gc_profile_args,
+                                    'galaxy_Re': 6.0,
+                                    'host_halo_mass': 10 ** 13.3,
+                                    'f': 3.4e-5,
+                                    'center_x': [0.5, 1.0],
+                                    'center_y': [-1, 0]}
+        wdm = WDM(0.5, 1.5,
+                  7.7,
+                  add_globular_clusters=True,
+                  kwargs_globular_clusters=kwargs_globular_clusters)
+        _ = wdm.lensing_quantities()
 
     def test_SIDM(self):
 
@@ -93,6 +127,22 @@ def test_WDM_general(self):
         _ = wdm.lensing_quantities()
         wdm = func(0.5, 1.5, 7.7, -2.0, truncation_model_subhalos='TRUNCATION_GALACTICUS')
         self._test_default_infall_model(wdm, 'hybrid')
+        gc_profile_args = {'gamma': 2.5,
+                           'gc_size_lightyear': 250,
+                           'r_core_fraction': 0.05}
+        kwargs_globular_clusters = {'log10_mgc_mean': 5.0,
+                                    'log10_mgc_sigma': 0.5,
+                                    'rendering_radius_arcsec': 0.1,
+                                    'gc_profile_args': gc_profile_args,
+                                    'galaxy_Re': 6.0,
+                                    'host_halo_mass': 10 ** 13.3,
+                                    'f': 3.4e-5,
+                                    'center_x': [0.5, 1.0],
+                                    'center_y': [-1, 0]}
+        wdm = func(0.5, 1.5, 7.7, -2.5,
+                  add_globular_clusters=True,
+                  kwargs_globular_clusters=kwargs_globular_clusters)
+        _ = wdm.lensing_quantities()
 
     def test_CDM_emulator(self):
 

tests/test_realization_extensions.py

@@ -11,7 +11,6 @@
 from pyHalo.PresetModels.cdm import CDM
 from scipy.special import eval_chebyt
 
-
 class TestRealizationExtensions(object):
 
     def test_globular_clusters(self):