Setting up documentation for mflike (#121)

* first commit to set docs for mflike

* modifications

* adding figure

* reformatting

* docs for mflike and theoryforge

* small fixes

* adding documentation for bandpass

* adding foreground docs

* Update bandpass.py

Add link referring to bandpass shift bug.

* codestyle update bandpass.py

---------

Co-authored-by: Ian Harrison <itrharrison@gmail.com>

docs/bandpass.rst

@@ -0,0 +1,11 @@
+Bandpass
+========
+
+Bandpass computation
+--------------------
+
+.. automodule:: soliket.bandpass
+    :exclude-members: initialize
+    :members:
+    :show-inheritance:
+    :private-members:

docs/foreground.rst

@@ -0,0 +1,11 @@
+Foreground
+==========
+
+Foreground computation
+----------------------
+
+.. autoclass:: soliket.Foreground
+    :members:
+    :show-inheritance:
+    :private-members:
+

docs/index.rst

@@ -21,11 +21,15 @@ Table of contents
    bias
    ccl
    cosmopower
+   bandpass
+   foreground
 
 .. toctree::
    :caption: Likelihoods
    :maxdepth: 2
 
+   mflike
+
 .. toctree::
    :caption: Developers
    :maxdepth: 2

docs/mflike.rst

@@ -0,0 +1,24 @@
+MFLike
+======
+
+.. automodule:: soliket.mflike.mflike
+
+Multi Frequency Likelihood
+--------------------------
+
+.. autoclass:: soliket.mflike.MFLike
+    :exclude-members: initialize
+    :members:
+    :show-inheritance:
+    :private-members:
+
+.. automodule:: soliket.mflike.theoryforge_MFLike
+
+TheoryForge_MFLike
+------------------
+
+.. autoclass:: soliket.mflike.TheoryForge_MFLike
+    :exclude-members: initialize
+    :members:
+    :show-inheritance:
+    :private-members:

soliket/bandpass.py

@@ -13,14 +13,27 @@
 
 
 def _cmb2bb(nu):
+    r"""
+    Computes the conversion factor :math:`\frac{\partial B_{\nu}}{\partial T}`
+    from CMB thermodynamic units to antenna temperature units.
+    There is an additional :math:`\nu^2` factor to convert passbands from
+    Rayleigh-Jeans units to antenna temperature units.
 
+    :param nu: frequency array
+
+    :return: the array :math:`\frac{\partial B_{\nu}}{\partial T} \nu^2`
+    """
     # NB: numerical factors not included
     x = nu * h_Planck * 1e9 / k_Boltzmann / T_CMB
     return np.exp(x) * (nu * x / np.expm1(x))**2
 
 
 # Provides the frequency value given the bandpass name. To be modified - it is ACT based!!
 def _get_fr(array):
+    r"""
+    Provides the strings for the ACT frequency array. It will be removed in 
+    future versions.
+    """
 
     #a = array.split("_")[0]
     #if a == 'PA1' or a == 'PA2':
@@ -39,6 +52,10 @@ def _get_fr(array):
 
 
 def _get_arrays_weights(arrays, polarized_arrays, freqs):
+    """
+    Provides the array weights for the ACT frequency array. It could be removed in
+    future versions.
+    """
     array_weights = {}
     counter = []
     for array in arrays:
@@ -103,6 +120,12 @@ def must_provide(self, **requirements):
             self.freqs = requirements["bandint_freqs"]["freqs"]
 
     def calculate(self, state, want_derived=False, **params_values_dict):
+        r"""
+        Adds the bandpass transmission to the ``state`` dictionary of the
+        BandPass Theory class.
+
+        :param *params_values_dict: dictionary of nuisance parameters
+        """
 
         nuis_params = {k: params_values_dict[k] for k in self.expected_params_bp}
 
@@ -115,11 +138,34 @@ def calculate(self, state, want_derived=False, **params_values_dict):
         state["bandint_freqs"] = self.bandint_freqs
 
     def get_bandint_freqs(self):
+        """
+        Returns the ``state`` dictionary of bandpass transmissions
+        """
         return self.current_state["bandint_freqs"]
 
     # Takes care of the bandpass construction. It returns a list of nu-transmittance for
     # each frequency or an array with the effective freqs.
     def _bandpass_construction(self, **params):
+        r"""
+        Note: There is currently
+        `a known bug <https://github.com/simonsobs/LAT_MFLike/pull/58>`_ here where the
+        denominator does not depend on the bandpass shift as it should.
+        Results should be used with caution.
+        Builds the bandpass transmission 
+        :math:`\frac{\frac{\partial B_{\nu+\Delta \nu}}{\partial T} 
+        (\nu+\Delta \nu)^2 \tau(\nu+\Delta \nu)}{\int d\nu 
+        \frac{\partial B_{\nu}}{\partial T} (\nu)^2 \tau(\nu)}`  
+        using passbands :math:`\tau(\nu)` (in RJ units, not read from file)
+        and bandpass shift :math:`\Delta \nu`. :math:`\tau(\nu)` is built as a top-hat
+        with width ``bandint_width`` and number of samples ``nsteps``, 
+        read from the ``BandPass.yaml``.
+        If ``nstep = 1`` and ``bandint_width = 0``, the passband is a Dirac delta
+        centered at :math:`\nu+\Delta \nu`.
+
+        :param *params: dictionary of nuisance parameters 
+        :return: the list of [nu, transmission] in the multifrequency case  
+                 or just an array of frequencies in the single frequency one
+        """
 
         if not hasattr(self.bandint_width, "__len__"):
             self.bandint_width = np.full_like(self.freqs, self.bandint_width,
@@ -158,6 +204,12 @@ def _bandpass_construction(self, **params):
     #        trans_norm = np.trapz(bp * _cmb2bb(nu_ghz), nu_ghz)
     #        self.external_bandpass.append([fr, nu_ghz, bp / trans_norm])
     def _init_external_bandpass_construction(self, path, arrays):
+        """
+        Initializes the passband reading for ``_external_bandpass_construction``.
+
+        :param path: path of the passband txt file
+        :param arrays: list of arrays
+        """
         self.external_bandpass = []
         for array in arrays:
             fr, pa = _get_fr(array)
@@ -174,6 +226,22 @@ def _init_external_bandpass_construction(self, path, arrays):
     #        trans = bp * _cmb2bb(nub)
     #        bandint_freqs.append([nub, trans])
     def _external_bandpass_construction(self, **params):
+        r"""
+        Note: There is currently
+        `a known bug <https://github.com/simonsobs/LAT_MFLike/pull/58>`_ here where the
+        denominator does not depend on the bandpass shift as it should.
+        Results should be used with caution.
+        Builds bandpass transmission 
+        :math:`\frac{\frac{\partial B_{\nu+\Delta \nu}}{\partial T} 
+        (\nu+\Delta \nu)^2 \tau(\nu+\Delta \nu)}{\int d\nu 
+        \frac{\partial B_{\nu}}{\partial T} (\nu)^2 \tau(\nu)}`   
+        using passbands :math:`\tau(\nu)` (in RJ units) read from file and 
+        possible bandpass shift parameters :math:`\Delta \nu`.
+
+        :param *params: dictionary of nuisance parameters
+
+        :return: the list of [nu, transmission]
+        """
         bandint_freqs = []
         order = []
         for pa, fr, nu_ghz, bp in self.external_bandpass:

soliket/foreground.py

@@ -15,7 +15,11 @@ class Foreground(Theory):
 
     # Initializes the foreground model. It sets the SED and reads the templates
     def initialize(self):
-
+        """
+        Initializes the foreground models from ``fgspectra``. Sets the SED 
+        of kSZ, tSZ, dust, radio, CIB Poisson and clustered,
+        tSZxCIB, and reads the templates for CIB and tSZxCIB.
+        """
         from fgspectra import cross as fgc
         from fgspectra import frequency as fgf
         from fgspectra import power as fgp
@@ -68,6 +72,24 @@ def _get_foreground_model(self,
                               freqs=None,
                               bandint_freqs=None,
                               **fg_params):
+        r"""
+        Gets the foreground power spectra for each component computed by ``fgspectra``.
+        The computation assumes the bandpass transmissions from the ``BandPass`` class
+        and integration in frequency is performed if the passbands are not Dirac delta.
+
+        :param requested_cls: the fields required. If ``None``, 
+                              it uses the default ones in the 
+                              ``Foreground.yaml``
+        :param ell: ell range. If ``None`` the default range 
+                    set in ``Foreground.yaml`` is used
+        :param freqs: list of frequency channels. If ``None``, 
+                      it uses the default ones in the ``Foreground.yaml``
+        :param bandint_freqs: the bandpass transmissions. If ``None`` it uses ``freqs``, 
+                              otherwise the transmissions computed by the ``BandPass`` 
+                              class
+
+        :return: the foreground dictionary
+        """
 
         if not requested_cls:
             requested_cls = self.requested_cls
@@ -210,10 +232,18 @@ def must_provide(self, **requirements):
             return {"bandint_freqs": {"freqs": self.freqs}}
 
     def get_bandpasses(self, **params):
+        """
+        Gets bandpass transmissions from the ``BandPass`` class.
+        """
         return self.provider.get_bandint_freqs()
 
     def calculate(self, state, want_derived=False, **params_values_dict):
-
+        """
+        Fills the ``state`` dictionary of the ``Foreground`` Theory class
+        with the foreground spectra, computed using the bandpass 
+        transmissions from the ``BandPass`` class and the sampled foreground
+        parameters.
+        """
         self.bandint_freqs = self.get_bandpasses(**params_values_dict)
         fg_params = {k: params_values_dict[k] for k in self.expected_params_fg}
         state["fg_dict"] = self._get_foreground_model(requested_cls=self.requested_cls,
@@ -222,4 +252,7 @@ def calculate(self, state, want_derived=False, **params_values_dict):
                                                     **fg_params)
 
     def get_fg_dict(self):
+        """
+        Returns the ``state`` dictionary of fogreground spectra
+        """
         return self.current_state["fg_dict"]

soliket/mflike/mflike.py

@@ -1,10 +1,28 @@
-"""
+r"""
 .. module:: mflike
 
-:Synopsis: Definition of simplistic likelihood for Simons Observatory
+:Synopsis: Multi frequency likelihood for TTTEEE CMB power spectra for Simons Observatory
 :Authors: Thibaut Louis, Xavier Garrido, Max Abitbol,
           Erminia Calabrese, Antony Lewis, David Alonso.
 
+MFLike is a multi frequency likelihood code interfaced with the Cobaya 
+sampler and a theory Boltzmann code such as CAMB, CLASS or Cosmopower.
+The ``MFLike`` likelihood class reads the data file (in ``sacc`` format) 
+and all the settings 
+for the MCMC run (such as file paths, :math:`\ell` ranges, experiments 
+and frequencies to be used, parameters priors...)
+from the ``MFLike.yaml`` file. 
+
+The theory :math:`C_{\ell}` are then summed to the (possibly frequency 
+integrated) foreground power spectra and modified by systematic effects 
+in the ``TheoryForge_MFLike`` class. The foreground power spectra are 
+computed by the ``soliket.Foreground`` class, while the bandpasses from 
+the ``soliket.BandPass`` one; the ``Foreground`` class is required by 
+``TheoryForge_MFLike``, while ``BandPass`` is requires by ``Foreground``.
+This is a scheme of how ``MFLike`` and ``TheoryForge_MFLike`` are interfaced:
+
+.. image:: images/mflike_scheme.png
+   :width: 400
 """
 import os
 from typing import Optional
@@ -84,6 +102,14 @@ def logp(self, **params_values):
         return self.loglike(cmbfg_dict)
 
     def loglike(self, cmbfg_dict):
+        """
+        Computes the gaussian log-likelihood
+
+        :param cmbfg_dict: the dictionary of theory + foregrounds
+                           :math:`D_{\ell}`
+
+        :return: the exact loglikelihood :math:`\ln \mathcal{L}` 
+        """
         ps_vec = self._get_power_spectra(cmbfg_dict)
         delta = self.data_vec - ps_vec
         logp = -0.5 * (delta @ self.inv_cov @ delta)
@@ -94,6 +120,14 @@ def loglike(self, cmbfg_dict):
         return logp
 
     def prepare_data(self, verbose=False):
+        """
+        Reads the sacc data, extracts the data tracers,
+        trims the spectra and covariance according to the ell scales
+        set in the input file. It stores the ell vector, the deta vector
+        and the covariance in a GaussianData object.
+        If ``verbose=True``, it plots the tracer names, the spectrum name,
+        the shape of the indices array, lmin, lmax.
+        """
         import sacc
         data = self.data
         # Read data
@@ -134,10 +168,15 @@ def xp_nu(xp, nu):
             return f"{xp}_{nu}"
 
         def get_cl_meta(spec):
-            # For each of the entries of the `spectra` section of the
-            # yaml file, extract the relevant information: experiments,
-            # frequencies, polarization combinations, scale cuts and
-            # whether TE should be symmetrized.
+            """
+            Lower-level function of `prepare_data`.
+            For each of the entries of the `spectra` section of the
+            yaml file, extracts the relevant information: experiments,
+            frequencies, polarization combinations, scale cuts and
+            whether TE should be symmetrized.
+
+            :param spec: the dictionary ``data["spectra"]`` 
+            """
             # Experiments/frequencies
             exp_1, exp_2 = spec["experiments"]
             freq_1, freq_2 = spec["frequencies"]
@@ -167,10 +206,22 @@ def get_cl_meta(spec):
             return exp_1, exp_2, freq_1, freq_2, pols, scls, symm
 
         def get_sacc_names(pol, exp_1, exp_2, freq_1, freq_2):
-            # Translate the polarization combination, experiment
-            # and frequency names of a given entry in the `spectra`
-            # part of the input yaml file into the names expected
-            # in the SACC files.
+            """
+            Lower-level function of `prepare_data`.
+            Translates the polarization combination, experiment
+            and frequency names of a given entry in the `spectra`
+            part of the input yaml file into the names expected
+            in the SACC files.
+
+            :param pol: temperature or polarization fields, i.e. 'TT', 'TE'
+            :param exp_1: experiment of map 1
+            :param exp_2: experiment of map 2
+            :param freq_1: frequency of map 1
+            :param freq_2: frequency of map 2
+
+            :return: tracer name 1, tracer name 2, string for :math:`C_{\ell}`
+                     type
+            """
             p1, p2 = pol
             tname_1 = xp_nu(exp_1, freq_1)
             tname_2 = xp_nu(exp_2, freq_2)
@@ -340,7 +391,14 @@ def get_sacc_names(pol, exp_1, exp_2, freq_1, freq_2):
         self.data = GaussianData("mflike", self.ell_vec, self.data_vec, self.cov)
 
     def _get_power_spectra(self, cmbfg):
-        # Get Dl's from the theory component
+        """
+        Get :math:`D_{\ell}` from the theory component
+        already modified by ``theoryforge_MFLike``
+
+        :param cmbfg: the dictionary of theory+foreground :math:`D_{\ell}`
+
+        :return: the binned data vector
+        """
         ps_vec = np.zeros_like(self.data_vec)
         DlsObs = dict()
         # Note we rescale l_bpws because cmbfg spectra start from l=2