Initial type-hinting attempt

soliket/bandpass/bandpass.py

@@ -87,7 +87,7 @@
 """
 
 import os
-from typing import Optional
+from typing import List, Optional, Union
 
 import numpy as np
 from cobaya.log import LoggedError
@@ -100,7 +100,7 @@
 # Numerical factors not included, it needs proper normalization when used.
 
 
-def _cmb2bb(nu):
+def _cmb2bb(nu: np.ndarray) -> np.ndarray:
     r"""
     Computes the conversion factor :math:`\frac{\partial B_{\nu}}{\partial T}`
     from CMB thermodynamic units to differential source intensity.
@@ -132,9 +132,9 @@ def initialize(self):
                                    "bandint_shift_LAT_145",
                                    "bandint_shift_LAT_225"]
 
-        self.exp_ch = None
-        # self.eff_freqs = None
-        self.bands = None
+        self.exp_ch: Optional[List[str]] = None
+        # self.eff_freqs: Optional[list] = None
+        self.bands: dict = None
 
         # To read passbands stored in the sacc files
         # default for mflike
@@ -171,7 +171,7 @@ def initialize_with_params(self):
                 self.log, "Configuration error in parameters: %r.",
                 differences)
 
-    def must_provide(self, **requirements):
+    def must_provide(self, **requirements: dict):
         # bandint_freqs is required by Foreground
         # and requires some params to be computed
         # Assign those from Foreground
@@ -180,7 +180,7 @@ def must_provide(self, **requirements):
             self.exp_ch = [k.replace("_s0", "") for k in self.bands.keys()
                            if "_s0" in k]
 
-    def calculate(self, state, want_derived=False, **params_values_dict):
+    def calculate(self, state: dict, want_derived=False, **params_values_dict: dict):
         r"""
         Adds the bandpass transmission to the ``state`` dictionary of the
         BandPass Theory class.
@@ -198,15 +198,17 @@ def calculate(self, state, want_derived=False, **params_values_dict):
 
         state["bandint_freqs"] = self.bandint_freqs
 
-    def get_bandint_freqs(self):
+    def get_bandint_freqs(self) -> dict:
         """
         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):
+    def _bandpass_construction(
+        self, **params: dict
+    ) -> Union[np.ndarray, List[np.ndarray]]:
         r"""
         Builds the bandpass transmission
         :math:`\frac{\frac{\partial B_{\nu+\Delta \nu}}{\partial T}
@@ -277,7 +279,7 @@ def _bandpass_construction(self, **params):
 
         return bandint_freqs
 
-    def _init_external_bandpass_construction(self, path, exp_ch):
+    def _init_external_bandpass_construction(self, path: str, exp_ch: List[str]):
         """
         Initializes the passband reading for ``_external_bandpass_construction``.
 
@@ -290,7 +292,7 @@ def _init_external_bandpass_construction(self, path, exp_ch):
             nu_ghz, bp = np.loadtxt(path + "/" + expc, usecols=(0, 1), unpack=True)
             self.external_bandpass.append([expc, nu_ghz, bp])
 
-    def _external_bandpass_construction(self, **params):
+    def _external_bandpass_construction(self, **params: dict) -> List[np.ndarray]:
         r"""
         Builds bandpass transmission
         :math:`\frac{\frac{\partial B_{\nu+\Delta \nu}}{\partial T}

soliket/bias/bias.py

@@ -26,24 +26,27 @@
 function (have a look at the linear bias model for ideas).
 """
 
+from typing import Dict, Set, Tuple
+
 import numpy as np
-from cobaya.theory import Theory
+from cobaya.theory import Provider, Theory
 
 
 class Bias(Theory):
     """Parent class for bias models."""
 
+    provider: Provider
     _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()
+    def initialize(self) -> None:
+        self._var_pairs: Set[Tuple[str, str]] = set()
 
-    def get_requirements(self):
+    def get_requirements(self) -> Dict[str, dict]:
         return {}
 
-    def must_provide(self, **requirements):
+    def must_provide(self, **requirements: dict) -> Dict[str, dict]:
         options = requirements.get("linear_bias") or {}
 
         self.kmax = max(self.kmax, options.get("kmax", self.kmax))
@@ -52,7 +55,7 @@ def must_provide(self, **requirements):
              np.atleast_1d(self.z))))
 
         # Dictionary of the things needed from CAMB/CLASS
