Files
analysis_pdb/runner.py
2023-12-04 16:25:17 +08:00

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)