Merge pull request #28 from martinkilbinger/cosmo

cosmo functions

Author: martinkilbinger

cs_util/calc.py

@@ -43,3 +43,29 @@ def weighted_avg_and_std(values, weights, corrected=False):
         variance = variance * n / (n - 1)
 
     return average, np.sqrt(variance)
+
+
+def transform_nan(value):
+    """Transform Nan.
+
+    Transform a ``nan`` to a very large number.
+
+    Parameters
+    ----------
+    value : float
+        input value
+
+    Returns
+    -------
+    float
+        output value
+
+    """
+    large = 1e30
+
+    if np.isnan(value) or np.isinf(value):
+        res = 1e30
+    else:
+        res = value
+
+    return res

cs_util/cosmo.py

@@ -14,8 +14,40 @@
 from astropy import constants
 from astropy import units
 
+import pyccl as ccl
 
-def sigma_crit(z_lens, z_source, cosmo, d_lens=None, d_source=None):
+
+# Set ell_max to large value, for spline interpolation (in integral over
+# C_ell to get real-space correlation functions). Avoid aliasing
+# (oscillations)
+ccl.spline_params.ELL_MAX_CORR = 10_000_000
+
+ccl.spline_params.N_ELL_CORR = 5_000
+
+
+def get_cosmo_default():
+    """Get Cosmo Default.
+
+    Return default cosmology.
+
+    Returns
+    -------
+    Cosmology
+        pyccl cosmology object
+
+    """
+    cos = ccl.Cosmology(
+        Omega_c=0.27,
+        Omega_b=0.045,
+        h=0.67,
+        sigma8=0.83,
+        n_s=0.96,
+    )
+
+    return cos
+
+
+def sigma_crit(z_lens, z_source, cos, d_lens=None, d_source=None):
     """Sigma Crit.
 
     Critical surface mass density.
@@ -26,7 +58,7 @@ def sigma_crit(z_lens, z_source, cosmo, d_lens=None, d_source=None):
         lens redshift
     z_source : float
         source redshift
-    cosmo : pyccl.core.Cosmology
+    cos : pyccl.core.Cosmology
         cosmological parameters
     d_lens : astropy.units.Quantity, optional
         precomputed anguar diameter distance to lens, computed from z_lens
@@ -50,13 +82,11 @@ def sigma_crit(z_lens, z_source, cosmo, d_lens=None, d_source=None):
     a_lens = 1 / (1 + z_lens)
     a_source = 1 / (1 + z_source)
     if not d_lens:
-        d_lens = cosmo.angular_diameter_distance(a_lens) * units.Mpc
+        d_lens = cos.angular_diameter_distance(a_lens) * units.Mpc
     if not d_source:
-        d_source = cosmo.angular_diameter_distance(a_source) * units.Mpc
+        d_source = cos.angular_diameter_distance(a_source) * units.Mpc
 
-    d_lens_source = (
-        cosmo.angular_diameter_distance(a_lens, a_source) * units.Mpc
-    )
+    d_lens_source = cos.angular_diameter_distance(a_lens, a_source) * units.Mpc
 
     frac = d_source / (d_lens_source * d_lens)
     pref = constants.c**2 / (4 * np.pi * constants.G)
@@ -70,7 +100,7 @@ def sigma_crit_eff(
     z_lens,
     z_source_arr,
     nz_source_arr,
-    cosmo,
+    cos,
     d_lens=None,
     d_source_arr=None,
 ):
@@ -87,7 +117,7 @@ def sigma_crit_eff(
         source redshifts
     nz_source_arr : list
         number of galaxies at z_source
-    cosmo : pyccl.core.Cosmology
+    cos : pyccl.core.Cosmology
         cosmological parameters
     d_lens : astropy.units.Quantity, optional
         precomputed anguar diameter distance to lens;
@@ -124,7 +154,7 @@ def sigma_crit_eff(
         sigma_cr = sigma_crit(
             z_lens,
             z_source_arr[idx],
-            cosmo,
+            cos,
             d_lens=d_lens,
             d_source=d_source_arr[idx],
         )
@@ -146,7 +176,7 @@ def sigma_crit_m1_eff(
     z_lens,
     z_source_arr,
     nz_source_arr,
-    cosmo,
+    cos,
     d_lens=None,
     d_source_arr=None,
 ):
@@ -164,7 +194,7 @@ def sigma_crit_m1_eff(
         source redshifts
     nz_source_arr : list
         number of galaxies at z_source
-    cosmo : pyccl.core.Cosmology
+    cos : pyccl.core.Cosmology
         cosmological parameters
     d_lens : astropy.units.Quantity, optional
         precomputed anguar diameter distance to lens;
@@ -204,7 +234,7 @@ def sigma_crit_m1_eff(
         sigma_cr = sigma_crit(
             z_lens,
             z_source_arr[idx],
-            cosmo,
+            cos,
             d_lens=d_lens,
             d_source=d_source_arr[idx],
         )
@@ -225,3 +255,57 @@ def sigma_crit_m1_eff(
     sigma_cr_m1_eff = np.average(sigma_cr_m1_arr, weights=weights)
 
     return sigma_cr_m1_eff * unit
+
+
+def xipm_theo(
+    theta,
+    cos,
+    z,
+    dndz,
+):
+    """Xipm Theo.
+
+    Return theoretical prediction of the shear two-point correlation function.
+
+    Parameters
+    ----------
+    theta : list
+        angular scales, list of type astropy.units.Quantity
+    cos : pyccl.core.Cosmology
+        cosmological parameters
+    z : list
+        redshift centers
+    dndz : list
+        number of galaxies for each z (arbitrary normalisation)
+
+    Returns
+    -------
+    numpy.ndarray
+        xi_+
+    numpy.ndarray
+        xi_-
+
+    """
+    # Create objects to represent tracers of the weak lensing signal with this
+    # number density (with has_intrinsic_alignment=False)
+    lens_tr = ccl.WeakLensingTracer(cos, dndz=(z, dndz))
+
+    # Calculate the angular cross-spectrum of the two tracers as a function
+    # of ell
+    # MKDEBUG TODO: vary, use unions-shear-ustc-cea/unions_wl/defaults.py
+    ell = np.logspace(0, np.log10(10000), 1000)
+    cl = ccl.angular_cl(cos, lens_tr, lens_tr, ell)
+
+    method = "Bessel"
+
+    xipm = {}
+    for corr_type in ("GG+", "GG-"):
+        xipm[corr_type] = ccl.correlation(
+            cos,
+            ell,
+            cl,
+            theta.to("deg"),
+            type=corr_type,
+        )
+
+    return xipm["GG+"], xipm["GG-"]

cs_util/tests/test_cosmo.py

@@ -5,21 +5,21 @@
 """
 
 import os
