From a7d7465d3744d1d20e7afa1183c6a674c8a8cd2f Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Wed, 27 Nov 2024 16:50:06 +0000
Subject: [PATCH 01/14] ENH: add `cosmology` to CBC result object

---
 bilby/gw/result.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/bilby/gw/result.py b/bilby/gw/result.py
index cc9f97ed..d28ab1ff 100644
--- a/bilby/gw/result.py
+++ b/bilby/gw/result.py
@@ -117,6 +117,18 @@ def parameter_conversion(self):
         return self.__get_from_nested_meta_data(
             'likelihood', 'parameter_conversion')
 
+    @property
+    def cosmology(self):
+        """The global cosmology used in the analysis.
+
+        If the cosmology is not set, the default cosmology is returned.
+        See :py:func:`bilby.gw.cosmology.get_cosmology` for details.
+
+        .. versionadded:: 2.5.0
+        """
+        from .cosmology import get_cosmology
+        return get_cosmology()
+
     def detector_injection_properties(self, detector):
         """ Returns a dictionary of the injection properties for each detector
 

From 6248fca251c00d81fd09dc12538318dc659d3471 Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Wed, 27 Nov 2024 16:50:29 +0000
Subject: [PATCH 02/14] TST: add test for `cosmology`  in CBC result object

---
 test/gw/result_test.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/test/gw/result_test.py b/test/gw/result_test.py
index 73780174..a92cb104 100644
--- a/test/gw/result_test.py
+++ b/test/gw/result_test.py
@@ -3,6 +3,7 @@
 import unittest
 
 import pandas as pd
+from parameterized import parameterized
 
 import bilby
 
@@ -200,6 +201,15 @@ def test_detector_injection_properties_no_injection(self):
             self.result.detector_injection_properties("not_a_detector"), None
         )
 
+    @parameterized.expand(["Planck18", "Planck15", "Planck15_LAL"])
+    def test_cosmology(self, cosmology):
+        from bilby.gw.cosmology import get_cosmology, set_cosmology
+        set_cosmology(cosmology)
+        cosmo = get_cosmology()
+        self.assertEqual(self.result.cosmology, cosmo)
+        # Reset the cosmology to the default
+        set_cosmology("Planck15")
+
 
 if __name__ == "__main__":
     unittest.main()

From 496a80e5590c5e9a4acbe1fabd7a87a28e71c83e Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Thu, 28 Nov 2024 15:14:54 +0000
Subject: [PATCH 03/14] MAINT: specify cosmology in meta data when using
 `run_sampler`

---
 bilby/core/sampler/__init__.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/bilby/core/sampler/__init__.py b/bilby/core/sampler/__init__.py
index ef65ce1d..087dba54 100644
--- a/bilby/core/sampler/__init__.py
+++ b/bilby/core/sampler/__init__.py
@@ -251,6 +251,14 @@ def run_sampler(
     meta_data["likelihood"] = likelihood.meta_data
     meta_data["loaded_modules"] = loaded_modules_dict()
     meta_data["environment_packages"] = env_package_list(as_dataframe=True)
+    # Add the cosmology if available
+    if "cosmology" not in meta_data:
+        try:
+            from bilby.gw.cosmology import get_cosmology
+
+            meta_data["cosmology"] = get_cosmology()
+        except ImportError:
+            logger.info("Could not specify cosmology in meta data")
 
     if command_line_args.bilby_zero_likelihood_mode:
         from bilby.core.likelihood import ZeroLikelihood

From aa778f5589bbedc4a66205de3650a3909523536c Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Thu, 28 Nov 2024 15:15:11 +0000
Subject: [PATCH 04/14] MAINT: move cosmology to meta data

---
 bilby/gw/result.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/bilby/gw/result.py b/bilby/gw/result.py
index d28ab1ff..b642aa74 100644
--- a/bilby/gw/result.py
+++ b/bilby/gw/result.py
@@ -19,6 +19,14 @@ class CompactBinaryCoalescenceResult(CoreResult):
     of compact binaries.
     """
     def __init__(self, **kwargs):
