diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py index 33c272c..ae434ba 100644 --- a/src/cdtools/tools/data/data.py +++ b/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') diff --git a/tests/conftest.py b/tests/conftest.py index ca3a2c7..51a58dc 100644 --- a/tests/conftest.py +++ b/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] diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index 8e1751f..b8d964c 100644 --- a/tests/tools/test_data.py +++ b/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)