Fix WL functions in pyccl.halos.profiles (#943)

* Fix array_like a_source

* Update docstrings

* Add tests

* Replace np.iterable with builtin function

* Revert removal of a_source = np.array(a_source)

* Add docstrings on shape of array_like a_source

* Minor fix for docstrings

benchmarks/test_haloprofile.py

@@ -146,6 +146,8 @@ def test_weak_lensing_functions():
         np.log(rmax/rmin) * np.arange(data.shape[0]) / (data.shape[0]-1))
 
     r_al = r / a_lens
+    len_r = len(r)
+    a_source = [a_source]*len_r
 
     mdef = ccl.halos.MassDef(mDelta, 'matter')
     c = ccl.halos.ConcentrationConstant(c=concentration, mdef=mdef)

pyccl/halos/profiles.py

@@ -290,8 +290,9 @@ def convergence(self, cosmo, r, M, a_lens, a_source, mass_def=None):
             cosmo (:class:`~pyccl.core.Cosmology`): A Cosmology object.
             r (float or array_like): comoving radius in Mpc.
             M (float or array_like): halo mass in units of M_sun.
-            a_lens (float or array_like): scale factor of lens.
+            a_lens (float): scale factor of lens.
             a_source (float or array_like): scale factor of source.
+                If array_like, it must have the same shape as `r`.
             mass_def (:class:`~pyccl.halos.massdef.MassDef`):
                 a mass definition object.
 
@@ -300,6 +301,9 @@ def convergence(self, cosmo, r, M, a_lens, a_source, mass_def=None):
                 :math:`\\kappa`
         """
         Sigma = self.projected(cosmo, r, M, a_lens, mass_def) / a_lens**2
+        if hasattr(a_source, "__iter__"):
+            a_source = np.array(a_source)
+            a_lens = np.full_like(a_source, a_lens)
         Sigma_crit = sigma_critical(cosmo, a_lens, a_source)
         return Sigma / Sigma_crit
 
@@ -320,8 +324,9 @@ def shear(self, cosmo, r, M, a_lens, a_source, mass_def=None):
             cosmo (:class:`~pyccl.core.Cosmology`): A Cosmology object.
             r (float or array_like): comoving radius in Mpc.
             M (float or array_like): halo mass in units of M_sun.
-            a_lens (float or array_like): lens' scale factor.
+            a_lens (float): scale factor of lens.
             a_source (float or array_like): source's scale factor.
+                If array_like, it must have the same shape as `r`.
             mass_def (:class:`~pyccl.halos.massdef.MassDef`):
                 a mass definition object.
 
@@ -331,6 +336,9 @@ def shear(self, cosmo, r, M, a_lens, a_source, mass_def=None):
         """
         Sigma = self.projected(cosmo, r, M, a_lens, mass_def)
         Sigma_bar = self.cumul2d(cosmo, r, M, a_lens, mass_def)
+        if hasattr(a_source, "__iter__"):
+            a_source = np.array(a_source)
+            a_lens = np.full_like(a_source, a_lens)
         Sigma_crit = sigma_critical(cosmo, a_lens, a_source)
         return (Sigma_bar - Sigma) / (Sigma_crit * a_lens**2)
 
@@ -349,8 +357,9 @@ def reduced_shear(self, cosmo, r, M, a_lens, a_source, mass_def=None):
             cosmo (:class:`~pyccl.core.Cosmology`): A Cosmology object.
             r (float or array_like): comoving radius in Mpc.
             M (float or array_like): halo mass in units of M_sun.
-            a_lens (float or array_like): lens' scale factor.
+            a_lens (float): scale factor of lens.
             a_source (float or array_like): source's scale factor.
+                If array_like, it must have the same shape as `r`.
             mass_def (:class:`~pyccl.halos.massdef.MassDef`):
                 a mass definition object.
 
@@ -376,8 +385,9 @@ def magnification(self, cosmo, r, M, a_lens, a_source, mass_def=None):
             cosmo (:class:`~pyccl.core.Cosmology`): A Cosmology object.
             r (float or array_like): comoving radius in Mpc.
             M (float or array_like): halo mass in units of M_sun.
-            a_lens (float or array_like): lens' scale factor.
+            a_lens (float): scale factor of lens.
             a_source (float or array_like): source's scale factor.
+                If array_like, it must have the same shape as `r`.
             mass_def (:class:`~pyccl.halos.massdef.MassDef`):
                 a mass definition object.