diff --git a/soliket/bandpass/bandpass.py b/soliket/bandpass/bandpass.py index 35edf0e6..27d85580 100644 --- a/soliket/bandpass/bandpass.py +++ b/soliket/bandpass/bandpass.py @@ -173,7 +173,7 @@ def initialize_with_params(self): self.log, "Configuration error in parameters: %r.", differences) - def must_provide(self, **requirements: dict): + def must_provide(self, **requirements): # bandint_freqs is required by Foreground # and requires some params to be computed # Assign those from Foreground @@ -209,7 +209,7 @@ def get_bandint_freqs(self) -> dict: # 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: dict + self, **params ) -> Union[np.ndarray, List[np.ndarray]]: r""" Builds the bandpass transmission @@ -294,7 +294,9 @@ def _init_external_bandpass_construction(self, path: str, exp_ch: List[str]): 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: dict) -> List[np.ndarray]: + def _external_bandpass_construction( + self, **params + ) -> Union[np.ndarray, List[np.ndarray]]: r""" Builds bandpass transmission :math:`\frac{\frac{\partial B_{\nu+\Delta \nu}}{\partial T} diff --git a/soliket/bias/bias.py b/soliket/bias/bias.py index 5e67b8b7..c7d63aa1 100644 --- a/soliket/bias/bias.py +++ b/soliket/bias/bias.py @@ -26,7 +26,7 @@ function (have a look at the linear bias model for ideas). """ -from typing import Dict, List, Optional, Set, Tuple, Union +from typing import Any, Dict, List, Optional, Set, Tuple, Union import numpy as np from cobaya.theory import Theory @@ -47,13 +47,13 @@ class Bias(Theory): _default_z_sampling = 10 ** _logz _default_z_sampling[0] = 0 - def initialize(self) -> None: + def initialize(self): self._var_pairs: Set[Tuple[str, str]] = set() - def get_requirements(self) -> Dict[str, dict]: + def get_requirements(self) -> Dict[str, Any]: return {} - def must_provide(self, **requirements: dict) -> Dict[str, dict]: + def must_provide(self, **requirements) -> Dict[str, Any]: options = requirements.get("linear_bias") or {} self.kmax = max(self.kmax, options.get("kmax", self.kmax)) @@ -62,7 +62,7 @@ def must_provide(self, **requirements: dict) -> Dict[str, dict]: np.atleast_1d(self.z)))) # Dictionary of the things needed from CAMB/CLASS - needs: Dict[str, dict] = {} + needs = {} self.nonlinear = self.nonlinear or options.get("nonlinear", False) self._var_pairs.update( @@ -103,7 +103,7 @@ class Linear_bias(Bias): params: dict def calculate(self, state: dict, want_derived: bool = True, - **params_values_dict) -> None: + **params_values_dict): Pk_mm = self._get_Pk_mm() state["Pk_gg_grid"] = params_values_dict["b_lin"] ** 2. * Pk_mm diff --git a/soliket/cash/cash.py b/soliket/cash/cash.py index 4f8d571a..64b41c8c 100644 --- a/soliket/cash/cash.py +++ b/soliket/cash/cash.py @@ -13,7 +13,7 @@ class CashCLikelihood(Likelihood): enforce_types: bool = True - def initialize(self) -> None: + def initialize(self): x, N = self._get_data() self.data = CashCData(self.name, N) @@ -23,12 +23,12 @@ def _get_data(self) -> Tuple[np.ndarray, np.ndarray]: x = data[:, :-1] return x, N - def _get_theory(self, **kwargs: dict) -> np.ndarray: + def _get_theory(self, **kwargs) -> np.ndarray: if "cash_test_logp" in kwargs: return np.arange(kwargs["cash_test_logp"]) else: raise NotImplementedError - def logp(self, **params_values: dict) -> float: + def logp(self, **params_values) -> float: theory = self._get_theory(**params_values) return self.data.loglike(theory) diff --git a/soliket/cash/cash_data.py b/soliket/cash/cash_data.py index 55146b1a..9dae8e44 100644 --- a/soliket/cash/cash_data.py +++ b/soliket/cash/cash_data.py @@ -31,7 +31,7 @@ class CashCData: def __init__( self, name: str, N: Union[np.ndarray, float], usestirling: bool = True - ) -> None: + ): self.name = str(name) self.data = N self.usestirling = usestirling diff --git a/soliket/ccl/ccl.py b/soliket/ccl/ccl.py index 2c8c60e4..1b26254b 100644 --- a/soliket/ccl/ccl.py +++ b/soliket/ccl/ccl.py @@ -98,7 +98,7 @@ class CCL(Theory): _default_z_sampling[0] = 0 provider: Provider - def initialize(self) -> None: + def initialize(self): try: import pyccl as ccl except ImportError: @@ -161,7 +161,7 @@ def get_can_support_params(self) -> Sequence[str]: return [] def calculate(self, state: dict, want_derived: bool = True, - **params_values_dict) -> None: + **params_values_dict): # 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 diff --git a/soliket/cosmopower/cosmopower.py b/soliket/cosmopower/cosmopower.py index 33255a1b..bb569f54 100644 --- a/soliket/cosmopower/cosmopower.py +++ b/soliket/cosmopower/cosmopower.py @@ -91,7 +91,7 @@ information on how SOLikeT infers these values. """ import os -from typing import Dict, Iterable, List, Tuple +from typing import Any, Dict, List import numpy as np from cobaya.log import LoggedError @@ -112,7 +112,7 @@ class CosmoPower(BoltzmannBase): enforce_types: bool = True - def initialize(self) -> None: + def initialize(self): super().initialize() @@ -283,10 +283,10 @@ def cmb_unit_factor(self, spectra: str, return res - def get_can_support_parameters(self) -> Iterable[str]: + def get_can_support_parameters(self) -> List[str]: return self.all_parameters - def get_requirements(self) -> Iterable[Tuple[str, str]]: + def get_requirements(self) -> Dict[str, Any]: requirements = [] for k in self.all_parameters: if k in self.renames.values(): @@ -303,7 +303,7 @@ def get_requirements(self) -> Iterable[Tuple[str, str]]: class CosmoPowerDerived(Theory): """A theory class that can calculate derived parameters from CosmoPower networks.""" - def initialize(self) -> None: + def initialize(self): super().initialize() if self.network_settings is None: @@ -365,10 +365,10 @@ def calculate(self, state: dict, want_derived: bool = True, **params) -> bool: def get_param(self, p) -> float: return self.current_state["derived"][self.translate_param(p)] - def get_can_support_parameters(self) -> Iterable[str]: + def get_can_support_parameters(self) -> List[str]: return self.input_parameters - def get_requirements(self) -> Iterable[Tuple[str, str]]: + def get_requirements(self) -> Dict[str, Any]: requirements = [] for k in self.input_parameters: if k in self.renames.values(): @@ -381,6 +381,6 @@ def get_requirements(self) -> Iterable[Tuple[str, str]]: return requirements - def get_can_provide(self) -> Iterable[str]: + def get_can_provide(self) -> List[str]: return set([par for par in self.derived_parameters if (len(par) > 0 and not par == "_")]) diff --git a/soliket/cross_correlation/cross_correlation.py b/soliket/cross_correlation/cross_correlation.py index d40b94c0..d671ee4d 100644 --- a/soliket/cross_correlation/cross_correlation.py +++ b/soliket/cross_correlation/cross_correlation.py @@ -36,14 +36,14 @@ def initialize(self): self._get_sacc_data() self._check_tracers() - def get_requirements(self) -> Dict[str, dict]: + def get_requirements(self) -> Dict[str, Any]: return {"CCL": {"kmax": 10, "nonlinear": True}, "zstar": None} def _get_CCL_results(self) -> Tuple[CCL, dict]: cosmo_dict = self.provider.get_CCL() return cosmo_dict["ccl"], cosmo_dict["cosmo"] - def _check_tracers(self) -> None: + def _check_tracers(self): # check correct tracers for tracer_comb in self.sacc_data.get_tracer_combinations(): @@ -72,7 +72,7 @@ def _get_nz( z: np.ndarray, tracer: Any, tracer_name: str, - **params_values: dict + **params_values ) -> np.ndarray: if self.z_nuisance_mode == 'deltaz': bias = params_values[f'{tracer_name}_deltaz'] @@ -82,7 +82,7 @@ def _get_nz( return nz_biased - def _get_sacc_data(self, **params_values: dict) -> None: + def _get_sacc_data(self, **params_values): self.sacc_data = sacc.Sacc.load_fits(self.datapath) if self.use_spectra == 'all': @@ -109,7 +109,7 @@ def _construct_ell_bins(self) -> np.ndarray: return np.concatenate(ell_eff) def _get_data( - self, **params_values: dict + self, **params_values ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: data_auto = np.loadtxt(self.auto_file) data_cross = np.loadtxt(self.cross_file) @@ -138,7 +138,7 @@ def get_binning(self, tracer_comb: tuple) -> Tuple[np.ndarray, np.ndarray]: return ells_theory, w_bins - def logp(self, **params_values: dict) -> float: + def logp(self, **params_values) -> float: theory = self._get_theory(**params_values) return self.data.loglike(theory) @@ -150,7 +150,7 @@ class GalaxyKappaLikelihood(CrossCorrelationLikelihood): _allowable_tracers: ClassVar[List[str]] = ['cmb_convergence', 'galaxy_density'] params: dict - def _get_theory(self, **params_values: dict) -> np.ndarray: + def _get_theory(self, **params_values) -> np.ndarray: ccl, cosmo = self._get_CCL_results() tracer_comb = self.sacc_data.get_tracer_combinations() @@ -194,7 +194,7 @@ class ShearKappaLikelihood(CrossCorrelationLikelihood): ia_mode: Optional[str] params: dict - def _get_theory(self, **params_values: dict) -> np.ndarray: + def _get_theory(self, **params_values) -> np.ndarray: ccl, cosmo = self._get_CCL_results() cl_binned_list: List[np.ndarray] = [] diff --git a/soliket/foreground/foreground.py b/soliket/foreground/foreground.py index ab7c0db3..2b8faa97 100644 --- a/soliket/foreground/foreground.py +++ b/soliket/foreground/foreground.py @@ -306,7 +306,7 @@ def _get_foreground_model( fg_dict[s, "all", f1, f2] += fg_dict[s, comp, f1, f2] return fg_dict - def must_provide(self, **requirements: dict) -> dict: + def must_provide(self, **requirements) -> dict: # fg_dict is required by theoryforge # and requires some params to be computed # Assign those from theoryforge @@ -323,7 +323,7 @@ def must_provide(self, **requirements: dict) -> dict: return {"bandint_freqs": {"bands": self.bands}} return {} - def get_bandpasses(self, **params: dict) -> np.ndarray: + def get_bandpasses(self, **params) -> np.ndarray: """ Gets bandpass transmissions from the ``BandPass`` class. """ diff --git a/soliket/gaussian/gaussian.py b/soliket/gaussian/gaussian.py index 2b980c32..73c963af 100644 --- a/soliket/gaussian/gaussian.py +++ b/soliket/gaussian/gaussian.py @@ -21,7 +21,7 @@ class GaussianLikelihood(Likelihood): enforce_types: bool = True - def initialize(self) -> None: + def initialize(self): x, y = self._get_data() cov = self._get_cov() self.data = GaussianData(self.name, x, y, cov, self.ncovsims) @@ -34,16 +34,16 @@ def _get_cov(self) -> np.ndarray: cov = np.loadtxt(self.covpath) return cov - def _get_theory(self, **kwargs: dict) -> np.ndarray: + def _get_theory(self, **kwargs) -> np.ndarray: raise NotImplementedError - def logp(self, **params_values: dict) -> float: + def logp(self, **params_values) -> float: theory = self._get_theory(**params_values) return self.data.loglike(theory) class CrossCov(dict): - def save(self, path: str) -> None: + def save(self, path: str): np.savez(path, **{str(k): v for k, v in self.items()}) @classmethod @@ -58,7 +58,7 @@ class MultiGaussianLikelihood(GaussianLikelihood): options: Optional[Sequence] = None cross_cov_path: Optional[str] = None - def __init__(self, info: dict = empty_dict, **kwargs) -> None: + def __init__(self, info: dict = empty_dict, **kwargs): if 'components' in info: self.likelihoods: List[Likelihood] = [ @@ -72,7 +72,7 @@ def __init__(self, info: dict = empty_dict, **kwargs) -> None: super().__init__(info=default_info, **kwargs) - def initialize(self) -> None: + def initialize(self): self.cross_cov: Optional[CrossCov] = CrossCov.load(self.cross_cov_path) data_list = [like.data for like in self.likelihoods] @@ -80,7 +80,7 @@ def initialize(self) -> None: self.log.info('Initialized.') - def initialize_with_provider(self, provider: Provider) -> None: # pragma: no cover + def initialize_with_provider(self, provider: Provider): # pragma: no cover for like in self.likelihoods: like.initialize_with_provider(provider) # super().initialize_with_provider(provider) diff --git a/soliket/gaussian/gaussian_data.py b/soliket/gaussian/gaussian_data.py index 6fe74e73..4bf1011b 100644 --- a/soliket/gaussian/gaussian_data.py +++ b/soliket/gaussian/gaussian_data.py @@ -17,7 +17,7 @@ class GaussianData: _fast_chi_squared = staticmethod(functions.chi_squared) def __init__(self, name: str, x: Sequence[float], y: Sequence[float], cov: np.ndarray, - ncovsims: Optional[int] = None) -> None: + ncovsims: Optional[int] = None): self.name: str = str(name) self.ncovsims: Optional[int] = ncovsims @@ -64,7 +64,7 @@ def __init__( self, data_list: List[GaussianData], cross_covs: Optional[Dict[Tuple[str, str], np.ndarray]] = None, - ) -> None: + ): if cross_covs is None: cross_covs = {} @@ -148,7 +148,7 @@ def _slice(self, *names: str) -> slice: return np.s_[tuple(slice(*self._index_range(n)) for n in names)] - def _assemble_data(self) -> None: + def _assemble_data(self): x = np.concatenate([d.x for d in self.data_list]) y = np.concatenate([d.y for d in self.data_list]) @@ -161,7 +161,7 @@ def _assemble_data(self) -> None: self._data = GaussianData(" + ".join(self.names), x, y, cov) - def plot_cov(self, **kwargs) -> None: + def plot_cov(self, **kwargs): import holoviews as hv data = [ diff --git a/soliket/halo_model/halo_model.py b/soliket/halo_model/halo_model.py index d8cb8c9f..97c49482 100644 --- a/soliket/halo_model/halo_model.py +++ b/soliket/halo_model/halo_model.py @@ -32,7 +32,7 @@ function (have a look at the simple pyhalomodel model for ideas). """ -from typing import Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union import numpy as np import pyhalomodel as halo from cobaya.theory import Provider, Theory @@ -54,7 +54,7 @@ class HaloModel(Theory): _default_z_sampling[0] = 0 provider: Provider - def initialize(self) -> None: + def initialize(self): self._var_pairs = set() self._required_results = {} @@ -92,11 +92,11 @@ class HaloModel_pyhm(HaloModel): Mmax: float nM: int - def initialize(self) -> None: + def initialize(self): super().initialize() self.Ms = np.logspace(np.log10(self.Mmin), np.log10(self.Mmax), self.nM) - def get_requirements(self) -> Dict[str, Optional[None]]: + def get_requirements(self) -> Dict[str, Any]: return {"omegam": None} def must_provide(self, **requirements) -> dict: @@ -128,7 +128,7 @@ def must_provide(self, **requirements) -> dict: return needs def calculate(self, state: dict, want_derived: bool = True, - **params_values_dict) -> None: + **params_values_dict): Pk_mm_lin: np.ndarray = self._get_Pk_mm_lin() diff --git a/soliket/lensing/lensing.py b/soliket/lensing/lensing.py index c1e77768..f9264aeb 100644 --- a/soliket/lensing/lensing.py +++ b/soliket/lensing/lensing.py @@ -211,7 +211,7 @@ def _get_binning_matrix(self) -> np.ndarray: self.binning_matrix = binning_matrix return binning_matrix - def _get_theory(self, **params_values: dict) -> np.ndarray: + def _get_theory(self, **params_values) -> np.ndarray: r""" Generate binned theory vector of :math:`\kappa \kappa` with correction terms. diff --git a/soliket/mflike/mflike.py b/soliket/mflike/mflike.py index 9c956d68..49f41056 100644 --- a/soliket/mflike/mflike.py +++ b/soliket/mflike/mflike.py @@ -112,11 +112,11 @@ def get_requirements(self) -> dict: "bands": self.bands} return reqs - def _get_theory(self, **params_values: dict) -> np.ndarray: + def _get_theory(self, **params_values) -> np.ndarray: cmbfg_dict = self.provider.get_cmbfg_dict() return self._get_power_spectra(cmbfg_dict) - def logp(self, **params_values: dict) -> float: + def logp(self, **params_values) -> float: cmbfg_dict = self.provider.get_cmbfg_dict() return self.loglike(cmbfg_dict) diff --git a/soliket/mflike/theoryforge_MFLike.py b/soliket/mflike/theoryforge_MFLike.py index 0de83453..fe7e3900 100644 --- a/soliket/mflike/theoryforge_MFLike.py +++ b/soliket/mflike/theoryforge_MFLike.py @@ -132,7 +132,7 @@ def initialize_with_params(self): self.log, "Configuration error in parameters: %r.", differences) - def must_provide(self, **requirements: dict) -> dict: + def must_provide(self, **requirements) -> dict: # cmbfg_dict is required by mflike # and requires some params to be computed # Assign required params from mflike @@ -159,10 +159,10 @@ def must_provide(self, **requirements: dict) -> dict: "exp_ch": self.exp_ch, "bands": self.bands} return reqs - def get_cmb_theory(self, **params: dict) -> dict: + def get_cmb_theory(self, **params) -> dict: return self.provider.get_Cl(ell_factor=True) - def get_foreground_theory(self, **params: dict) -> dict: + def get_foreground_theory(self, **params) -> dict: return self.provider.get_fg_dict() def calculate(self, state, want_derived=False, **params_values_dict): @@ -178,7 +178,7 @@ def calculate(self, state, want_derived=False, **params_values_dict): def get_cmbfg_dict(self) -> dict: return self.current_state["cmbfg_dict"] - def get_modified_theory(self, Dls: dict, fg_dict: dict, **params: dict) -> dict: + def get_modified_theory(self, Dls: dict, fg_dict: dict, **params) -> dict: r""" Takes the theory :math:`D_{\ell}`, sums it to the total foreground power spectrum (possibly computed with bandpass diff --git a/soliket/poisson/poisson.py b/soliket/poisson/poisson.py index 3ffb11f7..997f5a4b 100644 --- a/soliket/poisson/poisson.py +++ b/soliket/poisson/poisson.py @@ -12,7 +12,7 @@ class PoissonLikelihood(Likelihood): enforce_types: bool = True - def initialize(self) -> None: + def initialize(self): catalog = self._get_catalog() if self.columns is None: self.columns = catalog.columns @@ -25,17 +25,17 @@ def _get_catalog(self) -> pd.DataFrame: catalog = pd.read_csv(self.data_path) return catalog - def _get_rate_fn(self, **kwargs: dict) -> Callable: + def _get_rate_fn(self, **kwargs) -> Callable: """Returns a callable rate function that takes each of 'columns' as kwargs. """ raise NotImplementedError - def _get_n_expected(self, **kwargs: dict) -> float: + def _get_n_expected(self, **kwargs) -> float: """Computes and returns the integral of the rate function """ raise NotImplementedError - def logp(self, **params_values: dict) -> float: + def logp(self, **params_values) -> float: rate_fn = self._get_rate_fn(**params_values) n_expected = self._get_n_expected(**params_values) return self.data.loglike(rate_fn, n_expected) diff --git a/soliket/poisson/poisson_data.py b/soliket/poisson/poisson_data.py index 676b7c38..97c08bfe 100644 --- a/soliket/poisson/poisson_data.py +++ b/soliket/poisson/poisson_data.py @@ -10,7 +10,7 @@ def __init__( catalog: pd.DataFrame, columns: List[str], samples: Optional[Dict[str, np.ndarray]] = None, - ) -> None: + ): self.name = str(name) self.catalog = pd.DataFrame(catalog)[columns] diff --git a/soliket/ps/ps.py b/soliket/ps/ps.py index a6b38d19..41e885a6 100644 --- a/soliket/ps/ps.py +++ b/soliket/ps/ps.py @@ -1,4 +1,4 @@ -from typing import Dict, Tuple +from typing import Any, Dict, Tuple from cobaya.theory import Provider import numpy as np @@ -12,13 +12,13 @@ class PSLikelihood(GaussianLikelihood): lmax: int = 6000 provider: Provider - def get_requirements(self) -> Dict[str, Dict[str, int]]: + def get_requirements(self) -> Dict[str, Dict[str, Any]]: return {"Cl": {self.kind: self.lmax}} def _get_Cl(self) -> Dict[str, np.ndarray]: return self.provider.get_Cl(ell_factor=True) - def _get_theory(self, **params_values: dict) -> np.ndarray: + def _get_theory(self, **params_values) -> np.ndarray: cl_theory = self._get_Cl() return cl_theory[self.kind][:self.lmax] @@ -26,7 +26,7 @@ def _get_theory(self, **params_values: dict) -> np.ndarray: class BinnedPSLikelihood(PSLikelihood): binning_matrix_path: str = "" - def initialize(self) -> None: + def initialize(self): self.binning_matrix = self._get_binning_matrix() self.bin_centers = self.binning_matrix.dot( np.arange(self.binning_matrix.shape[1]) @@ -45,6 +45,6 @@ def _get_binning_matrix(self) -> np.ndarray: def _get_data(self) -> Tuple[np.ndarray, np.ndarray]: return self.bin_centers, np.loadtxt(self.datapath) - def _get_theory(self, **params_values: dict) -> np.ndarray: + def _get_theory(self, **params_values) -> np.ndarray: cl_theory: Dict[str, np.ndarray] = self._get_Cl() return self.binning_matrix.dot(cl_theory[self.kind][:self.lmax]) diff --git a/soliket/xcorr/xcorr.py b/soliket/xcorr/xcorr.py index 7ef1788d..8e2ac094 100644 --- a/soliket/xcorr/xcorr.py +++ b/soliket/xcorr/xcorr.py @@ -167,7 +167,7 @@ def _bin( & (self.ell_range < lmax[i])]) return binned_theory_cl - def _get_sacc_data(self, **params_values: dict) -> dict: + def _get_sacc_data(self, **params_values) -> dict: data_sacc = sacc.Sacc.load_fits(self.datapath) # TODO: would be better to use keep_selection @@ -198,7 +198,7 @@ def _get_sacc_data(self, **params_values: dict) -> dict: def _get_data( - self, **params_values: dict + self, **params_values ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: data_auto = np.loadtxt(self.auto_file) @@ -245,7 +245,7 @@ def _setup_chi(self) -> dict: return chi_result - def _get_theory(self, **params_values: dict) -> np.ndarray: + def _get_theory(self, **params_values) -> np.ndarray: setup_chi_out = self._setup_chi()