1"""This module defines an ASE interface to GROMACS.
2
3http://www.gromacs.org/
4It is VERY SLOW compared to standard Gromacs
5(due to slow formatted io required here).
6
7Mainly intended to be the MM part in the ase QM/MM
8
9Markus.Kaukonen@iki.fi
10
11To be done:
121) change the documentation for the new file-io-calculator (test works now)
132) change gromacs program names
14-now:     hard coded
15-future:  set as dictionary in params_runs
16
17"""
18
19import os
20import subprocess
21from glob import glob
22from shutil import which
23
24import numpy as np
25
26from ase import units
27from ase.calculators.calculator import (EnvironmentError,
28                                        FileIOCalculator,
29                                        all_changes)
30from ase.io.gromos import read_gromos, write_gromos
31
32
33def parse_gromacs_version(output):
34    import re
35    match = re.search(r'GROMACS version\:\s*(\S+)', output, re.M)
36    return match.group(1)
37
38
39def get_gromacs_version(executable):
40    output = subprocess.check_output([executable, '--version'],
41                                     encoding='utf-8')
42    return parse_gromacs_version(output)
43
44
45def do_clean(name='#*'):
46    """ remove files matching wildcards """
47    myfiles = glob(name)
48    for myfile in myfiles:
49        try:
50            os.remove(myfile)
51        except OSError:
52            pass
53
54
55class Gromacs(FileIOCalculator):
56    """Class for doing GROMACS calculations.
57    Before running a gromacs calculation you must prepare the input files
58    separately (pdb2gmx and grompp for instance.)
59
60    Input parameters for gromacs runs (the .mdp file)
61    are given in self.params and can be set when initializing the calculator
62    or by method set_own.
63    for example::
64
65        CALC_MM_RELAX = Gromacs()
66        CALC_MM_RELAX.set_own_params('integrator', 'steep',
67                                     'use steepest descent')
68
69    Run command line arguments for gromacs related programs:
70    pdb2gmx, grompp, mdrun, energy, traj.  These can be given as::
71
72        CALC_MM_RELAX = Gromacs()
73        CALC_MM_RELAX.set_own_params_runs('force_field', 'oplsaa')
74    """
75
76    implemented_properties = ['energy', 'forces']
77    discard_results_on_any_change = True
78
79    default_parameters = dict(
80        define='-DFLEXIBLE',
81        integrator='cg',
82        nsteps='10000',
83        nstfout='10',
84        nstlog='10',
85        nstenergy='10',
86        nstlist='10',
87        ns_type='grid',
88        pbc='xyz',
89        rlist='1.15',
90        coulombtype='PME-Switch',
91        rcoulomb='0.8',
92        vdwtype='shift',
93        rvdw='0.8',
94        rvdw_switch='0.75',
95        DispCorr='Ener')
96
97    def __init__(self, restart=None,
98                 ignore_bad_restart_file=FileIOCalculator._deprecated,
99                 label='gromacs', atoms=None,
100                 do_qmmm=False, clean=True,
101                 water_model='tip3p', force_field='oplsaa', command=None,
102                 **kwargs):
103        """Construct GROMACS-calculator object.
104
105        Parameters
106        ==========
107        label: str
108            Prefix to use for filenames (label.in, label.txt, ...).
109            Default is 'gromacs'.
110
111        do_qmmm : bool
112            Is gromacs used as mm calculator for a qm/mm calculation
113
114        clean :     bool
115            Remove gromacs backup files
116            and old gormacs.* files
117
118        water_model: str
119            Water model to be used in gromacs runs (see gromacs manual)
120
121        force_field: str
122            Force field to be used in gromacs runs
123
124        command : str
125            Gromacs executable; if None (default), choose available one from
126            ('gmx', 'gmx_d', 'gmx_mpi', 'gmx_mpi_d')
127        """
128
129        gmxes = ('gmx', 'gmx_d', 'gmx_mpi', 'gmx_mpi_d')
130        if command is not None:
131            self.command = command
132        else:
133            for command in gmxes:
134                if which(command):
135                    self.command = command
136                    break
137            else:
138                self.command = None
139                self.missing_gmx = 'missing gromacs executable {}'.format(gmxes)
140
141        self.do_qmmm = do_qmmm
142        self.water_model = water_model
143        self.force_field = force_field
144        self.clean = clean
145        self.params_doc = {}
146        # add comments for gromacs input file
147        self.params_doc['define'] = \
148            'flexible/ rigid water'
149        self.params_doc['integrator'] = \
150            'md: molecular dynamics(Leapfrog), \n' + \
151            '; md-vv: molecular dynamics(Velocity Verlet), \n' + \
152            '; steep: steepest descent minimization, \n' + \
153            '; cg: conjugate cradient minimization \n'
154
155        self.positions = None
156        self.atoms = None
157        # storage for energy and forces
158        #self.energy = None
159        #self.forces = None
160
161        FileIOCalculator.__init__(self, restart, ignore_bad_restart_file,
162                                  label, atoms, **kwargs)
163        self.set(**kwargs)
164        # default values for runtime parameters
165        # can be changed by self.set_own_params_runs('key', 'value')
166        self.params_runs = {}
167        self.params_runs['index_filename'] = 'index.ndx'
168        self.params_runs['init_structure'] = self.label + '.pdb'
169        self.params_runs['water'] = self.water_model
170        self.params_runs['force_field'] = self.force_field
171
172        # these below are required by qm/mm
173        self.topology_filename = self.label + '.top'
174        self.name = 'Gromacs'
175
176        # clean up gromacs backups
177        if self.clean:
178            do_clean('gromacs.???')
179
180        # write input files for gromacs program energy
181        self.write_energy_files()
182
183        if self.do_qmmm:
184            self.parameters['integrator'] = 'md'
185            self.parameters['nsteps'] = '0'
186
187    def _execute_gromacs(self, command):
188        """ execute gmx command
189        Parameters
190        ----------
191        command : str
192        """
193        if self.command:
194            subprocess.check_call(self.command + ' ' + command, shell=True)
195        else:
196            raise EnvironmentError(self.missing_gmx)
197
198    def generate_g96file(self):
199        """ from current coordinates (self.structure_file)
200            write a structure file in .g96 format
201        """
202        # generate structure file in g96 format
203        write_gromos(self.label + '.g96', self.atoms)
204
205    def run_editconf(self):
206        """ run gromacs program editconf, typically to set a simulation box
207        writing to the input structure"""
208        subcmd = 'editconf'
209        command = ' '.join([
210            subcmd,
211            '-f', self.label + '.g96',
212            '-o', self.label + '.g96',
213            self.params_runs.get('extra_editconf_parameters', ''),
214            '> {}.{}.log 2>&1'.format(self.label, subcmd)])
215        self._execute_gromacs(command)
216
217    def run_genbox(self):
218        """Run gromacs program genbox, typically to solvate the system
219        writing to the input structure
220        as extra parameter you need to define the file containing the solvent
221
222        for instance::
223
224           CALC_MM_RELAX = Gromacs()
225           CALC_MM_RELAX.set_own_params_runs(
226                'extra_genbox_parameters', '-cs spc216.gro')
227        """
228        subcmd = 'genbox'
229        command = ' '.join([
230            subcmd,
231            '-cp', self.label + '.g96',
232            '-o', self.label + '.g96',
233            '-p', self.label + '.top',
234            self.params_runs.get('extra_genbox_parameters', ''),
235            '> {}.{}.log 2>&1'.format(self.label, subcmd)])
236        self._execute_gromacs(command)
237
238    def run(self):
239        """ runs a gromacs-mdrun with the
240        current atom-configuration """
241
242        # clean up gromacs backups
243        if self.clean:
244            do_clean('#*')
245
246        subcmd = 'mdrun'
247        command = [subcmd]
248        if self.do_qmmm:
249            command += [
250                '-s', self.label + '.tpr',
251                '-o', self.label + '.trr',
252                '-e', self.label + '.edr',
253                '-g', self.label + '.log',
254                '-rerun', self.label + '.g96',
255                self.params_runs.get('extra_mdrun_parameters', ''),
256                '> QMMM.log 2>&1']
257            command = ' '.join(command)
258            self._execute_gromacs(command)
259        else:
260            command += [
261                '-s', self.label + '.tpr',
262                '-o', self.label + '.trr',
263                '-e', self.label + '.edr',
264                '-g', self.label + '.log',
265                '-c', self.label + '.g96',
266                self.params_runs.get('extra_mdrun_parameters', ''),
267                '> MM.log 2>&1']
268            command = ' '.join(command)
269            self._execute_gromacs(command)
270
271            atoms = read_gromos(self.label + '.g96')
272            self.atoms = atoms.copy()
273
274    def generate_topology_and_g96file(self):
275        """ from coordinates (self.label.+'pdb')
276            and gromacs run input file (self.label + '.mdp)
277            generate topology (self.label+'top')
278            and structure file in .g96 format (self.label + '.g96')
279        """
280        #generate structure and topology files
281        # In case of predefinded topology file this is not done
282        subcmd = 'pdb2gmx'
283        command = ' '.join([
284            subcmd,
285            '-f', self.params_runs['init_structure'],
286            '-o', self.label + '.g96',
287            '-p', self.label + '.top',
288            '-ff', self.params_runs['force_field'],
289            '-water', self.params_runs['water'],
290            self.params_runs.get('extra_pdb2gmx_parameters', ''),
291            '> {}.{}.log 2>&1'.format(self.label, subcmd)])
292        self._execute_gromacs(command)
293
294        atoms = read_gromos(self.label + '.g96')
295        self.atoms = atoms.copy()
296
297    def generate_gromacs_run_file(self):
298        """ Generates input file for a gromacs mdrun
299        based on structure file and topology file
300        resulting file is self.label + '.tpr
301        """
302
303        #generate gromacs run input file (gromacs.tpr)
304        try:
305            os.remove(self.label + '.tpr')
306        except OSError:
307            pass
308
309        subcmd = 'grompp'
310        command = ' '.join([
311            subcmd,
312            '-f', self.label + '.mdp',
313            '-c', self.label + '.g96',
314            '-p', self.label + '.top',
315            '-o', self.label + '.tpr',
316            '-maxwarn', '100',
317            self.params_runs.get('extra_grompp_parameters', ''),
318            '> {}.{}.log 2>&1'.format(self.label, subcmd)])
319        self._execute_gromacs(command)
320
321    def write_energy_files(self):
322        """write input files for gromacs force and energy calculations
323        for gromacs program energy"""
324        filename = 'inputGenergy.txt'
325        with open(filename, 'w') as output:
326            output.write('Potential  \n')
327            output.write('   \n')
328            output.write('   \n')
329
330        filename = 'inputGtraj.txt'
331        with open(filename, 'w') as output:
332            output.write('System  \n')
333            output.write('   \n')
334            output.write('   \n')
335
336    def set_own_params(self, key, value, docstring=""):
337        """Set own gromacs parameter with doc strings."""
338        self.parameters[key] = value
339        self.params_doc[key] = docstring
340
341    def set_own_params_runs(self, key, value):
342        """Set own gromacs parameter for program parameters
343        Add spaces to avoid errors """
344        self.params_runs[key] = ' ' + value + ' '
345
346    def write_input(self, atoms=None, properties=None, system_changes=None):
347        """Write input parameters to input file."""
348
349        FileIOCalculator.write_input(self, atoms, properties, system_changes)
350        #print self.parameters
351        with open(self.label + '.mdp', 'w') as myfile:
352            for key, val in self.parameters.items():
353                if val is not None:
354                    docstring = self.params_doc.get(key, '')
355                    myfile.write('%-35s = %s ; %s\n'
356                                 % (key, val, ';' + docstring))
357
358    def update(self, atoms):
359        """ set atoms and do the calculation """
360        # performs an update of the atoms
361        self.atoms = atoms.copy()
362        #must be g96 format for accuracy, alternatively binary formats
363        write_gromos(self.label + '.g96', atoms)
364        # does run to get forces and energies
365        self.calculate()
366
367    def calculate(self, atoms=None, properties=['energy', 'forces'],
368                  system_changes=all_changes):
369        """ runs a gromacs-mdrun and
370        gets energy and forces
371        rest below is to make gromacs calculator
372        compactible with ase-Calculator class
373
374        atoms: Atoms object
375            Contains positions, unit-cell, ...
376        properties: list of str
377            List of what needs to be calculated.  Can be any combination
378            of 'energy', 'forces'
379        system_changes: list of str
380            List of what has changed since last calculation.  Can be
381            any combination of these five: 'positions', 'numbers', 'cell',
382            'pbc', 'initial_charges' and 'initial_magmoms'.
383
384        """
385
386        self.run()
387        if self.clean:
388            do_clean('#*')
389        # get energy
390        try:
391            os.remove('tmp_ene.del')
392        except OSError:
393            pass
394
395        subcmd = 'energy'
396        command = ' '.join([
397            subcmd,
398            '-f', self.label + '.edr',
399            '-o', self.label + '.Energy.xvg',
400            '< inputGenergy.txt',
401            '> {}.{}.log 2>&1'.format(self.label, subcmd)])
402        self._execute_gromacs(command)
403        with open(self.label + '.Energy.xvg') as fd:
404            lastline = fd.readlines()[-1]
405            energy = float(lastline.split()[1])
406        #We go for ASE units !
407        #self.energy = energy * units.kJ / units.mol
408        self.results['energy'] = energy * units.kJ / units.mol
409        # energies are about 100 times bigger in Gromacs units
410        # when compared to ase units
411
412        subcmd = 'traj'
413        command = ' '.join([
414            subcmd,
415            '-f', self.label + '.trr',
416            '-s', self.label + '.tpr',
417            '-of', self.label + '.Force.xvg',
418            '< inputGtraj.txt',
419            '> {}.{}.log 2>&1'.format(self.label, subcmd)])
420        self._execute_gromacs(command)
421        with open(self.label + '.Force.xvg', 'r') as fd:
422            lastline = fd.readlines()[-1]
423            forces = np.array([float(f) for f in lastline.split()[1:]])
424        #We go for ASE units !gromacsForce.xvg
425        #self.forces = np.array(forces)/ units.nm * units.kJ / units.mol
426        #self.forces = np.reshape(self.forces, (-1, 3))
427        tmp_forces = forces / units.nm * units.kJ / units.mol
428        tmp_forces = np.reshape(tmp_forces, (-1, 3))
429        self.results['forces'] = tmp_forces
430        #self.forces = np.array(forces)
431