Add simple halo model (simonsobs#150)

* remove depracated likelihood.theory calls

* codestyle

* sphinx compatability update

* sphinx compatability update 2204

* add build tools

* move rtd python ver

* rtd try different pyver

* rtd try different order

* conda on rtd

* update rtd mock import

* add rtd theme

* add rtd theme correctly

* add rtd theme in setup.cfg

* Testing yaml instead of yml.

* added docs conda env

* add full requirements in docs

* pin tensorflow dep

* started basic halo model

* started basic pyhm implementation and tests

* pkmm computation with z loop

* remove empty functions

* test against precomputed value

* codestyle and fix test

* add missing default files

* added docs

* line length

* remove refernce to bias model in docs

* improve docs

---------

Co-authored-by: Hidde Jense <JenseH@cardiff.ac.uk>

docs/index.rst

@@ -35,6 +35,7 @@ The pages here describe how to install and run SOLikeT, and document the functio
    bandpass
    foreground
    bias
+   halo_model
 
 .. toctree::
    :caption: Likelihood codes

requirements.txt

@@ -17,3 +17,4 @@ pyccl
 sacc
 fgspectra>=1.1.0
 syslibrary
+pyhalomodel

setup.cfg

@@ -24,9 +24,10 @@ install_requires =
     sacc
     fgspectra>=1.1.0
     syslibrary
+    pyhalomodel
 
 [options.package_data]
-soliket = *.yaml,*.bibtex,clusters/data/*,clusters/data/selFn_equD56/*,lensing/data/*.txt,lensing/data/*.fits,mflike/*.yaml,tests/*.yaml,data/xcorr_simulated/*.txt,data/CosmoPower/CP_paper/CMB/*.pkl,tests/data/test_bandpass/*,cross_correlation/*.yaml,tests/data/*.fits
+soliket = *.yaml,*.bibtex,clusters/data/*,clusters/data/selFn_equD56/*,lensing/data/*.txt,lensing/data/*.fits,mflike/*.yaml,tests/*.yaml,data/xcorr_simulated/*.txt,data/CosmoPower/CP_paper/CMB/*.pkl,tests/data/test_bandpass/*,cross_correlation/*.yaml,tests/data/*.fits,halo_model/*.yaml
 testpaths = "soliket"
 text_file_format = rst
 

soliket/__init__.py

@@ -3,12 +3,12 @@
 # from .studentst import StudentstLikelihood  # noqa: F401
 from .ps import PSLikelihood, BinnedPSLikelihood  # noqa: F401
 from .mflike import MFLike  # noqa: F401
-from .mflike import TheoryForge_MFLike
+from .mflike import TheoryForge_MFLike  # noqa: F401
 from .cross_correlation import CrossCorrelationLikelihood, GalaxyKappaLikelihood, ShearKappaLikelihood  # noqa: F401, E501
 from .xcorr import XcorrLikelihood  # noqa: F401
-from .foreground import Foreground
-from .bandpass import BandPass
-from .cosmopower import CosmoPower, CosmoPowerDerived
+from .foreground import Foreground  # noqa: F401
+from .bandpass import BandPass  # noqa: F401
+from .cosmopower import CosmoPower, CosmoPowerDerived  # noqa: F401
 from .ccl import CCL  # noqa: F401
 
 try:

soliket/halo_model/HaloModel.yaml

@@ -0,0 +1,10 @@
+# Default parameters for the Bias theory.
+
+# Maximum k for P(k) (units 1/Mpc). Set to zero if not needed.
+kmax: 10.0
+
+# Redshift sampling.
+z: []
+
+# Extra (non-parameter) arguments passed to ccl.Cosmology()
+extra_args:

soliket/halo_model/HaloModel_pyhm.yaml

@@ -0,0 +1,15 @@
+# Default parameters for the pyhalomodel HaloModel theory.
+hmf_name: 'Tinker et al. (2010)'
+hmf_Dv: 330.
+Mmin: 1.e9
+Mmax: 1.e17
+nM: 256
+
+# Maximum k for P(k) (units 1/Mpc). Set to zero if not needed.
+kmax: 10.0
+
+# Redshift sampling.
+z: []
+
+# Extra (non-parameter) arguments passed to ccl.Cosmology()
+extra_args:

soliket/halo_model/__init__.py

@@ -0,0 +1 @@
+from .halo_model import HaloModel, HaloModel_pyhm  # noqa: F401

soliket/halo_model/halo_model.py

@@ -0,0 +1,174 @@
+"""
+.. module:: soliket.halo_model
+
+:Synopsis: Class to calculate Halo Models for non-linear power spectra.
+:Author: Ian Harrison
+
+.. |br| raw:: html
+
+   <br />
+
+
+Usage
+-----
+
+Halo Models for calculating non-linear power spectra for use in large scale structure 
+and lensing likelihoods. The abstract HaloModel base class should be built on with 
+specific model implementations. HaloModels can be added as theory codes alongside others 
+in your run settings. e.g.:
+
+.. code-block:: yaml
+
+  theory:
+    camb:
+    soliket.halo_model.HaloModel_pyhm:
+
+
+Implementing your own halo model
+--------------------------------
+
+If you want to add your own halo model, you can do so by inheriting from the
+``soliket.HaloModel`` theory class and implementing your own custom ``calculate()``
+function (have a look at the simple pyhalomodel model for ideas).
+"""
+
+import numpy as np
+# from cobaya.theories.cosmo.boltzmannbase import PowerSpectrumInterpolator
+from scipy.interpolate import RectBivariateSpline
+from typing import Optional, Sequence
+from cobaya.theory import Theory
+from cobaya.typing import InfoDict
+import pyhalomodel as halo
+
+
+class HaloModel(Theory):
+    """Abstract parent class for implementing Halo Models."""
+
+    _logz = np.linspace(-3, np.log10(1100), 150)
+    _default_z_sampling = 10**_logz
+    _default_z_sampling[0] = 0
+
+    def initialize(self):
+        self._var_pairs = set()
+        self._required_results = {}
+
+    # def must_provide(self, **requirements):
+    #     options = requirements.get("halo_model") or {}
+
+    def _get_Pk_mm_lin(self):
+        for pair in self._var_pairs:
+            self.k, self.z, Pk_mm = \
+                self.provider.get_Pk_grid(var_pair=pair, nonlinear=False)
+
+        return Pk_mm
+
+    def get_Pk_mm_grid(self):
+
+        return self.current_state["Pk_mm_grid"]
+
+    def get_Pk_gg_grid(self):
+
+        return self.current_state["Pk_gg_grid"]
+
+    def get_Pk_gm_grid(self):
+
+        return self.current_state["Pk_gm_grid"]
+
+
+class HaloModel_pyhm(HaloModel):
+    """Halo Model wrapping the simple pyhalomodel code of Asgari, Mead & Heymans (2023)
+
+    We include this simple halo model for the non-linear matter-matter power spectrum with 
+    NFW profiles. This is calculated via the `pyhalomodel
+    <https://github.com/alexander-mead/pyhalomodel>`_ code.
+    """
+
+    def initialize(self):
+        super().initialize()
+        self.Ms = np.logspace(np.log10(self.Mmin), np.log10(self.Mmax), self.nM)
+
+    def get_requirements(self):
+
+        return {"omegam": None}
+
+    def must_provide(self, **requirements):
+
+        options = requirements.get("halo_model") or {}
+        self._var_pairs.update(
+            set((x, y) for x, y in
+                options.get("vars_pairs", [("delta_tot", "delta_tot")])))
+
+        self.kmax = max(self.kmax, options.get("kmax", self.kmax))
+        self.z = np.unique(np.concatenate(
+                            (np.atleast_1d(options.get("z", self._default_z_sampling)),
+                            np.atleast_1d(self.z))))
+
+        needs = {}
+
+        needs["Pk_grid"] = {
+                "vars_pairs": self._var_pairs,
+                "nonlinear": (False, False),
+                "z": self.z,
+                "k_max": self.kmax
+            }
+
+        needs["sigma_R"] = {"vars_pairs": self._var_pairs,
+                           "z": self.z,
+                           "k_max": self.kmax,
+                           "R": np.linspace(0.14, 66, 256) # list of radii required
+                           }
+
+
+        return needs
+
+    def calculate(self, state: dict, want_derived: bool = True,
+                  **params_values_dict):
+
+        Pk_mm_lin = self._get_Pk_mm_lin()
+
+        # now wish to interpolate sigma_R to these Rs
+        zinterp, rinterp, sigmaRinterp = self.provider.get_sigma_R()
+        # sigmaRs = PowerSpectrumInterpolator(zinterp, rinterp, sigma_R)
+        sigmaRs = RectBivariateSpline(zinterp, rinterp, sigmaRinterp)
+
+        output_Pk_hm_mm = np.empty([len(self.z), len(self.k)])
+
+        # for sure we could avoid the for loop with some thought
+        for iz, zeval in enumerate(self.z):
+            hmod = halo.model(zeval, self.provider.get_param('omegam'),
+                              name=self.hmf_name, Dv=self.hmf_Dv)
+
+            Rs = hmod.Lagrangian_radius(self.Ms)
+            rvs = hmod.virial_radius(self.Ms)
+
+            cs = 7.85 * (self.Ms / 2e12)**-0.081 * (1. + zeval)**-0.71
+            Uk = self._win_NFW(self.k, rvs, cs)
+            matter_profile = halo.profile.Fourier(self.k, self.Ms, Uk,
+                                                  amplitude=self.Ms,
+                                                  normalisation=hmod.rhom,
+                                                  mass_tracer=True)
+
+            Pk_2h, Pk_1h, Pk_hm = hmod.power_spectrum(self.k, Pk_mm_lin[iz],
+                                                      self.Ms, sigmaRs(zeval, Rs)[0],
+                                                      {'m': matter_profile},
+                                                      verbose=False)
+
+            output_Pk_hm_mm[iz] = Pk_hm['m-m']
+
+        state['Pk_mm_grid'] = output_Pk_hm_mm
+        # state['Pk_gm_grid'] = Pk_hm['g-m']
+        # state['Pk_gg_grid'] = Pk_hm['g-g']
+
+    def _win_NFW(self, k, rv, c):
+        from scipy.special import sici
+        rs = rv / c
+        kv = np.outer(k, rv)
+        ks = np.outer(k, rs)
+        Sisv, Cisv = sici(ks + kv)
+        Sis, Cis = sici(ks)
+        f1 = np.cos(ks) * (Cisv - Cis)
+        f2 = np.sin(ks) * (Sisv - Sis)
+        f3 = np.sin(kv) / (ks + kv)
+        f4 = np.log(1. + c) - c / (1. + c)
+        Wk = (f1 + f2 - f3) / f4
+        return Wk

soliket/tests/test_bias.py

@@ -1,10 +1,6 @@
-# pytest -k bias -v .
-
-import pytest
 import numpy as np
 
 from cobaya.model import get_model
-from cobaya.run import run
 
 info = {"params": {
                    "b_lin": 1.1,
@@ -22,11 +18,11 @@
 
 
 def test_bias_import():
-    from soliket.bias import Bias
+    from soliket.bias import Bias # noqa F401
 
 
 def test_linear_bias_import():
-    from soliket.bias import Linear_bias
+    from soliket.bias import Linear_bias # noqa F401
 
 
 def test_linear_bias_model():

soliket/tests/test_cash.py

@@ -23,7 +23,7 @@ def toy_data():
 
 
 def test_cash_import():
-    from soliket.cash import CashCLikelihood
+    from soliket.cash import CashCLikelihood # noqa F401
 
 
 def test_cash_read_data(request):

soliket/tests/test_ccl.py

@@ -1,7 +1,6 @@
 """
 Check that CCL works correctly.
 """
-import pytest
 import numpy as np
 from cobaya.model import get_model
 from cobaya.likelihood import Likelihood
@@ -49,7 +48,7 @@ def test_ccl_import(request):
     """
     Test whether we can import pyCCL.
     """
-    import pyccl
+    import pyccl # noqa F401
 
 
 def test_ccl_cobaya(request):

soliket/tests/test_clusters.py

@@ -1,5 +1,4 @@
 import numpy as np
-import pytest
 
 from cobaya.model import get_model
 

soliket/tests/test_cosmopower.py

@@ -2,10 +2,8 @@
 Check that CosmoPower gives the correct Planck CMB power spectrum.
 """
 import os
-import tempfile
 import pytest
 import numpy as np
-import matplotlib.pyplot as plt
 
 from cobaya.model import get_model
 from soliket.cosmopower import HAS_COSMOPOWER

soliket/tests/test_cross_correlation.py

@@ -1,6 +1,5 @@
 import numpy as np
 import os
-import pytest
 from soliket.ccl import CCL
 from cobaya.model import get_model
 
@@ -28,12 +27,12 @@
 
 def test_galaxykappa_import(request):
 
-    from soliket.cross_correlation import GalaxyKappaLikelihood
+    from soliket.cross_correlation import GalaxyKappaLikelihood # noqa F401
 
 
 def test_shearkappa_import(request):
 
-    from soliket.cross_correlation import ShearKappaLikelihood
+    from soliket.cross_correlation import ShearKappaLikelihood # noqa F401
 
 
 def test_galaxykappa_model(request):

soliket/tests/test_foreground.py

@@ -1,11 +1,7 @@
-# pytest -k bandpass -v .
-
-import pytest
 import numpy as np
 import os
 
 from cobaya.model import get_model
-from cobaya.run import run
 
 info = {"params": {
                    "a_tSZ": 3.3044404448917724,
@@ -30,7 +26,7 @@
 
 
 def test_foreground_import():
-    from soliket.foreground import Foreground
+    from soliket.foreground import Foreground # noqa F401
 
 
 def test_foreground_model():