update

runner.py

@@ -22,7 +22,6 @@ logging.basicConfig(level=logging.INFO, filename='simulation_log.log', filemode=
                     format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger()
 
-@dataclass
 class SimulationRunner:
     pdb_file: Path
     nsteps: int
@@ -30,18 +29,73 @@ class SimulationRunner:
     base_folder: Path
     bash_script: Path = None  # Bash脚本路径作为可选参数
     gmxrc_path: Path = None  # GMXRC文件路径作为可选参数
+    temp_folder: Path  # 用于存储临时数据和结果数据的路径(轨迹数据处理中间数据等)
     runner_folder: Path = field(init=False)
+    tpr_file: Path = field(init=False)
+    xtc_file: Path = field(init=False)
 
     def __post_init__(self):
+        # 初始化文件夹和文件路径
         self.runner_folder = self.base_folder / f"runner_{self.pdb_file.stem}"
+        self.tpr_file = self.runner_folder / "md.tpr"
+        self.xtc_file = self.runner_folder / "md.xtc"
         self.runner_folder.mkdir(exist_ok=True)
-        self.bash_script = self.bash_script.absolute() or Path(__file__).resolve().parent / "md_gromacs.sh"
+        self.temp_folder = self.runner_folder / "temp"
+        self.temp_folder.mkdir(exist_ok=True)
+        self.bash_script = self.bash_script.absolute() if self.bash_script else Path(__file__).resolve().parent / "md_gromacs.sh"
         # 设置 GMXRC_PATH 环境变量
-        self.gmxrc_path = self.gmxrc_path or Path("$HOME/software/gmx2023.2/bin/GMXRC")
+        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)
 
+    @staticmethod
+    def read_ndx_file(filename):
+        ndx_dict = {}
+        current_section = None
+
+        with open(filename, 'r') as file:
+            for line in file:
+                line = line.strip()
+
+                if line.startswith('[') and line.endswith(']'):
+                    current_section = line[1:-1].strip()
+                    ndx_dict[current_section] = []
+                else:
+                    if current_section is not None:
+                        ndx_dict[current_section].extend(map(int, line.split()))
+
+        return ndx_dict
+
+    @staticmethod
+echo "Protein" | gmx_mpi trjconv -dt {extract_interval} -s {tpr_file} -f {xtc_file} -n {temp_folder}/tarj_show.ndx -pbc mol -o {temp_folder}/temp.xtc
+
+    # 新增处理轨迹的方法
+    def process_trajectory(self, extract_interval):
+        # 根据提供的脚本逻辑读取和保存索引文件
+        ndx_dict = self.read_ndx_file(f'{self.runner_folder}/index.ndx')
+        
+        # 根据索引文件内容决定如何处理轨迹
+        if any(key.startswith("LG") for key in ndx_dict):
+            # 处理含有LG组的情况
+            new_ndx_dict = {key: value for key, value in ndx_dict.items() if key.startswith("LG") or key in ["Protein", "Protein_LIG"]}
+            self.save_ndx_file(f"{self.temp_folder}/tarj_show.ndx", new_ndx_dict)
+            # 构建处理轨迹的命令
+            command_1 = f'echo "Protein_LIG" | gmx trjconv -dt {extract_interval} -s {self.tpr_file} -f {self.xtc_file} -n {self.temp_folder}/tarj_show.ndx -pbc mol -o {self.temp_folder}/temp.xtc'
+            command_2 = f'echo "Protein\nProtein\nProtein_LIG" | gmx trjconv -s {self.tpr_file} -f {self.temp_folder}/temp.xtc -n {self.temp_folder}/tarj_show.ndx -center -fit rot+trans -o {self.output_folder}/traj_show.xtc'
+            command_3 = f'echo "Protein\nProtein\nProtein_LIG" | gmx trjconv -s {self.tpr_file} -f {self.temp_folder}/temp.xtc -n {self.temp_folder}/tarj_show.ndx -center -fit rot+trans -b 0 -e 0 -o {self.output_folder}/tarj_show.pdb'
+        else:
+            # 处理只含有蛋白质组的情况
+            new_ndx_dict = {key: value for key, value in ndx_dict.items() if key in ["Protein"]}
+            self.save_ndx_file(f"{self.temp_folder}/tarj_show.ndx", new_ndx_dict)
+            # 构建处理轨迹的命令
+            command_1 = f'echo "Protein" | gmx trjconv -dt {extract_interval} -s {self.tpr_file} -f {self.xtc_file} -n {self.temp_folder}/tarj_show.ndx -pbc mol -o {self.temp_folder}/temp.xtc'
+            command_2 = f'echo "Protein\nProtein\nProtein" | gmx trjconv -s {self.tpr_file} -f {self.temp_folder}/temp.xtc -n {self.temp_folder}/tarj_show.ndx -center -fit rot+trans -o {self.output_folder}/traj_show.xtc'
+            command_3 = f'echo "Protein\nProtein\nProtein" | gmx trjconv -s {self.tpr_file} -f {self.temp_folder}/temp.xtc -n {self.temp_folder}/tarj_show.ndx -center -fit rot+trans -b 0 -e 0 -o {self.output_folder}/tarj_show.pdb'
+        subprocess.run(command_1, shell=True, check=True)
+        subprocess.run(command_2, shell=True, check=True)
+        subprocess.run(command_3, shell=True, check=True)
+
     def run_simulation(self):
         start_time = time.time()
         result = None
@@ -84,6 +138,8 @@ def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_p
         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}.")
+        runner.process_trajectory(extract_interval=100)  # 例如，每100ps抽取一次轨迹
+        logger.info(f"Finished processing trajectory for {pdb_file.name}. per 100ps")
 
 if __name__ == "__main__":
     NSTEPS = 50000000  # Example: 50000000 steps
@@ -92,7 +148,7 @@ if __name__ == "__main__":
     # 传入自定义的bash脚本路径
     CUSTOM_BASH_SCRIPT_PATH = Path('md_gromacs.sh')
     # 传入 GMXRC 文件的路径
-    GMXRC_PATH = Path('/home/zenglingyu/software/gmx2023.2/bin/GMXRC')
+    GMXRC_PATH = Path('/root/software/gmx2023.2/bin/GMXRC')
     main(NSTEPS, DT, PDB_FOLDER_PATH, CUSTOM_BASH_SCRIPT_PATH, GMXRC_PATH)