From 95a8cd5b6a54c4e2071546b016d24f8c9eed85a6 Mon Sep 17 00:00:00 2001
From: Joe Zuntz <joezuntz@googlemail.com>
Date: Wed, 20 Nov 2024 12:50:31 +0000
Subject: [PATCH] Add utility function test_datablock_to_astropy

---
 cosmosis/test/test_utils.py | 149 ++++++++++++++++++++++++++++++++++++
 cosmosis/utils.py           |  55 +++++++++++++
 2 files changed, 204 insertions(+)

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’."""