adding h5

pyproject.toml

@@ -14,6 +14,7 @@ flake8 = "^7.0.0"
 black = "^24.1.1"
 pyarrow = "^15.0.0"
 getdist = "^1.4.7"
+h5py = "^3.10.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.3.2"

src/scripts/io.py

@@ -1,5 +1,7 @@
 import pickle
 import h5py
+import numpy as np
+import torch
 
 class ModelLoader:
     def save_model_pkl(self, path, model_name, posterior):
@@ -82,13 +84,15 @@ def save_data_h5(self,
         :param data_name: Name of the data
         :param data: Data to be saved
         """
+        data_arrays = {key: np.asarray(value) for key, value in data.items()}
+
         file_name = path + data_name + ".h5"
         with h5py.File(file_name, "w") as file:
-            file.create_dataset(data_name, data=data)
+            # Save each array as a dataset in the HDF5 file
+            for key, value in data_arrays.items():
+                file.create_dataset(key, data=value)
 
-    def load_data_h5(self,
-                     data_name,
-                     path='../saveddata/'):
+    def load_data_h5(self, data_name, path='../saveddata/'):
         """
         Load data from an h5 file.
 
@@ -97,6 +101,8 @@ def load_data_h5(self,
         :return: Loaded data
         """
         file_name = path + data_name + ".h5"
+        loaded_data = {}
         with h5py.File(file_name, "r") as file:
-            data = file[data_name][...]
-        return data
+            for key in file.keys():
+                loaded_data[key] = torch.Tensor(file[key][...])
+        return loaded_data