diff --git a/dpdata/plugins/pymatgen.py b/dpdata/plugins/pymatgen.py index e9c91fc8..3a17f5ab 100644 --- a/dpdata/plugins/pymatgen.py +++ b/dpdata/plugins/pymatgen.py @@ -66,9 +66,7 @@ def to_system(self, data, **kwargs): except ModuleNotFoundError as e: raise ImportError("No module pymatgen.Molecule") from e - species = [] - for name, numb in zip(data["atom_names"], data["atom_numbs"]): - species.extend([name] * numb) + species = [data["atom_names"][tt] for tt in data["atom_types"]] data = dpdata.system.remove_pbc(data) for ii in range(np.array(data["coords"]).shape[0]): molecule = Molecule(species, data["coords"][ii]) diff --git a/tests/test_pymatgen_molecule.py b/tests/test_pymatgen_molecule.py index e6a1b5ee..27e63c11 100644 --- a/tests/test_pymatgen_molecule.py +++ b/tests/test_pymatgen_molecule.py @@ -36,6 +36,27 @@ def test_poscar_to_molecule(self): max_dist_1 = np.max(np.linalg.norm(dist, axis=1)) self.assertAlmostEqual(max_dist_0, max_dist_1) + def test_ungrouped_atom_types_species_match_coords(self): + # Regression for gh-994: species must follow atom_types order, not the + # grouped atom_names/atom_numbs order, otherwise sites are assigned the + # wrong element when atoms are not grouped by type. + system = dpdata.System( + data={ + "atom_names": ["H", "O"], + "atom_numbs": [2, 1], + "atom_types": np.array([0, 1, 0]), + "orig": np.zeros(3), + "cells": np.eye(3).reshape(1, 3, 3) * 20, + "coords": np.array( + [[[0.0, 0.0, 0.0], [9.0, 0.0, 0.0], [1.0, 0.0, 0.0]]] + ), + } + ) + mol = system.to("pymatgen/molecule")[0] + self.assertEqual([str(site.specie) for site in mol], ["H", "O", "H"]) + for site, coord in zip(mol, system["coords"][0]): + np.testing.assert_allclose(site.coords, coord) + if __name__ == "__main__": unittest.main()