diff --git a/dpgen/data/gen.py b/dpgen/data/gen.py index 27134ef64..003e7fda8 100644 --- a/dpgen/data/gen.py +++ b/dpgen/data/gen.py @@ -237,12 +237,24 @@ def poscar_scale_cartesian(str_in, scale): 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]: + elif "C" == lines[coord_type_line][0] or "c" == lines[coord_type_line][0]: lines = poscar_scale_cartesian(lines, scale) 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))