93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
"""
|
|
为指定案例添加次紧邻键 (k2, k=100, r0=1.41421356)。
|
|
|
|
用法: python add_k2.py case16
|
|
python add_k2.py case16 case17 case18
|
|
python add_k2.py --all
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
def add_k2(case_dir):
|
|
coord_path = os.path.join(case_dir, "input", "coord.txt")
|
|
conn_path = os.path.join(case_dir, "input", "connection.txt")
|
|
bond_path = os.path.join(case_dir, "input", "bond.txt")
|
|
|
|
if not os.path.exists(coord_path):
|
|
print(f" [跳过] {case_dir}: 找不到 coord.txt")
|
|
return False
|
|
|
|
# 读取 coord.txt 获取网格尺寸
|
|
with open(coord_path, "r", encoding="utf-8") as f:
|
|
lines = f.readlines()
|
|
n_atoms = len(lines) - 1 # 去掉表头
|
|
N = int(n_atoms ** 0.5)
|
|
if N * N != n_atoms:
|
|
print(f" [跳过] {case_dir}: 非正方形网格 (n_atoms={n_atoms})")
|
|
return False
|
|
|
|
print(f" {case_dir}: {N}x{N} 网格")
|
|
|
|
# 读取现有 connection.txt，检查是否已有 k2
|
|
has_k2 = False
|
|
if os.path.exists(conn_path):
|
|
with open(conn_path, "r") as f:
|
|
for line in f:
|
|
if "k2" in line:
|
|
has_k2 = True
|
|
break
|
|
|
|
if has_k2:
|
|
print(f" k2 已存在，跳过 connection.txt")
|
|
else:
|
|
# 追加 k2 键到 connection.txt
|
|
with open(conn_path, "a", encoding="utf-8") as f:
|
|
cnt = 0
|
|
for row in range(N):
|
|
for col in range(N):
|
|
id1 = row * N + col + 1
|
|
if col + 1 < N and row + 1 < N:
|
|
f.write(f"{id1} {(row + 1) * N + (col + 1) + 1} k2\n")
|
|
cnt += 1
|
|
if col - 1 >= 0 and row + 1 < N:
|
|
f.write(f"{id1} {(row + 1) * N + (col - 1) + 1} k2\n")
|
|
cnt += 1
|
|
print(f" connection.txt: 追加 {cnt} 条 k2 键")
|
|
|
|
# 检查 bond.txt 是否有 k2
|
|
has_bond = False
|
|
if os.path.exists(bond_path):
|
|
with open(bond_path, "r") as f:
|
|
for line in f:
|
|
if line.startswith("k2"):
|
|
has_bond = True
|
|
break
|
|
|
|
if has_bond:
|
|
print(f" bond.txt: k2 已存在")
|
|
else:
|
|
with open(bond_path, "a", encoding="utf-8") as f:
|
|
f.write("k2 100.0 1.41421356\n")
|
|
print(f" bond.txt: 追加 k2 定义")
|
|
|
|
return True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
targets = []
|
|
if "--all" in sys.argv:
|
|
base = os.path.dirname(os.path.abspath(__file__))
|
|
for d in sorted(os.listdir(base)):
|
|
if d.startswith("case") and os.path.isdir(os.path.join(base, d)):
|
|
targets.append(os.path.join(base, d))
|
|
else:
|
|
for arg in sys.argv[1:]:
|
|
if arg.startswith("--"):
|
|
continue
|
|
p = arg if os.path.isabs(arg) else os.path.join(os.path.dirname(os.path.abspath(__file__)), arg)
|
|
targets.append(p)
|
|
|
|
for t in targets:
|
|
add_k2(t)
|