From 0e5d72b5a72be9724035a0ee97a0b42af96e8e8d Mon Sep 17 00:00:00 2001 From: Thomas Warford Date: Sat, 9 Nov 2024 23:58:41 +0000 Subject: [PATCH] add head_key=head_name to info dict of atoms regardless of isolated_atom_config_value --- mace/data/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mace/data/utils.py b/mace/data/utils.py index bb8e5448..59b868ed 100644 --- a/mace/data/utils.py +++ b/mace/data/utils.py @@ -265,7 +265,6 @@ def load_from_xyz( atoms_without_iso_atoms = [] for idx, atoms in enumerate(atoms_list): - atoms.info[head_key] = head_name isolated_atom_config = ( len(atoms) == 1 and atoms.info.get("config_type") == "IsolatedAtom" ) @@ -288,6 +287,9 @@ def load_from_xyz( if not keep_isolated_atoms: atoms_list = atoms_without_iso_atoms + for atoms in atoms_list: + atoms.info[head_key] = head_name + configs = config_from_atoms_list( atoms_list, config_type_weights=config_type_weights,