99 lines
3.9 KiB
Python
99 lines
3.9 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
'''
|
|
@file :runner.py
|
|
@Description: :
|
|
@Date :2023/12/04 14:34:36
|
|
@Author :lyzeng
|
|
@Email :pylyzeng@gmail.com
|
|
@version :1.0
|
|
'''
|
|
|
|
import time
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
import subprocess
|
|
import shutil
|
|
import os
|
|
|
|
# 设置日志记录
|
|
logging.basicConfig(level=logging.INFO, filename='simulation_log.log', filemode='a',
|
|
format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger()
|
|
|
|
@dataclass
|
|
class SimulationRunner:
|
|
pdb_file: Path
|
|
nsteps: int
|
|
dt: float
|
|
base_folder: Path
|
|
bash_script: Path = None # Bash脚本路径作为可选参数
|
|
gmxrc_path: Path = None # GMXRC文件路径作为可选参数
|
|
runner_folder: Path = field(init=False)
|
|
|
|
def __post_init__(self):
|
|
self.runner_folder = self.base_folder / f"runner_{self.pdb_file.stem}"
|
|
self.runner_folder.mkdir(exist_ok=True)
|
|
self.bash_script = self.bash_script.absolute() or Path(__file__).resolve().parent / "md_gromacs.sh"
|
|
# 设置 GMXRC_PATH 环境变量
|
|
self.gmxrc_path = self.gmxrc_path or Path("/home/lingyuzeng/software/gmx2023.2/bin/GMXRC")
|
|
|
|
def copy_pdb(self):
|
|
shutil.copy(self.pdb_file, self.runner_folder / self.pdb_file.name)
|
|
|
|
def run_simulation(self):
|
|
start_time = time.time()
|
|
result = None
|
|
try:
|
|
env_vars = {
|
|
"NAME": self.pdb_file.stem, # 首先将 NAME 设置为文件的 stem
|
|
"NSTEPS": str(self.nsteps),
|
|
"DT": str(self.dt),
|
|
"GMXRC_PATH": str(self.gmxrc_path)
|
|
}
|
|
env_vars["HOME"] = os.environ["HOME"]
|
|
logger.info(f"pdb_file: {self.pdb_file.name}")
|
|
logger.info(f"Executing script at: {self.bash_script}")
|
|
result = subprocess.run(["bash", str(self.bash_script)], env=env_vars, cwd=self.runner_folder,
|
|
capture_output=True, text=True, check=True)
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(f"Error in simulation for {self.pdb_file.name}: {e}")
|
|
if e.stdout:
|
|
logger.error(f"Standard Output:\n{e.stdout}")
|
|
if e.stderr:
|
|
logger.error(f"Standard Error:\n{e.stderr}")
|
|
|
|
end_time = time.time()
|
|
duration = end_time - start_time
|
|
|
|
if result:
|
|
logger.info(f"Simulation for {self.pdb_file.name} completed successfully in {duration:.2f} seconds.")
|
|
if result.stdout:
|
|
logger.info(f"Shell Script Output:\n{result.stdout}")
|
|
if result.stderr:
|
|
logger.error(f"Shell Script Error Output:\n{result.stderr}")
|
|
else:
|
|
logger.error(f"Simulation for {self.pdb_file.name} failed in {duration:.2f} seconds.")
|
|
|
|
def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_path):
|
|
pdb_folder = Path(pdb_folder_path).resolve()
|
|
for pdb_file in pdb_folder.glob("*.pdb"):
|
|
runner = SimulationRunner(pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path)
|
|
runner.copy_pdb()
|
|
logger.info(f"Running simulation for {pdb_file.name} in {runner.runner_folder}...")
|
|
runner.run_simulation()
|
|
logger.info(f"Finished simulation for {pdb_file.name}.")
|
|
|
|
if __name__ == "__main__":
|
|
NSTEPS = 50000000 # Example: 50000000 steps
|
|
DT = 0.002 # Example: 2 fs time step
|
|
PDB_FOLDER_PATH = Path("./pdb_files") # Assuming the PDB files are in a folder named 'pdb_files' in the current directory
|
|
# 传入自定义的bash脚本路径
|
|
CUSTOM_BASH_SCRIPT_PATH = Path('analysis_pdb/md_gromacs.sh')
|
|
# 传入 GMXRC 文件的路径
|
|
GMXRC_PATH = Path('/home/zenglingyu/software/gmx2023.2/bin/GMXRC')
|
|
main(NSTEPS, DT, PDB_FOLDER_PATH, CUSTOM_BASH_SCRIPT_PATH, GMXRC_PATH)
|
|
|
|
|