diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py index 556aa38..b69d0c0 100644 --- a/src/cdtools/tools/data/data.py +++ b/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 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 bf048e5..e637509 100644 --- a/tests/tools/test_data.py +++ b/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)