diff --git a/cosmosis/test/test_utils.py b/cosmosis/test/test_utils.py index 50e81e0c..c23e1298 100644 --- a/cosmosis/test/test_utils.py +++ b/cosmosis/test/test_utils.py @@ -89,3 +89,152 @@ def test_git_nosub(): finally: cosmosis.utils.dulwich = dulwich_orig del os.environ["COSMOSIS_NO_SUBPROCESS"] + + +def test_datablock_to_astropy(): + try: + import astropy.cosmology + except ImportError: + pytest.skip("astropy not installed") + block = cosmosis.DataBlock() + + # implicitly should be FlatLambdaCDM + test_params = [ + (astropy.cosmology.FlatLambdaCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.7, + "omega_b": 0.05, + "mnu": 0.0, + }), + (astropy.cosmology.FlatLambdaCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.7, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.0, + }), + (astropy.cosmology.FlatLambdaCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.7, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.0, + "wa": 0.0, + }), + (astropy.cosmology.FlatLambdaCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.7, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.0, + "wa": 0.0, + "omega_k": 0.0, + }), + (astropy.cosmology.FlatwCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.7, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.50, + "wa": 0.0, + "omega_k": 0.0, + }), + (astropy.cosmology.FlatwCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.7, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.50, + }), + (astropy.cosmology.Flatw0waCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.7, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.50, + "wa": 0.5, + }), + (astropy.cosmology.wCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.65, + "omega_k": 0.05, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.50, + "wa": 0.0, + }), + (astropy.cosmology.wCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.65, + "omega_k": 0.05, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.50, + }), + (astropy.cosmology.w0waCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.65, + "omega_k": 0.05, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.50, + "wa": -0.5, + }), + (astropy.cosmology.LambdaCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.65, + "omega_k": 0.05, + "omega_b": 0.05, + "mnu": 0.0, + }), + (astropy.cosmology.LambdaCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.65, + "omega_k": 0.05, + "omega_b": 0.05, + "mnu": 0.0, + "w": -1.0, + }), + (astropy.cosmology.LambdaCDM, + { + "h0": 0.7, + "omega_m": 0.3, + "omega_lambda": 0.65, + "omega_k": 0.05, + "omega_b": 0.05, + "mnu": 0.0, + "wa": 0.0, + }), + + ] + + for expected_class, params in test_params: + block = cosmosis.DataBlock() + for k, v in params.items(): + block["cosmological_parameters", k] = v + c = cosmosis.utils.datablock_to_astropy(block) + assert isinstance(c, expected_class) diff --git a/cosmosis/utils.py b/cosmosis/utils.py index bf3f4c16..38fe2385 100644 --- a/cosmosis/utils.py +++ b/cosmosis/utils.py @@ -34,6 +34,61 @@ def __getitem__(self, param): +def datablock_to_astropy(block): + from .datablock import names + cosmo = names.cosmological_parameters + from astropy.cosmology import FlatLambdaCDM, LambdaCDM, FlatwCDM, wCDM, w0waCDM, Flatw0waCDM + + # Required parameters that must be in the data block + # for this to work + omega_lambda = block[cosmo, 'omega_lambda'] + omega_m = block[cosmo, 'omega_m'] + omega_b = block[cosmo, 'omega_b'] + H0 = block[cosmo, 'h0'] * 100 + m_nu = block[cosmo, 'mnu'] + + # Optional parameters with simple default values + omega_k = block.get_double(cosmo, 'omega_k', 0.0) + w0 = block.get_double(cosmo, 'w', -1.0) + wa = block.get_double(cosmo, 'wa', 0.0) + + kw = { + "H0": H0, + "Ob0": omega_b, + "Om0": omega_m, + "m_nu": m_nu, + } + + if omega_k == 0: + if wa == 0: + if w0 == -1.0: + model = FlatLambdaCDM + else: + model = FlatwCDM + kw["w0"] = w0 + else: + model = Flatw0waCDM + kw["w0"] = w0 + kw["wa"] = wa + else: + if wa == 0: + if w0 == -1.0: + model = LambdaCDM + kw["Ode0"] = omega_lambda + else: + model = wCDM + kw["Ode0"] = omega_lambda + kw["w0"] = w0 + else: + model = w0waCDM + kw["Ode0"] = omega_lambda + kw["w0"] = w0 + kw["wa"] = wa + + return model(**kw) + + + class ParseExtraParameters(argparse.Action): u"""Extended command-line argument parser :class:`Action` which knows how to read arguments of the form ‘section.name=value’."""