-        needs = {}
+        needs: Dict[str, dict] = {}
 
         self.nonlinear = self.nonlinear or options.get("nonlinear", False)
         self._var_pairs.update(
@@ -69,7 +72,7 @@ def must_provide(self, **requirements):
         assert len(self._var_pairs) < 2, "Bias doesn't support other Pk yet"
         return needs
 
-    def _get_Pk_mm(self):
+    def _get_Pk_mm(self) -> np.ndarray:
         self.k, self.z, Pk_mm = \
             self.provider.get_Pk_grid(var_pair=list(self._var_pairs)[0],
                                       nonlinear=self.nonlinear)
@@ -90,7 +93,7 @@ class Linear_bias(Bias):
     """
 
     def calculate(self, state: dict, want_derived: bool = True,
-                  **params_values_dict):
+                  **params_values_dict) -> None:
         Pk_mm = self._get_Pk_mm()
 
         state["Pk_gg_grid"] = params_values_dict["b_lin"] ** 2. * Pk_mm

soliket/cash/cash.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Tuple
 import numpy as np
 from cobaya.likelihood import Likelihood
 from .cash_data import CashCData
@@ -9,25 +9,24 @@
 
 class CashCLikelihood(Likelihood):
     name: str = "Cash-C"
-    datapath = Optional[str]
-
-    def initialize(self):
+    datapath: Optional[str] = None
 
+    def initialize(self) -> None:
         x, N = self._get_data()
         self.data = CashCData(self.name, N)
 
-    def _get_data(self):
+    def _get_data(self) -> Tuple[np.ndarray, np.ndarray]:
         data = np.loadtxt(self.datapath, unpack=False)
         N = data[:, -1]  # assume data stored like column_stack([z, q, N])
         x = data[:, :-1]
         return x, N
 
-    def _get_theory(self, **kwargs):
+    def _get_theory(self, **kwargs: dict) -> np.ndarray:
         if "cash_test_logp" in kwargs:
             return np.arange(kwargs["cash_test_logp"])
         else:
             raise NotImplementedError
 
-    def logp(self, **params_values):
+    def logp(self, **params_values: dict) -> float:
         theory = self._get_theory(**params_values)
         return self.data.loglike(theory)

soliket/cash/cash_data.py

@@ -1,8 +1,13 @@
+from typing import Union
 import numpy as np
 from scipy.special import factorial
 
 
-def cash_c_logpdf(theory, data, usestirling=True):
+def cash_c_logpdf(
+    theory: Union[np.ndarray, float],
+    data: Union[np.ndarray, float],
+    usestirling: bool = True
+) -> float:
     data = np.asarray(data, dtype=int)
 
     ln_fac = np.zeros_like(data, dtype=float)
@@ -24,13 +29,15 @@ class CashCData:
     """Named multi-dimensional Cash-C distributed data
     """
 
-    def __init__(self, name, N, usestirling=True):
+    def __init__(
+            self, name: str, N: Union[np.ndarray, float], usestirling: bool = True
+    ) -> None:
         self.name = str(name)
         self.data = N
         self.usestirling = usestirling
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.data)
 
-    def loglike(self, theory):
+    def loglike(self, theory: Union[np.ndarray, float]) -> float:
         return cash_c_logpdf(theory, self.data)

soliket/ccl/ccl.py

@@ -79,8 +79,8 @@
 # https://cobaya.readthedocs.io/en/devel/theories_and_dependencies.html
 
 import numpy as np
-from typing import Sequence
-from cobaya.theory import Theory
+from typing import Dict, Sequence
+from cobaya.theory import Provider, Theory
 from cobaya.tools import LoggedError
 
 
@@ -92,6 +92,7 @@ class CCL(Theory):
     kmax: float
     z: np.ndarray
     nonlinear: bool
+    provider: Provider
 
     def initialize(self) -> None:
         try:
@@ -125,7 +126,7 @@ def must_provide(self, **requirements) -> dict:
              np.atleast_1d(self.z))))
 
         # Dictionary of the things CCL needs from CAMB/CLASS
-        needs = {}
+        needs: Dict[str, dict] = {}
 
         if self.kmax:
             self.nonlinear = self.nonlinear or options.get('nonlinear', False)
@@ -156,7 +157,7 @@ def get_can_support_params(self) -> Sequence[str]:
         return []
 
     def calculate(self, state: dict, want_derived: bool = True,
-                  **params_values_dict):
+                  **params_values_dict) -> None:
         # calculate the general CCL cosmo object which likelihoods can then use to get
         # what they need (likelihoods should cache results appropriately)
         # get our requirements from self.provider
@@ -231,5 +232,5 @@ def calculate(self, state: dict, want_derived: bool = True,
         for required_result, method in self._required_results.items():
             state['CCL'][required_result] = method(cosmo)
 
-    def get_CCL(self):
+    def get_CCL(self) -> dict:
         return self._current_state['CCL']