改成在服务器上直接用的python脚本

将给定的 Bash 脚本转换为 Python 脚本，并使其能够在服务器上直接运行。以下是转换后的 Python 脚本：

import os
import subprocess
import time

# 设置 SLURM 参数
job_name = "AIMD-job"
node = "node01"
num_nodes = 1
num_tasks = 36
time_limit = "72:00:00"
output_file = "vasp.out"
partition = "maxtime"

# 获取当前工作目录
submit_dir = os.getcwd()

# 打印作业信息
print(f"JOBID: {os.environ.get('SLURM_JOBID', '')}")
print(f"Job Name: {job_name}")
print(f"Node list: {os.environ.get('SLURM_JOB_NODELIST', '')}")
print(f"Requested total cores#: {num_tasks}")
print(f"Cores on node: {os.environ.get('SLURM_CPUS_ON_NODE', '')}")
print(f"OpenMP thread: {os.environ.get('OMP_NUM_THREADS', '')}")
print(f"Path: {os.environ.get('PATH', '')}")
print(f"Current working directory: {submit_dir}")

# 获取 CPU 和内存信息
cpu_info = subprocess.check_output("grep MHz /proc/cpuinfo | head -1", shell=True).decode().strip()
core_info = subprocess.check_output("grep cores /proc/cpuinfo | head -1", shell=True).decode().strip()
processor_count = int(subprocess.check_output("grep processor /proc/cpuinfo | wc -l", shell=True).decode().strip())
cpu_name = subprocess.check_output("grep 'model name' /proc/cpuinfo | head -1", shell=True).decode().strip()
memory_info = subprocess.check_output("grep MemTotal /proc/meminfo", shell=True).decode().strip()

print(cpu_info)
print(core_info)
print(f"Processor count: {processor_count}")
print(cpu_name)
print(memory_info)

# 记录开始时间
start_time = time.time()

# 加热过程
while not os.path.exists('Heating_success'):
    print("Running Heating...")
    os.system("cp INCAR.heat INCAR")
    os.system("ulimit -s unlimited")
    with open("vasp.out.heat", "w") as f:
        subprocess.run(["mpirun", "-np", str(num_tasks), "/public/soft/vasp.6.4.2/bin/vasp_gam"], stdout=f, stderr=subprocess.STDOUT)
    
    success = subprocess.check_output("grep '1000 T=' vasp.out.heat", shell=True).decode().strip()
    time.sleep(10)
    
    if success:
        os.system("touch Heating_success")
        os.makedirs("RUN_0", exist_ok=True)
        os.system("cp POSCAR CONTCAR INCAR* KPOINTS POTCAR RUN_0")
        os.system("mv DOSCAR EIGENVAL IBZKPT OSZICAR OUTCAR PCDAT PROCAR REPORT XDATCAR vasp.out.heat vasprun.xml RUN_0")
        
        heating_time = time.time() - start_time
        print(f"Heating time was {heating_time:.2f} seconds")
        break
    else:
        print("Run heating again")

# 分子动力学过程
os.system("cp INCAR.MD INCAR")

for i in range(1, 151):
    while not os.path.exists(f"RUN_{i}_success"):
        print(f"Running {i}...")
        run_start = time.time()
        final = i - 1
        os.system(f"cp RUN_{final}/CONTCAR ./POSCAR")
        os.system("ulimit -s unlimited")
        with open(f"vasp.out.{i}", "w") as f:
            subprocess.run(["mpirun", "-np", str(num_tasks), "/public/soft/vasp.6.4.2/bin/vasp_gam"], stdout=f, stderr=subprocess.STDOUT)
        
        success = subprocess.check_output(f"grep '1000 T=' vasp.out.{i}", shell=True).decode().strip()
        time.sleep(10)
        
        if success:
            os.system(f"touch RUN_{i}_success")
            os.makedirs(f"RUN_{i}", exist_ok=True)
            os.system(f"cp POSCAR CONTCAR INCAR* KPOINTS POTCAR RUN_{i}")
            os.system(f"mv DOSCAR EIGENVAL IBZKPT OSZICAR OUTCAR PCDAT PROCAR REPORT XDATCAR vasp.out.{i} vasprun.xml RUN_{i}")
            
            run_end = time.time()
            runtime = run_end - run_start
            print(f"Running time of this run was {runtime:.2f} seconds")
            break
        else:
            print(f"Run {i} again")

### 说明：
1. **环境变量**：使用 `os.environ.get` 获取环境变量。
2. **命令执行**：使用 `subprocess.run` 和 `os.system` 执行系统命令。
3. **文件操作**：使用 `os.makedirs` 创建目录，使用 `os.system` 复制和移动文件。
4. **时间记录**：使用 `time.time()` 记录时间和计算运行时间。

这个脚本可以在服务器上直接运行，假设你已经安装了必要的依赖项（如 `subprocess` 模块）。你可以通过以下命令运行该脚本：

python3 your_script.py

请确保将 `your_script.py` 替换为你保存脚本的实际文件名。