feat(lammps): generate masses for lammps/lmp export

Prefer explicitly stored masses when available, and otherwise infer
per-type masses from atom_names when all names are valid element
symbols. Keep the previous behavior for unknown type names so exports
do not emit unsafe Masses sections.

Add regression tests covering both known-element and unknown-type cases,
and make the new tests independent of the current working directory.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.4)

Author: njzjz-bot

dpdata/lammps/lmp.py

@@ -3,6 +3,8 @@
 
 import numpy as np
 
+from dpdata.periodic_table import Element, ELEMENTS
+
 ptr_float_fmt = "%15.10f"
 ptr_int_fmt = "%6d"
 ptr_key_fmt = "%15s"
@@ -484,6 +486,34 @@ def rotate_to_lower_triangle(
     return cell, coord
 
 
+def _get_lammps_masses(system) -> np.ndarray | None:
+    """Get masses for the LAMMPS ``Masses`` section.
+
+    Prefer explicitly stored masses when available. Otherwise, infer masses from
+    ``atom_names`` when all names are valid chemical element symbols.
+
+    Parameters
+    ----------
+    system : dict
+        System data dictionary
+
+    Returns
+    -------
+    np.ndarray or None
+        Per-type masses aligned with ``atom_names``. Returns ``None`` when the
+        masses cannot be determined safely.
+    """
+    atom_names = system["atom_names"]
+    masses = system.get("masses")
+    if masses is not None and len(masses) >= len(atom_names):
+        return np.asarray(masses[: len(atom_names)], dtype=float)
+
+    if not all(name in ELEMENTS for name in atom_names):
+        return None
+
+    return np.array([Element(name).mass for name in atom_names], dtype=float)
+
+
 def from_system_data(system, f_idx=0):
     ret = ""
     ret += "\n"
@@ -514,6 +544,16 @@ def from_system_data(system, f_idx=0):
         cell[2][1],
     )  # noqa: UP031
     ret += "\n"
+
+    masses = _get_lammps_masses(system)
+    if masses is not None:
+        ret += "Masses\n"
+        ret += "\n"
+        mass_fmt = ptr_int_fmt + " " + ptr_float_fmt + " # %s\n"  # noqa: UP031
+        for ii, (mass, atom_name) in enumerate(zip(masses, system["atom_names"])):
+            ret += mass_fmt % (ii + 1, mass, atom_name)
+        ret += "\n"
+
     ret += "Atoms # atomic\n"
     ret += "\n"
     coord_fmt = (

tests/test_lammps_lmp_dump.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import tempfile
 import unittest
 
 import numpy as np
@@ -9,6 +10,9 @@
 
 from dpdata.lammps.lmp import rotate_to_lower_triangle
 
+TEST_DIR = os.path.dirname(__file__)
+POSCAR_CONF_LMP = os.path.join(TEST_DIR, "poscars", "conf.lmp")
+
 
 class TestLmpDump(unittest.TestCase, TestPOSCARoh):
     def setUp(self):
@@ -100,5 +104,30 @@ def test_negative_diagonal(self):
         )
 
 
+class TestLmpDumpMasses(unittest.TestCase):
+    def test_dump_known_elements_writes_masses(self):
+        system = dpdata.System(POSCAR_CONF_LMP, type_map=["O", "H"])
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output = os.path.join(tmpdir, "tmp_masses.lmp")
+            system.to_lammps_lmp(output)
+            with open(output) as f:
+                content = f.read()
+
+        self.assertIn("Masses\n", content)
+        self.assertIn("     1   15.9994000000 # O", content)
+        self.assertIn("     2    1.0079400000 # H", content)
+        self.assertLess(content.index("Masses\n"), content.index("Atoms # atomic\n"))
+
+    def test_dump_unknown_types_skips_masses(self):
+        system = dpdata.System(POSCAR_CONF_LMP)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output = os.path.join(tmpdir, "tmp_unknown_types.lmp")
+            system.to_lammps_lmp(output)
+            with open(output) as f:
+                content = f.read()
+
+        self.assertNotIn("Masses\n", content)
+
+
 if __name__ == "__main__":
     unittest.main()