From b94f88bd43f490837a204a26a00bd60b37ee4a60 Mon Sep 17 00:00:00 2001
From: Celine Combet <celine.combet@lpsc.in2p3.fr>
Date: Fri, 11 Jun 2021 17:01:52 +0200
Subject: [PATCH] Adding Einasto and Hernquist profiles from CCL. Note: the
 FFTLog integration may still be optimized

---
 clmm/theory/ccl.py | 35 +++++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/clmm/theory/ccl.py b/clmm/theory/ccl.py
index 6a1e233df..215e8638e 100644
--- a/clmm/theory/ccl.py
+++ b/clmm/theory/ccl.py
@@ -3,6 +3,7 @@
 import pyccl as ccl
 
 import numpy as np
+from scipy.interpolate import interp1d
 import warnings
 from packaging import version
 
@@ -29,18 +30,20 @@ def __init__(self, massdef='mean', delta_mdef=200, halo_profile_model='nfw'):
             'mean': 'matter',
             'critical': 'critical',
             'virial': 'critical'}
-        self.hdpm_dict = {'nfw': ccl.halos.HaloProfileNFW}
+        self.hdpm_dict = {'nfw': ccl.halos.HaloProfileNFW,
+                          'einasto': ccl.halos.HaloProfileEinasto,
+                          'hernquist': ccl.halos.HaloProfileHernquist}
         # Uncomment lines below when CCL einasto and hernquist profiles are stable (also add version number)
         #if version.parse(ccl.__version__) >= version.parse('???'):
-        #    self.hdpm_dict.update({
+        # self.hdpm_dict.update({
         #        'einasto': ccl.halos.HaloProfileEinasto,
         #        'hernquist': ccl.halos.HaloProfileHernquist})
         # Attributes exclusive to this class
         self.hdpm_opts = {'nfw': {'truncated': False,
                                   'projected_analytic': True,
                                   'cumul2d_analytic': True},
-                          'einasto': {},
-                          'hernquist': {}}
+                          'einasto': {'truncated': False},
+                          'hernquist': {'truncated': False}}
         self.MDelta = 0.0
         self.cor_factor = _patch_rho_crit_to_cd2018(ccl.physical_constants.RHO_CRITICAL)
         # Set halo profile and cosmology
@@ -85,18 +88,34 @@ def eval_3d_density(self, r3d, z_cl):
 
     def eval_surface_density(self, r_proj, z_cl):
         a_cl = self.cosmo._get_a_from_z(z_cl)
-        return self.hdpm.projected(self.cosmo.be_cosmo, r_proj/a_cl, self.MDelta, a_cl, self.mdef)*self.cor_factor/a_cl**2
+        if self.halo_profile_model == 'nfw':
+            return self.hdpm.projected(self.cosmo.be_cosmo, r_proj/a_cl, self.MDelta, a_cl, self.mdef)*self.cor_factor/a_cl**2
+        else:
+            rtmp = np.geomspace(np.min(r_proj)/10., np.max(r_proj)*10., 1000)
+            tmp = self.hdpm.projected (self.cosmo.be_cosmo, rtmp/a_cl, self.MDelta, a_cl, self.mdef)*self.cor_factor/a_cl**2
+            ptf = interp1d(np.log(rtmp), np.log(tmp), bounds_error=False, fill_value=-100)
+            return np.exp(ptf(np.log(r_proj)))
+
 
     def eval_mean_surface_density(self, r_proj, z_cl):
         a_cl = self.cosmo._get_a_from_z(z_cl)
-        return self.hdpm.cumul2d(self.cosmo.be_cosmo, r_proj/a_cl, self.MDelta, self.cosmo._get_a_from_z(z_cl), self.mdef)*self.cor_factor/a_cl**2
+        if self.halo_profile_model =='nfw':
+            return self.hdpm.cumul2d(self.cosmo.be_cosmo, r_proj/a_cl, self.MDelta, self.cosmo._get_a_from_z(z_cl), self.mdef)*self.cor_factor/a_cl**2
+        else:
+            rtmp = np.geomspace(np.min(r_proj)/10., np.max(r_proj)*10., 1000)
+            tmp = self.hdpm.cumul2d (self.cosmo.be_cosmo, rtmp/a_cl, self.MDelta, a_cl, self.mdef)*self.cor_factor/a_cl**2
+            ptf = interp1d(np.log(rtmp), np.log(tmp), bounds_error=False, fill_value=-100)
+            return np.exp(ptf(np.log(r_proj)))
 
     def eval_excess_surface_density(self, r_proj, z_cl):
         a_cl = self.cosmo._get_a_from_z(z_cl)
         r_cor = r_proj/a_cl
 
-        return (self.hdpm.cumul2d(self.cosmo.be_cosmo, r_cor, self.MDelta, self.cosmo._get_a_from_z(z_cl), self.mdef)-
-                self.hdpm.projected(self.cosmo.be_cosmo, r_cor, self.MDelta, self.cosmo._get_a_from_z(z_cl), self.mdef))*self.cor_factor/a_cl**2
+        if self.halo_profile_model =='nfw':
+            return (self.hdpm.cumul2d(self.cosmo.be_cosmo, r_cor, self.MDelta, self.cosmo._get_a_from_z(z_cl), self.mdef)-
+                    self.hdpm.projected(self.cosmo.be_cosmo, r_cor, self.MDelta, self.cosmo._get_a_from_z(z_cl), self.mdef))*self.cor_factor/a_cl**2
+        else:
+            return self.eval_mean_surface_density(r_proj, z_cl) - self.eval_surface_density(r_proj, z_cl)
 
     def eval_tangential_shear(self, r_proj, z_cl, z_src):
         sigma_excess = self.eval_excess_surface_density(r_proj, z_cl)