+import pickle
 
+from numpy import testing as npt
 import numpy as np
-import pyccl as ccl
 
 from astropy import units
-
-from numpy import testing as npt
+import pyccl as ccl
 
 from unittest import TestCase
 
 from cs_util import cosmo
 
 
 class CosmoTestCase(TestCase):
-    """Test case for the ``cosmo` module."""
+    """Test case for the ``cosmo`` module."""
 
     def setUp(self):
         """Set test parameter values."""
@@ -45,6 +45,20 @@ def setUp(self):
             1678.82870081,
         ] * units.Mpc
 
+        self._cos_def = cosmo = ccl.Cosmology(
+            Omega_c=0.27,
+            Omega_b=0.045,
+            h=0.67,
+            sigma8=0.83,
+            n_s=0.96,
+        )
+
+        self._theta = [1, 10, 100] * units.arcmin
+        self._z = np.linspace(0.2, 1.2, 50)
+        self._nz = (self._z / 0.8) ** 2 * np.exp(-((self._z / 0.8) ** 1.5))
+        self._xip = [1.33045991e-04, 2.13181640e-05, 2.13598131e-06]
+        self._xim = [1.97627462e-05, 1.23127046e-05, 2.17498675e-06]
+
     def tearDown(self):
         """Unset test parameter values."""
         self._z_source = None
@@ -55,6 +69,12 @@ def tearDown(self):
         self._d_source = None
         self._d_lens = None
         self._ds_cosmo = None
+        self._cos_def = None
+        self._theta = None
+        self._z = None
+        self._nz = None
+        self._xip = None
+        self._xim = None
 
     def test_sigma_crit(self):
         """Test ``cs_util.cosmo.sigma_crit`` method."""
@@ -250,3 +270,22 @@ def test_sigma_crit_m1_eff(self):
             self._sigma_crit_value_eff_m1,
             decimal=3,
         )
+
+    def test_get_cosmo_default(self):
+        """Test ``cs_util.get_cosmo_default`` method."""
+
+        cos_def = cosmo.get_cosmo_default()
+
+        npt.assert_equal(pickle.dumps(cos_def), pickle.dumps(self._cos_def))
+
+    def test_xipm_theo(self):
+        xip, xim = cosmo.xipm_theo(
+            self._theta,
+            self._cosmo,
+            self._z,
+            self._nz,
+        )
+        npt.assert_equal(len(xip), len(self._theta))
+        for idx in range(len(self._theta)):
+            npt.assert_almost_equal(xip[idx], self._xip[idx], decimal=4)
+            npt.assert_almost_equal(xim[idx], self._xim[idx], decimal=4)