Enhance poscar_scale function (/data/gen.py)to detect Selective dynamics

dpgen/data/gen.py

@@ -216,7 +216,8 @@ def poscar_scale_direct(str_in, scale):
     return lines
 
 
-def poscar_scale_cartesian(str_in, scale):
+def poscar_scale_cartesian(str_in, scale, coordinate_start=8):
+    """Scale Cartesian coordinates while preserving Selective Dynamics flags."""
     lines = str_in.copy()
     numb_atoms = poscar_natoms(lines)
     # scale box
@@ -226,23 +227,36 @@ def poscar_scale_cartesian(str_in, scale):
         boxv = np.array(boxv) * scale
         lines[ii] = f"{boxv[0]:.16e} {boxv[1]:.16e} {boxv[2]:.16e}\n"
     # scale coord
-    for ii in range(8, 8 + numb_atoms):
-        cl = lines[ii].split()
-        cv = [float(ii) for ii in cl]
+    for ii in range(coordinate_start, coordinate_start + numb_atoms):
+        fields = lines[ii].split()
+        cv = [float(value) for value in fields[:3]]
         cv = np.array(cv) * scale
-        lines[ii] = f"{cv[0]:.16e} {cv[1]:.16e} {cv[2]:.16e}\n"
+        suffix = f" {' '.join(fields[3:])}" if len(fields) > 3 else ""
+        lines[ii] = f"{cv[0]:.16e} {cv[1]:.16e} {cv[2]:.16e}{suffix}\n"
     return lines
 
 
 def poscar_scale(poscar_in, poscar_out, scale):
     with open(poscar_in) as fin:
         lines = list(fin)
-    if "D" == lines[7][0] or "d" == lines[7][0]:
+
+    # Determine if "Selective dynamics" is present
+    if lines[7].strip().lower().startswith("s"):
+        coord_type_line = 8  # If present, coordinates type is on line 9
+    else:
+        coord_type_line = 7  # If not, coordinates type is on line 8
+
+    # Process according to the coordinate type
+    if "D" == lines[coord_type_line][0] or "d" == lines[coord_type_line][0]:
         lines = poscar_scale_direct(lines, scale)
-    elif "C" == lines[7][0] or "c" == lines[7][0]:
-        lines = poscar_scale_cartesian(lines, scale)
+    elif "C" == lines[coord_type_line][0] or "c" == lines[coord_type_line][0]:
+        lines = poscar_scale_cartesian(lines, scale, coord_type_line + 1)
     else:
-        raise RuntimeError("Unknow poscar style at line 7: %s" % lines[7])
+        raise RuntimeError(
+            f"Unknown poscar style at line {coord_type_line + 1}: {lines[coord_type_line]}"
+        )
+
+    # Write scaled positions back to output file
     with open(poscar_out, "w") as fout:
         fout.write("".join(lines))
 

tests/data/test_poscar_scale.py

@@ -0,0 +1,39 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+from dpgen.data.gen import poscar_scale
+
+
+class TestPoscarScale(unittest.TestCase):
+    def test_selective_cartesian_coordinates_and_flags(self):
+        """Selective Dynamics moves the coordinate start and appends T/F flags."""
+        poscar = """Mg
+1.0
+1 0 0
+0 1 0
+0 0 1
+Mg
+1
+Selective dynamics
+Cartesian
+0.25 0.50 0.75 T F T
+"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            source = Path(tmpdir) / "POSCAR.in"
+            target = Path(tmpdir) / "POSCAR.out"
+            source.write_text(poscar)
+
+            poscar_scale(source, target, 2.0)
+            scaled = target.read_text().splitlines()
+
+        self.assertEqual(scaled[8], "Cartesian")
+        self.assertEqual(
+            scaled[9],
+            "5.0000000000000000e-01 1.0000000000000000e+00 "
+            "1.5000000000000000e+00 T F T",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()