Merge pull request #152 from SWIFTSIM/cosmology-fixes

Added neutrinos and Tcmb0 read from snapshots

swiftsimio/conversions.py

@@ -14,6 +14,47 @@
     import astropy.units as astropy_units
     import numpy as np
 
+    def swift_neutrinos_to_astropy(N_eff, N_ur, M_nu_eV, deg_nu):
+        """
+        Parameters
+        ----------
+
+        N_eff: float
+            Fractional number of effective massless neutrinos at high redshift
+
+        N_ur: float
+            Fractional number of massless neutrino species
+
+        M_nu_eV: array of floats
+            Masses in eV of massive species only, up to degeneracy
+
+        deg_nu: array of floats
+            Fractional degeneracies of the massive neutrino species
+
+        Returns
+        -------
+
+        ap_m_nu
+            Array of neutrino masses in eV, replicated according to degeneracy,
+            including massless species, as desired by astropy
+        """
+        if np.isscalar(deg_nu):
+            deg_nu = np.array([deg_nu])
+        if np.isscalar(M_nu_eV):
+            M_nu_eV = np.array([M_nu_eV])
+        if not (deg_nu == deg_nu.astype(int)).all():
+            raise AttributeError(
+                "SWIFTsimIO uses astropy, which cannot handle this cosmological model."
+            )
+        if not int(N_eff) == deg_nu.astype(int).sum() + int(N_ur):
+            raise AttributeError(
+                "SWIFTsimIO uses astropy, which cannot handle this cosmological model."
+            )
+        ap_m_nu = [[m] * int(d) for m, d in zip(M_nu_eV, deg_nu)]  # replicate
+        ap_m_nu = sum(ap_m_nu, []) + [0.0] * int(N_ur)  # flatten + add massless
+        ap_m_nu = np.array(ap_m_nu) * astropy_units.eV
+        return ap_m_nu
+
     def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology:
         """
         Parameters
@@ -42,19 +83,52 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology:
         w_0 = cosmo["w_0"][0]
         w_a = cosmo["w_a"][0]
 
-        # expressions taken directly from astropy, since they do no longer
-        # allow access to these attributes (since version 5.1+)
-        critdens_const = (3.0 / (8.0 * np.pi * const.G)).cgs.value
-        a_B_c2 = (4.0 * const.sigma_sb / const.c ** 3).cgs.value
-
-        # SWIFT provides Omega_r, but we need a consistent Tcmb0 for astropy.
-        # This is an exact inversion of the procedure performed in astropy.
-        critical_density_0 = astropy_units.Quantity(
-            critdens_const * H0.to("1/s").value ** 2,
-            astropy_units.g / astropy_units.cm ** 3,
-        )
-
-        Tcmb0 = (Omega_r * critical_density_0.value / a_B_c2) ** (1.0 / 4.0)
+        # For backwards compatibility with previous cosmology constructs
+        # in snapshots
+        Tcmb0 = None
+        Neff = None
+        N_ur = None
+        M_nu_eV = None
+        deg_nu = None
+
+        try:
+            Tcmb0 = cosmo["T_CMB_0 [K]"][0]
+        except (IndexError, KeyError, AttributeError):
+            # expressions taken directly from astropy, since they do no longer
+            # allow access to these attributes (since version 5.1+)
+            critdens_const = (3.0 / (8.0 * np.pi * const.G)).cgs.value
+            a_B_c2 = (4.0 * const.sigma_sb / const.c ** 3).cgs.value
+
+            # SWIFT provides Omega_r, but we need a consistent Tcmb0 for astropy.
+            # This is an exact inversion of the procedure performed in astropy.
+            critical_density_0 = astropy_units.Quantity(
+                critdens_const * H0.to("1/s").value ** 2,
+                astropy_units.g / astropy_units.cm ** 3,
+            )
+
+            Tcmb0 = (Omega_r * critical_density_0.value / a_B_c2) ** (1.0 / 4.0)
+
+        try:
+            Neff = cosmo["N_eff"][0]
+        except (IndexError, KeyError, AttributeError):
+            Neff = 3.04  # Astropy default
+
+        try:
+            M_nu_eV = cosmo["M_nu_eV"]
+        except (IndexError, KeyError, AttributeError):
+            M_nu_eV = 0.0
+
+        try:
+            deg_nu = cosmo["deg_nu"]
+        except (IndexError, KeyError, AttributeError):
+            deg_nu = 0.0
+
+        try:
+            N_ur = cosmo["N_ur"]
+        except (IndexError, KeyError, AttributeError):
+            N_ur = 3.04  # Astropy default
+
+        ap_m_nu = swift_neutrinos_to_astropy(Neff, N_ur, M_nu_eV, deg_nu)
 
         return w0waCDM(
             H0=H0.to_astropy(),
@@ -64,6 +138,8 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology:
             wa=w_a,
             Tcmb0=Tcmb0,
             Ob0=Omega_b,
+            Neff=Neff,
+            m_nu=ap_m_nu,
         )