Coverage for biobb_gromacs/gromacs/mdrun.py: 56%

1#!/usr/bin/env python3 

2 

3"""Module containing the MDrun class and the command line interface.""" 

4from typing import Optional 

5from pathlib import PurePath 

6from biobb_common.generic.biobb_object import BiobbObject 

7from biobb_common.tools import file_utils as fu 

8from biobb_common.tools.file_utils import launchlogger 

9from biobb_gromacs.gromacs.common import get_gromacs_version 

10 

11 

12class Mdrun(BiobbObject): 

13 """ 

14 | biobb_gromacs Mdrun 

15 | Wrapper of the `GROMACS mdrun <http://manual.gromacs.org/current/onlinehelp/gmx-mdrun.html>`_ module. 

16 | MDRun is the main computational chemistry engine within GROMACS. It performs Molecular Dynamics simulations, but it can also perform Stochastic Dynamics, Energy Minimization, test particle insertion or (re)calculation of energies. 

17 

18 Args: 

19 input_tpr_path (str): Path to the portable binary run input file TPR. File type: input. `Sample file <https://github.com/bioexcel/biobb_gromacs/raw/master/biobb_gromacs/test/data/gromacs/mdrun.tpr>`_. Accepted formats: tpr (edam:format_2333). 

20 output_gro_path (str): Path to the output GROMACS structure GRO file. File type: output. `Sample file <https://github.com/bioexcel/biobb_gromacs/raw/master/biobb_gromacs/test/reference/gromacs/ref_mdrun.gro>`_. Accepted formats: gro (edam:format_2033). 

21 output_edr_path (str): Path to the output GROMACS portable energy file EDR. File type: output. `Sample file <https://github.com/bioexcel/biobb_gromacs/raw/master/biobb_gromacs/test/reference/gromacs/ref_mdrun.edr>`_. Accepted formats: edr (edam:format_2330). 

22 output_log_path (str): Path to the output GROMACS trajectory log file LOG. File type: output. `Sample file <https://github.com/bioexcel/biobb_gromacs/raw/master/biobb_gromacs/test/reference/gromacs/ref_mdrun.log>`_. Accepted formats: log (edam:format_2330). 

23 output_trr_path (str) (Optional): Path to the GROMACS uncompressed raw trajectory file TRR. File type: output. `Sample file <https://github.com/bioexcel/biobb_gromacs/raw/master/biobb_gromacs/test/reference/gromacs/ref_mdrun.trr>`_. Accepted formats: trr (edam:format_3910). 

24 input_cpt_path (str) (Optional): Path to the input GROMACS checkpoint file CPT. File type: input. Accepted formats: cpt (edam:format_2333). 

25 output_xtc_path (str) (Optional): Path to the GROMACS compressed trajectory file XTC. File type: output. Accepted formats: xtc (edam:format_3875). 

26 output_cpt_path (str) (Optional): Path to the output GROMACS checkpoint file CPT. File type: output. Accepted formats: cpt (edam:format_2333). 

27 output_dhdl_path (str) (Optional): Path to the output dhdl.xvg file only used when free energy calculation is turned on. File type: output. Accepted formats: xvg (edam:format_2033). 

28 properties (dict - Python dictionary object containing the tool parameters, not input/output files): 

29 * **mpi_bin** (*str*) - (None) Path to the MPI runner. Usually "mpirun" or "srun". 

30 * **mpi_np** (*int*) - (0) [0~1000|1] Number of MPI processes. Usually an integer bigger than 1. 

31 * **mpi_flags** (*str*) - (None) Path to the MPI hostlist file. 

32 * **checkpoint_time** (*int*) - (15) [0~1000|1] Checkpoint writing interval in minutes. Only enabled if an output_cpt_path is provided. 

33 * **noappend** (*bool*) - (False) Include the noappend flag to open new output files and add the simulation part number to all output file names 

34 * **num_threads** (*int*) - (0) [0~1000|1] Let GROMACS guess. The number of threads that are going to be used. 

35 * **num_threads_mpi** (*int*) - (0) [0~1000|1] Let GROMACS guess. The number of GROMACS MPI threads that are going to be used. 

36 * **num_threads_omp** (*int*) - (0) [0~1000|1] Let GROMACS guess. The number of GROMACS OPENMP threads that are going to be used. 

37 * **num_threads_omp_pme** (*int*) - (0) [0~1000|1] Let GROMACS guess. The number of GROMACS OPENMP_PME threads that are going to be used. 

38 * **use_gpu** (*bool*) - (False) Use settings appropriate for GPU. Adds: -nb gpu -pme gpu 

39 * **gpu_id** (*str*) - (None) list of unique GPU device IDs available to use. 

40 * **gpu_tasks** (*str*) - (None) list of GPU device IDs, mapping each PP task on each node to a device. 

41 * **gmx_lib** (*str*) - (None) Path set GROMACS GMXLIB environment variable. 

42 * **binary_path** (*str*) - ("gmx") Path to the GROMACS executable binary. 

43 * **remove_tmp** (*bool*) - (True) [WF property] Remove temporal files. 

44 * **restart** (*bool*) - (False) [WF property] Do not execute if output files exist. 

45 * **sandbox_path** (*str*) - ("./") [WF property] Parent path to the sandbox directory. 

46 * **container_path** (*str*) - (None) Path to the binary executable of your container. 

47 * **container_image** (*str*) - (None) Container Image identifier. 

48 * **container_volume_path** (*str*) - ("/data") Path to an internal directory in the container. 

49 * **container_working_dir** (*str*) - (None) Path to the internal CWD in the container. 

50 * **container_user_id** (*str*) - (None) User number id to be mapped inside the container. 

51 * **container_shell_path** (*str*) - ("/bin/bash") Path to the binary executable of the container shell. 

52 

53 Examples: 

54 This is a use example of how to use the building block from Python:: 

55 

56 from biobb_gromacs.gromacs.mdrun import mdrun 

57 prop = { 'num_threads': 0, 

58 'binary_path': 'gmx' } 

59 mdrun(input_tpr_path='/path/to/myPortableBinaryRunInputFile.tpr', 

60 output_trr_path='/path/to/newTrajectory.trr', 

61 output_gro_path='/path/to/newStructure.gro', 

62 output_edr_path='/path/to/newEnergy.edr', 

63 output_log_path='/path/to/newSimulationLog.log', 

64 properties=prop) 

65 

66 Info: 

67 * wrapped_software: 

68 * name: GROMACS Mdrun 

69 * version: 2025.2 

70 * license: LGPL 2.1 

71 * multinode: mpi 

72 * ontology: 

73 * name: EDAM 

74 * schema: http://edamontology.org/EDAM.owl 

75 """ 

