fix: build pymatgen Molecule species from atom_types order (#994)

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])

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()