From 7e073a437cc2a5164a90116e93cbafcf3bad896b Mon Sep 17 00:00:00 2001
From: Cedric Conday <277679649+CedricConday@users.noreply.github.com>
Date: Mon, 6 Jul 2026 23:43:12 +0000
Subject: [PATCH] fix: build pymatgen Molecule species from atom_types order
 (#994)

PyMatgenMoleculeFormat.to_system() built the species list by expanding
grouped atom_names/atom_numbs, but passed coords in atom_types order. When
atoms are not grouped by type, species were assigned to the wrong
coordinates. Build species per-atom from atom_types, matching the coord
order and the sibling Structure exporter (line 37).

Adds a regression test that fails on master (['H','H','O'] vs ['H','O','H'])
and passes with this change.
---
 dpdata/plugins/pymatgen.py      |  4 +---
 tests/test_pymatgen_molecule.py | 21 +++++++++++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/dpdata/plugins/pymatgen.py b/dpdata/plugins/pymatgen.py
index e9c91fc8a..3a17f5aba 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 e6a1b5ee5..27e63c111 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()