76 

77 def __init__(self, input_tpr_path: str, output_gro_path: str, output_edr_path: str, 

78 output_log_path: str, output_trr_path: Optional[str] = None, input_cpt_path: Optional[str] = None, 

79 output_xtc_path: Optional[str] = None, output_cpt_path: Optional[str] = None, 

80 output_dhdl_path: Optional[str] = None, properties: Optional[dict] = None, **kwargs) -> None: 

81 properties = properties or {} 

82 

83 # Call parent class constructor 

84 super().__init__(properties) 

85 self.locals_var_dict = locals().copy() 

86 

87 # Input/Output files 

88 self.io_dict = { 

89 "in": {"input_tpr_path": input_tpr_path, "input_cpt_path": input_cpt_path}, 

90 "out": {"output_trr_path": output_trr_path, "output_gro_path": output_gro_path, 

91 "output_edr_path": output_edr_path, "output_log_path": output_log_path, 

92 "output_xtc_path": output_xtc_path, "output_cpt_path": output_cpt_path, 

93 "output_dhdl_path": output_dhdl_path} 

94 } 

95 

96 # Properties specific for BB 

97 # general mpi properties 

98 self.mpi_bin = properties.get('mpi_bin') 

99 self.mpi_np = properties.get('mpi_np') 

100 self.mpi_flags = properties.get('mpi_flags') 

101 # gromacs cpu mpi/openmp properties 

102 self.num_threads = str(properties.get('num_threads', '')) 

103 self.num_threads_mpi = str(properties.get('num_threads_mpi', '')) 

104 self.num_threads_omp = str(properties.get('num_threads_omp', '')) 

105 self.num_threads_omp_pme = str( 

106 properties.get('num_threads_omp_pme', '')) 

107 # gromacs gpus 

108 self.use_gpu = properties.get( 

109 'use_gpu', False) # Adds: -nb gpu -pme gpu 

110 self.gpu_id = str(properties.get('gpu_id', '')) 

111 self.gpu_tasks = str(properties.get('gpu_tasks', '')) 

112 # gromacs 

113 self.checkpoint_time = properties.get('checkpoint_time') 

114 self.noappend = properties.get('noappend', False) 

115 

116 # Properties common in all GROMACS BB 

