diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py index ae434ba..b69d0c0 100644 --- a/src/cdtools/tools/data/data.py +++ b/src/cdtools/tools/data/data.py @@ -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,7 +832,7 @@ 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 ---------- @@ -842,7 +842,7 @@ def h5_to_nested_dict(h5_file): Returns ------- d : dict - 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 + 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 @@ -884,19 +884,15 @@ def nested_dict_to_numpy(d): Returns ------- - d_out : dict + 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() @@ -927,19 +923,15 @@ def nested_dict_to_torch(d, device=None): Returns ------- - d_out : dict + 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, device=device) # bools are an instance of number, but not np.bool_... - elif isinstance(value, np.bool_): - new_dict[key] = t.as_tensor(value, device=device) - elif isinstance(value, np.ndarray): + 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.to(device=device) diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index b8d964c..e637509 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -296,17 +296,15 @@ def test_nested_dict_to_h5(tmp_path, example_nested_dicts): def check_dict_equality(truth, to_test): for key in truth.keys(): - if type(truth[key]) == type({'a': 1}): + if isinstance(truth[key], dict): check_dict_equality(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])) + elif t.is_tensor(truth[key]): + assert isinstance(to_test[key], np.ndarray) 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])) + elif isinstance(truth[key], np.ndarray): + assert isinstance(to_test[key], np.ndarray) 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'): + elif isinstance(truth[key], (float, int, str)): assert truth[key] == to_test[key] else: assert 0 @@ -328,17 +326,15 @@ 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}): + if isinstance(truth[key], dict): 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])) + elif t.is_tensor(truth[key]): + assert isinstance(to_test[key], np.ndarray) 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])) + elif isinstance(truth[key], np.ndarray): + assert isinstance(to_test[key], np.ndarray) 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'): + elif isinstance(truth[key], (float, int, str)): assert truth[key] == to_test[key] else: assert 0 @@ -352,17 +348,15 @@ 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}): + if isinstance(truth[key], dict): 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]))) + elif t.is_tensor(truth[key]): + assert t.is_tensor(to_test[key]) 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]))) + 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 type(truth[key]) == type(1.3) or \ - type(truth[key]) == type(1) or \ - type(truth[key]) == type('1'): + elif isinstance(truth[key], (float, int, str)): assert truth[key] == to_test[key] else: assert 0