Merge pull request #25 from cdtools-developers/bugfix/h5_strings

Bugfix/h5 strings

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
     ----------
@@ -813,7 +813,7 @@ def nested_dict_to_h5(h5_file, d):
     
     for key in d.keys():
         value = d[key]
-        if isinstance(value, numbers.Number):
+        if isinstance(value, (numbers.Number, np.bool_)):
             arr = np.array(value)
             h5_file.create_dataset(key, data=arr)
         elif isinstance(value, np.ndarray):
@@ -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
+    """Loads a nested dictionary from 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 loaded into 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,16 +875,24 @@ 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
+    -------
+    new_dict : dict
+        A new dictionary with all array like objects sent to numpy 
+    """
+    
     new_dict = {}
     for key in d.keys():
         value = d[key]
-        if isinstance(value, numbers.Number):
-            new_dict[key] = value
         # bools are an instance of number, but not np.bool_...
-        elif isinstance(value, np.bool_):
-            new_dict[key] = value
-        elif isinstance(value, np.ndarray):
+        if isinstance(value, (numbers.Number, np.bool_, np.ndarray)):
             new_dict[key] = value
         elif t.is_tensor(value):
             new_dict[key] = value.cpu().numpy()
@@ -888,29 +901,45 @@ 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
+    -------
+    new_dict : 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)
         # bools are an instance of number, but not np.bool_...
-        elif isinstance(value, np.bool_):
-            new_dict[key] = t.as_tensor(value)
-        elif isinstance(value, np.ndarray):
-            new_dict[key] = t.as_tensor(value)
+        if isinstance(value, (numbers.Number, np.bool_, np.ndarray)):
+            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, device=device)
         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/conftest.py

@@ -1,4 +1,5 @@
 import numpy as np
+import torch as t
 import h5py
 import pytest
 import datetime
@@ -365,4 +366,26 @@ def lab_ptycho_cxi(pytestconfig):
         '/examples/example_data/lab_ptycho_data.cxi'
 
 
+@pytest.fixture(scope='module')
+def example_nested_dicts(pytestconfig):
+    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
+    }
+
+    return [test_dict_1, test_dict_2, test_dict_3]

tests/tools/test_data.py

@@ -289,3 +289,78 @@ 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, example_nested_dicts):
+    ### Tests both nested_dict_to_h5 and h5_to_nested_dict
+
+    def check_dict_equality(truth, to_test):
+        for key in truth.keys():
+            if isinstance(truth[key], dict):
+                check_dict_equality(truth[key], to_test[key])
+            elif t.is_tensor(truth[key]):
+                assert isinstance(to_test[key], np.ndarray)
+                assert np.allclose(truth[key].numpy(), to_test[key])
+            elif isinstance(truth[key], np.ndarray):
+                assert isinstance(to_test[key], np.ndarray)
+                assert np.allclose(truth[key], to_test[key])
+            elif isinstance(truth[key], (float, int, str)):
+                assert truth[key] == to_test[key]
+            else:
+                assert 0
+            
+    for test_dict in example_nested_dicts:
+        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)
+
+def test_nested_dict_to_numpy(example_nested_dicts):
+
+    def check_dict_numpyness(truth, to_test): 
+        for key in truth.keys():
+            if isinstance(truth[key], dict):
+                check_dict_numpyness(truth[key], to_test[key])
+            elif t.is_tensor(truth[key]):
+                assert isinstance(to_test[key], np.ndarray)
+                assert np.allclose(truth[key].numpy(), to_test[key])
+            elif isinstance(truth[key], np.ndarray):
+                assert isinstance(to_test[key], np.ndarray)
+                assert np.allclose(truth[key], to_test[key])
+            elif isinstance(truth[key], (float, int, str)):
+                assert truth[key] == to_test[key]
+            else:
+                assert 0
+
+    for test_dict in example_nested_dicts:
+        numpy_dict = data.nested_dict_to_numpy(test_dict)
+        check_dict_numpyness(test_dict, numpy_dict)            
+        
+    
+def test_nested_dict_to_torch(example_nested_dicts):
+
+    def check_dict_torchiness(truth, to_test): 
+        for key in truth.keys():
+            if isinstance(truth[key], dict):
+                check_dict_torchiness(truth[key], to_test[key])
+            elif t.is_tensor(truth[key]):
+                assert t.is_tensor(to_test[key])
+                assert t.allclose(truth[key], to_test[key])
+            elif isinstance(truth[key], np.ndarray):
+                assert t.is_tensor(to_test[key])
+                assert t.allclose(t.as_tensor(truth[key]), to_test[key])
+            elif isinstance(truth[key], (float, int, str)):
+                assert truth[key] == to_test[key]
+            else:
+                assert 0
+
+    for test_dict in example_nested_dicts:
+        torch_dict = data.nested_dict_to_torch(test_dict)
+        check_dict_torchiness(test_dict, torch_dict)