117 self.gmx_lib = properties.get('gmx_lib', None) 

118 self.binary_path: str = properties.get('binary_path', 'gmx') 

119 self.gmx_nobackup = properties.get('gmx_nobackup', True) 

120 self.gmx_nocopyright = properties.get('gmx_nocopyright', True) 

121 if self.gmx_nobackup: 

122 self.binary_path += ' -nobackup' 

123 if self.gmx_nocopyright: 

124 self.binary_path += ' -nocopyright' 

125 if (not self.mpi_bin) and (not self.container_path): 

126 self.gmx_version = get_gromacs_version(self.binary_path) 

127 

128 # Check the properties 

129 self.check_properties(properties) 

130 self.check_arguments() 

131 

132 @launchlogger 

133 def launch(self) -> int: 

134 """Execute the :class:`Mdrun <gromacs.mdrun.Mdrun>` object.""" 

135 

136 # Setup Biobb 

137 if self.check_restart(): 

138 return 0 

139 

140 # Optional output files (if not added mrun will create them using a generic name) 

141 if not self.stage_io_dict["out"].get("output_trr_path"): 

142 self.stage_io_dict["out"]["output_trr_path"] = fu.create_name( 

143 prefix=self.prefix, step=self.step, name='trajectory.trr') 

144 self.tmp_files.append(self.stage_io_dict["out"]["output_trr_path"]) 

145 

146 self.stage_files() 

147 

148 if self.container_path: 

149 working_dir = self.container_volume_path if self.container_volume_path else "/data" 

150 else: 

151 working_dir = self.stage_io_dict.get('unique_dir', '') 

152 

153 self.cmd = [self.binary_path, 'mdrun', 

154 '-o', PurePath(self.stage_io_dict["out"]["output_trr_path"]).name, 

155 '-s', PurePath(self.stage_io_dict["in"]["input_tpr_path"]).name, 

156 '-c', PurePath(self.stage_io_dict["out"]["output_gro_path"]).name, 

157 '-e', PurePath(self.stage_io_dict["out"]["output_edr_path"]).name, 

158 '-g', PurePath(self.stage_io_dict["out"]["output_log_path"]).name] 

159 

160 if self.stage_io_dict["in"].get("input_cpt_path"): 

161 self.cmd.append('-cpi') 

162 self.cmd.append(PurePath(self.stage_io_dict["in"]["input_cpt_path"]).name) 

163 if self.stage_io_dict["out"].get("output_xtc_path"): 

164 self.cmd.append('-x') 

165 self.cmd.append(PurePath(self.stage_io_dict["out"]["output_xtc_path"]).name) 

166 else: 

167 self.tmp_files.append('traj_comp.xtc') 

168 if self.stage_io_dict["out"].get("output_cpt_path"): 

169 self.cmd.append('-cpo') 

170 self.cmd.append(PurePath(self.stage_io_dict["out"]["output_cpt_path"]).name) 

171 if self.checkpoint_time: 

172 self.cmd.append('-cpt') 

173 self.cmd.append(str(self.checkpoint_time)) 

174 if self.stage_io_dict["out"].get("output_dhdl_path"): 

175 self.cmd.append('-dhdl') 

176 self.cmd.append(PurePath(self.stage_io_dict["out"]["output_dhdl_path"]).name) 

177 

178 # general mpi properties 

179 if self.mpi_bin: 

180 mpi_cmd = [self.mpi_bin] 

181 if self.mpi_np: 

182 mpi_cmd.append('-n') 

183 mpi_cmd.append(str(self.mpi_np)) 

184 if self.mpi_flags: 

185 mpi_cmd.extend(self.mpi_flags) 

186 self.cmd = mpi_cmd + self.cmd 

187 

188 self.cmd = ["cd", working_dir, ";"] + self.cmd 

189 

190 # gromacs cpu mpi/openmp properties 

191 if self.num_threads: 

192 fu.log( 

193 f'User added number of gmx threads: {self.num_threads}', self.out_log) 

194 self.cmd.append('-nt') 

195 self.cmd.append(self.num_threads) 

196 if self.num_threads_mpi: 

197 fu.log( 

198 f'User added number of gmx mpi threads: {self.num_threads_mpi}', self.out_log) 

199 self.cmd.append('-ntmpi') 

200 self.cmd.append(self.num_threads_mpi) 

201 if self.num_threads_omp: 

202 fu.log( 

203 f'User added number of gmx omp threads: {self.num_threads_omp}', self.out_log) 

