Fix free Einasto alpha dimension (#989)

* Remove unneeded code from last PR

* Fix free alpha dimension

* Update docstrings; Change default alpha to 'cosmo'

* Update docstrings for update_parameters

* Remove unneeded code

* Use np.full_like; Update tests

pyccl/halos/profiles.py

@@ -839,10 +839,12 @@ class HaloProfileEinasto(HaloProfile):
         truncated (bool): set to `True` if the profile should be
             truncated at :math:`r = R_\\Delta` (i.e. zero at larger
             radii.
+        alpha (float, 'cosmo'): Set the Einasto alpha parameter or set to
+            'cosmo' to calculate the value from cosmology. Default: 'cosmo'
     """
     name = 'Einasto'
 
-    def __init__(self, c_M_relation, truncated=True, alpha=None):
+    def __init__(self, c_M_relation, truncated=True, alpha='cosmo'):
         if not isinstance(c_M_relation, Concentration):
             raise TypeError("c_M_relation must be of type `Concentration`)")
 
@@ -861,8 +863,9 @@ def update_parameters(self, alpha=None):
 
         Arguments
         ---------
-        alpha : float
-            Profile shape parameter.
+        alpha : float, 'cosmo'
+            Profile shape parameter. Set to
+            'cosmo' to calculate the value from cosmology
         """
         if alpha is not None and alpha != self.alpha:
             self.alpha = alpha
@@ -871,15 +874,14 @@ def _get_cM(self, cosmo, M, a, mdef=None):
         return self.cM.get_concentration(cosmo, M, a, mdef_other=mdef)
 
     def _get_alpha(self, cosmo, M, a, mdef):
-        if self.alpha is None:
+        if self.alpha == 'cosmo':
             mdef_vir = MassDef('vir', 'matter')
             Mvir = mdef.translate_mass(cosmo, M, a, mdef_vir)
-            print(cosmo, M, a, mdef_vir)
             sM = sigmaM(cosmo, Mvir, a)
             nu = 1.686 / sM
             alpha = 0.155 + 0.0095 * nu * nu
         else:
-            alpha = self.alpha
+            alpha = np.full_like(M, self.alpha)
         return alpha
 
     def _norm(self, M, Rs, c, alpha):

pyccl/tests/test_profiles.py

@@ -115,8 +115,17 @@ def test_cib_2pt_raises():
 def test_einasto_smoke():
     c = ccl.halos.ConcentrationDuffy08(M200)
     p = ccl.halos.HaloProfileEinasto(c)
-    p.update_parameters(alpha=1.)
-    assert p._get_alpha(COSMO, 1E14, 1., M200) == 1.
+    for M in [1E14, [1E14, 1E15]]:
+        alpha_from_cosmo = p._get_alpha(COSMO, M, 1., M200)
+
+        p.update_parameters(alpha=1.)
+        alpha = p._get_alpha(COSMO, M, 1., M200)
+        assert np.ndim(M) == np.ndim(alpha)
+        assert np.all(p._get_alpha(COSMO, M, 1., M200) == np.full_like(M, 1.))
+
+        p.update_parameters(alpha='cosmo')
+        assert np.ndim(M) == np.ndim(alpha_from_cosmo)
+        assert np.all(p._get_alpha(COSMO, M, 1., M200) == alpha_from_cosmo)
 
 
 def test_gnfw_smoke():