Use the global "cosmo" instance for simplification

Signed-off-by: Aaron LI <aly@aaronly.me>

4 files changed, 34 insertions, 53 deletions

diff --git a/fg21sim/extragalactic/clusters/emission.py b/fg21sim/extragalactic/clusters/emission.py
index 93e04ff..4366025 100644
--- a/fg21sim/extragalactic/clusters/emission.py
+++ b/fg21sim/extragalactic/clusters/emission.py
@@ -18,11 +18,11 @@ import numpy as np
 import scipy.integrate
 import scipy.special
 
+from ...utils import cosmo
 from ...utils.units import (Units as AU,
                             UnitConversions as AUC,
                             Constants as AC)
 from ...utils.convert import Fnu_to_Tb_fast
-from ...utils.cosmology import Cosmology
 
 
 logger = logging.getLogger(__name__)
@@ -57,7 +57,6 @@ class SynchrotronEmission:
         self.n_e = n_e
         self.z = z
         self.radius = radius  # [kpc]
-        self.cosmo = Cosmology()
 
     @property
     def frequency_larmor(self):
@@ -165,7 +164,7 @@ class SynchrotronEmission:
             Synchrotron flux at frequency ``nu``.
             Unit: [Jy] = 1e-23 [erg/s/cm^2/Hz]
         """
-        DL = self.cosmo.DL(self.z) * AUC.Mpc2cm  # [cm]
+        DL = cosmo.DL(self.z) * AUC.Mpc2cm  # [cm]
         P_nu = self.power(nu)
         F_nu = 1e23 * P_nu / (4*np.pi * DL*DL)  # [Jy]
         return F_nu
@@ -192,7 +191,7 @@ class SynchrotronEmission:
             Synchrotron surface brightness at frequency ``nu``.
             Unit: [K] <-> [Jy/pixel]
         """
-        DA = self.cosmo.DL(self.z) * AUC.Mpc2cm  # [cm]
+        DA = cosmo.DL(self.z) * AUC.Mpc2cm  # [cm]
         radius = self.radius * AUC.kpc2cm  # [cm]
         omega = (np.pi * radius**2 / DA**2) * AUC.rad2deg**2  # [deg^2]
         pixelarea = (pixelsize * AUC.arcsec2deg) ** 2  # [deg^2]
diff --git a/fg21sim/extragalactic/clusters/formation.py b/fg21sim/extragalactic/clusters/formation.py
index 827c089..3bb44a9 100644
--- a/fg21sim/extragalactic/clusters/formation.py
+++ b/fg21sim/extragalactic/clusters/formation.py
@@ -21,7 +21,7 @@ import scipy.special
 import scipy.optimize
 
 from .mergertree import MergerTree
-from ...utils.cosmology import Cosmology
+from ...utils import cosmo
 
 
 logger = logging.getLogger(__name__)
@@ -54,8 +54,6 @@ class ClusterFormation:
         the merger is a major event or a  minor one.
         If ``M_main/M_sub < ratio_major``, then it is a major merger event.
         Default: 3.0
-    cosmo : `~Cosmology`, optional
-        Adopted cosmological model with custom utility functions.
     merger_mass_min : float, optional
         Minimum mass change to be regarded as a merger event instead of
         accretion.
@@ -70,12 +68,11 @@ class ClusterFormation:
         major merger event, or None if not found.
     """
     def __init__(self, M0, z0, zmax=3.0, ratio_major=3.0,
-                 cosmo=Cosmology(), merger_mass_min=1e12):
+                 merger_mass_min=1e12):
         self.M0 = M0  # [Msun]
         self.z0 = z0
         self.zmax = zmax
         self.ratio_major = ratio_major
-        self.cosmo = cosmo
         self.merger_mass_min = merger_mass_min
 
     @property
@@ -101,7 +98,7 @@ class ClusterFormation:
         References: Ref.[1],Eq.(2)
         """
         alpha = self.sigma_index
-        sigma = self.cosmo.sigma8 * (mass / self.cosmo.M8) ** (-alpha)
+        sigma = cosmo.sigma8 * (mass / cosmo.M8) ** (-alpha)
         return sigma
 
     def f_delta_c(self, z):
