Add to_dict utility method

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,25 +329,20 @@ 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.
 
         Args:
             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()

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]))

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

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]