fix: use total atom counts for defect energies

Fixes #1901

Coding-Agent: Codex

Codex-Version: codex-cli 0.149.0

Model: gpt-5.6-sol

Reasoning-Effort: xhigh

Author: njzjz-bot

dpgen/auto_test/Interstitial.py

@@ -8,7 +8,7 @@
 
 import dpgen.auto_test.lib.abacus as abacus
 import dpgen.auto_test.lib.lammps as lammps
-from dpgen.auto_test.Property import Property
+from dpgen.auto_test.Property import Property, _total_atom_count
 from dpgen.auto_test.refine import make_refine
 from dpgen.auto_test.reproduce import make_repro, post_repro
 
@@ -514,14 +514,14 @@ def _compute_lower(self, output_file, all_tasks, all_res):
                 idid += 1
                 structure_dir = os.path.basename(ii)
                 task_result = loadfn(all_res[idid])
-                natoms = task_result["atom_numbs"][0]
+                natoms = _total_atom_count(task_result)
                 equi_path = os.path.abspath(
                     os.path.join(
                         os.path.dirname(output_file), "../relaxation/relax_task"
                     )
                 )
                 equi_result = loadfn(os.path.join(equi_path, "result.json"))
-                equi_epa = equi_result["energies"][-1] / equi_result["atom_numbs"][0]
+                equi_epa = equi_result["energies"][-1] / _total_atom_count(equi_result)
                 evac = task_result["energies"][-1] - equi_epa * natoms
 
                 supercell_index = loadfn(os.path.join(ii, "supercell.json"))

dpgen/auto_test/Property.py

@@ -8,6 +8,11 @@
 from dpgen.auto_test.calculator import make_calculator
 
 
+def _total_atom_count(result):
+    """Return the total number of atoms across all species in a result."""
+    return sum(result["atom_numbs"])
+
+
 class Property(ABC):
     @abstractmethod
     def __init__(self, parameter):

dpgen/auto_test/Vacancy.py

@@ -8,7 +8,7 @@
 
 import dpgen.auto_test.lib.abacus as abacus
 from dpgen import dlog
-from dpgen.auto_test.Property import Property
+from dpgen.auto_test.Property import Property, _total_atom_count
 from dpgen.auto_test.refine import make_refine
 from dpgen.auto_test.reproduce import make_repro, post_repro
 
@@ -231,14 +231,14 @@ def _compute_lower(self, output_file, all_tasks, all_res):
                 idid += 1
                 structure_dir = os.path.basename(ii)
                 task_result = loadfn(all_res[idid])
-                natoms = task_result["atom_numbs"][0]
+                natoms = _total_atom_count(task_result)
                 equi_path = os.path.abspath(
                     os.path.join(
                         os.path.dirname(output_file), "../relaxation/relax_task"
                     )
                 )
                 equi_result = loadfn(os.path.join(equi_path, "result.json"))
-                equi_epa = equi_result["energies"][-1] / equi_result["atom_numbs"][0]
+                equi_epa = equi_result["energies"][-1] / _total_atom_count(equi_result)
                 evac = task_result["energies"][-1] - equi_epa * natoms
 
                 supercell_index = loadfn(os.path.join(ii, "supercell.json"))

tests/auto_test/test_property_atom_count.py

@@ -0,0 +1,12 @@
+import unittest
+
+from dpgen.auto_test.Property import _total_atom_count
+
+
+class TestPropertyAtomCount(unittest.TestCase):
+    def test_counts_all_species(self):
+        self.assertEqual(_total_atom_count({"atom_numbs": [2, 3, 4]}), 9)
+
+
+if __name__ == "__main__":
+    unittest.main()