@@ -113,7 +110,7 @@ class ClusterFormation:
 
         References: Ref.[1],App.A,Eq.(A1)
         """
-        return self.cosmo.overdensity_crit(z)
+        return cosmo.overdensity_crit(z)
 
     def f_dw_max(self, mass):
         """
@@ -148,7 +145,7 @@ class ClusterFormation:
         References: Ref.[1],Sec.(3)
         """
         alpha = self.sigma_index
-        mass = self.cosmo.M8 * (S / self.cosmo.sigma8**2)**(-1/(2*alpha))
+        mass = cosmo.M8 * (S / cosmo.sigma8**2)**(-1/(2*alpha))
         return mass
 
     @staticmethod
@@ -254,7 +251,7 @@ class ClusterFormation:
         Mc = self.M0
         mtree_root = MergerTree(data={"mass": Mc,
                                       "z": zc,
-                                      "age": self.cosmo.age(zc)})
+                                      "age": cosmo.age(zc)})
         logger.debug("[main] z=%.4f : mass=%g [Msun]" % (zc, Mc))
 
         mtree = mtree_root
@@ -275,7 +272,7 @@ class ClusterFormation:
             dS = self.gen_dS(dw)
             # Progenitor properties
             z1 = self.calc_z(w2 + dw)
-            age1 = self.cosmo.age(z1)
+            age1 = cosmo.age(z1)
             S1 = S2 + dS
             M1 = self.calc_mass(S1)
             dM = Mc - M1
@@ -313,7 +310,7 @@ class ClusterFormation:
         merger tree.
         """
         z = 0.0 if _z is None else _z
-        node_data = {"mass": M, "z": z, "age": self.cosmo.age(z)}
+        node_data = {"mass": M, "z": z, "age": cosmo.age(z)}
 
         # Whether to stop the trace
         if self.zmax is not None and z > self.zmax:
diff --git a/fg21sim/extragalactic/clusters/halo.py b/fg21sim/extragalactic/clusters/halo.py
index d7f351f..c78c2c3 100644
--- a/fg21sim/extragalactic/clusters/halo.py
+++ b/fg21sim/extragalactic/clusters/halo.py
@@ -26,7 +26,7 @@ import scipy.optimize
 
 from .formation import ClusterFormation
 from .solver import FokkerPlanckSolver
-from ...utils.cosmology import Cosmology
+from ...utils import cosmo
 from ...utils.units import (Units as AU,
                             UnitConversions as AUC,
                             Constants as AC)
