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

1#!/usr/bin/env python3 

2 

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

4from typing import Optional 

5from biobb_common.generic.biobb_object import BiobbObject 

6from biobb_common.tools import file_utils as fu 

7from biobb_common.tools.file_utils import launchlogger 

8from biobb_gromacs.gromacs.common import get_gromacs_version 

9 

10 

11class Mdrun(BiobbObject): 

12 """ 

13 | biobb_gromacs Mdrun 

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

15 | 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. 

16 

17 Args: 

18 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). 

19 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). 

20 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). 

21 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). 

22 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). 

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

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

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

26 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). 

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

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

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

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

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

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

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

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

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

36 * **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. 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

51 

52 Examples: 

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

54 

55 from biobb_gromacs.gromacs.mdrun import mdrun 

56 prop = { 'num_threads': 0, 

57 'binary_path': 'gmx' } 

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

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

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

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

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

63 properties=prop) 

64 

65 Info: 

66 * wrapped_software: 

67 * name: GROMACS Mdrun 

68 * version: 2025.2 

69 * license: LGPL 2.1 

70 * multinode: mpi 

71 * ontology: 

72 * name: EDAM 

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

74 """ 

75 

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

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

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

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

80 properties = properties or {} 

81 

82 # Call parent class constructor 

83 super().__init__(properties) 

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

85 

86 # Input/Output files 

87 self.io_dict = { 

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

89 "out": {"output_trr_path": output_trr_path, "output_gro_path": output_gro_path, 

90 "output_edr_path": output_edr_path, "output_log_path": output_log_path, 

91 "output_xtc_path": output_xtc_path, "output_cpt_path": output_cpt_path, 

92 "output_dhdl_path": output_dhdl_path} 

93 } 

94 

95 # Properties specific for BB 

96 # general mpi properties 

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

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

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

100 # gromacs cpu mpi/openmp properties 

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

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

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

104 self.num_threads_omp_pme = str( 

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

106 # gromacs gpus 

107 self.use_gpu = properties.get( 

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

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

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

111 # gromacs 

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

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

114 

115 # Properties common in all GROMACS BB 

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

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

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

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

120 if self.gmx_nobackup: 

121 self.binary_path += ' -nobackup' 

122 if self.gmx_nocopyright: 

123 self.binary_path += ' -nocopyright' 

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

125 self.gmx_version = get_gromacs_version(self.binary_path) 

126 

127 # Check the properties 

128 self.check_properties(properties) 

129 self.check_arguments() 

130 

131 @launchlogger 

132 def launch(self) -> int: 

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

134 

135 # Setup Biobb 

136 if self.check_restart(): 

137 return 0 

138 

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

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

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

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

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

144 

145 self.stage_files() 

146 

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

148 '-o', self.stage_io_dict["out"]["output_trr_path"], 

149 '-s', self.stage_io_dict["in"]["input_tpr_path"], 

150 '-c', self.stage_io_dict["out"]["output_gro_path"], 

151 '-e', self.stage_io_dict["out"]["output_edr_path"], 

152 '-g', self.stage_io_dict["out"]["output_log_path"]] 

153 

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

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

156 self.cmd.append(self.stage_io_dict["in"]["input_cpt_path"]) 

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

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

159 self.cmd.append(self.stage_io_dict["out"]["output_xtc_path"]) 

160 else: 

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

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

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

164 self.cmd.append(self.stage_io_dict["out"]["output_cpt_path"]) 

165 if self.checkpoint_time: 

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

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

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

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

170 self.cmd.append(self.stage_io_dict["out"]["output_dhdl_path"]) 

171 

172 # general mpi properties 

173 if self.mpi_bin: 

174 mpi_cmd = [self.mpi_bin] 

175 if self.mpi_np: 

176 mpi_cmd.append('-n') 

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

178 if self.mpi_flags: 

179 mpi_cmd.extend(self.mpi_flags) 

180 self.cmd = mpi_cmd + self.cmd 

181 

182 # gromacs cpu mpi/openmp properties 

183 if self.num_threads: 

184 fu.log( 

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

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

187 self.cmd.append(self.num_threads) 

188 if self.num_threads_mpi: 

189 fu.log( 

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

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

192 self.cmd.append(self.num_threads_mpi) 

193 if self.num_threads_omp: 

194 fu.log( 

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

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

197 self.cmd.append(self.num_threads_omp) 

198 if self.num_threads_omp_pme: 

199 fu.log( 

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

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

202 self.cmd.append(self.num_threads_omp_pme) 

203 # GMX gpu properties 

204 if self.use_gpu: 

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

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

207 if self.gpu_id: 

208 fu.log( 

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

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

211 self.cmd.append(self.gpu_id) 

212 if self.gpu_tasks: 

213 fu.log( 

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

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

216 self.cmd.append(self.gpu_tasks) 

217 

218 if self.noappend: 

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

220 

221 if self.gmx_lib: 

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

223 

224 # Run Biobb block 

225 self.run_biobb() 

226 

227 # Copy files to host 

228 self.copy_to_host() 

229 

230 # Remove temporal files 

231 self.remove_tmp_files() 

232 

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

234 return self.return_code 

235 

236 def copy_to_host(self): 

237 """ 

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

239 to catch changes due to noappend restart. 

240 

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

242 if the noappend flag is used. 

243 """ 

244 import pathlib 

245 

246 def capture_part_pattern(filename): 

247 """ 

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

249 """ 

250 import re 

251 pattern = r'part\d+' 

252 

253 match = re.search(pattern, filename) 

254 if match: 

255 return match.group(0) 

256 else: 

257 return None 

258 

259 if self.noappend: 

260 # List files in the staging directory 

261 staging_path = self.stage_io_dict["unique_dir"] 

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

263 

264 # Find the part000x pattern in the output files 

265 for file in files_in_staging: 

266 part_pattern = capture_part_pattern(file.name) 

267 if part_pattern: 

268 break 

269 

270 # Update expected output files 

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

272 if stage_file_path: 

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

274 parent_path = pathlib.Path(stage_file_path).parent 

275 file_stem = pathlib.Path(stage_file_path).stem 

276 file_suffix = pathlib.Path(stage_file_path).suffix 

277 

278 # Rename all output files except checkpoint files 

279 if file_suffix != '.cpt': 

280 # Create the new file name with the part pattern 

281 if part_pattern: 

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

283 new_file_path = parent_path / new_file_name 

284 # Update the stage_io_dict with the new file path 

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

286 new_file_path) 

287 return super().copy_to_host() 

288 

289 

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

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

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

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

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

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

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

297 

298 

299mdrun.__doc__ = Mdrun.__doc__ 

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

301 

302 

303if __name__ == '__main__': 

304 main()