204 self.cmd.append('-ntomp') 

205 self.cmd.append(self.num_threads_omp) 

206 if self.num_threads_omp_pme: 

207 fu.log( 

208 f'User added number of gmx omp_pme threads: {self.num_threads_omp_pme}', self.out_log) 

209 self.cmd.append('-ntomp_pme') 

210 self.cmd.append(self.num_threads_omp_pme) 

211 # GMX gpu properties 

212 if self.use_gpu: 

213 fu.log('Adding GPU specific settings adds: -nb gpu -pme gpu', self.out_log) 

214 self.cmd += ["-nb", "gpu", "-pme", "gpu"] 

215 if self.gpu_id: 

216 fu.log( 

217 f'list of unique GPU device IDs available to use: {self.gpu_id}', self.out_log) 

218 self.cmd.append('-gpu_id') 

219 self.cmd.append(self.gpu_id) 

220 if self.gpu_tasks: 

221 fu.log( 

222 f'list of GPU device IDs, mapping each PP task on each node to a device: {self.gpu_tasks}', self.out_log) 

223 self.cmd.append('-gputasks') 

224 self.cmd.append(self.gpu_tasks) 

225 

226 if self.noappend: 

227 self.cmd.append('-noappend') 

228 

229 if self.gmx_lib: 

230 self.env_vars_dict['GMXLIB'] = self.gmx_lib 

231 

232 # Run Biobb block 

233 self.run_biobb() 

234 

235 # Copy files to host 

236 self.copy_to_host() 

237 

238 # Remove temporal files 

239 self.remove_tmp_files() 

240 

241 self.check_arguments(output_files_created=True, raise_exception=False) 

242 return self.return_code 

243 

244 def copy_to_host(self): 

245 """ 

246 Updates the path to the original output files in the sandbox, 

247 to catch changes due to noappend restart. 

248 

249 GROMACS mdrun will change the output file names from md.gro to md.part0001.gro 

250 if the noappend flag is used. 

251 """ 

252 import pathlib 

253 

254 def capture_part_pattern(filename): 

255 """ 

256 Captures the 'part' pattern followed by digits from a string. 

257 """ 

258 import re 

259 pattern = r'part\d+' 

260 

261 match = re.search(pattern, filename) 

262 if match: 

263 return match.group(0) 

264 else: 

265 return None 

266 

267 if self.noappend: 

268 # List files in the staging directory 

269 staging_path = self.stage_io_dict["unique_dir"] 

270 files_in_staging = list(pathlib.Path(staging_path).glob('*')) 

271 

272 # Find the part000x pattern in the output files 

273 for file in files_in_staging: 

274 part_pattern = capture_part_pattern(file.name) 

275 if part_pattern: 

276 break 

277 

278 # Update expected output files 

279 for file_ref, stage_file_path in self.stage_io_dict["out"].items(): 

280 if stage_file_path: 

281 # Find the parent and the file name in the sandbox 

282 parent_path = pathlib.Path(stage_file_path).parent 

283 file_stem = pathlib.Path(stage_file_path).stem 

284 file_suffix = pathlib.Path(stage_file_path).suffix 

285 

286 # Rename all output files except checkpoint files 

287 if file_suffix != '.cpt': 

288 # Create the new file name with the part pattern 

289 if part_pattern: 

290 new_file_name = f"{file_stem}.{part_pattern}{file_suffix}" 

291 new_file_path = parent_path / new_file_name 

292 # Update the stage_io_dict with the new file path 

293 self.stage_io_dict["out"][file_ref] = str( 

294 new_file_path) 

295 return super().copy_to_host() 

296 

297 

298def mdrun(input_tpr_path: str, output_gro_path: str, output_edr_path: str, 

299 output_log_path: str, output_trr_path: Optional[str] = None, input_cpt_path: Optional[str] = None, 

300 output_xtc_path: Optional[str] = None, output_cpt_path: Optional[str] = None, 

301 output_dhdl_path: Optional[str] = None, properties: Optional[dict] = None, **kwargs) -> int: 

302 """Create :class:`Mdrun <gromacs.mdrun.Mdrun>` class and 

303 execute the :meth:`launch() <gromacs.mdrun.Mdrun.launch>` method.""" 

304 return Mdrun(**dict(locals())).launch() 

305 

306 

307mdrun.__doc__ = Mdrun.__doc__ 

308main = Mdrun.get_main(mdrun, "Wrapper for the GROMACS mdrun module.") 

309 

310 

311if __name__ == '__main__': 

312 main()