@@ -35,7 +35,7 @@ from ...utils.units import (Units as AU,
 logger = logging.getLogger(__name__)
 
 
-class HaloSingle:
+class RadioHalo:
     """
     Simulate a single (giant) radio halos following the "statistical
     magneto-turbulent model" proposed by Cassano & Brunetti (2005).
@@ -63,8 +63,6 @@ class HaloSingle:
     mec : float
         Unit for electron momentum (p): mec = m_e * c, p = gamma * mec,
         therefore value of p is the Lorentz factor.
-    cosmo : `~Cosmology`
-        Adopted cosmological model with custom utility functions.
     mtree : `~MergerTree`
         Merging history of this cluster.
     """
@@ -99,10 +97,6 @@ class HaloSingle:
         self.buffer_np = self.configs.getn(comp+"/buffer_np")
         self.time_step = self.configs.getn(comp+"/time_step")
         self.injection_index = self.configs.getn(comp+"/injection_index")
-        # Cosmology model
-        self.H0 = self.configs.getn("cosmology/H0")
-        self.OmegaM0 = self.configs.getn("cosmology/OmegaM0")
-        self.cosmo = Cosmology(H0=self.H0, Om0=self.OmegaM0)
         logger.info("Loaded and set up configurations")
 
     @property
@@ -185,13 +179,13 @@ class HaloSingle:
             Unit: [cm^-3 mec^-1]
         """
         if zbegin is None:
-            tstart = self.cosmo.age(self.zmax)
+            tstart = cosmo.age(self.zmax)
         else:
-            tstart = self.cosmo.age(zbegin)
+            tstart = cosmo.age(zbegin)
         if zend is None:
-            tstop = self.cosmo.age(self.z0)
+            tstop = cosmo.age(self.z0)
         else:
-            tstop = self.cosmo.age(zend)
+            tstop = cosmo.age(zend)
 
         fpsolver = FokkerPlanckSolver(
             xmin=self.pmin, xmax=self.pmax,
@@ -252,8 +246,8 @@ class HaloSingle:
         Rvir : float
             Virial radius (unit: kpc) of the cluster at given redshift
         """
-        Dc = self.cosmo.overdensity_virial(z)
-        rho = self.cosmo.rho_crit(z)  # [g/cm^3]
+        Dc = cosmo.overdensity_virial(z)
+        rho = cosmo.rho_crit(z)  # [g/cm^3]
         R_vir = (3*mass*AUC.Msun2g / (4*np.pi * Dc * rho))**(1/3)  # [cm]
         R_vir *= AUC.cm2kpc  # [kpc]
         return R_vir
@@ -316,7 +310,7 @@ class HaloSingle:
             http://adsabs.harvard.edu/abs/2005MNRAS.357.1313C
             Eq.(12)
         """
-        f_baryon = self.cosmo.Ob0 / self.cosmo.Om0
+        f_baryon = cosmo.Ob0 / cosmo.Om0
         Rv = self._radius_virial(mass, z) * AUC.kpc2cm  # [cm]
         V = (4*np.pi / 3) * Rv**3  # [cm^3]
         rho = f_baryon * mass*AUC.Msun2g / V  # [g/cm^3]
@@ -346,7 +340,7 @@ class HaloSingle:
             http://adsabs.harvard.edu/abs/2005MNRAS.357.1313C
             Eq.(13)
         """
-        f_baryon = self.cosmo.Ob0 / self.cosmo.Om0
+        f_baryon = cosmo.Ob0 / cosmo.Om0
         M_ICM = mass * f_baryon * AUC.Msun2g  # [g]
         r *= AUC.kpc2cm  # [cm]
         Rv = self._radius_virial(mass, z) * AUC.kpc2cm  # [cm]
@@ -477,12 +471,12 @@ class HaloSingle:
         time : float
             Elapsing time (unit: Gyr)
         """
-        t_begin = self.cosmo.age(z_begin)  # [Gyr]
+        t_begin = cosmo.age(z_begin)  # [Gyr]
         t_end = t_begin + time
-        if t_end >= self.cosmo.age(0):
+        if t_end >= cosmo.age(0):
             z_end = 0.0
         else:
-            z_end = self.cosmo.redshift(t_end)
+            z_end = cosmo.redshift(t_end)
         return z_end
 
     @property
@@ -542,7 +536,7 @@ class HaloSingle:
         """
         zgrid, chigrid = self._chi_data
         # XXX: force a minimal value instead of zero or too small
-        chi_min = 1 / (10 * self.cosmo.age0)
+        chi_min = 1 / (10 * cosmo.age0)
 
         try:
             zi = np.where(z < zgrid)[0][0]
@@ -677,7 +671,7 @@ class HaloSingle:
         """
         if not hasattr(self, "_electron_injection_rate"):
             e_th = self.e_thermal  # [erg/cm^3]
-            age = self.cosmo.age(self.z0)
+            age = cosmo.age(self.z0)
             term1 = (self.injection_index-2) * self.eta_e
             term2 = e_th / (self.pmin * self.mec * AC.c)  # [cm^-3]
             term3 = 1.0 / (age * self.pmin)  # [Gyr^-1 mec^-1]
@@ -714,7 +708,7 @@ class HaloSingle:
             http://adsabs.harvard.edu/abs/2013AN....334..515D
             Eq.(15)
         """
-        z = self.cosmo.redshift(t)
+        z = cosmo.redshift(t)
         chi = self._coef_acceleration(z)  # [Gyr^-1]
         # NOTE: Cassano & Brunetti's formula misses a factor of 2.
         Dpp = chi * p**2 / 4  # [mec^2/Gyr]
@@ -760,7 +754,7 @@ class HaloSingle:
             http://adsabs.harvard.edu/abs/2005MNRAS.357.1313C
             Eq.(38)
         """
-        z = self.cosmo.redshift(t)
+        z = cosmo.redshift(t)
         n_th = self._n_thermal(self.M0, z)
         coef = -3.3e-29 * AUC.Gyr2s / self.mec  # [mec/Gyr]
         dpdt = coef * n_th * (1 + np.log(p/n_th) / 75)
@@ -776,7 +770,7 @@ class HaloSingle:
             http://adsabs.harvard.edu/abs/2005MNRAS.357.1313C
             Eq.(39)
         """
-        z = self.cosmo.redshift(t)
+        z = cosmo.redshift(t)
         coef = -4.8e-4 * AUC.Gyr2s / self.mec  # [mec/Gyr]
         dpdt = (coef * (p*self.mec)**2 *
                 ((self.magnetic_field/3.2)**2 + (1+z)**4))
@@ -793,7 +787,7 @@ class HaloSingle:
             Energy density of the ICM (unit: erg/cm^3)
         """
         mass = self.M0
-        f_baryon = self.cosmo.Ob0 / self.cosmo.Om0
+        f_baryon = cosmo.Ob0 / cosmo.Om0
         kT = self.kT_mass(mass)  # [keV]
         N = mass * AUC.Msun2g * f_baryon / (AC.mu * AC.u)
         E_th = kT*AUC.keV2erg * N  # [erg]
@@ -820,7 +814,7 @@ class HaloSingle:
             Number density of the ICM
             Unit: [cm^-3]
         """
-        f_baryon = self.cosmo.Ob0 / self.cosmo.Om0
+        f_baryon = cosmo.Ob0 / cosmo.Om0
         N = mass * AUC.Msun2g * f_baryon / (AC.mu * AC.u)
         Rv = self._radius_virial(mass, z) * AUC.kpc2cm  # [cm]
         V = (4*np.pi / 3) * Rv**3  # [cm^3]
diff --git a/fg21sim/extragalactic/clusters/main.py b/fg21sim/extragalactic/clusters/main.py
index 438b860..b7c2d03 100644
--- a/fg21sim/extragalactic/clusters/main.py
+++ b/fg21sim/extragalactic/clusters/main.py
@@ -19,7 +19,7 @@ import pandas as pd
 from .formation import ClusterFormation
 from .halo import RadioHalo
 from ...sky import get_sky
-from ...utils.cosmology import Cosmology
+from ...utils import cosmo
 from ...utils.io import dataframe_to_csv
 
 
@@ -63,14 +63,6 @@ class GalaxyClusters:
         self.checksum = self.configs.getn("output/checksum")
         self.clobber = self.configs.getn("output/clobber")
 
-        # Cosmology model
-        self.H0 = self.configs.getn("cosmology/H0")
-        self.OmegaM0 = self.configs.getn("cosmology/OmegaM0")
-        self.Omegab0 = self.configs.getn("cosmology/Omegab0")
-        self.sigma8 = self.configs.getn("cosmology/sigma8")
-        self.cosmo = Cosmology(H0=self.H0, Om0=self.OmegaM0,
-                               Ob0=self.Omegab0, sigma8=self.sigma8)
-
         # Sky and resolution
         if self.sky.type_ == "patch":
             self.resolution = self.sky.pixelsize  # [arcsec]
@@ -184,11 +176,10 @@ class GalaxyClusters:
             if ii % 50 == 0:
                 logger.info("[%d/%d] %.1f%% ..." % (ii, num, 100*ii/num))
             z0, M0 = row.z, row.mass
-            age0 = self.cosmo.age(z0)
-            zmax = self.cosmo.redshift(age0 - self.tau_merger)
+            age0 = cosmo.age(z0)
+            zmax = cosmo.redshift(age0 - self.tau_merger)
             clform = ClusterFormation(M0=M0, z0=z0, zmax=zmax,
                                       ratio_major=self.ratio_major,
-                                      cosmo=self.cosmo,
                                       merger_mass_min=self.merger_mass_min)
             clform.simulate_mergertree(main_only=True)
             mmev = clform.last_major_merger