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