一些小修改

GPUMD/swap_li.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+# -*- coding: ascii -*-
+"""
+Randomly swap one Li-Y pair in a VASP5 POSCAR and write N new files.
+- Keeps coordinate mode (Direct/Cartesian), Selective Dynamics flags, and Velocities.
+- Requires VASP5+ POSCAR (with element symbols line).
+"""
+
+import random
+from pathlib import Path
+
+
+
+def _is_ints(tokens):
+    try:
+        _ = [int(t) for t in tokens]
+        return True
+    except ValueError:
+        return False
+
+def _find_species_index(species, target):
+    t = target.lower()
+    for i, s in enumerate(species):
+        if s.lower() == t:
+            return i
+    raise ValueError("Element '%s' not found in species line: %s" % (target, " ".join(species)))
+
+def parse_poscar(lines):
+    if len(lines) < 8:
+        raise ValueError("POSCAR too short")
+
+    comment = lines[0].rstrip("\n")
+    scale = lines[1].rstrip("\n")
+    lv = [lines[2].rstrip("\n"), lines[3].rstrip("\n"), lines[4].rstrip("\n")]
+
+    i = 5
+    tokens = lines[i].split()
+    if _is_ints(tokens):
+        raise ValueError("VASP4 format (no element symbols line) is not supported.")
+    species = tokens
+    i += 1
+
+    counts_line = lines[i].rstrip("\n")
+    counts = [int(x) for x in counts_line.split()]
+    i += 1
+
+    selective = False
+    sel_line = None
+    if i < len(lines) and lines[i].strip().lower().startswith("s"):
+        selective = True
+        sel_line = lines[i].rstrip("\n")
+        i += 1
+
+    coord_line = lines[i].rstrip("\n")
+    i += 1
+
+    natoms = sum(counts)
+    pos_start = i
+    pos_end = i + natoms
+    if pos_end > len(lines):
+        raise ValueError("Atom count exceeds file length.")
+    pos_lines = [lines[j].rstrip("\n") for j in range(pos_start, pos_end)]
+
+    # Optional Velocities section
+    k = pos_end
+    while k < len(lines) and lines[k].strip() == "":
+        k += 1
+
+    vel_header = None
+    vel_lines = None
+    vel_end = k
+    if k < len(lines) and lines[k].strip().lower().startswith("veloc"):
+        vel_header = lines[k].rstrip("\n")
+        vel_start = k + 1
+        vel_end = vel_start + natoms
+        if vel_end > len(lines):
+            raise ValueError("Velocities section length inconsistent with atom count.")
+        vel_lines = [lines[j].rstrip("\n") for j in range(vel_start, vel_end)]
+
+    tail_lines = [lines[j].rstrip("\n") for j in range(vel_end, len(lines))] if vel_end < len(lines) else []
+
+    # Species index ranges (by order in species list)
+    starts = []
+    acc = 0
+    for c in counts:
+        starts.append(acc)
+        acc += c
+    species_ranges = []
+    for idx, sp in enumerate(species):
+        s, e = starts[idx], starts[idx] + counts[idx]
+        species_ranges.append((sp, s, e))
+
+    return {
+        "comment": comment,
+        "scale": scale,
+        "lv": lv,
+        "species": species,
+        "counts": counts,
+        "counts_line": counts_line,
+        "selective": selective,
+        "sel_line": sel_line,
+        "coord_line": coord_line,
+        "natoms": natoms,
+        "pos_lines": pos_lines,
+        "vel_header": vel_header,
+        "vel_lines": vel_lines,
+        "tail_lines": tail_lines,
+        "species_ranges": species_ranges,
+    }
+
+def build_poscar(data, pos_lines, vel_lines=None):
+    out = []
+    out.append(data["comment"])
+    out.append(data["scale"])
+    out.extend(data["lv"])
+    out.append(" ".join(data["species"]))
+    out.append(data["counts_line"])
+    if data["selective"]:
+        out.append(data["sel_line"] if data["sel_line"] is not None else "Selective dynamics")
+    out.append(data["coord_line"])
+    out.extend(pos_lines)
+    if data["vel_header"] is not None and vel_lines is not None:
+        out.append(data["vel_header"])
+        out.extend(vel_lines)
+    if data["tail_lines"]:
+        out.extend(data["tail_lines"])
+    return "\n".join(out) + "\n"
+
+def _swap_once(data, rng, li_label="Li", y_label="Y"):
+    si_li = _find_species_index(data["species"], li_label)
+    si_y  = _find_species_index(data["species"], y_label)
+    _, li_start, li_end = data["species_ranges"][si_li]
+    _, y_start,  y_end  = data["species_ranges"][si_y]
+
+    li_pick = rng.randrange(li_start, li_end)
+    y_pick  = rng.randrange(y_start, y_end)
+
+    new_pos = list(data["pos_lines"])
+    new_pos[li_pick], new_pos[y_pick] = new_pos[y_pick], new_pos[li_pick]
+
+    new_vel = None
+    if data["vel_lines"] is not None:
+        new_vel = list(data["vel_lines"])
+        new_vel[li_pick], new_vel[y_pick] = new_vel[y_pick], new_vel[li_pick]
+
+    return new_pos, new_vel, (li_pick, y_pick)
+
+def swap(n, input_file, output_dir):
+    """
+    Generate n POSCAR files, each with one random Li-Y swap.
+
+    Returns: list of Path to written files.
+    """
+    input_path = Path(input_file)
+    out_dir = Path(output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    lines = input_path.read_text().splitlines()
+    data = parse_poscar(lines)
+
+    rng = random.Random()
+    base = input_path.name
+
+    out_paths = []
+    for k in range(1, n + 1):
+        new_pos, new_vel, picked = _swap_once(data, rng)
+        txt = build_poscar(data, new_pos, new_vel)
+        out_path = out_dir / f"swap_{k}_{base}"
+        out_path.write_text(txt)
+        out_paths.append(out_path)
+        print(f"Wrote {out_path}  (swapped Li idx {picked[0]} <-> Y idx {picked[1]})")
+    return out_paths
+# --------- Editable defaults for direct run ---------
+INPUT_FILE = "data_POSCAR/origin/p3m1.vasp"   # path to input POSCAR
+OUTPUT_DIR = "data_POSCAR/p3m1"    # output directory
+N = 5                   # number of files to generate
+# ----------------------------------------------------
+if __name__ == "__main__":
+    # Direct-run entry: edit INPUT_FILE/OUTPUT_DIR/N above to change behavior.
+    swap(n=N, input_file=INPUT_FILE, output_dir=OUTPUT_DIR)