Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

hotfix for string types #146

Merged
merged 1 commit into from
Oct 6, 2024
Merged

hotfix for string types #146

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions tests/test_data_types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,17 @@ def test_dict_data(tmp_path):
io.append(molecule)
molecule.info["b"] = {"a": 1, "b": 2, "c": 3, "d": 4}
io.append(molecule)
# molecule.info["elements"] = {"H": "H", 2: "O"}
# NOTE! Mixed types in dict keys are currently not supported
# TODO json dump / load?
molecule.info["elements"] = {"H": "H", "O": "O"}
io.append(molecule)

assert io[0].info["test"] == {"a": 1, "b": 2}
assert io[1].info["test"] == {"a": 1, "b": 2, "c": 3}
assert io[2].info["b"] == {"a": 1, "b": 2, "c": 3, "d": 4}
# assert io[3].info["elements"] == {"H": "H", 2: "O"}
assert io[3].info["elements"] == {"H": "H", "O": "O"}


def test_list_data(tmp_path):
Expand All @@ -81,10 +88,13 @@ def test_list_data(tmp_path):
io.append(molecule)
molecule.info["b"] = [1, 2, 3, 4]
io.append(molecule)
molecule.info["elements"] = ["H", "O"]
io.append(molecule)

npt.assert_array_equal(io[0].info["test"], [1, 2])
npt.assert_array_equal(io[1].info["test"], [1, 2, 3])
npt.assert_array_equal(io[2].info["b"], [1, 2, 3, 4])
npt.assert_array_equal(io[3].info["elements"], ["H", "O"])


def test_multiple_molecules_with_diff_length_dicts(tmp_path):
Expand Down
11 changes: 10 additions & 1 deletion znh5md/format.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,7 @@ def extract_atoms_data(atoms: Atoms, use_ase_calc: bool = True) -> ASEData: # n
raise ValueError(f"Key {key} is reserved for ASE calculator results.")
if key not in ASE_TO_H5MD and key not in CustomINFOData.__members__:
if isinstance(value, str):
# TODO: needs to be tested!
if len(value) > NUMPY_STRING_DTYPE.itemsize:
raise ValueError(f"String {key} is too long to be stored.")
info_data[key] = np.array(value.encode(), dtype=NUMPY_STRING_DTYPE)
Expand Down Expand Up @@ -298,7 +299,15 @@ def _combine_dicts(dicts: List[Dict[str, np.ndarray]]) -> Dict[str, np.ndarray]:
data = []
for d in dicts:
if key in d:
data.append(np.array(d[key]))
# import warnings
# warnings.warn(f"Key {key}:{type(d[key])}")
if isinstance(d[key], list):
if any(isinstance(x, str) for x in d[key]):
data.append(np.array(d[key], dtype=NUMPY_STRING_DTYPE))
else:
data.append(np.array(d[key]))
else:
data.append(np.array(d[key]))
else:
dims = dicts[0][key].ndim
# Create an array with the appropriate number of dimensions.
Expand Down
Loading