+
+        # Ensure cosmology is always stored in the meta_data
+        if "meta_data" not in kwargs:
+            kwargs["meta_data"] = dict()
+        if "cosmology" not in kwargs["meta_data"]:
+            from .cosmology import get_cosmology
+            kwargs["meta_data"]["cosmology"] = get_cosmology()
+
         super(CompactBinaryCoalescenceResult, self).__init__(**kwargs)
 
     def __get_from_nested_meta_data(self, *keys):
@@ -121,13 +129,9 @@ def parameter_conversion(self):
     def cosmology(self):
         """The global cosmology used in the analysis.
 
-        If the cosmology is not set, the default cosmology is returned.
-        See :py:func:`bilby.gw.cosmology.get_cosmology` for details.
-
         .. versionadded:: 2.5.0
         """
-        from .cosmology import get_cosmology
-        return get_cosmology()
+        return self.__get_from_nested_meta_data('cosmology')
 
     def detector_injection_properties(self, detector):
         """ Returns a dictionary of the injection properties for each detector

From c4e51b212a5205191ab9454adb09b3e0a4ad71a3 Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Thu, 28 Nov 2024 15:31:33 +0000
Subject: [PATCH 05/14] TST: update cosmology test in result tests

---
 test/gw/result_test.py | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/test/gw/result_test.py b/test/gw/result_test.py
index a92cb104..d0f8ca56 100644
--- a/test/gw/result_test.py
+++ b/test/gw/result_test.py
@@ -3,7 +3,6 @@
 import unittest
 
 import pandas as pd
-from parameterized import parameterized
 
 import bilby
 
@@ -201,14 +200,11 @@ def test_detector_injection_properties_no_injection(self):
             self.result.detector_injection_properties("not_a_detector"), None
         )
 
-    @parameterized.expand(["Planck18", "Planck15", "Planck15_LAL"])
-    def test_cosmology(self, cosmology):
-        from bilby.gw.cosmology import get_cosmology, set_cosmology
-        set_cosmology(cosmology)
-        cosmo = get_cosmology()
-        self.assertEqual(self.result.cosmology, cosmo)
-        # Reset the cosmology to the default
-        set_cosmology("Planck15")
+    def test_cosmology(self):
+        self.assertEqual(
+            self.result.cosmology,
+            self.meta_data["cosmology"],
+        )
 
 
 if __name__ == "__main__":

From 6b9e64be230e48da1e6d3fe22dfb81a967c0b515 Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Fri, 29 Nov 2024 15:53:45 +0000
Subject: [PATCH 06/14] ENH: add support for saving astropy cosmologies in HDF5
 files

---
 bilby/core/utils/io.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/bilby/core/utils/io.py b/bilby/core/utils/io.py
index cc8329e1..9179e126 100644
--- a/bilby/core/utils/io.py
+++ b/bilby/core/utils/io.py
@@ -11,7 +11,6 @@
 import pandas as pd
 
 from .log import logger
-from .introspection import infer_args_from_method
 
 
 def check_directory_exists_and_if_not_mkdir(directory):
@@ -87,11 +86,7 @@ def default(self, obj):
 
 
 def encode_astropy_cosmology(obj):
-    cls_name = obj.__class__.__name__
-    dct = {key: getattr(obj, key) for key in infer_args_from_method(obj.__init__)}
-    dct["__cosmology__"] = True
-    dct["__name__"] = cls_name
-    return dct
+    return {"__cosmology__": True, **obj.to_format("mapping")}
 
 
 def encode_astropy_quantity(dct):
@@ -103,11 +98,10 @@ def encode_astropy_quantity(dct):
 
 def decode_astropy_cosmology(dct):
     try:
-        from astropy import cosmology as cosmo
+        from astropy.cosmology import Cosmology
 
-        cosmo_cls = getattr(cosmo, dct["__name__"])
-        del dct["__cosmology__"], dct["__name__"]
-        return cosmo_cls(**dct)
+        del dct["__cosmology__"]
+        return Cosmology.from_format(dct, format="mapping")
     except ImportError:
         logger.debug(
             "Cannot import astropy, cosmological priors may not be " "properly loaded."
@@ -241,6 +235,8 @@ def decode_from_hdf5(item):
             output = item
     elif isinstance(item, np.bool_):
         output = bool(item)
+    elif isinstance(item, dict) and "__cosmology__" in item:
+        output = decode_astropy_cosmology(item)
     else:
         output = item
     return output
@@ -313,7 +309,14 @@ def encode_for_hdf5(key, item):
     elif isinstance(item, datetime.timedelta):
         output = item.total_seconds()
     else:
-        raise ValueError(f'Cannot save {key}: {type(item)} type')
+        try:
+            from astropy import cosmology as cosmo
+
+            if isinstance(item, cosmo.FLRW):
+                output = encode_astropy_cosmology(item)
+        except ImportError:
+            logger.debug("Cannot import astropy, cannot write cosmological priors")
+            raise ValueError(f'Cannot save {key}: {type(item)} type')
     return output
 
 

From 170b5b3d0bffb4e37f84db0cc16ba9c5caf84051 Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Mon, 2 Dec 2024 10:13:24 +0000
Subject: [PATCH 07/14] ENH: improve how astropy objects are encoded/decoded

---
 bilby/core/utils/io.py | 118 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 108 insertions(+), 10 deletions(-)

diff --git a/bilby/core/utils/io.py b/bilby/core/utils/io.py
index 9179e126..4b4362bd 100644
--- a/bilby/core/utils/io.py
+++ b/bilby/core/utils/io.py
@@ -52,8 +52,8 @@ def default(self, obj):
                 return encode_astropy_cosmology(obj)
             if isinstance(obj, units.Quantity):
                 return encode_astropy_quantity(obj)
-            if isinstance(obj, units.PrefixUnit):
-                return str(obj)
+            if isinstance(obj, (units.PrefixUnit, units.UnitBase, units.FunctionUnitBase)):
+                return encode_astropy_unit(obj)
         except ImportError:
             logger.debug("Cannot import astropy, cannot write cosmological priors")
         if isinstance(obj, np.ndarray):
@@ -86,33 +86,98 @@ def default(self, obj):
 
 
 def encode_astropy_cosmology(obj):
+    """Encode an astropy cosmology object to a dictionary.
+
+    Adds the key :code:`__cosmology__` to the dictionary to indicate that the
+    object is a cosmology object.
+
+    .. versionchange:: 2.5.0
+        Now uses the :code:`to_format("mapping")` method to encode the
+        cosmology object.
+    """
     return {"__cosmology__": True, **obj.to_format("mapping")}
 
 
 def encode_astropy_quantity(dct):
-    dct = dict(__astropy_quantity__=True, value=dct.value, unit=str(dct.unit))
+    """Encode an astropy quantity object to a dictionary.
+
+    Adds the key :code:`__astropy_quantity__` to the dictionary to indicate that
+    the object is a quantity object.
+    """
+    dct = dict(__astropy_quantity__=True, value=dct.value, unit=dct.unit.to_string())
     if isinstance(dct["value"], np.ndarray):
-        dct["value"] = list(dct["value"])
+        dct["value"] = dct["value"].tolist()
     return dct
 
 
+def encode_astropy_unit(obj):
+    """Encode an astropy unit object to a dictionary.
+
+    Adds the key :code:`__astropy_unit__` to the dictionary to indicate that the
+    object is a unit object.
+
+    .. versionadded:: 2.5.0
+    """
+    try:
+        from astropy import units
+
+        # Based on the JsonCustomEncoder in astropy.units.misc
+        if obj == units.dimensionless_unscaled:
+            return dict(__astropy_unit__=True, unit="dimensionless_unit")
+        return dict(__astropy_unit__=True, unit=obj.to_string())
+
+    except ImportError:
+        logger.debug(
+            "Cannot import astropy, cosmological priors may not be properly loaded."
+        )
+        return obj
+
+
 def decode_astropy_cosmology(dct):
+    """Decode an astropy cosmology from a dictionary.
+
+    The dictionary should have been encoded using
+    :py:func:`~bibly.core.utils.io.encode_astropy_cosmology` and should have the
+    key :code:`__cosmology__`.
+
+    .. versionchange:: 2.5.0
+        Now uses the :code:`from_format` method to decode the cosmology object.
+        Still supports decoding result files that used the previous encoding.
+    """
     try:
-        from astropy.cosmology import Cosmology
+        from astropy import cosmology as cosmo
 
         del dct["__cosmology__"]
-        return Cosmology.from_format(dct, format="mapping")
+        return cosmo.Cosmology.from_format(dct, format="mapping")
     except ImportError:
         logger.debug(
-            "Cannot import astropy, cosmological priors may not be " "properly loaded."
+            "Cannot import astropy, cosmological priors may not be properly loaded."
         )
         return dct
+    except KeyError:
+        # Support decoding result files that used the previous encoding
+        logger.warning(
+            "Failed to decode cosmology, falling back to legacy decoding."
+            "Support for legacy decoding will be removed in a future release."
+        )
+        cosmo_cls = getattr(cosmo, dct["__name__"])
+        del dct["__cosmology__"], dct["__name__"]
+        return cosmo_cls(**dct)
 
 
 def decode_astropy_quantity(dct):
+    """Decode an astropy quantity from a dictionary.
+
+    The dictionary should have been encoded using
+    :py:func:`~bilby.core.utils.io.encode_astropy_quantity` and should have the
+    key :code:`__astropy_quantity__`.
+    """
     try:
         from astropy import units
+        from astropy.cosmology import units as cosmo_units
 
+        # Enable cosmology units such as redshift
+        units.add_enabled_units(cosmo_units)
         if dct["value"] is None:
             return None
         else:
@@ -120,7 +185,34 @@ def decode_astropy_quantity(dct):
             return units.Quantity(**dct)
     except ImportError:
         logger.debug(
-            "Cannot import astropy, cosmological priors may not be " "properly loaded."
+            "Cannot import astropy, cosmological priors may not be properly loaded."
+        )
+        return dct
+
+
+def decode_astropy_unit(dct):
+    """Decode an astropy unit from a dictionary.
+
+    The dictionary should have been encoded using
+    :py:func:`~bilby.core.utils.io.encode_astropy_unit` and should have the
+    key :code:`__astropy_unit__`.
+
+    .. versionadded:: 2.5.0
+    """
+    try:
+        from astropy import units
+        from astropy.cosmology import units as cosmo_units
+
+        # Enable cosmology units such as redshift
+        units.add_enabled_units(cosmo_units)
+        if dct["unit"] == "dimensionless_unit":
+            return units.dimensionless_unscaled
+        else:
+            del dct["__astropy_unit__"]
+            return units.Unit(dct["unit"])
+    except ImportError:
+        logger.debug(
+            "Cannot import astropy, cosmological priors may not be properly loaded."
         )
         return dct
 
@@ -164,6 +256,8 @@ def decode_bilby_json(dct):
         return decode_astropy_cosmology(dct)
     if dct.get("__astropy_quantity__", False):
         return decode_astropy_quantity(dct)
+    if dct.get("__astropy_unit__", False):
+        return decode_astropy_unit(dct)
     if dct.get("__array__", False):
         return np.asarray(dct["content"])
     if dct.get("__complex__", False):
@@ -310,10 +404,14 @@ def encode_for_hdf5(key, item):
         output = item.total_seconds()
     else:
         try:
-            from astropy import cosmology as cosmo
+            from astropy import cosmology as cosmo, units
 
             if isinstance(item, cosmo.FLRW):
-                output = encode_astropy_cosmology(item)
+                return encode_astropy_cosmology(item)
+            if isinstance(item, units.Quantity):
+                return encode_astropy_quantity(item)
+            if isinstance(item, (units.PrefixUnit, units.UnitBase, units.FunctionUnitBase)):
+                return encode_astropy_unit(item)
         except ImportError:
             logger.debug("Cannot import astropy, cannot write cosmological priors")
             raise ValueError(f'Cannot save {key}: {type(item)} type')

From f8999c619323aa75214f71ee2ec7e0f799227dd1 Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Mon, 2 Dec 2024 10:14:37 +0000
Subject: [PATCH 08/14] TST: switch to using pytest's tmp_path fixture

---
 test/core/result_test.py | 35 ++++++++++++++++++++++++++++-------
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/test/core/result_test.py b/test/core/result_test.py
index 2a876109..c96217e9 100644
--- a/test/core/result_test.py
+++ b/test/core/result_test.py
@@ -4,6 +4,7 @@
 import shutil
 import os
 import json
+import pytest
 from unittest.mock import patch
 
 import bilby
@@ -11,6 +12,7 @@
 
 
 class TestJson(unittest.TestCase):
+
     def setUp(self):
         self.encoder = bilby.core.utils.BilbyJsonEncoder
         self.decoder = bilby.core.utils.decode_bilby_json
@@ -50,6 +52,12 @@ def test_dataframe_encoding(self):
 
 
 class TestResult(unittest.TestCase):
+
+    @pytest.fixture(autouse=True)
+    def init_outdir(self, tmp_path):
+        # Use pytest's tmp_path fixture to create a temporary directory
+        self.outdir = str(tmp_path / "test")
+
     def setUp(self):
         np.random.seed(7)
         bilby.utils.command_line_args.bilby_test_mode = False
@@ -61,7 +69,6 @@ def setUp(self):
                 d=2,
             )
         )
-        self.outdir = "test_outdir"
         result = bilby.core.result.Result(
             label="label",
             outdir=self.outdir,
@@ -543,18 +550,32 @@ class NotAResult(bilby.core.result.Result):
             pass
 
         result = bilby.run_sampler(
-            likelihood, priors, sampler='bilby_mcmc', nsamples=10, L1steps=1,
-            proposal_cycle="default_noGMnoKD", printdt=1,
+            likelihood,
+            priors,
+            sampler='bilby_mcmc',
+            outdir=self.outdir,
+            nsamples=10,
+            L1steps=1,
+            proposal_cycle="default_noGMnoKD",
+            printdt=1,
             check_point_plot=False,
-            result_class=NotAResult)
+            result_class=NotAResult
+        )
         # result should be specified result_class
         assert isinstance(result, NotAResult)
 
         cached_result = bilby.run_sampler(
-            likelihood, priors, sampler='bilby_mcmc', nsamples=10, L1steps=1,
-            proposal_cycle="default_noGMnoKD", printdt=1,
+            likelihood,
+            priors,
+            sampler='bilby_mcmc',
+            outdir=self.outdir,
+            nsamples=10,
+            L1steps=1,
+            proposal_cycle="default_noGMnoKD",
+            printdt=1,
             check_point_plot=False,
-            result_class=NotAResult)
+            result_class=NotAResult
+        )
 
         # so should a result loaded from cache
         assert isinstance(cached_result, NotAResult)

From cd902d65fa8b6889b009391799efdaed352a1a76 Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Mon, 2 Dec 2024 11:13:47 +0000
Subject: [PATCH 09/14] ENH: add support for deconding astropy and bilby types
 from hdf5 files

---
 bilby/core/utils/io.py | 52 +++++++++++++++++++++++++++++++-----------
 1 file changed, 39 insertions(+), 13 deletions(-)

diff --git a/bilby/core/utils/io.py b/bilby/core/utils/io.py
index 4b4362bd..e37b5d5a 100644
--- a/bilby/core/utils/io.py
+++ b/bilby/core/utils/io.py
@@ -329,8 +329,6 @@ def decode_from_hdf5(item):
             output = item
     elif isinstance(item, np.bool_):
         output = bool(item)
-    elif isinstance(item, dict) and "__cosmology__" in item:
-        output = decode_astropy_cosmology(item)
     else:
         output = item
     return output
@@ -354,6 +352,13 @@ def encode_for_hdf5(key, item):
     """
     from ..prior.dict import PriorDict
 
