TST: add tests for saving and loading numpy random generators

Author: mj-will

test/core/utils_test.py

@@ -1,10 +1,14 @@
 import unittest
 import os
 
+import dill
 import numpy as np
 from astropy import constants
 import lal
 import matplotlib.pyplot as plt
+import h5py
+import json
+import pytest
 
 import bilby
 from bilby.core import utils
@@ -412,5 +416,56 @@ def test_integral_func2_irregular_steps(self):
         self.assertTrue(np.abs((np.exp(res) - self.func2int) / self.func2int) < 1e-2)
 
 
+class TestSavingNumpyRandomGenerator(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 = tmp_path / "test"
+        self.outdir.mkdir()
+
+    def setUp(self):
+        self.filename = "test_random_state.npy"
+        self.data = {
+            "rng": np.random.default_rng(),
+            "seed": 1234,
+        }
+
+    def test_hdf5(self):
+        with h5py.File(self.outdir / "test.h5", "w") as f:
+            bilby.core.utils.recursively_save_dict_contents_to_group(
+                f, "/", self.data
+            )
+        a = self.data["rng"].random()
+
+        with h5py.File(self.outdir / "test.h5", "r") as f:
+            data = bilby.core.utils.recursively_load_dict_contents_from_group(f, "/")
+
+        b = data["rng"].random()
+        self.assertEqual(a, b)
+
+    def test_json(self):
+        with open(self.outdir / "test.json", 'w') as file:
+            json.dump(self.data, file, indent=2, cls=bilby.core.utils.BilbyJsonEncoder)
+
+        a = self.data["rng"].random()
+
+        with open(self.outdir / "test.json", 'r') as file:
+            data = json.load(file, object_hook=bilby.core.utils.decode_bilby_json)
+
+        b = data["rng"].random()
+        self.assertEqual(a, b)
+
+    def test_pickle(self):
+        with open(self.outdir / "test.pkl", 'wb') as file:
+            dill.dump(self.data, file)
+        a = self.data["rng"].random()
+
+        with open(self.outdir / "test.pkl", 'rb') as file:
+            data = dill.load(file)
+        b = data["rng"].random()
+        self.assertEqual(a, b)
+
+
 if __name__ == "__main__":
     unittest.main()