Add tests for nested_dict_to_numpy and nested_dict_to_torch

src/cdtools/tools/data/data.py

@@ -946,7 +946,7 @@ def nested_dict_to_torch(d, device=None):
         elif isinstance(value, str):
             new_dict[key] = value
         elif isinstance(value, Mapping):
-            new_dict[key] = nested_dict_to_torch(value, dtype=dtype)
+            new_dict[key] = nested_dict_to_torch(value, device=device)
         else:
             raise ValueError(f'{value} is not a number, numpy array, torch tensor, string, or mapping')
 

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

@@ -291,45 +291,27 @@ def test_add_ptycho_translations(tmp_path):
     assert np.allclose(-translations, read_translations_3)
 
 
-def test_nested_dict_to_h5(tmp_path):
+def test_nested_dict_to_h5(tmp_path, example_nested_dicts):
     ### 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):
+            if type(truth[key]) == type({'a': 1}):
                 check_dict_equality(truth[key], to_test[key])
-            elif type(truth[key]) == type(example_tensor):
-                assert type(to_test[key]) == type(example_array)
+            elif type(truth[key]) == type(t.as_tensor(np.array([1]))):
+                assert type(to_test[key]) == type(np.array([1]))
                 assert np.allclose(truth[key].numpy(), to_test[key])
-            elif type(truth[key]) == type(example_array):
-                assert type(to_test[key]) == type(example_array)
+            elif type(truth[key]) == type(np.array([1])):
+                assert type(to_test[key]) == type(np.array([1]))
                 assert np.allclose(truth[key], to_test[key])
-            elif type(truth[key]) == type(example_scalar):
+            elif type(truth[key]) == type(1.3) or \
+                 type(truth[key]) == type(1) or \
+                 type(truth[key]) == type('1'):
                 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]:
+            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)
@@ -341,3 +323,50 @@ def test_h5_to_nested_dict(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 type(truth[key]) == type({'a': 1}):
+                check_dict_numpyness(truth[key], to_test[key])
+            elif type(truth[key]) == type(t.as_tensor(np.array([1]))):
+                assert type(to_test[key]) == type(np.array([1]))
+                assert np.allclose(truth[key].numpy(), to_test[key])
+            elif type(truth[key]) == type(np.array([1])):
+                assert type(to_test[key]) == type(np.array([1]))
+                assert np.allclose(truth[key], to_test[key])
+            elif type(truth[key]) == type(1.3) or \
+                 type(truth[key]) == type(1) or \
+                 type(truth[key]) == type('1'):
+                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 type(truth[key]) == type({'a': 1}):
+                check_dict_torchiness(truth[key], to_test[key])
+            elif type(truth[key]) == type(t.as_tensor(np.array([1]))):
+                assert type(to_test[key]) == type(t.as_tensor(np.array([1])))
+                assert t.allclose(truth[key], to_test[key])
+            elif type(truth[key]) == type(np.array([1])):
+                assert type(to_test[key]) == type(t.as_tensor(np.array([1])))
+                assert t.allclose(t.as_tensor(truth[key]), to_test[key])
+            elif type(truth[key]) == type(1.3) or \
+                 type(truth[key]) == type(1) or \
+                 type(truth[key]) == type('1'):
+                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)