diff --git a/pyccl/cosmology.py b/pyccl/cosmology.py index 63eaee46b..276339eb7 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, (int, float, 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 @@ -235,13 +253,23 @@ def __init__( self.lin_pk_emu = None if isinstance(transfer_function, emulators.EmulatorPk): self.lin_pk_emu = transfer_function - transfer_function = 'emulator' + 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 - matter_power_spectrum = 'emulator' + 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): @@ -277,6 +305,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 +329,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. @@ -306,18 +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, dict): - make_yaml_friendly(v) - - params = {**self._params_init_kwargs, **self._config_init_kwargs} - make_yaml_friendly(params) + params = _make_yaml_friendly(self.to_dict()) if isinstance(filename, str): with open(filename, "w") as fp: @@ -337,12 +362,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 @@ -360,9 +387,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 @@ -519,7 +546,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 @@ -549,7 +576,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 " @@ -583,7 +610,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.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_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 diff --git a/pyccl/tests/test_yaml.py b/pyccl/tests/test_yaml.py index 1e4af5d3e..2e525b2e6 100644 --- a/pyccl/tests/test_yaml.py +++ b/pyccl/tests/test_yaml.py @@ -1,13 +1,19 @@ import tempfile +import pytest import filecmp import io + +import numpy as np + import pyccl as ccl +from pyccl.cosmology import _make_yaml_friendly 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", + ) # Make temporary files with tempfile.NamedTemporaryFile(delete=True) as tmpfile1, \ @@ -19,6 +25,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 +41,39 @@ def test_yaml(): cosmo2.write_yaml(stream2) 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(), + 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()) + + +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]