Fix bug with loading h5 files including strings, and add test coverage for h5_to_nested_dict and nested_dict_to_h5

src/cdtools/tools/data/data.py

@@ -796,7 +796,7 @@ def add_ptycho_translations(cxi_file, translations):
 
 
 def nested_dict_to_h5(h5_file, d):
-    """saves a nested dictionary to an h5 file object
+    """Saves a nested dictionary to an h5 file object
 
     Parameters
     ----------
@@ -832,14 +832,17 @@ def nested_dict_to_h5(h5_file, d):
 
 
 def h5_to_nested_dict(h5_file):
-    """saves a nested dictionary to an h5 file object
+    """Saves a nested dictionary to an h5 file object
 
     Parameters
     ----------
     h5_file : h5py.File
         A file object, or path to a file, to load from
+
+    Returns
+    -------
     d : dict
-        A mapping whose keys are all strings and whose values are only numpy arrays, pytorch tensors, scalars, python strings, or other mappings meeting the same conditions
+        A dictionary whose keys are all strings and whose values are numpy arrays, scalars, or python strings. Will raise an error if the data cannot be loadedinto this format
     """
     
     # If a bare string is passed
@@ -852,11 +855,13 @@ def h5_to_nested_dict(h5_file):
         value = h5_file[key]
         if isinstance(value, h5py.Dataset):
             arr = value[()]
-            if arr.dtype == object:
+            # Strings stored via nested_dict_to_h5 will wind up as bytes objs
+            if type(arr) == type(b''):
+                d[key] = arr.decode('utf-8')
+            # Some strings in h5 files seem to be stored this way
+            elif hasattr(arr, 'dtype') and arr.dtype == object:
                 d[key] = arr.ravel()[0].decode('utf-8')
-            elif arr.ndim == 0:
-                # TODO is this needed with arr = value[()]?
-                d[key] = arr.ravel()[0] 
+            # This is the default case: it's an array of numbers
             else:
                 d[key] = arr
             
@@ -870,7 +875,19 @@ def h5_to_nested_dict(h5_file):
 
 
 def nested_dict_to_numpy(d):
+    """Sends all array like objects in a nested dict to numpy arrays
 
+    Parameters
+    ----------
+    d : dict
+        A mapping whose keys are all strings and whose values are only numpy arrays, pytorch tensors, scalars, or other mappings meeting the same conditions
+
+    Returns
+    -------
+    d_out : dict
+        A new dictionary with all array like objects sent to numpy 
+    """
+    
     new_dict = {}
     for key in d.keys():
         value = d[key]
@@ -888,29 +905,49 @@ def nested_dict_to_numpy(d):
         elif isinstance(value, Mapping):
             new_dict[key] = nested_dict_to_numpy(value)
         else:
-            raise ValueError(f'{value} is not a number, numpy array, torch tensor, or mapping')
+            raise ValueError(f'{value} is not a number, numpy array, torch tensor, string, or mapping')
 
     return new_dict
 
-def nested_dict_to_torch(d):
+def nested_dict_to_torch(d, device=None):
+    """Sends all array like objects in a nested dict to pytorch tensors
+
+    This will also send all the tensors to a specific device, if specified.
+    There is no option to send all tensors to a specific dtype, as tensors
+    are often a mixture of integer, floating point, and complex types. In
+    the future, this may support a "precision" option to send all tensors to
+    a specified precision.
     
+    Parameters
+    ----------
+    d : dict
+        A mapping whose keys are all strings and whose values are only numpy arrays, pytorch tensors, scalars, or other mappings meeting the same conditions
+    device : torch.device
+        A valid device argument for torch.Tensor.to
+    
+    Returns
+    -------
+    d_out : dict
+        A new dictionary with all array like objects sent to torch tensors 
+    """
+
     new_dict = {}
     for key in d.keys():
         value = d[key]
         if isinstance(value, numbers.Number):
-            new_dict[key] = t.as_tensor(value)
+            new_dict[key] = t.as_tensor(value, device=device)
         # bools are an instance of number, but not np.bool_...
         elif isinstance(value, np.bool_):
-            new_dict[key] = t.as_tensor(value)
+            new_dict[key] = t.as_tensor(value, device=device)
         elif isinstance(value, np.ndarray):
-            new_dict[key] = t.as_tensor(value)
+            new_dict[key] = t.as_tensor(value, device=device)
         elif t.is_tensor(value):
-            new_dict[key] = value
+            new_dict[key] = value.to(device=device)
         elif isinstance(value, str):
             new_dict[key] = value
         elif isinstance(value, Mapping):
-            new_dict[key] = nested_dict_to_numpy(value)
+            new_dict[key] = nested_dict_to_torch(value, dtype=dtype)
         else:
-            raise ValueError(f'{value} is not a number, numpy array, torch tensor, or mapping')
+            raise ValueError(f'{value} is not a number, numpy array, torch tensor, string, or mapping')
 
     return new_dict

tests/tools/test_data.py

@@ -289,3 +289,55 @@ def test_add_ptycho_translations(tmp_path):
     assert np.allclose(-translations, read_translations_1)
     assert np.allclose(-translations, read_translations_2)
     assert np.allclose(-translations, read_translations_3)
+
+
+def test_nested_dict_to_h5(tmp_path):
+    ### Tests both nested_dict_to_h5 and h5_to_nested_dict
+    example_tensor = t.as_tensor(np.array([1,4.5,7]))
+    example_array = np.ones([10,20,30])
+    example_scalar = 4.5
+    example_single_element_array = np.array([0.3])
+    example_string = 'testing'
+
+    test_dict_1 = {}
+    test_dict_2 = {
+        'example_tensor': example_tensor,
+        'example_array': example_array,
+        'example_scalar': example_scalar,
+        'example_single_element_array': example_single_element_array,
+        'example_string': example_string
+    }
+    test_dict_3 = {
+        'example_array': example_array,
+        'example_string': example_string,
+        'example_dict': test_dict_2
+    }
+
+    def check_dict_equality(truth, to_test):
+        for key in truth.keys():
+            if type(truth[key]) == type(test_dict_2):
+                check_dict_equality(truth[key], to_test[key])
+            elif type(truth[key]) == type(example_tensor):
+                assert type(to_test[key]) == type(example_array)
+                assert np.allclose(truth[key].numpy(), to_test[key])
+            elif type(truth[key]) == type(example_array):
+                assert type(to_test[key]) == type(example_array)
+                assert np.allclose(truth[key], to_test[key])
+            elif type(truth[key]) == type(example_scalar):
+                assert truth[key] == to_test[key]
+            elif type(truth[key]) == type(example_string):
+                assert truth[key] == to_test[key]
+
+
+    for test_dict in [test_dict_1, test_dict_2, test_dict_3]:
+        filename = tmp_path / 'example_dataset.h5'
+        data.nested_dict_to_h5(filename, test_dict)
+        roundtrip = data.h5_to_nested_dict(filename)
+        check_dict_equality(test_dict, roundtrip)
+                
+    
+def test_h5_to_nested_dict(test_ptycho_cxis):
+    for cxi, expected in test_ptycho_cxis:
+        # Just test that it runs without errors for these ones.
+        # A round-trip test is in test_nested_dict_to_h5
+        d = data.h5_to_nested_dict(cxi)