From 0de451791e7ae2da85ff86465550b8fd721328d1 Mon Sep 17 00:00:00 2001 From: Ian Harrison Date: Fri, 4 Nov 2022 15:13:11 +0000 Subject: [PATCH 1/3] added filenames --- soliket/cosmopower.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/soliket/cosmopower.py b/soliket/cosmopower.py index 29353360..531ff3f5 100755 --- a/soliket/cosmopower.py +++ b/soliket/cosmopower.py @@ -19,9 +19,11 @@ class CosmoPower(BoltzmannBase): soliket_data_path: str = "soliket/data/CosmoPower" network_path: str = "CP_paper/CMB" + network_path_pk: str = "CP_paper/PK" cmb_tt_nn_filename: str = "cmb_TT_NN" cmb_te_pcaplusnn_filename: str = "cmb_TE_PCAplusNN" cmb_ee_nn_filename: str = "cmb_EE_NN" + pk_lin_nn_filename: str = "PKLIN_NN" extra_args: InfoDict = {} From 0b2ef6041c8c7e7b3e554cf0309c2cd65bdd1230 Mon Sep 17 00:00:00 2001 From: Ian Harrison Date: Fri, 4 Nov 2022 17:09:34 +0000 Subject: [PATCH 2/3] added Pk methods and provides --- soliket/cosmopower.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/soliket/cosmopower.py b/soliket/cosmopower.py index 531ff3f5..8f9059cf 100755 --- a/soliket/cosmopower.py +++ b/soliket/cosmopower.py @@ -76,6 +76,10 @@ def calculate(self, state, want_derived=True, **params): state["ee"] = self.cp_ee_nn.ten_to_predictions_np(cmb_params)[0, :] state["ell"] = self.cp_tt_nn.modes + state["Pk_grid"] = self.get_Pk_grid() + state["k"] = self.cp_pk_nn.modes + state["z"] = self.z + def get_Cl(self, ell_factor=False, units="FIRASmuK2"): cls_old = self.current_state.copy() @@ -104,5 +108,44 @@ def get_Cl(self, ell_factor=False, units="FIRASmuK2"): return cls + def must_provide(self, **requirements): + + if 'Pk_interpolator' not in requirements and 'Pk_grid' not in requirements: + return {} + + self.kmax = max(self.k_list) + self.z = np.unique(np.concatenate( + (np.atleast_1d(options.get("z", self._default_z_sampling)), + np.atleast_1d(self.z)))) + + self.nonlinear = self.nonlinear or options.get('nonlinear', False) + + self._var_pairs.update(set((x, y) for x, y in + options.get('vars_pairs', [('delta_tot', 'delta_tot')]))) + + needs['Pk_grid'] = { + 'vars_pairs': self._var_pairs or [('delta_tot', 'delta_tot')], + 'nonlinear': (True, False) if self.nonlinear else False, + } + + return needs + + + def get_Pk_grid(self, var_pair=("delta_tot", "delta_tot"), nonlinear=True, + extrap_kmin=None, extrap_kmax=None): + + if var_pair != ("delta_tot", "delta_tot") or nonlinear + raise LoggedError(self.log, + 'COSMOPOWER P(k, z) only trained for linear delta_tot pk') + + if self.kmax > max(self.cp_pk_nn.modes): + raise LoggedError(self.log, + 'COSMOPOWER P(k, z) only trained up to {}'.format(max(self.cp_pk_nn.modes)) + 'but you have requested kmax {}.'.format(self.kmax)) + + k = self.cp_pk_nn.modes + pk = self.cp_pk_nn.predictions_np(params) + + def get_can_support_params(self): return ["omega_b", "omega_cdm", "h", "logA", "ns", "tau_reio"] From 05642be38afa47b21a112030e33831618dff0aca Mon Sep 17 00:00:00 2001 From: Ian Harrison Date: Tue, 29 Nov 2022 16:44:03 +0000 Subject: [PATCH 3/3] added pk_grid saved to state --- soliket/cosmopower.py | 110 ++++++++++++++++++------------- soliket/tests/test_cosmopower.py | 56 ++++++++++++++++ 2 files changed, 121 insertions(+), 45 deletions(-) diff --git a/soliket/cosmopower.py b/soliket/cosmopower.py index 8f9059cf..fe25c4e4 100755 --- a/soliket/cosmopower.py +++ b/soliket/cosmopower.py @@ -9,6 +9,7 @@ from cobaya.theories.cosmo import BoltzmannBase from cobaya.typing import InfoDict +from cobaya.log import LoggedError """ Simple CosmoPower theory wrapper for Cobaya. @@ -24,6 +25,8 @@ class CosmoPower(BoltzmannBase): cmb_te_pcaplusnn_filename: str = "cmb_TE_PCAplusNN" cmb_ee_nn_filename: str = "cmb_EE_NN" pk_lin_nn_filename: str = "PKLIN_NN" + k_max: int = 1. + z: float = np.linspace(0.0, 2, 128) extra_args: InfoDict = {} @@ -53,6 +56,10 @@ def initialize(self): restore=True, restore_filename=os.path.join(base_path, self.cmb_ee_nn_filename), ) + self.cp_pk_nn = cp.cosmopower_NN( + restore=True, + restore_filename=os.path.join(self.soliket_data_path, self.network_path_pk, self.pk_lin_nn_filename), + ) if "lmax" not in self.extra_args: self.extra_args["lmax"] = None @@ -60,25 +67,34 @@ def initialize(self): self.log.info(f"Loaded CosmoPower from directory {self.network_path}") def calculate(self, state, want_derived=True, **params): - cmb_params = {} + network_params = {} for par in self.renames: if par in params: - cmb_params[par] = [params[par]] + network_params[par] = [params[par]] else: for r in self.renames[par]: if r in params: - cmb_params[par] = [params[r]] + network_params[par] = [params[r]] break - state["tt"] = self.cp_tt_nn.ten_to_predictions_np(cmb_params)[0, :] - state["te"] = self.cp_te_nn.predictions_np(cmb_params)[0, :] - state["ee"] = self.cp_ee_nn.ten_to_predictions_np(cmb_params)[0, :] + state["tt"] = self.cp_tt_nn.ten_to_predictions_np(network_params)[0, :] + state["te"] = self.cp_te_nn.predictions_np(network_params)[0, :] + state["ee"] = self.cp_ee_nn.ten_to_predictions_np(network_params)[0, :] state["ell"] = self.cp_tt_nn.modes - state["Pk_grid"] = self.get_Pk_grid() - state["k"] = self.cp_pk_nn.modes - state["z"] = self.z + # if "Pk_grid" in self.requested() or "Pk_interpolator" in self.requested(): + + network_params_z = {} + + for k in network_params.keys(): + network_params_z[k] = np.repeat(network_params[k], len(self.z)) + network_params_z['z'] = self.z + + var_pair=("delta_tot", "delta_tot") + nonlinear=True + + state[("Pk_grid", bool(nonlinear)) + tuple(sorted(var_pair))] = self.cp_pk_nn.modes, self.z, self.cp_pk_nn.predictions_np(network_params_z) def get_Cl(self, ell_factor=False, units="FIRASmuK2"): cls_old = self.current_state.copy() @@ -108,44 +124,48 @@ def get_Cl(self, ell_factor=False, units="FIRASmuK2"): return cls - def must_provide(self, **requirements): - - if 'Pk_interpolator' not in requirements and 'Pk_grid' not in requirements: - return {} - - self.kmax = max(self.k_list) - self.z = np.unique(np.concatenate( - (np.atleast_1d(options.get("z", self._default_z_sampling)), - np.atleast_1d(self.z)))) - - self.nonlinear = self.nonlinear or options.get('nonlinear', False) - - self._var_pairs.update(set((x, y) for x, y in - options.get('vars_pairs', [('delta_tot', 'delta_tot')]))) - - needs['Pk_grid'] = { - 'vars_pairs': self._var_pairs or [('delta_tot', 'delta_tot')], - 'nonlinear': (True, False) if self.nonlinear else False, - } - - return needs - - - def get_Pk_grid(self, var_pair=("delta_tot", "delta_tot"), nonlinear=True, - extrap_kmin=None, extrap_kmax=None): + # def must_provide(self, **requirements): + + # super().must_provide(**requirements) + + # for k, v in self._must_provide.items(): + # if isinstance(k, tuple) and k[0] == "Pk_grid": + # v = deepcopy(v) + # self.add_P_k_max(v.pop("k_max"), units="1/Mpc") + # # NB: Actually, only the max z is used, and the actual sampling in z + # # for computing P(k,z) is controlled by `perturb_sampling_stepsize` + # # (default: 0.1). But let's leave it like this in case this changes + # # in the future. + # self.add_z_for_matter_power(v.pop("z")) + # if v["nonlinear"]: + # if "non_linear" not in self.extra_args: + # # this is redundant with initialisation, but just in case + # self.extra_args["non_linear"] = non_linear_default_code + # elif self.extra_args["non_linear"] == non_linear_null_value: + # raise LoggedError( + # self.log, ("Non-linear Pk requested, but `non_linear: " + # f"{non_linear_null_value}` imposed in " + # "`extra_args`")) + # pair = k[2:] + # if pair == ("delta_tot", "delta_tot"): + # self.collectors[k] = Collector( + # method="get_pk_and_k_and_z", + # kwargs=v, + # post=(lambda P, kk, z: (kk, z, np.array(P).T))) + # else: + # raise LoggedError(self.log, "NotImplemented in cosmopower: %r", pair) + + # return needs + + + # def get_Pk_grid(self, params, var_pair=("delta_tot", "delta_tot"), nonlinear=False, + # extrap_kmin=None, extrap_kmax=None): - if var_pair != ("delta_tot", "delta_tot") or nonlinear - raise LoggedError(self.log, - 'COSMOPOWER P(k, z) only trained for linear delta_tot pk') - - if self.kmax > max(self.cp_pk_nn.modes): - raise LoggedError(self.log, - 'COSMOPOWER P(k, z) only trained up to {}'.format(max(self.cp_pk_nn.modes)) - 'but you have requested kmax {}.'.format(self.kmax)) - - k = self.cp_pk_nn.modes - pk = self.cp_pk_nn.predictions_np(params) + # return self.current_state['k'], self.current_state['pk_z'], self.current_state['pk'] def get_can_support_params(self): return ["omega_b", "omega_cdm", "h", "logA", "ns", "tau_reio"] + + def get_can_provide(self): + return ['Cl', 'Pk_grid', 'Pk_interpolator'] diff --git a/soliket/tests/test_cosmopower.py b/soliket/tests/test_cosmopower.py index e0259053..5702d007 100644 --- a/soliket/tests/test_cosmopower.py +++ b/soliket/tests/test_cosmopower.py @@ -22,6 +22,30 @@ "H0": {"value": "lambda h: h * 100.0"}, } +des_params = { + 'DES_DzL1': 0.001, + 'DES_DzL2': 0.002, + 'DES_DzL3': 0.001, + 'DES_DzL4': 0.003, + 'DES_DzL5': 0, + 'DES_b1': 1.45, + 'DES_b2': 1.55, + 'DES_b3': 1.65, + 'DES_b4': 1.8, + 'DES_b5': 2.0, + 'DES_DzS1': -0.001, + 'DES_DzS2': -0.019, + 'DES_DzS3': 0.009, + 'DES_DzS4': -0.018, + 'DES_m1': 0.012, + 'DES_m2': 0.012, + 'DES_m3': 0.012, + 'DES_m4': 0.012, + 'DES_AIA': 1, + 'DES_alphaIA': 1, + 'DES_z0IA': 0.62, +} + info_dict = { "params": fiducial_params, "likelihood": { @@ -39,6 +63,27 @@ }, } +def lss_part_likelihood(_self): + results = _self.provider.get_Pk_interpolator().P(0.1, 1.0) + # results = _self.provider.get_Pk_grid() + return 1 + + +info_dict_pk = { + "params": fiducial_params, + "likelihood": {'lss': {'external': lss_part_likelihood, 'requires': {'Pk_interpolator': {'z' : np.linspace(0.0, 2, 128), + 'k_max': 1.}}} + }, + "theory": { + "soliket.CosmoPower": { + "soliket_data_path": "soliket/data/CosmoPower", + "stop_at_error": True, + "provides": 'Pk_grid', + } + # 'camb': {"stop_at_error": True,} + } + } + @pytest.mark.skipif(not HAS_COSMOPOWER, reason='test requires cosmopower') def test_cosmopower_theory(): @@ -77,3 +122,14 @@ def test_cosmopower_against_camb(): assert np.allclose(cp_cls['tt'][nanmask], camb_cls['tt'][nanmask], rtol=1.e-2) assert np.isclose(logL_camb, logL_cp, rtol=1.e-1) + + +@pytest.mark.skipif(not HAS_COSMOPOWER, reason='test requires cosmopower') +def test_cosmopower_pkgrid(): + + model_cp = get_model(info_dict_pk) + + logL_cp = float(model_cp.loglikes({})[0]) + + assert np.isfinite(logL_cp) +