+    try:
+        from astropy import cosmology as cosmo, units
+    except ImportError:
+        logger.debug("Cannot import astropy, cannot write cosmological priors")
+        cosmo = None
+        units = None
+
     if isinstance(item, np.int_):
         item = int(item)
     elif isinstance(item, np.float64):
@@ -392,6 +397,12 @@ def encode_for_hdf5(key, item):
         output = json.dumps(item._get_json_dict())
     elif isinstance(item, pd.DataFrame):
         output = item.to_dict(orient="list")
+    elif cosmo is not None and isinstance(item, cosmo.FLRW):
+        output = encode_astropy_cosmology(item)
+    elif units is not None and isinstance(item, units.Quantity):
+        output = encode_astropy_quantity(item)
+    elif units is not None and isinstance(item, (units.PrefixUnit, units.UnitBase, units.FunctionUnitBase)):
+        output = encode_astropy_unit(item)
     elif inspect.isfunction(item) or inspect.isclass(item):
         output = dict(
             __module__=item.__module__, __name__=item.__name__, __class__=True
@@ -403,18 +414,26 @@ def encode_for_hdf5(key, item):
     elif isinstance(item, datetime.timedelta):
         output = item.total_seconds()
     else:
-        try:
-            from astropy import cosmology as cosmo, units
+        raise ValueError(f'Cannot save {key}: {type(item)} type')
+    return output
 
-            if isinstance(item, cosmo.FLRW):
-                return encode_astropy_cosmology(item)
-            if isinstance(item, units.Quantity):
-                return encode_astropy_quantity(item)
-            if isinstance(item, (units.PrefixUnit, units.UnitBase, units.FunctionUnitBase)):
-                return encode_astropy_unit(item)
-        except ImportError:
-            logger.debug("Cannot import astropy, cannot write cosmological priors")
-            raise ValueError(f'Cannot save {key}: {type(item)} type')
+
+def decode_hdf5_dict(output):
+    """Decode a dictionary constructed from a HDF5 file.
+
+    This handles decoding of Bilby types and astropy types from the dictionary.
+
+    .. versionadded:: 2.5.0
+    """
+    if ("__function__" in output) or ("__class__" in output):
+        default = ".".join([output["__module__"], output["__name__"]])
+        output = getattr(import_module(output["__module__"]), output["__name__"], default)
+    elif "__cosmology__" in output:
+        output = decode_astropy_cosmology(output)
+    elif "__astropy_quantity__" in output:
+        output = decode_astropy_quantity(output)
+    elif "__astropy_unit__" in output:
+        output = decode_astropy_unit(output)
     return output
 
 
@@ -424,6 +443,9 @@ def recursively_load_dict_contents_from_group(h5file, path):
 
     .. versionadded:: 1.1.0
 
+    .. versionchanged: 2.5.0
+        Now decodes astropy and bilby types
+
     Parameters
     ----------
     h5file: h5py.File
@@ -446,6 +468,10 @@ def recursively_load_dict_contents_from_group(h5file, path):
             output[key] = recursively_load_dict_contents_from_group(
                 h5file, path + key + "/"
             )
+    # Some items may be encoded as dictionaries, so we need to decode them
+    # after the dictionary has been constructed.
+    # This includes decoding astropy and bilby types
+    output = decode_hdf5_dict(output)
     return output
 
 

From d84d656d584fd97b34fb9717b7542beab1d96d0e Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Mon, 2 Dec 2024 11:32:53 +0000
Subject: [PATCH 10/14] TST: add save and load test for CBC result

---
 test/gw/result_test.py | 60 +++++++++++++++++++++++++++++++++++++-----
 1 file changed, 53 insertions(+), 7 deletions(-)

diff --git a/test/gw/result_test.py b/test/gw/result_test.py
index d0f8ca56..6b886201 100644
--- a/test/gw/result_test.py
+++ b/test/gw/result_test.py
@@ -1,13 +1,20 @@
 import os
-import shutil
 import unittest
 
 import pandas as pd
+from parameterized import parameterized, parameterized_class
+import pytest
 
 import bilby
 
 
-class TestCBCResult(unittest.TestCase):
+class BaseCBCResultTest(unittest.TestCase):
+
+    @pytest.fixture(autouse=True)
+    def init_outdir(self, tmp_path):
+        # Use pytest's tmp_path fixture to create a temporary directory
+        self.outdir = str(tmp_path / "test")
+
     def setUp(self):
         bilby.utils.command_line_args.bilby_test_mode = False
         priors = bilby.gw.prior.BBHPriorDict()
@@ -35,7 +42,7 @@ def setUp(self):
         )
         self.result = bilby.gw.result.CBCResult(
             label="label",
-            outdir="outdir",
+            outdir=self.outdir,
             sampler="nestle",
             search_parameter_keys=list(priors.keys()),
             fixed_parameter_keys=list(),
@@ -50,12 +57,11 @@ def setUp(self):
 
     def tearDown(self):
         bilby.utils.command_line_args.bilby_test_mode = True
-        try:
-            shutil.rmtree(self.result.outdir)
-        except OSError:
-            pass
         del self.result
 
+
+class TestCBCResult(BaseCBCResultTest):
+
     def test_phase_marginalization(self):
         self.assertEqual(
             self.result.phase_marginalization,
@@ -207,5 +213,45 @@ def test_cosmology(self):
         )
 
 
+@parameterized_class(
+    ["cosmology_name"],
+    [["Planck15"], ["Planck15_LAL"]]
+)
+class TestCBCResultSaveAndLoad(BaseCBCResultTest):
+
+    def setUp(self):
+        self.orig_cosmology = bilby.gw.cosmology.get_cosmology()
+        bilby.gw.cosmology.set_cosmology(self.cosmology_name)
+        return super().setUp()
+
+    def tearDown(self):
+        super().tearDown()
+        bilby.gw.cosmology.set_cosmology(self.orig_cosmology)
+
+    @parameterized.expand(
+        [
+            ("json", False),
+            ("json", True),
+            ("pkl", False),
+            ("hdf5", False),
+        ]
+    )
+    def test_save_and_load(self, extension, gzip):
+        self.result.save_to_file(extension=extension, gzip=gzip)
+        loaded_result = bilby.core.result.read_in_result(
+            outdir=self.result.outdir,
+            label=self.result.label,
+            extension=extension,
+            gzip=gzip,
+            result_class=bilby.gw.result.CBCResult,
+        )
+        self.assertEqual(self.result.cosmology, loaded_result.cosmology)
+        self.assertEqual(self.result.parameter_conversion, loaded_result.parameter_conversion)
+        self.assertEqual(self.result.frequency_domain_source_model, loaded_result.frequency_domain_source_model)
+        self.assertEqual(self.result.waveform_generator_class, loaded_result.waveform_generator_class)
+        self.assertEqual(self.result.waveform_arguments, loaded_result.waveform_arguments)
+        self.assertEqual(self.result.interferometers, loaded_result.interferometers)
+
+
 if __name__ == "__main__":
     unittest.main()

From 66479bed6de6a1589af2b7320ab4df3dcd4111e8 Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Fri, 6 Dec 2024 17:42:52 +0000
Subject: [PATCH 11/14] MAINT: remove cosmology from meta data in `run_sampler`

---
 bilby/core/sampler/__init__.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/bilby/core/sampler/__init__.py b/bilby/core/sampler/__init__.py
index 087dba54..ef65ce1d 100644
--- a/bilby/core/sampler/__init__.py
+++ b/bilby/core/sampler/__init__.py
@@ -251,14 +251,6 @@ def run_sampler(
     meta_data["likelihood"] = likelihood.meta_data
     meta_data["loaded_modules"] = loaded_modules_dict()
     meta_data["environment_packages"] = env_package_list(as_dataframe=True)
-    # Add the cosmology if available
-    if "cosmology" not in meta_data:
-        try:
-            from bilby.gw.cosmology import get_cosmology
-
-            meta_data["cosmology"] = get_cosmology()
-        except ImportError:
-            logger.info("Could not specify cosmology in meta data")
 
     if command_line_args.bilby_zero_likelihood_mode:
         from bilby.core.likelihood import ZeroLikelihood

From 8b0d541493f48482ac2ee0231534c3b664af07cb Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Mon, 9 Dec 2024 13:32:05 +0000
Subject: [PATCH 12/14] MAINT: move cosmology to global meta data

---
 bilby/gw/result.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/bilby/gw/result.py b/bilby/gw/result.py
index b642aa74..00f71b8d 100644
--- a/bilby/gw/result.py
+++ b/bilby/gw/result.py
@@ -20,12 +20,14 @@ class CompactBinaryCoalescenceResult(CoreResult):
     """
     def __init__(self, **kwargs):
 
-        # Ensure cosmology is always stored in the meta_data
         if "meta_data" not in kwargs:
             kwargs["meta_data"] = dict()
-        if "cosmology" not in kwargs["meta_data"]:
+        if "global_meta_data" not in kwargs:
+            kwargs["meta_data"]["global_meta_data"] = dict()
+        # Ensure cosmology is always stored in the meta_data
+        if "cosmology" not in kwargs["meta_data"]["global_meta_data"]:
             from .cosmology import get_cosmology
-            kwargs["meta_data"]["cosmology"] = get_cosmology()
+            kwargs["meta_data"]["global_meta_data"]["cosmology"] = get_cosmology()
 
         super(CompactBinaryCoalescenceResult, self).__init__(**kwargs)
 
@@ -131,7 +133,9 @@ def cosmology(self):
 
         .. versionadded:: 2.5.0
         """
-        return self.__get_from_nested_meta_data('cosmology')
+        return self.__get_from_nested_meta_data(
+            'global_meta_data', 'cosmology'
+        )
 
     def detector_injection_properties(self, detector):
         """ Returns a dictionary of the injection properties for each detector

From 3fc3289516355391ef894dd1d1e6585a6ebc1d81 Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Mon, 9 Dec 2024 13:32:21 +0000
Subject: [PATCH 13/14] TST: update cosmology test

---
 test/gw/result_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/gw/result_test.py b/test/gw/result_test.py
index 6b886201..98a7ce52 100644
--- a/test/gw/result_test.py
+++ b/test/gw/result_test.py
@@ -209,7 +209,7 @@ def test_detector_injection_properties_no_injection(self):
     def test_cosmology(self):
         self.assertEqual(
             self.result.cosmology,
-            self.meta_data["cosmology"],
+            self.meta_data["global_meta_data"]["cosmology"],
         )
 
 

From 1a274bb4c4afbb4a8f03b21bd275080f56c4be3f Mon Sep 17 00:00:00 2001
From: mj-will <michaeljw1@googlemail.com>
Date: Thu, 9 Jan 2025 10:12:53 +0100
Subject: [PATCH 14/14] MAINT: fix `decode_astropy_cosmology` for legacy result
 files

---
 bilby/core/utils/io.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/bilby/core/utils/io.py b/bilby/core/utils/io.py
index e37b5d5a..96df21bf 100644
--- a/bilby/core/utils/io.py
+++ b/bilby/core/utils/io.py
@@ -157,11 +157,11 @@ def decode_astropy_cosmology(dct):
     except KeyError:
         # Support decoding result files that used the previous encoding
         logger.warning(
-            "Failed to decode cosmology, falling back to legacy decoding."
+            "Failed to decode cosmology, falling back to legacy decoding. "
             "Support for legacy decoding will be removed in a future release."
         )
         cosmo_cls = getattr(cosmo, dct["__name__"])
-        del dct["__cosmology__"], dct["__name__"]
+        del dct["__name__"]
         return cosmo_cls(**dct)