biobb_gromacs/gromacs/mdrun.py source
1 #!/usr/bin/env python3
2
3 """Module containing the MDrun class and the command line interface."""
4 import argparse
5 from typing import Optional
6 from biobb_common.generic.biobb_object import BiobbObject
7 from biobb_common.configuration import settings
8 from biobb_common.tools import file_utils as fu
9 from biobb_common.tools.file_utils import launchlogger
10 from biobb_gromacs.gromacs.common import get_gromacs_version
11
12
13 class Mdrun(BiobbObject):
14 """
15 | biobb_gromacs Mdrun
16 | Wrapper of the `GROMACS mdrun <http://manual.gromacs.org/current/onlinehelp/gmx-mdrun.html>`_ module.
17 | 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.
18
19 Args:
20 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).
21 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).
22 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).
23 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).
24 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).
25 input_cpt_path (str) (Optional): Path to the input GROMACS checkpoint file CPT. File type: input. Accepted formats: cpt (edam:format_2333).
26 output_xtc_path (str) (Optional): Path to the GROMACS compressed trajectory file XTC. File type: output. Accepted formats: xtc (edam:format_3875).
27 output_cpt_path (str) (Optional): Path to the output GROMACS checkpoint file CPT. File type: output. Accepted formats: cpt (edam:format_2333).
28 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).
29 properties (dict - Python dictionary object containing the tool parameters, not input/output files):
30 * **mpi_bin** (*str*) - (None) Path to the MPI runner. Usually "mpirun" or "srun".
31 * **mpi_np** (*int*) - (0) [0~1000|1] Number of MPI processes. Usually an integer bigger than 1.
32 * **mpi_flags** (*str*) - (None) Path to the MPI hostlist file.
33 * **checkpoint_time** (*int*) - (15) [0~1000|1] Checkpoint writing interval in minutes. Only enabled if an output_cpt_path is provided.
34 * **noappend** (*bool*) - (False) Include the noappend flag to open new output files and add the simulation part number to all output file names
35 * **num_threads** (*int*) - (0) [0~1000|1] Let GROMACS guess. The number of threads that are going to be used.
36 * **num_threads_mpi** (*int*) - (0) [0~1000|1] Let GROMACS guess. The number of GROMACS MPI threads that are going to be used.
37 * **num_threads_omp** (*int*) - (0) [0~1000|1] Let GROMACS guess. The number of GROMACS OPENMP threads that are going to be used.
38 * **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.
39 * **use_gpu** (*bool*) - (False) Use settings appropriate for GPU. Adds: -nb gpu -pme gpu
40 * **gpu_id** (*str*) - (None) list of unique GPU device IDs available to use.
41 * **gpu_tasks** (*str*) - (None) list of GPU device IDs, mapping each PP task on each node to a device.
42 * **gmx_lib** (*str*) - (None) Path set GROMACS GMXLIB environment variable.
43 * **binary_path** (*str*) - ("gmx") Path to the GROMACS executable binary.
44 * **remove_tmp** (*bool*) - (True) [WF property] Remove temporal files.
45 * **restart** (*bool*) - (False) [WF property] Do not execute if output files exist.
46 * **sandbox_path** (*str*) - ("./") [WF property] Parent path to the sandbox directory.
47 * **container_path** (*str*) - (None) Path to the binary executable of your container.
48 * **container_image** (*str*) - (None) Container Image identifier.
49 * **container_volume_path** (*str*) - ("/data") Path to an internal directory in the container.
50 * **container_working_dir** (*str*) - (None) Path to the internal CWD in the container.
51 * **container_user_id** (*str*) - (None) User number id to be mapped inside the container.
52 * **container_shell_path** (*str*) - ("/bin/bash") Path to the binary executable of the container shell.
53
54 Examples:
55 This is a use example of how to use the building block from Python::
56
57 from biobb_gromacs.gromacs.mdrun import mdrun
58 prop = { 'num_threads': 0,
59 'binary_path': 'gmx' }
60 mdrun(input_tpr_path='/path/to/myPortableBinaryRunInputFile.tpr',
61 output_trr_path='/path/to/newTrajectory.trr',
62 output_gro_path='/path/to/newStructure.gro',
63 output_edr_path='/path/to/newEnergy.edr',
64 output_log_path='/path/to/newSimulationLog.log',
65 properties=prop)
66
67 Info:
68 * wrapped_software:
69 * name: GROMACS Mdrun
70 * version: 2025.2
71 * license: LGPL 2.1
72 * multinode: mpi
73 * ontology:
74 * name: EDAM
75 * schema: http://edamontology.org/EDAM.owl
76 """
77
78 def __init__(self, input_tpr_path: str, output_gro_path: str, output_edr_path: str,
79 output_log_path: str, output_trr_path: Optional[str] = None, input_cpt_path: Optional[str] = None,
80 output_xtc_path: Optional[str] = None, output_cpt_path: Optional[str] = None,
81 output_dhdl_path: Optional[str] = None, properties: Optional[dict] = None, **kwargs) -> None:
82 properties = properties or {}
83
84 # Call parent class constructor
85 super().__init__(properties)
86 self.locals_var_dict = locals().copy()
87
88 # Input/Output files
89 self.io_dict = {
90 "in": {"input_tpr_path": input_tpr_path, "input_cpt_path": input_cpt_path},
91 "out": {"output_trr_path": output_trr_path, "output_gro_path": output_gro_path,
92 "output_edr_path": output_edr_path, "output_log_path": output_log_path,
93 "output_xtc_path": output_xtc_path, "output_cpt_path": output_cpt_path,
94 "output_dhdl_path": output_dhdl_path}
95 }
96
97 # Properties specific for BB
98 # general mpi properties
99 self.mpi_bin = properties.get('mpi_bin')
100 self.mpi_np = properties.get('mpi_np')
101 self.mpi_flags = properties.get('mpi_flags')
102 # gromacs cpu mpi/openmp properties
103 self.num_threads = str(properties.get('num_threads', ''))
104 self.num_threads_mpi = str(properties.get('num_threads_mpi', ''))
105 self.num_threads_omp = str(properties.get('num_threads_omp', ''))
106 self.num_threads_omp_pme = str(properties.get('num_threads_omp_pme', ''))
107 # gromacs gpus
108 self.use_gpu = properties.get('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(prefix=self.prefix, step=self.step, name='trajectory.trr')
142 self.tmp_files.append(self.stage_io_dict["out"]["output_trr_path"])
143
144 self.stage_files()
145
146 self.cmd = [self.binary_path, 'mdrun',
147 '-o', self.stage_io_dict["out"]["output_trr_path"],
148 '-s', self.stage_io_dict["in"]["input_tpr_path"],
149 '-c', self.stage_io_dict["out"]["output_gro_path"],
150 '-e', self.stage_io_dict["out"]["output_edr_path"],
151 '-g', self.stage_io_dict["out"]["output_log_path"]]
152
153 if self.stage_io_dict["in"].get("input_cpt_path"):
154 self.cmd.append('-cpi')
155 self.cmd.append(self.stage_io_dict["in"]["input_cpt_path"])
156 if self.stage_io_dict["out"].get("output_xtc_path"):
157 self.cmd.append('-x')
158 self.cmd.append(self.stage_io_dict["out"]["output_xtc_path"])
159 else:
160 self.tmp_files.append('traj_comp.xtc')
161 if self.stage_io_dict["out"].get("output_cpt_path"):
162 self.cmd.append('-cpo')
163 self.cmd.append(self.stage_io_dict["out"]["output_cpt_path"])
164 if self.checkpoint_time:
165 self.cmd.append('-cpt')
166 self.cmd.append(str(self.checkpoint_time))
167 if self.stage_io_dict["out"].get("output_dhdl_path"):
168 self.cmd.append('-dhdl')
169 self.cmd.append(self.stage_io_dict["out"]["output_dhdl_path"])
170
171 # general mpi properties
172 if self.mpi_bin:
173 mpi_cmd = [self.mpi_bin]
174 if self.mpi_np:
175 mpi_cmd.append('-n')
176 mpi_cmd.append(str(self.mpi_np))
177 if self.mpi_flags:
178 mpi_cmd.extend(self.mpi_flags)
179 self.cmd = mpi_cmd + self.cmd
180
181 # gromacs cpu mpi/openmp properties
182 if self.num_threads:
183 fu.log(f'User added number of gmx threads: {self.num_threads}', self.out_log)
184 self.cmd.append('-nt')
185 self.cmd.append(self.num_threads)
186 if self.num_threads_mpi:
187 fu.log(f'User added number of gmx mpi threads: {self.num_threads_mpi}', self.out_log)
188 self.cmd.append('-ntmpi')
189 self.cmd.append(self.num_threads_mpi)
190 if self.num_threads_omp:
191 fu.log(f'User added number of gmx omp threads: {self.num_threads_omp}', self.out_log)
192 self.cmd.append('-ntomp')
193 self.cmd.append(self.num_threads_omp)
194 if self.num_threads_omp_pme:
195 fu.log(f'User added number of gmx omp_pme threads: {self.num_threads_omp_pme}', self.out_log)
196 self.cmd.append('-ntomp_pme')
197 self.cmd.append(self.num_threads_omp_pme)
198 # GMX gpu properties
199 if self.use_gpu:
200 fu.log('Adding GPU specific settings adds: -nb gpu -pme gpu', self.out_log)
201 self.cmd += ["-nb", "gpu", "-pme", "gpu"]
202 if self.gpu_id:
203 fu.log(f'list of unique GPU device IDs available to use: {self.gpu_id}', self.out_log)
204 self.cmd.append('-gpu_id')
205 self.cmd.append(self.gpu_id)
206 if self.gpu_tasks:
207 fu.log(f'list of GPU device IDs, mapping each PP task on each node to a device: {self.gpu_tasks}', self.out_log)
208 self.cmd.append('-gputasks')
209 self.cmd.append(self.gpu_tasks)
- W293 Blank line contains whitespace
211 if self.noappend:
212 self.cmd.append('-noappend')
213
214 if self.gmx_lib:
215 self.env_vars_dict['GMXLIB'] = self.gmx_lib
216
217 # Run Biobb block
218 self.run_biobb()
219
220 # Copy files to host
221 self.copy_to_host()
222
223 # Remove temporal files
224 self.remove_tmp_files()
225
226 self.check_arguments(output_files_created=True, raise_exception=False)
227 return self.return_code
228
229 def copy_to_host(self):
230 """
231 Updates the path to the original output files in the sandbox,
232 to catch changes due to noappend restart.
233
234 GROMACS mdrun will change the output file names from md.gro to md.part0001.gro
235 if the noappend flag is used.
236 """
237 import pathlib
- W293 Blank line contains whitespace
239 def capture_part_pattern(filename):
240 """
241 Captures the 'part' pattern followed by digits from a string.
242 """
243 import re
244 pattern = r'part\d+'
- W293 Blank line contains whitespace
246 match = re.search(pattern, filename)
247 if match:
248 return match.group(0)
249 else:
250 return None
251
252 if self.noappend:
253 # List files in the staging directory
254 staging_path = self.stage_io_dict["unique_dir"]
255 files_in_staging = list(pathlib.Path(staging_path).glob('*'))
256
257 # Find the part000x pattern in the output files
258 for file in files_in_staging:
259 part_pattern = capture_part_pattern(file.name)
260 if part_pattern:
261 break
262
263 # Update expected output files
264 for file_ref, stage_file_path in self.stage_io_dict["out"].items():
265 if stage_file_path:
266 # Find the parent and the file name in the sandbox
267 parent_path = pathlib.Path(stage_file_path).parent
268 file_stem = pathlib.Path(stage_file_path).stem
269 file_suffix = pathlib.Path(stage_file_path).suffix
- W293 Blank line contains whitespace
271 # Rename all output files except checkpoint files
272 if file_suffix != '.cpt':
273 # Create the new file name with the part pattern
274 if part_pattern:
275 new_file_name = f"{file_stem}.{part_pattern}{file_suffix}"
276 new_file_path = parent_path / new_file_name
277 # Update the stage_io_dict with the new file path
278 self.stage_io_dict["out"][file_ref] = str(new_file_path)
279
280 return super().copy_to_host()
281
- E302 Expected 2 blank lines, found 1
283 output_log_path: str, output_trr_path: Optional[str] = None, input_cpt_path: Optional[str] = None,
284 output_xtc_path: Optional[str] = None, output_cpt_path: Optional[str] = None,
285 output_dhdl_path: Optional[str] = None, properties: Optional[dict] = None, **kwargs) -> int:
286 """Create :class:`Mdrun <gromacs.mdrun.Mdrun>` class and
287 execute the :meth:`launch() <gromacs.mdrun.Mdrun.launch>` method."""
288
289 return Mdrun(input_tpr_path=input_tpr_path, output_trr_path=output_trr_path,
290 output_gro_path=output_gro_path, output_edr_path=output_edr_path,
291 output_log_path=output_log_path, input_cpt_path=input_cpt_path,
292 output_xtc_path=output_xtc_path, output_cpt_path=output_cpt_path,
293 output_dhdl_path=output_dhdl_path, properties=properties,
294 **kwargs).launch()
295
296
297 mdrun.__doc__ = Mdrun.__doc__
298
299
300 def main():
301 """Command line execution of this building block. Please check the command line documentation."""
302 parser = argparse.ArgumentParser(description="Wrapper for the GROMACS mdrun module.",
303 formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, width=99999))
304 parser.add_argument('-c', '--config', required=False, help="This file can be a YAML file, JSON file or JSON string")
305
306 # Specific args of each building block
307 required_args = parser.add_argument_group('required arguments')
308 required_args.add_argument('--input_tpr_path', required=True)
309 required_args.add_argument('--output_gro_path', required=True)
310 required_args.add_argument('--output_edr_path', required=True)
311 required_args.add_argument('--output_log_path', required=True)
312 parser.add_argument('--output_trr_path', required=False)
313 parser.add_argument('--input_cpt_path', required=False)
314 parser.add_argument('--output_xtc_path', required=False)
315 parser.add_argument('--output_cpt_path', required=False)
316 parser.add_argument('--output_dhdl_path', required=False)
317
318 args = parser.parse_args()
319 config = args.config if args.config else None
320 properties = settings.ConfReader(config=config).get_prop_dic()
321
322 # Specific call of each building block
323 mdrun(input_tpr_path=args.input_tpr_path, output_trr_path=args.output_trr_path,
324 output_gro_path=args.output_gro_path, output_edr_path=args.output_edr_path,
325 output_log_path=args.output_log_path, input_cpt_path=args.input_cpt_path,
326 output_xtc_path=args.output_xtc_path, output_cpt_path=args.output_cpt_path,
327 output_dhdl_path=args.output_dhdl_path, properties=properties)
328
329
330 if __name__ == '__main__':
331 main()