From e49ab74a1c1d5bbeb01a41bf943962a3fdbd709f Mon Sep 17 00:00:00 2001 From: Tilman Troester Date: Wed, 21 Feb 2024 17:50:17 +0100 Subject: [PATCH 1/5] add to_dict method --- pyccl/cosmology.py | 26 +++++++++++++++++++------- pyccl/tests/test_yaml.py | 22 +++++++++++++++++++++- 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/pyccl/cosmology.py b/pyccl/cosmology.py index 63eaee46b..1a49678bc 100644 --- a/pyccl/cosmology.py +++ b/pyccl/cosmology.py @@ -235,13 +235,11 @@ def __init__( self.lin_pk_emu = None if isinstance(transfer_function, emulators.EmulatorPk): self.lin_pk_emu = transfer_function - transfer_function = 'emulator' # initialise nonlinear Pk emulators if needed self.nl_pk_emu = None if isinstance(matter_power_spectrum, emulators.EmulatorPk): self.nl_pk_emu = matter_power_spectrum - matter_power_spectrum = 'emulator' self.baryons = baryonic_effects if not isinstance(self.baryons, baryons.Baryons): @@ -277,6 +275,8 @@ def __init__( self._config_init_kwargs = dict( transfer_function=transfer_function, matter_power_spectrum=matter_power_spectrum, + baryonic_effects=baryonic_effects, + mg_parametrization=mg_parametrization, extra_parameters=extra_parameters) self._build_cosmo() @@ -299,6 +299,12 @@ def _build_cosmo(self): if self.cosmo.status != 0: raise CCLError(f"{self.cosmo.status}: {self.cosmo.status_message}") + def to_dict(self): + """Returns a dictionary of the arguments used to create the Cosmology + object such that ``cosmo == pyccl.Cosmology(**cosmo.to_dict())`` + is ``True``.""" + return {**self._params_init_kwargs, **self._config_init_kwargs} + def write_yaml(self, filename, *, sort_keys=False): """Write a YAML representation of the parameters to file. @@ -316,7 +322,7 @@ def make_yaml_friendly(d): elif isinstance(v, dict): make_yaml_friendly(v) - params = {**self._params_init_kwargs, **self._config_init_kwargs} + params = self.to_dict() make_yaml_friendly(params) if isinstance(filename, str): @@ -337,12 +343,14 @@ def read_yaml(cls, filename, **kwargs): loader = yaml.Loader if isinstance(filename, str): with open(filename, 'r') as fp: - return cls(**{**yaml.load(fp, Loader=loader), **kwargs}) - return cls(**{**yaml.load(filename, Loader=loader), **kwargs}) + params = yaml.load(fp, Loader=loader) + else: + params = yaml.load(filename, Loader=loader) + return cls(**{**params, **kwargs}) def _build_config( - self, transfer_function=None, matter_power_spectrum=None, - extra_parameters=None): + self, *, transfer_function=None, matter_power_spectrum=None, + **kwargs): """Build a ccl_configuration struct. This function builds C ccl_configuration struct. This structure @@ -353,6 +361,10 @@ def _build_config( It also does some error checking on the inputs to make sure they are valid and physically consistent. """ + if isinstance(transfer_function, emulators.EmulatorPk): + transfer_function = 'emulator' + if isinstance(matter_power_spectrum, emulators.EmulatorPk): + matter_power_spectrum = 'emulator' if (matter_power_spectrum == "camb" and transfer_function != "boltzmann_camb"): raise CCLError( diff --git a/pyccl/tests/test_yaml.py b/pyccl/tests/test_yaml.py index 1e4af5d3e..18944cc88 100644 --- a/pyccl/tests/test_yaml.py +++ b/pyccl/tests/test_yaml.py @@ -7,7 +7,9 @@ def test_yaml(): cosmo = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.7, A_s=2.1e-9, n_s=0.97, m_nu=[0.01, 0.2, 0.3], - transfer_function="boltzmann_camb") + transfer_function="boltzmann_camb", + baryonic_effects=ccl.baryons.BaryonsvanDaalen19() + ) # Make temporary files with tempfile.NamedTemporaryFile(delete=True) as tmpfile1, \ @@ -19,6 +21,8 @@ def test_yaml(): # Compare the contents of the two files assert filecmp.cmp(tmpfile1.name, tmpfile2.name, shallow=False) + # Compare the two Cosmology objects + assert cosmo == cosmo2 cosmo = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.7, A_s=2.1e-9, n_s=0.97, m_nu=0.1, mass_split="equal", @@ -33,3 +37,19 @@ def test_yaml(): cosmo2.write_yaml(stream2) assert stream.getvalue() == stream2.getvalue() + + +def test_to_dict(): + cosmo = ccl.CosmologyVanillaLCDM( + transfer_function=ccl.emulators.EmulatorPk(), + matter_power_spectrum=ccl.emulators.CosmicemuMTIIPk(), + baryonic_effects=ccl.baryons.BaryonsvanDaalen19(), + mg_parametrization=ccl.modified_gravity.MuSigmaMG() + ) + + assert cosmo == ccl.Cosmology(**cosmo.to_dict()) + + # Check that all arguments to Cosmology are stored + init_params = {k: v for k, v in cosmo.__signature__.parameters.items() + if k != "self"} + assert set(cosmo.to_dict().keys()) == set(init_params.keys()) From 84408b7e0a5a6f358f13428ef8d7e863c359e909 Mon Sep 17 00:00:00 2001 From: Tilman Troester Date: Wed, 21 Feb 2024 23:01:40 +0100 Subject: [PATCH 2/5] fix a bug --- pyccl/cosmology.py | 26 ++++++++++++++++-------- pyccl/tests/test_cosmology_parameters.py | 2 +- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/pyccl/cosmology.py b/pyccl/cosmology.py index 1a49678bc..769fc4430 100644 --- a/pyccl/cosmology.py +++ b/pyccl/cosmology.py @@ -235,11 +235,23 @@ def __init__( self.lin_pk_emu = None if isinstance(transfer_function, emulators.EmulatorPk): self.lin_pk_emu = transfer_function + self.transfer_function_type = "emulator" + elif isinstance(transfer_function, str): + self.transfer_function_type = transfer_function + else: + raise ValueError(f"transfer_function={transfer_function} not " + f"supported.") # initialise nonlinear Pk emulators if needed self.nl_pk_emu = None if isinstance(matter_power_spectrum, emulators.EmulatorPk): self.nl_pk_emu = matter_power_spectrum + self.matter_power_spectrum_type = "emulator" + elif isinstance(matter_power_spectrum, str): + self.matter_power_spectrum_type = matter_power_spectrum + else: + raise ValueError(f"matter_power_spectrum={matter_power_spectrum} " + f"not supported.") self.baryons = baryonic_effects if not isinstance(self.baryons, baryons.Baryons): @@ -361,10 +373,6 @@ def _build_config( It also does some error checking on the inputs to make sure they are valid and physically consistent. """ - if isinstance(transfer_function, emulators.EmulatorPk): - transfer_function = 'emulator' - if isinstance(matter_power_spectrum, emulators.EmulatorPk): - matter_power_spectrum = 'emulator' if (matter_power_spectrum == "camb" and transfer_function != "boltzmann_camb"): raise CCLError( @@ -372,9 +380,9 @@ def _build_config( "the transfer function should be 'boltzmann_camb'.") config = lib.configuration() - tf = transfer_function_types[transfer_function] + tf = transfer_function_types[self.transfer_function_type] config.transfer_function_method = tf - mps = matter_power_spectrum_types[matter_power_spectrum] + mps = matter_power_spectrum_types[self.matter_power_spectrum_type] config.matter_power_spectrum_method = mps # Store ccl_configuration for later access @@ -531,7 +539,7 @@ def _compute_linear_power(self): self.compute_growth() # Populate power spectrum splines - trf = self._config_init_kwargs['transfer_function'] + trf = self.transfer_function_type pk = None rescale_s8 = True rescale_mg = True @@ -561,7 +569,7 @@ def _compute_linear_power(self): # we set the nonlin power spectrum first, but keep the linear via a # status variable to use it later if the transfer function is CAMB too. pkl = None - if self._config_init_kwargs["matter_power_spectrum"] == "camb": + if self.matter_power_spectrum_type == "camb": rescale_mg = False if self.mg_parametrization.mu_0 != 0: raise ValueError("Can't rescale non-linear power spectrum " @@ -595,7 +603,7 @@ def _compute_nonlin_power(self): self.compute_distances() # Populate power spectrum splines - mps = self._config_init_kwargs['matter_power_spectrum'] + mps = self.matter_power_spectrum_type # needed for halofit, and linear options if (mps not in ['emulator']) and (mps is not None): self.compute_linear_power() diff --git a/pyccl/tests/test_cosmology_parameters.py b/pyccl/tests/test_cosmology_parameters.py index bc965f9bd..f24fca32e 100644 --- a/pyccl/tests/test_cosmology_parameters.py +++ b/pyccl/tests/test_cosmology_parameters.py @@ -323,7 +323,7 @@ def test_parameters_read_write(): # check new parameters and config correctly updated assert params3["n_s"] == 1.1 assert params["sum_nu_masses"] == params3["sum_nu_masses"] - assert params3._config_init_kwargs['matter_power_spectrum'] == 'linear' + assert params3.matter_power_spectrum_type == 'linear' # Now make a file that will be deleted so it does not exist # and check the right error is raise From 5e2feb4dd7ae2d5fda56574369b16622d8801508 Mon Sep 17 00:00:00 2001 From: Tilman Troester Date: Wed, 21 Feb 2024 23:21:24 +0100 Subject: [PATCH 3/5] Disallow complex types for yaml --- pyccl/cosmology.py | 6 ++++++ pyccl/tests/test_yaml.py | 11 ++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/pyccl/cosmology.py b/pyccl/cosmology.py index 769fc4430..16f98dd24 100644 --- a/pyccl/cosmology.py +++ b/pyccl/cosmology.py @@ -331,8 +331,14 @@ def make_yaml_friendly(d): d[k] = int(v) elif isinstance(v, float): d[k] = float(v) + elif isinstance(v, tuple): + d[k] = list(v) + elif isinstance(v, np.ndarray): + d[k] = v.tolist() elif isinstance(v, dict): make_yaml_friendly(v) + elif not (isinstance(v, (str, list)) or v is None): + raise ValueError(f"{k}={v} cannot be serialised to YAML.") params = self.to_dict() make_yaml_friendly(params) diff --git a/pyccl/tests/test_yaml.py b/pyccl/tests/test_yaml.py index 18944cc88..052debec6 100644 --- a/pyccl/tests/test_yaml.py +++ b/pyccl/tests/test_yaml.py @@ -1,4 +1,5 @@ import tempfile +import pytest import filecmp import io import pyccl as ccl @@ -8,7 +9,6 @@ def test_yaml(): cosmo = ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.7, A_s=2.1e-9, n_s=0.97, m_nu=[0.01, 0.2, 0.3], transfer_function="boltzmann_camb", - baryonic_effects=ccl.baryons.BaryonsvanDaalen19() ) # Make temporary files @@ -39,6 +39,15 @@ def test_yaml(): assert stream.getvalue() == stream2.getvalue() +def test_write_yaml_complex_types(): + cosmo = ccl.CosmologyVanillaLCDM( + baryonic_effects=ccl.baryons.BaryonsvanDaalen19() + ) + with pytest.raises(ValueError): + with tempfile.NamedTemporaryFile(delete=True) as tmpfile: + cosmo.write_yaml(tmpfile) + + def test_to_dict(): cosmo = ccl.CosmologyVanillaLCDM( transfer_function=ccl.emulators.EmulatorPk(), From bab970aa92813e4b640d688606645e1158edf090 Mon Sep 17 00:00:00 2001 From: Tilman Troester Date: Mon, 26 Feb 2024 13:01:27 +0100 Subject: [PATCH 4/5] more tests for coverage --- pyccl/cosmology.py | 37 ++++++++++++++++++----------------- pyccl/tests/test_cosmology.py | 6 ++++++ pyccl/tests/test_yaml.py | 15 ++++++++++++++ 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/pyccl/cosmology.py b/pyccl/cosmology.py index 16f98dd24..6b1512adb 100644 --- a/pyccl/cosmology.py +++ b/pyccl/cosmology.py @@ -19,6 +19,7 @@ "Cosmology", "CosmologyVanillaLCDM", "CosmologyCalculator",) import yaml +from copy import deepcopy from enum import Enum from inspect import getmembers, isfunction, signature from numbers import Real @@ -113,6 +114,23 @@ class _CosmologyBackgroundData: age0: float = None +def _make_yaml_friendly(d): + """Turn python objects into yaml types where possible.""" + + d = deepcopy(d) + for k, v in d.items(): + if isinstance(v, tuple): + d[k] = list(v) + elif isinstance(v, np.ndarray): + d[k] = v.tolist() + elif isinstance(v, dict): + d[k] = _make_yaml_friendly(v) + elif not (isinstance(v, (str, list)) or v is None): + raise ValueError(f"{k}={v} cannot be serialised to YAML.") + + return d + + @_make_methods(modules=("", "halos", "nl_pt",), name="cosmo") class Cosmology(CCLObject): """Stores information about cosmological parameters and associated data @@ -324,24 +342,7 @@ def write_yaml(self, filename, *, sort_keys=False): filename (:obj:`str`): file name, file pointer, or stream to write parameters to. """ - def make_yaml_friendly(d): - # serialize numpy types and dicts - for k, v in d.items(): - if isinstance(v, int): - d[k] = int(v) - elif isinstance(v, float): - d[k] = float(v) - elif isinstance(v, tuple): - d[k] = list(v) - elif isinstance(v, np.ndarray): - d[k] = v.tolist() - elif isinstance(v, dict): - make_yaml_friendly(v) - elif not (isinstance(v, (str, list)) or v is None): - raise ValueError(f"{k}={v} cannot be serialised to YAML.") - - params = self.to_dict() - make_yaml_friendly(params) + params = _make_yaml_friendly(self.to_dict()) if isinstance(filename, str): with open(filename, "w") as fp: diff --git a/pyccl/tests/test_cosmology.py b/pyccl/tests/test_cosmology.py index ca6a93a8d..25a2a44c9 100644 --- a/pyccl/tests/test_cosmology.py +++ b/pyccl/tests/test_cosmology.py @@ -120,6 +120,12 @@ def test_cosmology_init(): with pytest.raises(KeyError): ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.7, A_s=2.1e-9, n_s=0.96, transfer_function='x') + with pytest.raises(ValueError): + ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.7, A_s=2.1e-9, n_s=0.96, + matter_power_spectrum=None) + with pytest.raises(ValueError): + ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.7, A_s=2.1e-9, n_s=0.96, + transfer_function=None) with pytest.raises(ValueError): ccl.Cosmology(Omega_c=0.25, Omega_b=0.05, h=0.7, A_s=2.1e-9, n_s=0.96, m_nu=np.array([0.1, 0.1, 0.1, 0.1])) diff --git a/pyccl/tests/test_yaml.py b/pyccl/tests/test_yaml.py index 052debec6..2e525b2e6 100644 --- a/pyccl/tests/test_yaml.py +++ b/pyccl/tests/test_yaml.py @@ -2,7 +2,11 @@ import pytest import filecmp import io + +import numpy as np + import pyccl as ccl +from pyccl.cosmology import _make_yaml_friendly def test_yaml(): @@ -62,3 +66,14 @@ def test_to_dict(): init_params = {k: v for k, v in cosmo.__signature__.parameters.items() if k != "self"} assert set(cosmo.to_dict().keys()) == set(init_params.keys()) + + +def test_yaml_types(): + d = { + "tuple": (1, 2, 3), + "array": np.array([1.0, 42.0]) + } + + d_out = _make_yaml_friendly(d) + assert d_out["tuple"] == [1, 2, 3] + assert d_out["array"] == [1.0, 42.0] From da4fe7a69f0335dcbd4221c49043d6b96b13aa9c Mon Sep 17 00:00:00 2001 From: Tilman Troester Date: Mon, 26 Feb 2024 13:23:02 +0100 Subject: [PATCH 5/5] fix bug --- pyccl/cosmology.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyccl/cosmology.py b/pyccl/cosmology.py index 6b1512adb..276339eb7 100644 --- a/pyccl/cosmology.py +++ b/pyccl/cosmology.py @@ -125,7 +125,7 @@ def _make_yaml_friendly(d): d[k] = v.tolist() elif isinstance(v, dict): d[k] = _make_yaml_friendly(v) - elif not (isinstance(v, (str, list)) or v is None): + elif not (isinstance(v, (int, float, str, list)) or v is None): raise ValueError(f"{k}={v} cannot be serialised to YAML.") return d