Fix some minor problems in boost calculation and improve docstring

clmm/utils/__init__.py

@@ -13,8 +13,8 @@
 from .boost import (
     compute_nfw_boost,
     compute_powerlaw_boost,
-    correct_sigma_with_boost_values,
-    correct_sigma_with_boost_model,
+    correct_with_boost_values,
+    correct_with_boost_model,
     boost_models,
 )
 

clmm/utils/boost.py

@@ -1,24 +1,39 @@
 """General utility functions that are used in multiple modules"""
+
 import numpy as np
 
 
-def compute_nfw_boost(rvals, rscale=1000, boost0=0.1):
-    """Given a list of `rvals`, and optional `rscale` and `boost0`, return the corresponding
-    boost factor at each rval
+def compute_nfw_boost(rvals, rscale, boost0=0.1):
+    r"""Computes the boost factor at radii rvals using a parametric form 
+    following that of the analytical NFW excess surface density.
+
+    Setting :math:`x = R/R_s` 
+
+    .. math::
+        B(x) = 1 + B_0\frac{1-F(x)}{(x)^2 -1}
+
+    where
+
+    .. math::
+        F(x) = \begin{cases} \frac{\arctan \sqrt{x^2 -1 }}{\sqrt{x^2 -1 }} & \text{if $x>1$}, \\
+        \text{1} & \text{if x = 1}, \\
+        \frac{\text{arctanh} \sqrt{1-x^2 }}{\sqrt{1- x^2 }} & \text{if $x<1$}.
+        \end{cases}
+
 
     Parameters
     ----------
     rvals : array_like
         Radii
-    rscale : float, optional
-        scale radius for NFW in same units as rvals (default 2000 kpc)
-    boost0 : float, optional
-        Boost factor at each value of rvals
+    rscale : float
+        Scale radius in same units as rvals
+    boost0: float, optional
+        Boost factor normalisation
 
     Returns
     -------
     array
-        Boost factor
+        Boost factor at each value of rvals
     """
 
     r_norm = np.array(rvals) / rscale
@@ -39,25 +54,27 @@ def _calc_finternal(r_norm):
     return 1.0 + boost0 * (1 - _calc_finternal(r_norm)) / (r_norm**2 - 1)
 
 
-def compute_powerlaw_boost(rvals, rscale=1000, boost0=0.1, alpha=-1.0):
-    """Given a list of `rvals`, and optional `rscale` and `boost0`, and `alpha`,
-    return the corresponding boost factor at each `rval`
+def compute_powerlaw_boost(rvals, rscale, boost0=0.1, alpha=-1):
+    r"""Computes the boost factor at radii rvals using a power-law parametric form
+
+    .. math::
+        B(R) = 1 + B_0 \left(\frac{R}{R_s}\right)^\alpha
 
     Parameters
     ----------
     rvals : array_like
         Radii
-    rscale : float, optional
-        Scale radius for NFW in same units as rvals (default 2000 kpc)
+    rscale : float
+        Scale radius in same units as rvals
     boost0 : float, optional
-        Boost factor at each value of rvals
+        Boost factor normalisation
     alpha : float, optional
-        Exponent from Melchior+16. Default: -1.0
+        Exponent for the power-law parametrisation. NB: -1.0 (as in Melchior+16)
 
     Returns
     -------
     array
-        Boost factor
+        Boost factor at each value of rvals
     """
 
     r_norm = np.array(rvals) / rscale
@@ -71,35 +88,36 @@ def compute_powerlaw_boost(rvals, rscale=1000, boost0=0.1, alpha=-1.0):
 }
 
 
-def correct_sigma_with_boost_values(sigma_vals, boost_factors):
-    """Given a list of boost values and sigma profile, compute corrected sigma
+def correct_with_boost_values(profile_vals, boost_factors):
+    """Given a list of profile values (e.g., shear or DeltaSigma) and boost values,
+    computes corrected profile
 
     Parameters
     ----------
-    sigma_vals : array_like
-        uncorrected sigma with cluster member dilution
+    profile_vals : array_like
+        Uncorrected profile values
     boost_factors : array_like
-        Boost values pre-computed
+        Boost values, pre-computed
 
     Returns
     -------
-    sigma_corrected : numpy.ndarray
-        correted radial profile
+    profile_vals_corrected : numpy.ndarray
+        Corrected radial profile
     """
 
-    sigma_corrected = np.array(sigma_vals) / np.array(boost_factors)
-    return sigma_corrected
+    profile_vals_corrected = np.array(profile_vals) / np.array(boost_factors)
+    return profile_vals_corrected
 
 
-def correct_sigma_with_boost_model(rvals, sigma_vals, boost_model="nfw_boost", **boost_model_kw):
+def correct_with_boost_model(rvals, profile_vals, boost_model, boost_rscale, **boost_model_kw):
     """Given a boost model and sigma profile, compute corrected sigma
 
     Parameters
     ----------
     rvals : array_like
-        radii
-    sigma_vals : array_like
-        uncorrected sigma with cluster member dilution
+        Radii
+    profile_vals : array_like
+        Uncorrected profile values
     boost_model : str, optional
         Boost model to use for correcting sigma
 
@@ -108,14 +126,14 @@ def correct_sigma_with_boost_model(rvals, sigma_vals, boost_model="nfw_boost", *
 
     Returns
     -------
-    sigma_corrected : numpy.ndarray
-        correted radial profile
+    profile_vals_corrected : numpy.ndarray
+        Correted radial profile
     """
     boost_model_func = boost_models[boost_model]
-    boost_factors = boost_model_func(rvals, **boost_model_kw)
+    boost_factors = boost_model_func(rvals, boost_rscale, **boost_model_kw)
 
-    sigma_corrected = np.array(sigma_vals) / boost_factors
-    return sigma_corrected
+    profile_vals_corrected = np.array(profile_vals) / boost_factors
+    return profile_vals_corrected
 
 
 boost_models = {