feat(abacus): support STRU files with repeated atom type labels

dpdata/formats/abacus/stru.py

@@ -441,6 +441,46 @@ def right_hand_rule(
     return cell, coord
 
 
+def validate_duplicate_species(atom_names, masses, pp_files, orb_files):
+    """Check that duplicate atom species definitions are consistent.
+
+    get_frame_from_stru merges repeated labels and keeps the metadata of the
+    first matching row for masses, pp_files, and orb_files. If a duplicate
+    row conflicts with the first one, raise an error instead of silently
+    discarding the conflicting metadata.
+
+    Args:
+        atom_names (list): list of atom names.
+        masses (list): list of atomic masses.
+        pp_files (list): list of pseudo potential files.
+        orb_files (list): list of orbital files.
+
+    Raises
+    ------
+        RuntimeError: if duplicate species have conflicting metadata.
+    """
+    for name in dict.fromkeys(atom_names):
+        indices = [j for j in range(len(atom_names)) if atom_names[j] == name]
+        if len(indices) < 2:
+            continue
+        ref_mass = masses[indices[0]]
+        ref_pp = pp_files[indices[0]]
+        ref_orb = orb_files[indices[0]] if orb_files else None
+        for j in indices[1:]:
+            if not np.isclose(masses[j], ref_mass):
+                raise RuntimeError(
+                    f"Conflicting duplicate species '{name}': mass {masses[j]} != {ref_mass}"
+                )
+            if pp_files[j] != ref_pp:
+                raise RuntimeError(
+                    f"Conflicting duplicate species '{name}': pp_file {pp_files[j]} != {ref_pp}"
+                )
+            if ref_orb is not None and orb_files[j] != ref_orb:
+                raise RuntimeError(
+                    f"Conflicting duplicate species '{name}': orb_file {orb_files[j]} != {ref_orb}"
+                )
+
+
 def get_frame_from_stru(stru):
     """Read the ABACUS STRU file and return the dpdata frame.
 
@@ -492,27 +532,46 @@ def get_frame_from_stru(stru):
         blocks["ATOMIC_POSITIONS"], atom_names, celldm, cell
     )
 
+    validate_duplicate_species(atom_names, masses, pp_files, orb_files)
+
     cell, coords = right_hand_rule(cell, coords)
+    uniq_name = []
+    uniq_atom_num = []
+    for i in atom_names:
+        if i not in uniq_name:
+            uniq_name.append(i)
+            uniq_atom_num.append(
+                sum(
+                    [
+                        atom_numbs[j]
+                        for j in range(len(atom_names))
+                        if atom_names[j] == i
+                    ]
+                )
+            )
     data = {
-        "atom_names": atom_names,
-        "atom_numbs": atom_numbs,
+        "atom_names": uniq_name,
+        "atom_numbs": uniq_atom_num,
         "atom_types": np.array(
-            [i for i in range(len(atom_numbs)) for j in range(atom_numbs[i])]
+            [
+                uniq_name.index(atom_names[i])
+                for i in range(len(atom_numbs))
+                for j in range(atom_numbs[i])
+            ]
         ),
-        "masses": np.array(masses),
-        "pp_files": pp_files,
+        "masses": np.array([masses[atom_names.index(i)] for i in uniq_name]),
+        "pp_files": [pp_files[atom_names.index(i)] for i in uniq_name],
         "cells": np.array([cell]),
         "coords": np.array([coords]),
     }
     if len(mags) > 0:
         data["spins"] = np.array([mags])
     if len(orb_files) > 0:
-        data["orb_files"] = orb_files
+        data["orb_files"] = [orb_files[atom_names.index(i)] for i in uniq_name]
     if len(dpks_descriptor) > 0:
         data["dpks_descriptor"] = dpks_descriptor[0].strip()
     if len(move) > 0:
         data["move"] = np.array([move])
-
     return data
 
 

tests/abacus.scf/STRU-conflict-atomtype.ch4

@@ -0,0 +1,37 @@
+#This is the atom card containing all the information
+#about the lattice structure.
+
+ATOMIC_SPECIES
+C 1.000 C_ONCV_PBE-1.0.upf 	#Element, Mass, Pseudopotential
+H 1.000 H_ONCV_PBE-1.0.upf 
+H 2.000 H_ONCV_PBE-1.0.upf 
+
+NUMERICAL_ORBITAL
+c.orb
+h.orb
+h.orb
+
+LATTICE_CONSTANT
+10  					#Lattice constant
+
+LATTICE_VECTORS
+1 0.0 0.0 			#Lattice vector 1
+0.0 1 0.0 			#Lattice vector 2
+0.0 0.0 1 			#Lattice vector 3
+
+ATOMIC_POSITIONS
+Cartesian 				#Cartesian(Unit is LATTICE_CONSTANT)
+C 						#Name of element	
+0.0						#Magnetic for this element.
+1				#Number of atoms
+0.981274803    0.861285385    0.838442496 1 1 1
+H
+0.0
+2
+1.023557202    0.758025625    0.66351336 0 0 0
+0.78075702    0.889445935    0.837363468 1 0 1
+H
+0.0
+2
+1.064091613    1.043438905    0.840995502 1 0 1
+1.039321214    0.756530859    1.009609207 0 1 1

tests/abacus.scf/STRU-repeat-atomtype.ch4

@@ -0,0 +1,37 @@
+#This is the atom card containing all the information
+#about the lattice structure.
+
+ATOMIC_SPECIES
+C 1.000 C_ONCV_PBE-1.0.upf 	#Element, Mass, Pseudopotential
+H 1.000 H_ONCV_PBE-1.0.upf 
+H 1.000 H_ONCV_PBE-1.0.upf 
+
+NUMERICAL_ORBITAL
+c.orb
+h.orb
+h.orb
+
+LATTICE_CONSTANT
+10  					#Lattice constant
+
+LATTICE_VECTORS
+1 0.0 0.0 			#Lattice vector 1
+0.0 1 0.0 			#Lattice vector 2
+0.0 0.0 1 			#Lattice vector 3
+
+ATOMIC_POSITIONS
+Cartesian 				#Cartesian(Unit is LATTICE_CONSTANT)
+C 						#Name of element	
+0.0						#Magnetic for this element.
+1				#Number of atoms
+0.981274803    0.861285385    0.838442496 1 1 1
+H
+0.0
+2
+1.023557202    0.758025625    0.66351336 0 0 0
+0.78075702    0.889445935    0.837363468 1 0 1
+H
+0.0
+2
+1.064091613    1.043438905    0.840995502 1 0 1
+1.039321214    0.756530859    1.009609207 0 1 1

tests/test_abacus_stru_dump.py

@@ -262,6 +262,70 @@ def test_dump_chaotic_atomic_species(self):
         self.assertTrue(ref_c in lines)
 
 
+class TestStruRepeatAtomtype(unittest.TestCase):
+    def test_read_stru_with_repeat_atomtype(self):
+        sys_tmp = dpdata.System("abacus.scf/STRU-repeat-atomtype.ch4", fmt="stru")
+        self.assertEqual(sys_tmp.data["atom_names"], ["C", "H"])
+        self.assertEqual(sys_tmp.data["atom_numbs"], [1, 4])
+        self.assertEqual(sys_tmp.data["atom_types"].tolist(), [0, 1, 1, 1, 1])
+        self.assertEqual(sys_tmp.data["masses"].tolist(), [1.0, 1.0])
+        self.assertEqual(
+            sys_tmp.data["pp_files"], ["C_ONCV_PBE-1.0.upf", "H_ONCV_PBE-1.0.upf"]
+        )
+        self.assertEqual(sys_tmp.data["orb_files"], ["c.orb", "h.orb"])
+        self.assertEqual(sys_tmp.data["coords"].shape, (1, 5, 3))
+
+    def test_dump_stru_with_repeat_atomtype(self):
+        sys_tmp = dpdata.System("abacus.scf/STRU-repeat-atomtype.ch4", fmt="stru")
+        sys_tmp.to("stru", "STRU_tmp", mass=[12, 1])
+        with open("STRU_tmp") as f:
+            c = f.read()
+        self.assertTrue(
+            "ATOMIC_SPECIES\nC 12.000 C_ONCV_PBE-1.0.upf\nH 1.000 H_ONCV_PBE-1.0.upf"
+            in c
+        )
+        self.assertTrue("H\n0.0\n4\n" in c)
+        os.remove("STRU_tmp")
+
+    def test_read_stru_with_conflicting_duplicate_atomtype(self):
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Conflicting duplicate species 'H'.*mass 2.0 != 1.0",
+        ):
+            dpdata.System("abacus.scf/STRU-conflict-atomtype.ch4", fmt="stru")
+
+    def test_validate_duplicate_species(self):
+        from dpdata.formats.abacus.stru import validate_duplicate_species
+
+        validate_duplicate_species(
+            ["C", "H", "H"],
+            [1.0, 1.0, 1.0],
+            ["C.upf", "H.upf", "H.upf"],
+            ["c.orb", "h.orb", "h.orb"],
+        )
+        with self.assertRaisesRegex(RuntimeError, "mass"):
+            validate_duplicate_species(
+                ["C", "H", "H"],
+                [1.0, 1.0, 2.0],
+                ["C.upf", "H.upf", "H.upf"],
+                ["c.orb", "h.orb", "h.orb"],
+            )
+        with self.assertRaisesRegex(RuntimeError, "pp_file"):
+            validate_duplicate_species(
+                ["C", "H", "H"],
+                [1.0, 1.0, 1.0],
+                ["C.upf", "H.upf", "O.upf"],
+                ["c.orb", "h.orb", "h.orb"],
+            )
+        with self.assertRaisesRegex(RuntimeError, "orb_file"):
+            validate_duplicate_species(
+                ["C", "H", "H"],
+                [1.0, 1.0, 1.0],
+                ["C.upf", "H.upf", "H.upf"],
+                ["c.orb", "h.orb", "o.orb"],
+            )
+
+
 class TestABACUSParseStru(unittest.TestCase):
     def test_parse_pos_oneline(self):
         pos, move, velocity, magmom, angle1, angle2, constrain, lambda1 = (