Angular distance for one lens and multiple sources (#950)

* implemented and tested

* extra code in profiles.py not needed

pyccl/background.py

@@ -117,10 +117,15 @@ def comoving_angular_distance(cosmo, a):
 def angular_diameter_distance(cosmo, a1, a2=None):
     """Angular diameter distance.
 
-    .. note:: The angular diameter distance in Mpc from scale factor
-              a1 to scale factor a2. If a2 is not provided, it is assumed that
-              the distance will be calculated between 1 and a1. Note that a2
-              has to be smaller than a1.
+    The angular diameter distance in Mpc from scale factor
+    `a1` to scale factor `a2`. If `a2` is not provided, it is
+    assumed that the distance will be calculated between 1 and
+    `a1`.
+
+    .. note:: `a2` has to be smaller than `a1` (i.e. a source at
+              `a2` is behind one at `a1`). You can compute the
+              distance between a single lens at `a1` and multiple
+              sources at `a2` by passing a scalar `a1`.
 
     Args:
         cosmo (:class:`~pyccl.core.Cosmology`): Cosmological parameters.
@@ -133,6 +138,9 @@ def angular_diameter_distance(cosmo, a1, a2=None):
     """
     cosmo.compute_distances()
     if(a2 is not None):
+        # One lens, multiple sources
+        if (np.ndim(a1) == 0) and (np.ndim(a2) != 0):
+            a1 = np.full(len(a2), a1)
         return _vectorize_fn5(lib.angular_diameter_distance,
                               lib.angular_diameter_distance_vec,
                               cosmo, a1, a2)

pyccl/halos/profiles.py

@@ -299,9 +299,6 @@ def convergence(self, cosmo, r, M, a_lens, a_source, mass_def):
                 :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
 
@@ -334,9 +331,6 @@ def shear(self, cosmo, r, M, a_lens, a_source, mass_def):
         """
         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)
 

pyccl/tests/test_background.py

@@ -28,6 +28,31 @@
 input_fgrowth = ccl.background.growth_rate(COSMO, input_a_array)
 
 
+def test_angular_distance_arrs():
+    a_lens = 1.0
+    a_source = 0.5
+
+    # Pass no lens, one source
+    d0 = ccl.angular_diameter_distance(COSMO, a_source)
+
+    # Pass no lens, many sources
+    d = ccl.angular_diameter_distance(COSMO, np.full(10, a_source))
+    assert np.all(np.fabs(d-d0) < 1E-10)
+
+    # One source, one lens
+    d = ccl.angular_diameter_distance(COSMO, a_lens, a_source)
+    assert np.fabs(d0-d) < 1E-10
+
+    # Many sources, one lens
+    d = ccl.angular_diameter_distance(COSMO, a_lens, np.full(10, a_source))
+    assert np.all(np.fabs(d-d0) < 1E-10)
+
+    # Many sources, many lenses
+    d = ccl.angular_diameter_distance(COSMO, np.full(10, a_lens),
+                                      np.full(10, a_source))
+    assert np.all(np.fabs(d-d0) < 1E-10)
+
+
 @pytest.mark.parametrize('a', AVALS)
 @pytest.mark.parametrize('func', [
     ccl.growth_factor,