forcefield.py 244 KB
Newer Older
1
2
"""
forcefield.py: Constructs OpenMM System objects based on a Topology and an XML force field description
3
4
5
6
7
8

This is part of the OpenMM molecular simulation toolkit originating from
Simbios, the NIH National Center for Physics-Based Simulation of
Biological Structures at Stanford, funded under the NIH Roadmap for
Medical Research, grant U54 GM072970. See https://simtk.org.

9
Portions copyright (c) 2012-2018 Stanford University and the Authors.
10
11
12
Authors: Peter Eastman, Mark Friedrichs
Contributors:

Justin MacCallum's avatar
Justin MacCallum committed
13
Permission is hereby granted, free of charge, to any person obtaining a
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
USE OR OTHER DEALINGS IN THE SOFTWARE.
30
"""
31
32
from __future__ import absolute_import, print_function

33
34
35
36
37
38
__author__ = "Peter Eastman"
__version__ = "1.0"

import os
import itertools
import xml.etree.ElementTree as etree
39
import math
40
from math import sqrt, cos
41
from copy import deepcopy
peastman's avatar
peastman committed
42
from collections import defaultdict
43
44
import simtk.openmm as mm
import simtk.unit as unit
45
from . import element as elem
46
from simtk.openmm.app import Topology
47
from simtk.openmm.app.internal.singleton import Singleton
peastman's avatar
peastman committed
48
from simtk.openmm.app.internal import compiled
49

50
51
# Directories from which to load built in force fields.

52
53
54
55
56
57
58
59
60
61
62
63
64
_dataDirectories = None

def _getDataDirectories():
    global _dataDirectories
    if _dataDirectories is None:
        _dataDirectories = [os.path.join(os.path.dirname(__file__), 'data')]
        try:
            from pkg_resources import iter_entry_points
            for entry in iter_entry_points(group='openmm.forcefielddir'):
                _dataDirectories.append(entry.load()())
        except:
            pass # pkg_resources is not installed
    return _dataDirectories
65

66
67
def _convertParameterToNumber(param):
    if unit.is_quantity(param):
68
69
70
        if param.unit.is_compatible(unit.bar):
            return param / unit.bar
        return param.value_in_unit_system(unit.md_unit_system)
71
72
    return float(param)

73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
def _parseFunctions(element):
    """Parse the attributes on an XML tag to find any tabulated functions it defines."""
    functions = []
    for function in element.findall('Function'):
        values = [float(x) for x in function.text.split()]
        if 'type' in function.attrib:
            functionType = function.attrib['type']
        else:
            functionType = 'Continuous1D'
        params = {}
        for key in function.attrib:
            if key.endswith('size'):
                params[key] = int(function.attrib[key])
            elif key.endswith('min') or key.endswith('max'):
                params[key] = float(function.attrib[key])
        functions.append((function.attrib['name'], functionType, values, params))
    return functions

def _createFunctions(force, functions):
    """Add TabulatedFunctions to a Force based on the information that was recorded by _parseFunctions()."""
    for (name, type, values, params) in functions:
        if type == 'Continuous1D':
            force.addTabulatedFunction(name, mm.Continuous1DFunction(values, params['min'], params['max']))
        elif type == 'Continuous2D':
            force.addTabulatedFunction(name, mm.Continuous2DFunction(params['xsize'], params['ysize'], values, params['xmin'], params['xmax'], params['ymin'], params['ymax']))
        elif type == 'Continuous3D':
            force.addTabulatedFunction(name, mm.Continuous2DFunction(params['xsize'], params['ysize'], params['zsize'], values, params['xmin'], params['xmax'], params['ymin'], params['ymax'], params['zmin'], params['zmax']))
        elif type == 'Discrete1D':
            force.addTabulatedFunction(name, mm.Discrete1DFunction(values))
        elif type == 'Discrete2D':
            force.addTabulatedFunction(name, mm.Discrete2DFunction(params['xsize'], params['ysize'], values))
        elif type == 'Discrete3D':
            force.addTabulatedFunction(name, mm.Discrete2DFunction(params['xsize'], params['ysize'], params['zsize'], values))

107
108
# Enumerated values for nonbonded method

109
class NoCutoff(Singleton):
110
111
112
113
    def __repr__(self):
        return 'NoCutoff'
NoCutoff = NoCutoff()

114
class CutoffNonPeriodic(Singleton):
115
116
117
118
    def __repr__(self):
        return 'CutoffNonPeriodic'
CutoffNonPeriodic = CutoffNonPeriodic()

119
class CutoffPeriodic(Singleton):
120
121
122
123
    def __repr__(self):
        return 'CutoffPeriodic'
CutoffPeriodic = CutoffPeriodic()

124
class Ewald(Singleton):
125
126
127
128
    def __repr__(self):
        return 'Ewald'
Ewald = Ewald()

129
class PME(Singleton):
130
131
132
    def __repr__(self):
        return 'PME'
PME = PME()
133

Peter Eastman's avatar
Peter Eastman committed
134
135
136
137
138
class LJPME(Singleton):
    def __repr__(self):
        return 'LJPME'
LJPME = LJPME()

139
140
# Enumerated values for constraint type

141
class HBonds(Singleton):
142
143
144
145
    def __repr__(self):
        return 'HBonds'
HBonds = HBonds()

146
class AllBonds(Singleton):
147
148
149
150
    def __repr__(self):
        return 'AllBonds'
AllBonds = AllBonds()

151
class HAngles(Singleton):
152
153
154
    def __repr__(self):
        return 'HAngles'
HAngles = HAngles()
155
156
157
158
159
160
161
162
163
164

# A map of functions to parse elements of the XML file.

parsers = {}

class ForceField(object):
    """A ForceField constructs OpenMM System objects based on a Topology."""

    def __init__(self, *files):
        """Load one or more XML files and create a ForceField object based on them.
Justin MacCallum's avatar
Justin MacCallum committed
165

Robert McGibbon's avatar
Robert McGibbon committed
166
167
168
169
170
171
172
173
        Parameters
        ----------
        files : list
            A list of XML files defining the force field.  Each entry may
            be an absolute file path, a path relative to the current working
            directory, a path relative to this module's data subdirectory
            (for built in force fields), or an open file-like object with a
            read() method from which the forcefield XML data can be loaded.
174
175
176
        """
        self._atomTypes = {}
        self._templates = {}
177
178
        self._patches = {}
        self._templatePatches = {}
179
        self._templateSignatures = {None:[]}
180
        self._atomClasses = {'':set()}
181
        self._forces = []
182
        self._scripts = []
183
        self._templateGenerators = []
184
        self.loadFile(files)
185

186
    def loadFile(self, files):
187
        """Load an XML file and add the definitions from it to this ForceField.
188

Robert McGibbon's avatar
Robert McGibbon committed
189
190
        Parameters
        ----------
191
192
193
        files : string or file or tuple
            An XML file or tuple of XML files containing force field definitions.
            Each entry may be either an absolute file path, a path relative to the current working
Robert McGibbon's avatar
Robert McGibbon committed
194
195
196
            directory, a path relative to this module's data subdirectory (for
            built in force fields), or an open file-like object with a read()
            method from which the forcefield XML data can be loaded.
peastman's avatar
peastman committed
197
        """
198

199
200
201
202
        if isinstance(files, tuple):
            files = list(files)
        else:
            files = [files]
203
204
205

        trees = []

206
207
208
        i = 0
        while i < len(files):
            file = files[i]
209
            tree = None
210
211
212
213
            try:
                # this handles either filenames or open file-like objects
                tree = etree.parse(file)
            except IOError:
214
                for dataDir in _getDataDirectories():
215
216
217
218
                    f = os.path.join(dataDir, file)
                    if os.path.isfile(f):
                        tree = etree.parse(f)
                        break
219
220
221
222
223
224
225
226
227
228
229
            except Exception as e:
                # Fail with an error message about which file could not be read.
                # TODO: Also handle case where fallback to 'data' directory encounters problems,
                # but this is much less worrisome because we control those files.
                msg  = str(e) + '\n'
                if hasattr(file, 'name'):
                    filename = file.name
                else:
                    filename = str(file)
                msg += "ForceField.loadFile() encountered an error reading file '%s'\n" % filename
                raise Exception(msg)
230
231
            if tree is None:
                raise ValueError('Could not locate file "%s"' % file)
232
233

            trees.append(tree)
234
            i += 1
235

236
            # Process includes in this file.
237

238
239
            if isinstance(file, str):
                parentDir = os.path.dirname(file)
peastman's avatar
peastman committed
240
241
            else:
                parentDir = ''
peastman's avatar
peastman committed
242
243
            for included in tree.getroot().findall('Include'):
                includeFile = included.attrib['file']
peastman's avatar
peastman committed
244
245
246
                joined = os.path.join(parentDir, includeFile)
                if os.path.isfile(joined):
                    includeFile = joined
247
248
                if includeFile not in files:
                    files.append(includeFile)
peastman's avatar
peastman committed
249
250
251

        # Load the atom types.

252
253
254
255
        for tree in trees:
            if tree.getroot().find('AtomTypes') is not None:
                for type in tree.getroot().find('AtomTypes').findall('Type'):
                    self.registerAtomType(type.attrib)
peastman's avatar
peastman committed
256
257
258

        # Load the residue templates.

259
260
261
262
263
        for tree in trees:
            if tree.getroot().find('Residues') is not None:
                for residue in tree.getroot().find('Residues').findall('Residue'):
                    resName = residue.attrib['name']
                    template = ForceField._TemplateData(resName)
264
265
                    if 'override' in residue.attrib:
                        template.overrideLevel = int(residue.attrib['override'])
266
267
                    atomIndices = template.atomIndices
                    for ia, atom in enumerate(residue.findall('Atom')):
268
                        params = {}
269
270
271
272
                        for key in atom.attrib:
                            if key not in ('name', 'type'):
                                params[key] = _convertParameterToNumber(atom.attrib[key])
                        atomName = atom.attrib['name']
273
274
                        if atomName in atomIndices:
                            raise ValueError('Residue '+resName+' contains multiple atoms named '+atomName)
275
                        typeName = atom.attrib['type']
276
                        atomIndices[atomName] = ia
277
278
279
280
281
282
283
284
285
286
287
288
289
                        template.atoms.append(ForceField._TemplateAtomData(atomName, typeName, self._atomTypes[typeName].element, params))
                    for site in residue.findall('VirtualSite'):
                        template.virtualSites.append(ForceField._VirtualSiteData(site, atomIndices))
                    for bond in residue.findall('Bond'):
                        if 'atomName1' in bond.attrib:
                            template.addBondByName(bond.attrib['atomName1'], bond.attrib['atomName2'])
                        else:
                            template.addBond(int(bond.attrib['from']), int(bond.attrib['to']))
                    for bond in residue.findall('ExternalBond'):
                        if 'atomName' in bond.attrib:
                            template.addExternalBondByName(bond.attrib['atomName'])
                        else:
                            template.addExternalBond(int(bond.attrib['from']))
290
                    for patch in residue.findall('AllowPatch'):
291
                        patchName = patch.attrib['name']
292
293
                        if ':' in patchName:
                            colonIndex = patchName.find(':')
294
295
296
                            self.registerTemplatePatch(resName, patchName[:colonIndex], int(patchName[colonIndex+1:])-1)
                        else:
                            self.registerTemplatePatch(resName, patchName, 0)
297
                    self.registerResidueTemplate(template)
peastman's avatar
peastman committed
298

299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
        # Load the patch defintions.

        for tree in trees:
            if tree.getroot().find('Patches') is not None:
                for patch in tree.getroot().find('Patches').findall('Patch'):
                    patchName = patch.attrib['name']
                    if 'residues' in patch.attrib:
                        numResidues = int(patch.attrib['residues'])
                    else:
                        numResidues = 1
                    patchData = ForceField._PatchData(patchName, numResidues)
                    for atom in patch.findall('AddAtom'):
                        params = {}
                        for key in atom.attrib:
                            if key not in ('name', 'type'):
                                params[key] = _convertParameterToNumber(atom.attrib[key])
                        atomName = atom.attrib['name']
peastman's avatar
peastman committed
316
                        if atomName in patchData.allAtomNames:
317
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
318
                        patchData.allAtomNames.add(atomName)
319
320
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
321
                        patchData.addedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
322
323
324
325
326
327
                    for atom in patch.findall('ChangeAtom'):
                        params = {}
                        for key in atom.attrib:
                            if key not in ('name', 'type'):
                                params[key] = _convertParameterToNumber(atom.attrib[key])
                        atomName = atom.attrib['name']
peastman's avatar
peastman committed
328
                        if atomName in patchData.allAtomNames:
329
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
330
                        patchData.allAtomNames.add(atomName)
331
332
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
333
                        patchData.changedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
334
335
                    for atom in patch.findall('RemoveAtom'):
                        atomName = atom.attrib['name']
peastman's avatar
peastman committed
336
                        if atomName in patchData.allAtomNames:
337
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
338
                        patchData.allAtomNames.add(atomName)
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
                        atomDescription = ForceField._PatchAtomData(atomName)
                        patchData.deletedAtoms.append(atomDescription)
                    for bond in patch.findall('AddBond'):
                        atom1 = ForceField._PatchAtomData(bond.attrib['atomName1'])
                        atom2 = ForceField._PatchAtomData(bond.attrib['atomName2'])
                        patchData.addedBonds.append((atom1, atom2))
                    for bond in patch.findall('RemoveBond'):
                        atom1 = ForceField._PatchAtomData(bond.attrib['atomName1'])
                        atom2 = ForceField._PatchAtomData(bond.attrib['atomName2'])
                        patchData.deletedBonds.append((atom1, atom2))
                    for bond in patch.findall('AddExternalBond'):
                        atom = ForceField._PatchAtomData(bond.attrib['atomName'])
                        patchData.addedExternalBonds.append(atom)
                    for bond in patch.findall('RemoveExternalBond'):
                        atom = ForceField._PatchAtomData(bond.attrib['atomName'])
                        patchData.deletedExternalBonds.append(atom)
                    for residue in patch.findall('ApplyToResidue'):
                        name = residue.attrib['name']
                        if ':' in name:
                            colonIndex = name.find(':')
                            self.registerTemplatePatch(name[colonIndex+1:], patchName, int(name[:colonIndex])-1)
                        else:
                            self.registerTemplatePatch(name, patchName, 0)
                    self.registerPatch(patchData)

peastman's avatar
peastman committed
364
365
        # Load force definitions

366
367
368
369
        for tree in trees:
            for child in tree.getroot():
                if child.tag in parsers:
                    parsers[child.tag](child, self)
peastman's avatar
peastman committed
370
371
372

        # Load scripts

373
374
375
        for tree in trees:
            for node in tree.getroot().findall('Script'):
                self.registerScript(node.text)
376

377
378
379
    def getGenerators(self):
        """Get the list of all registered generators."""
        return self._forces
380

381
382
383
    def registerGenerator(self, generator):
        """Register a new generator."""
        self._forces.append(generator)
384

385
386
387
388
389
390
391
392
393
394
395
396
    def registerAtomType(self, parameters):
        """Register a new atom type."""
        name = parameters['name']
        if name in self._atomTypes:
            raise ValueError('Found multiple definitions for atom type: '+name)
        atomClass = parameters['class']
        mass = _convertParameterToNumber(parameters['mass'])
        element = None
        if 'element' in parameters:
            element = parameters['element']
            if not isinstance(element, elem.Element):
                element = elem.get_by_symbol(element)
397
        self._atomTypes[name] = ForceField._AtomType(name, atomClass, mass, element)
398
399
400
401
402
403
404
        if atomClass in self._atomClasses:
            typeSet = self._atomClasses[atomClass]
        else:
            typeSet = set()
            self._atomClasses[atomClass] = typeSet
        typeSet.add(name)
        self._atomClasses[''].add(name)
405

406
407
    def registerResidueTemplate(self, template):
        """Register a new residue template."""
408
409
        if template.name in self._templates:
            # There is already a template with this name, so check the override levels.
410

411
412
413
414
415
416
417
418
419
420
421
            existingTemplate = self._templates[template.name]
            if template.overrideLevel < existingTemplate.overrideLevel:
                # The existing one takes precedence, so just return.
                return
            if template.overrideLevel > existingTemplate.overrideLevel:
                # We need to delete the existing template.
                del self._templates[template.name]
                existingSignature = _createResidueSignature([atom.element for atom in existingTemplate.atoms])
                self._templateSignatures[existingSignature].remove(existingTemplate)
            else:
                raise ValueError('Residue template %s with the same override level %d already exists.' % (template.name, template.overrideLevel))
422

423
        # Register the template.
424

425
426
427
        self._templates[template.name] = template
        signature = _createResidueSignature([atom.element for atom in template.atoms])
        if signature in self._templateSignatures:
428
            self._templateSignatures[signature].append(template)
429
430
        else:
            self._templateSignatures[signature] = [template]
431

432
433
434
    def registerPatch(self, patch):
        """Register a new patch that can be applied to templates."""
        self._patches[patch.name] = patch
435

436
437
438
    def registerTemplatePatch(self, residue, patch, patchResidueIndex):
        """Register that a particular patch can be used with a particular residue."""
        if residue not in self._templatePatches:
439
440
            self._templatePatches[residue] = set()
        self._templatePatches[residue].add((patch, patchResidueIndex))
441

442
443
444
    def registerScript(self, script):
        """Register a new script to be executed after building the System."""
        self._scripts.append(script)
445

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
446
    def registerTemplateGenerator(self, generator):
447
448
449
450
        """Register a residue template generator that can be used to parameterize residues that do not match existing forcefield templates.

        This functionality can be used to add handlers to parameterize small molecules or unnatural/modified residues.

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
451
452
        .. CAUTION:: This method is experimental, and its API is subject to change.

453
454
        Parameters
        ----------
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
455
        generator : function
456
457
            A function that will be called when a residue is encountered that does not match an existing forcefield template.

John Chodera's avatar
John Chodera committed
458
        When a residue without a template is encountered, the ``generator`` function is called with:
459

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
460
461
        ::
           success = generator(forcefield, residue)
462

John Chodera's avatar
John Chodera committed
463
464
465
        where ``forcefield`` is the calling ``ForceField`` object and ``residue`` is a simtk.openmm.app.topology.Residue object.

        ``generator`` must conform to the following API:
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
466
467

        ::
John Chodera's avatar
John Chodera committed
468
469
470
           generator API

           Parameters
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
471
472
473
474
475
           ----------
           forcefield : simtk.openmm.app.ForceField
               The ForceField object to which residue templates and/or parameters are to be added.
           residue : simtk.openmm.app.Topology.Residue
               The residue topology for which a template is to be generated.
476

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
477
478
479
480
481
           Returns
           -------
           success : bool
               If the generator is able to successfully parameterize the residue, `True` is returned.
               If the generator cannot parameterize the residue, it should return `False` and not modify `forcefield`.
482

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
483
484
485
486
487
           The generator should either register a residue template directly with `forcefield.registerResidueTemplate(template)`
           or it should call `forcefield.loadFile(file)` to load residue definitions from an ffxml file.

           It can also use the `ForceField` programmatic API to add additional atom types (via `forcefield.registerAtomType(parameters)`)
           or additional parameters.
488
489

        """
490
        self._templateGenerators.append(generator)
491

492
    def _findAtomTypes(self, attrib, num):
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
        """Parse the attributes on an XML tag to find the set of atom types for each atom it involves.

        Parameters
        ----------
        attrib : dict of attributes
            The dictionary of attributes for an XML parameter tag.
        num : int
            The number of atom specifiers (e.g. 'class1' through 'class4') to extract.

        Returns
        -------
        types : list
            A list of atom types that match.

        """
508
509
510
511
512
513
514
        types = []
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
515
            typeAttrib = 'type'+suffix
516
            if classAttrib in attrib:
517
                if typeAttrib in attrib:
518
                    raise ValueError('Specified both a type and a class for the same atom: '+str(attrib))
519
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
520
521
522
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
523
524
            elif typeAttrib in attrib:
                if attrib[typeAttrib] == '':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
525
                    types.append(self._atomClasses[''])
526
                elif attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
527
528
529
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
530
531
            else:
                types.append(None) # Unknown atom type
532
533
        return types

534
    def _parseTorsion(self, attrib):
535
        """Parse the node defining a torsion."""
536
        types = self._findAtomTypes(attrib, 4)
peastman's avatar
peastman committed
537
        if None in types:
538
539
540
541
542
            return None
        torsion = PeriodicTorsion(types)
        index = 1
        while 'phase%d'%index in attrib:
            torsion.periodicity.append(int(attrib['periodicity%d'%index]))
543
544
            torsion.phase.append(_convertParameterToNumber(attrib['phase%d'%index]))
            torsion.k.append(_convertParameterToNumber(attrib['k%d'%index]))
545
            index += 1
Justin MacCallum's avatar
Justin MacCallum committed
546
547
        return torsion

548
    class _SystemData(object):
549
550
551
        """Inner class used to encapsulate data about the system being created."""
        def __init__(self):
            self.atomType = {}
552
            self.atomParameters = {}
553
            self.atomTemplateIndexes = {}
554
            self.atoms = []
555
            self.excludeAtomWith = []
556
            self.virtualSites = {}
557
558
559
560
561
562
            self.bonds = []
            self.angles = []
            self.propers = []
            self.impropers = []
            self.atomBonds = []
            self.isAngleConstrained = []
563
564
565
566
567
568
            self.constraints = {}

        def addConstraint(self, system, atom1, atom2, distance):
            """Add a constraint to the system, avoiding duplicate constraints."""
            key = (min(atom1, atom2), max(atom1, atom2))
            if key in self.constraints:
569
                if self.constraints[key] != distance:
570
571
572
573
                    raise ValueError('Two constraints were specified between atoms %d and %d with different distances' % (atom1, atom2))
            else:
                self.constraints[key] = distance
                system.addConstraint(atom1, atom2, distance)
574

575
576
577
578
579
580
        def recordMatchedAtomParameters(self, residue, template, matches):
            """Record parameters for atoms based on having matched a residue to a template."""
            matchAtoms = dict(zip(matches, residue.atoms()))
            for atom, match in zip(residue.atoms(), matches):
                self.atomType[atom] = template.atoms[match].type
                self.atomParameters[atom] = template.atoms[match].parameters
581
                self.atomTemplateIndexes[atom] = match
582
583
584
                for site in template.virtualSites:
                    if match == site.index:
                        self.virtualSites[atom] = (site, [matchAtoms[i].index for i in site.atoms], matchAtoms[site.excludeWith].index)
585

586
    class _TemplateData(object):
587
588
589
590
        """Inner class used to encapsulate data about a residue template definition."""
        def __init__(self, name):
            self.name = name
            self.atoms = []
591
            self.atomIndices = {}
592
            self.virtualSites = []
593
594
            self.bonds = []
            self.externalBonds = []
595
            self.overrideLevel = 0
596

597
598
        def getAtomIndexByName(self, atom_name):
            """Look up an atom index by atom name, providing a helpful error message if not found."""
599
600
601
            index = self.atomIndices.get(atom_name, None)
            if index is not None:
                return index
602
603

            # Provide a helpful error message if atom name not found.
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
604
            msg =  "Atom name '%s' not found in residue template '%s'.\n" % (atom_name, self.name)
605
            msg += "Possible atom names are: %s" % str(list(map(lambda x: x.name, self.atoms)))
606
607
            raise ValueError(msg)

608
609
        def addAtom(self, atom):
            self.atoms.append(atom)
610
            self.atomIndices[atom.name] = len(self.atoms)-1
611

612
        def addBond(self, atom1, atom2):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
613
            """Add a bond between two atoms in a template given their indices in the template."""
614
615
616
            self.bonds.append((atom1, atom2))
            self.atoms[atom1].bondedTo.append(atom2)
            self.atoms[atom2].bondedTo.append(atom1)
617

618
        def addBondByName(self, atom1_name, atom2_name):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
619
            """Add a bond between two atoms in a template given their atom names."""
620
621
622
623
624
            atom1 = self.getAtomIndexByName(atom1_name)
            atom2 = self.getAtomIndexByName(atom2_name)
            self.addBond(atom1, atom2)

        def addExternalBond(self, atom_index):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
625
            """Designate that an atom in a residue template has an external bond, using atom index within template."""
626
627
628
629
            self.externalBonds.append(atom_index)
            self.atoms[atom_index].externalBonds += 1

        def addExternalBondByName(self, atom_name):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
630
            """Designate that an atom in a residue template has an external bond, using atom name within template."""
631
632
633
            atom = self.getAtomIndexByName(atom_name)
            self.addExternalBond(atom)

634
    class _TemplateAtomData(object):
635
        """Inner class used to encapsulate data about an atom in a residue template definition."""
636
        def __init__(self, name, type, element, parameters={}):
637
638
639
            self.name = name
            self.type = type
            self.element = element
640
            self.parameters = parameters
641
642
643
            self.bondedTo = []
            self.externalBonds = 0

644
    class _BondData(object):
645
646
647
648
649
650
        """Inner class used to encapsulate data about a bond."""
        def __init__(self, atom1, atom2):
            self.atom1 = atom1
            self.atom2 = atom2
            self.isConstrained = False
            self.length = 0.0
Justin MacCallum's avatar
Justin MacCallum committed
651

652
    class _VirtualSiteData(object):
653
        """Inner class used to encapsulate data about a virtual site."""
654
        def __init__(self, node, atomIndices):
655
656
657
            attrib = node.attrib
            self.type = attrib['type']
            if self.type == 'average2':
658
                numAtoms = 2
659
660
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
661
                numAtoms = 3
662
663
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
664
                numAtoms = 3
665
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
666
            elif self.type == 'localCoords':
667
668
669
670
671
672
673
674
675
                numAtoms = 0
                self.originWeights = []
                self.xWeights = []
                self.yWeights = []
                while ('wo%d' % (numAtoms+1)) in attrib:
                    numAtoms += 1
                    self.originWeights.append(float(attrib['wo%d' % numAtoms]))
                    self.xWeights.append(float(attrib['wx%d' % numAtoms]))
                    self.yWeights.append(float(attrib['wy%d' % numAtoms]))
676
                self.localPos = [float(attrib['p1']), float(attrib['p2']), float(attrib['p3'])]
677
678
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
679
680
681
682
683
684
            if 'siteName' in attrib:
                self.index = atomIndices[attrib['siteName']]
                self.atoms = [atomIndices[attrib['atomName%d'%(i+1)]] for i in range(numAtoms)]
            else:
                self.index = int(attrib['index'])
                self.atoms = [int(attrib['atom%d'%(i+1)]) for i in range(numAtoms)]
685
686
687
688
            if 'excludeWith' in attrib:
                self.excludeWith = int(attrib['excludeWith'])
            else:
                self.excludeWith = self.atoms[0]
689

690
691
692
693
694
    class _PatchData(object):
        """Inner class used to encapsulate data about a patch definition."""
        def __init__(self, name, numResidues):
            self.name = name
            self.numResidues = numResidues
695
696
            self.addedAtoms = [[] for i in range(numResidues)]
            self.changedAtoms = [[] for i in range(numResidues)]
697
698
699
700
701
            self.deletedAtoms = []
            self.addedBonds = []
            self.deletedBonds = []
            self.addedExternalBonds = []
            self.deletedExternalBonds = []
peastman's avatar
peastman committed
702
            self.allAtomNames = set()
703

704
705
706
707
        def createPatchedTemplates(self, templates):
            """Apply this patch to a set of templates, creating new modified ones."""
            if len(templates) != self.numResidues:
                raise ValueError("Patch '%s' expected %d templates, received %d", (self.name, self.numResidues, len(templates)))
708

709
            # Construct a new version of each template.
710

711
712
713
714
            newTemplates = []
            for index, template in enumerate(templates):
                newTemplate = ForceField._TemplateData("%s-%s" % (template.name, self.name))
                newTemplates.append(newTemplate)
715

716
                # Build the list of atoms in it.
717

718
719
                for atom in template.atoms:
                    if not any(deleted.name == atom.name and deleted.residue == index for deleted in self.deletedAtoms):
720
                        newTemplate.addAtom(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
721
                for atom in self.addedAtoms[index]:
722
723
                    if any(a.name == atom.name for a in newTemplate.atoms):
                        raise ValueError("Patch '%s' adds an atom with the same name as an existing atom: %s" % (self.name, atom.name))
724
                    newTemplate.addAtom(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
725
726
727
728
729
                oldAtomIndex = dict([(atom.name, i) for i, atom in enumerate(template.atoms)])
                newAtomIndex = dict([(atom.name, i) for i, atom in enumerate(newTemplate.atoms)])
                for atom in self.changedAtoms[index]:
                    if atom.name not in newAtomIndex:
                        raise ValueError("Patch '%s' modifies nonexistent atom '%s' in template '%s'" % (self.name, atom.name, template.name))
730
                    newTemplate.atoms[newAtomIndex[atom.name]] = ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters)
731

732
                # Copy over the virtual sites, translating the atom indices.
733

734
735
736
737
738
739
740
                indexMap = dict([(oldAtomIndex[name], newAtomIndex[name]) for name in newAtomIndex if name in oldAtomIndex])
                for site in template.virtualSites:
                    if site.index in indexMap and all(i in indexMap for i in site.atoms):
                        newSite = deepcopy(site)
                        newSite.index = indexMap[site.index]
                        newSite.atoms = [indexMap[i] for i in site.atoms]
                        newTemplate.virtualSites.append(newSite)
741

742
                # Build the lists of bonds and external bonds.
743

744
745
746
747
748
                atomMap = dict([(template.atoms[i], indexMap[i]) for i in indexMap])
                deletedBonds = [(atom1.name, atom2.name) for atom1, atom2 in self.deletedBonds if atom1.residue == index and atom2.residue == index]
                for atom1, atom2 in template.bonds:
                    a1 = template.atoms[atom1]
                    a2 = template.atoms[atom2]
749
                    if a1 in atomMap and a2 in atomMap and (a1.name, a2.name) not in deletedBonds and (a2.name, a1.name) not in deletedBonds:
750
751
752
753
                        newTemplate.addBond(atomMap[a1], atomMap[a2])
                deletedExternalBonds = [atom.name for atom in self.deletedExternalBonds if atom.residue == index]
                for atom in template.externalBonds:
                    if template.atoms[atom].name not in deletedExternalBonds:
754
                        newTemplate.addExternalBond(indexMap[atom])
755
756
757
758
759
760
761
762
763
764
                for atom1, atom2 in self.addedBonds:
                    if atom1.residue == index and atom2.residue == index:
                        newTemplate.addBondByName(atom1.name, atom2.name)
                    elif atom1.residue == index:
                        newTemplate.addExternalBondByName(atom1.name)
                    elif atom2.residue == index:
                        newTemplate.addExternalBondByName(atom2.name)
                for atom in self.addedExternalBonds:
                    newTemplate.addExternalBondByName(atom.name)
            return newTemplates
765

766
767
768
769
770
771
772
773
774
775
776
    class _PatchAtomData(object):
        """Inner class used to encapsulate data about an atom in a patch definition."""
        def __init__(self, description):
            if ':' in description:
                colonIndex = description.find(':')
                self.residue = int(description[:colonIndex])-1
                self.name = description[colonIndex+1:]
            else:
                self.residue = 0
                self.name = description

777
    class _AtomType(object):
778
779
780
781
782
783
784
        """Inner class used to record atom types and associated properties."""
        def __init__(self, name, atomClass, mass, element):
            self.name = name
            self.atomClass = atomClass
            self.mass = mass
            self.element = element

785
    class _AtomTypeParameters(object):
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
        """Inner class used to record parameter values for atom types."""
        def __init__(self, forcefield, forceName, atomTag, paramNames):
            self.ff = forcefield
            self.forceName = forceName
            self.atomTag = atomTag
            self.paramNames = paramNames
            self.paramsForType = {}
            self.extraParamsForType = {}

        def registerAtom(self, parameters, expectedParams=None):
            if expectedParams is None:
                expectedParams = self.paramNames
            types = self.ff._findAtomTypes(parameters, 1)
            if None not in types:
                values = {}
                extraValues = {}
                for key in parameters:
                    if key in expectedParams:
                        values[key] = _convertParameterToNumber(parameters[key])
                    else:
                        extraValues[key] = parameters[key]
                if len(values) < len(expectedParams):
                    for key in expectedParams:
                        if key not in values:
                            raise ValueError('%s: No value specified for "%s"' % (self.forceName, key))
                for t in types[0]:
                    self.paramsForType[t] = values
                    self.extraParamsForType[t] = extraValues

        def parseDefinitions(self, element):
            """"Load the definitions from an XML element."""
            expectedParams = list(self.paramNames)
            excludedParams = [node.attrib['name'] for node in element.findall('UseAttributeFromResidue')]
            for param in excludedParams:
                if param not in expectedParams:
                    raise ValueError('%s: <UseAttributeFromResidue> specified an invalid attribute: %s' % (self.forceName, param))
                expectedParams.remove(param)
            for atom in element.findall(self.atomTag):
                for param in excludedParams:
                    if param in atom.attrib:
                        raise ValueError('%s: The attribute "%s" appeared in both <%s> and <UseAttributeFromResidue> tags' % (self.forceName, param, self.atomTag))
                self.registerAtom(atom.attrib, expectedParams)

        def getAtomParameters(self, atom, data):
            """Get the parameter values for a particular atom."""
            t = data.atomType[atom]
            p = data.atomParameters[atom]
            if t in self.paramsForType:
                values = self.paramsForType[t]
                result = [None]*len(self.paramNames)
                for i, name in enumerate(self.paramNames):
                    if name in values:
                        result[i] = values[name]
                    elif name in p:
                        result[i] = p[name]
                    else:
                        raise ValueError('%s: No value specified for "%s"' % (self.forceName, name))
                return result
            else:
                raise ValueError('%s: No parameters defined for atom type %s' % (self.forceName, t))

        def getExtraParameters(self, atom, data):
            """Get extra parameter values for an atom that appeared in the <Atom> tag but were not included in paramNames."""
            t = data.atomType[atom]
            if t in self.paramsForType:
                return self.extraParamsForType[t]
            else:
                raise ValueError('%s: No parameters defined for atom type %s' % (self.forceName, t))


856
    def _getResidueTemplateMatches(self, res, bondedToAtom, templateSignatures=None, ignoreExternalBonds=False):
857
858
859
860
861
862
        """Return the residue template matches, or None if none are found.

        Parameters
        ----------
        res : Topology.Residue
            The residue for which template matches are to be retrieved.
863
864
        bondedToAtom : list of set of int
            bondedToAtom[i] is the set of atoms bonded to atom index i
865
866
867
868
869
870
871
872
873
874
875
876

        Returns
        -------
        template : _ForceFieldTemplate
            The matching forcefield residue template, or None if no matches are found.
        matches : list
            a list specifying which atom of the template each atom of the residue
            corresponds to, or None if it does not match the template

        """
        template = None
        matches = None
877
878
        if templateSignatures is None:
            templateSignatures = self._templateSignatures
879
        signature = _createResidueSignature([atom.element for atom in res.atoms()])
880
        if signature in templateSignatures:
881
            allMatches = []
882
            for t in templateSignatures[signature]:
peastman's avatar
peastman committed
883
                match = compiled.matchResidueToTemplate(res, t, bondedToAtom, ignoreExternalBonds)
884
885
886
887
888
889
                if match is not None:
                    allMatches.append((t, match))
            if len(allMatches) == 1:
                template = allMatches[0][0]
                matches = allMatches[0][1]
            elif len(allMatches) > 1:
890
                raise Exception('Multiple matching templates found for residue %d (%s): %s.' % (res.index+1, res.name, ', '.join(match[0].name for match in allMatches)))
891
892
        return [template, matches]

893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
    def _buildBondedToAtomList(self, topology):
        """Build a list of which atom indices are bonded to each atom.

        Parameters
        ----------
        topology : Topology
            The Topology whose bonds are to be indexed.

        Returns
        -------
        bondedToAtom : list of set of int
            bondedToAtom[index] is the set of atom indices bonded to atom `index`

        """
        bondedToAtom = []
        for atom in topology.atoms():
            bondedToAtom.append(set())
910
911
912
        for (atom1, atom2) in topology.bonds():
            bondedToAtom[atom1.index].add(atom2.index)
            bondedToAtom[atom2.index].add(atom1.index)
913
        return bondedToAtom
914

915
916
917
    def getUnmatchedResidues(self, topology):
        """Return a list of Residue objects from specified topology for which no forcefield templates are available.

918
919
        .. CAUTION:: This method is experimental, and its API is subject to change.

920
921
922
        Parameters
        ----------
        topology : Topology
923
            The Topology whose residues are to be checked against the forcefield residue templates.
924
925
926
927

        Returns
        -------
        unmatched_residues : list of Residue
928
            List of Residue objects from `topology` for which no forcefield residue templates are available.
929
930
931
932
933
            Note that multiple instances of the same residue appearing at different points in the topology may be returned.

        This method may be of use in generating missing residue templates or diagnosing parameterization failures.
        """
        # Find the template matching each residue, compiling a list of residues for which no templates are available.
934
        bondedToAtom = self._buildBondedToAtomList(topology)
935
        unmatched_residues = list() # list of unmatched residues
936
937
938
939
940
941
        for res in topology.residues():
            # Attempt to match one of the existing templates.
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
            if matches is None:
                # No existing templates match.
                unmatched_residues.append(res)
942
943
944

        return unmatched_residues

945
    def getMatchingTemplates(self, topology, ignoreExternalBonds=False):
946
947
948
949
950
951
952
953
        """Return a list of forcefield residue templates matching residues in the specified topology.

        .. CAUTION:: This method is experimental, and its API is subject to change.

        Parameters
        ----------
        topology : Topology
            The Topology whose residues are to be checked against the forcefield residue templates.
954
955
        ignoreExternalBonds : bool=False
            If true, ignore external bonds when matching residues to templates.
956
957
958
959
960
961
962
963
964
965
966
967
968
        Returns
        -------
        templates : list of _TemplateData
            List of forcefield residue templates corresponding to residues in the topology.
            templates[index] is template corresponding to residue `index` in topology.residues()

        This method may be of use in debugging issues related to parameter assignment.
        """
        # Find the template matching each residue, compiling a list of residues for which no templates are available.
        bondedToAtom = self._buildBondedToAtomList(topology)
        templates = list() # list of templates matching the corresponding residues
        for res in topology.residues():
            # Attempt to match one of the existing templates.
969
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
970
971
972
973
974
975
976
977
            # Raise an exception if we have found no templates that match.
            if matches is None:
                raise ValueError('No template found for residue %d (%s).  %s' % (res.index+1, res.name, _findMatchErrors(self, res)))
            else:
                templates.append(template)

        return templates

978
979
    def generateTemplatesForUnmatchedResidues(self, topology):
        """Generate forcefield residue templates for residues in specified topology for which no forcefield templates are available.
980

981
982
        .. CAUTION:: This method is experimental, and its API is subject to change.

983
984
985
        Parameters
        ----------
        topology : Topology
986
            The Topology whose residues are to be checked against the forcefield residue templates.
987
988
989

        Returns
        -------
990
991
992
993
994
995
        templates : list of _TemplateData
            List of forcefield residue templates corresponding to residues in `topology` for which no forcefield templates are currently available.
            Atom types will be set to `None`, but template name, atom names, elements, and connectivity will be taken from corresponding Residue objects.
        residues : list of Residue
            List of Residue objects that were used to generate the templates.
            `residues[index]` is the Residue that was used to generate the template `templates[index]`
996
997
998
999
1000

        """
        # Get a non-unique list of unmatched residues.
        unmatched_residues = self.getUnmatchedResidues(topology)
        # Generate a unique list of unmatched residues by comparing fingerprints.
1001
        bondedToAtom = self._buildBondedToAtomList(topology)
1002
1003
        unique_unmatched_residues = list() # list of unique unmatched Residue objects from topology
        templates = list() # corresponding _TemplateData templates
1004
1005
1006
        signatures = set()
        for residue in unmatched_residues:
            signature = _createResidueSignature([ atom.element for atom in residue.atoms() ])
1007
            template = _createResidueTemplate(residue)
1008
1009
1010
1011
            is_unique = True
            if signature in signatures:
                # Signature is the same as an existing residue; check connectivity.
                for check_residue in unique_unmatched_residues:
peastman's avatar
peastman committed
1012
                    matches = compiled.matchResidueToTemplate(check_residue, template, bondedToAtom, False)
1013
1014
1015
1016
1017
1018
                    if matches is not None:
                        is_unique = False
            if is_unique:
                # Residue is unique.
                unique_unmatched_residues.append(residue)
                signatures.add(signature)
1019
                templates.append(template)
1020

1021
        return [templates, unique_unmatched_residues]
1022

1023
    def createSystem(self, topology, nonbondedMethod=NoCutoff, nonbondedCutoff=1.0*unit.nanometer,
1024
                     constraints=None, rigidWater=True, removeCMMotion=True, hydrogenMass=None, residueTemplates=dict(),
1025
                     ignoreExternalBonds=False, switchDistance=None, flexibleConstraints=False, **args):
1026
        """Construct an OpenMM System representing a Topology with this force field.
Justin MacCallum's avatar
Justin MacCallum committed
1027

Robert McGibbon's avatar
Robert McGibbon committed
1028
1029
1030
1031
1032
1033
        Parameters
        ----------
        topology : Topology
            The Topology for which to create a System
        nonbondedMethod : object=NoCutoff
            The method to use for nonbonded interactions.  Allowed values are
Peter Eastman's avatar
Peter Eastman committed
1034
            NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME, or LJPME.
Robert McGibbon's avatar
Robert McGibbon committed
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
        nonbondedCutoff : distance=1*nanometer
            The cutoff distance to use for nonbonded interactions
        constraints : object=None
            Specifies which bonds and angles should be implemented with constraints.
            Allowed values are None, HBonds, AllBonds, or HAngles.
        rigidWater : boolean=True
            If true, water molecules will be fully rigid regardless of the value
            passed for the constraints argument
        removeCMMotion : boolean=True
            If true, a CMMotionRemover will be added to the System
        hydrogenMass : mass=None
            The mass to use for hydrogen atoms bound to heavy atoms.  Any mass
            added to a hydrogen is subtracted from the heavy atom to keep
            their total mass the same.
1049
        residueTemplates : dict=dict()
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
            Key: Topology Residue object
            Value: string, name of _TemplateData residue template object to use for (Key) residue.
            This allows user to specify which template to apply to particular Residues
            in the event that multiple matching templates are available (e.g Fe2+ and Fe3+
            templates in the ForceField for a monoatomic iron ion in the topology).
        ignoreExternalBonds : boolean=False
            If true, ignore external bonds when matching residues to templates.  This is
            useful when the Topology represents one piece of a larger molecule, so chains are
            not terminated properly.  This option can create ambiguities where multiple
            templates match the same residue.  If that happens, use the residueTemplates
            argument to specify which one to use.
1061
1062
1063
        switchDistance : float=None
            The distance at which the potential energy switching function is turned on for
            Lennard-Jones interactions. If this is None, no switching function will be used.
1064
1065
        flexibleConstraints : boolean=False
            If True, parameters for constrained degrees of freedom will be added to the System
Robert McGibbon's avatar
Robert McGibbon committed
1066
        args
1067
1068
1069
            Arbitrary additional keyword arguments may also be specified.
            This allows extra parameters to be specified that are specific to
            particular force fields.
Robert McGibbon's avatar
Robert McGibbon committed
1070
1071
1072
1073
1074

        Returns
        -------
        system
            the newly created System
1075
        """
1076
1077
        args['switchDistance'] = switchDistance
        args['flexibleConstraints'] = flexibleConstraints
1078
        data = ForceField._SystemData()
1079
        data.atoms = list(topology.atoms())
1080
1081
        for atom in data.atoms:
            data.excludeAtomWith.append([])
1082
1083

        # Make a list of all bonds
Justin MacCallum's avatar
Justin MacCallum committed
1084

1085
        for bond in topology.bonds():
1086
            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
1087
1088

        # Record which atoms are bonded to each other atom
Justin MacCallum's avatar
Justin MacCallum committed
1089

1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
        bondedToAtom = []
        for i in range(len(data.atoms)):
            bondedToAtom.append(set())
            data.atomBonds.append([])
        for i in range(len(data.bonds)):
            bond = data.bonds[i]
            bondedToAtom[bond.atom1].add(bond.atom2)
            bondedToAtom[bond.atom2].add(bond.atom1)
            data.atomBonds[bond.atom1].append(i)
            data.atomBonds[bond.atom2].append(i)

        # Find the template matching each residue and assign atom types.
Justin MacCallum's avatar
Justin MacCallum committed
1102

1103
        unmatchedResidues = []
1104
1105
        for chain in topology.chains():
            for res in chain.residues():
1106
1107
1108
                if res in residueTemplates:
                    tname = residueTemplates[res]
                    template = self._templates[tname]
peastman's avatar
peastman committed
1109
                    matches = compiled.matchResidueToTemplate(res, template, bondedToAtom, ignoreExternalBonds)
1110
1111
1112
1113
                    if matches is None:
                        raise Exception('User-supplied template %s does not match the residue %d (%s)' % (tname, res.index+1, res.name))
                else:
                    # Attempt to match one of the existing templates.
1114
                    [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1115
                if matches is None:
1116
1117
1118
                    unmatchedResidues.append(res)
                else:
                    data.recordMatchedAtomParameters(res, template, matches)
1119

1120
        # Try to apply patches to find matches for any unmatched residues.
1121

1122
        if len(unmatchedResidues) > 0:
1123
            unmatchedResidues = _applyPatchesToMatchResidues(self, data, unmatchedResidues, bondedToAtom, ignoreExternalBonds)
1124

1125
1126
        # If we still haven't found a match for a residue, attempt to use residue template generators to create
        # new templates (and potentially atom types/parameters).
1127

1128
1129
        for res in unmatchedResidues:
            # A template might have been generated on an earlier iteration of this loop.
1130
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1131
1132
1133
1134
1135
            if matches is None:
                # Try all generators.
                for generator in self._templateGenerators:
                    if generator(self, res):
                        # This generator has registered a new residue template that should match.
1136
                        [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
                        if matches is None:
                            # Something went wrong because the generated template does not match the residue signature.
                            raise Exception('The residue handler %s indicated it had correctly parameterized residue %s, but the generated template did not match the residue signature.' % (generator.__class__.__name__, str(res)))
                        else:
                            # We successfully generated a residue template.  Break out of the for loop.
                            break
            if matches is None:
                raise ValueError('No template found for residue %d (%s).  %s' % (res.index+1, res.name, _findMatchErrors(self, res)))
            else:
                data.recordMatchedAtomParameters(res, template, matches)
1147
1148

        # Create the System and add atoms
Justin MacCallum's avatar
Justin MacCallum committed
1149

1150
1151
        sys = mm.System()
        for atom in topology.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1152
            # Look up the atom type name, returning a helpful error message if it cannot be found.
1153
1154
1155
1156
            if atom not in data.atomType:
                raise Exception("Could not identify atom type for atom '%s'." % str(atom))
            typename = data.atomType[atom]

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1157
            # Look up the type name in the list of registered atom types, returning a helpful error message if it cannot be found.
1158
1159
1160
1161
            if typename not in self._atomTypes:
                msg  = "Could not find typename '%s' for atom '%s' in list of known atom types.\n" % (typename, str(atom))
                msg += "Known atom types are: %s" % str(self._atomTypes.keys())
                raise Exception(msg)
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1162
1163

            # Add the particle to the OpenMM system.
1164
            mass = self._atomTypes[typename].mass
1165
            sys.addParticle(mass)
1166

1167
        # Adjust hydrogen masses if requested.
1168

1169
        if hydrogenMass is not None:
1170
1171
            if not unit.is_quantity(hydrogenMass):
                hydrogenMass *= unit.dalton
1172
            for atom1, atom2 in topology.bonds():
1173
                if atom1.element is elem.hydrogen:
1174
                    (atom1, atom2) = (atom2, atom1)
1175
                if atom2.element is elem.hydrogen and atom1.element not in (elem.hydrogen, None):
1176
1177
1178
                    transferMass = hydrogenMass-sys.getParticleMass(atom2.index)
                    sys.setParticleMass(atom2.index, hydrogenMass)
                    sys.setParticleMass(atom1.index, sys.getParticleMass(atom1.index)-transferMass)
Justin MacCallum's avatar
Justin MacCallum committed
1179

1180
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
1181

1182
1183
1184
        boxVectors = topology.getPeriodicBoxVectors()
        if boxVectors is not None:
            sys.setDefaultPeriodicBoxVectors(boxVectors[0], boxVectors[1], boxVectors[2])
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1185
        elif nonbondedMethod not in [NoCutoff, CutoffNonPeriodic]:
1186
1187
1188
            raise ValueError('Requested periodic boundary conditions for a Topology that does not specify periodic box dimensions')

        # Make a list of all unique angles
Justin MacCallum's avatar
Justin MacCallum committed
1189

1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
        uniqueAngles = set()
        for bond in data.bonds:
            for atom in bondedToAtom[bond.atom1]:
                if atom != bond.atom2:
                    if atom < bond.atom2:
                        uniqueAngles.add((atom, bond.atom1, bond.atom2))
                    else:
                        uniqueAngles.add((bond.atom2, bond.atom1, atom))
            for atom in bondedToAtom[bond.atom2]:
                if atom != bond.atom1:
                    if atom > bond.atom1:
                        uniqueAngles.add((bond.atom1, bond.atom2, atom))
                    else:
                        uniqueAngles.add((atom, bond.atom2, bond.atom1))
        data.angles = sorted(list(uniqueAngles))
Justin MacCallum's avatar
Justin MacCallum committed
1205

1206
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1207

1208
1209
1210
        uniquePropers = set()
        for angle in data.angles:
            for atom in bondedToAtom[angle[0]]:
pgrinaway's avatar
pgrinaway committed
1211
                if atom not in angle:
1212
1213
1214
1215
1216
                    if atom < angle[2]:
                        uniquePropers.add((atom, angle[0], angle[1], angle[2]))
                    else:
                        uniquePropers.add((angle[2], angle[1], angle[0], atom))
            for atom in bondedToAtom[angle[2]]:
pgrinaway's avatar
pgrinaway committed
1217
                if atom not in angle:
1218
1219
1220
1221
1222
                    if atom > angle[0]:
                        uniquePropers.add((angle[0], angle[1], angle[2], atom))
                    else:
                        uniquePropers.add((atom, angle[2], angle[1], angle[0]))
        data.propers = sorted(list(uniquePropers))
Justin MacCallum's avatar
Justin MacCallum committed
1223

1224
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1225

1226
1227
1228
1229
1230
        for atom in range(len(bondedToAtom)):
            bondedTo = bondedToAtom[atom]
            if len(bondedTo) > 2:
                for subset in itertools.combinations(bondedTo, 3):
                    data.impropers.append((atom, subset[0], subset[1], subset[2]))
Justin MacCallum's avatar
Justin MacCallum committed
1231

1232
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1233

1234
1235
1236
1237
1238
1239
1240
        if constraints == AllBonds or constraints == HAngles:
            for bond in data.bonds:
                bond.isConstrained = True
        elif constraints == HBonds:
            for bond in data.bonds:
                atom1 = data.atoms[bond.atom1]
                atom2 = data.atoms[bond.atom2]
1241
                bond.isConstrained = atom1.element is elem.hydrogen or atom2.element is elem.hydrogen
1242
1243
1244
1245
1246
1247
        if rigidWater:
            for bond in data.bonds:
                atom1 = data.atoms[bond.atom1]
                atom2 = data.atoms[bond.atom2]
                if atom1.residue.name == 'HOH' and atom2.residue.name == 'HOH':
                    bond.isConstrained = True
Justin MacCallum's avatar
Justin MacCallum committed
1248

1249
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1250

1251
1252
1253
1254
1255
1256
        if constraints == HAngles:
            for angle in data.angles:
                atom1 = data.atoms[angle[0]]
                atom2 = data.atoms[angle[1]]
                atom3 = data.atoms[angle[2]]
                numH = 0
1257
                if atom1.element is elem.hydrogen:
1258
                    numH += 1
1259
                if atom3.element is elem.hydrogen:
1260
                    numH += 1
1261
                data.isAngleConstrained.append(numH == 2 or (numH == 1 and atom2.element is elem.oxygen))
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
        else:
            data.isAngleConstrained = len(data.angles)*[False]
        if rigidWater:
            for i in range(len(data.angles)):
                angle = data.angles[i]
                atom1 = data.atoms[angle[0]]
                atom2 = data.atoms[angle[1]]
                atom3 = data.atoms[angle[2]]
                if atom1.residue.name == 'HOH' and atom2.residue.name == 'HOH' and atom3.residue.name == 'HOH':
                    data.isAngleConstrained[i] = True
Justin MacCallum's avatar
Justin MacCallum committed
1272

1273
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
1274

1275
        for atom in data.virtualSites:
1276
            (site, atoms, excludeWith) = data.virtualSites[atom]
1277
            index = atom.index
1278
            data.excludeAtomWith[excludeWith].append(index)
1279
            if site.type == 'average2':
1280
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(atoms[0], atoms[1], site.weights[0], site.weights[1]))
1281
            elif site.type == 'average3':
1282
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
1283
            elif site.type == 'outOfPlane':
1284
1285
                sys.setVirtualSite(index, mm.OutOfPlaneSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
            elif site.type == 'localCoords':
1286
                sys.setVirtualSite(index, mm.LocalCoordinatesSite(atoms, site.originWeights, site.xWeights, site.yWeights, site.localPos))
Justin MacCallum's avatar
Justin MacCallum committed
1287

1288
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
1289

1290
1291
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
1292
1293
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
1294

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1295
        # Let force generators do postprocessing
Justin MacCallum's avatar
Justin MacCallum committed
1296

peastman's avatar
peastman committed
1297
1298
1299
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
1300

1301
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
1302

1303
        for script in self._scripts:
1304
            exec(script, locals())
1305
1306
1307
        return sys


1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
def _findBondsForExclusions(data, sys):
    """Create a list of bonds to use when identifying exclusions."""
    bondIndices = []
    for bond in data.bonds:
        bondIndices.append((bond.atom1, bond.atom2))

    # If a virtual site does *not* share exclusions with another atom, add a bond between it and its first parent atom.

    for i in range(sys.getNumParticles()):
        if sys.isVirtualSite(i):
            (site, atoms, excludeWith) = data.virtualSites[data.atoms[i]]
            if excludeWith is None:
                bondIndices.append((i, site.getParticle(0)))

    # Certain particles, such as lone pairs and Drude particles, share exclusions with a parent atom.
    # If the parent atom does not interact with an atom, the child particle does not either.

    for atom1, atom2 in bondIndices:
        for child1 in data.excludeAtomWith[atom1]:
            bondIndices.append((child1, atom2))
            for child2 in data.excludeAtomWith[atom2]:
                bondIndices.append((child1, child2))
        for child2 in data.excludeAtomWith[atom2]:
            bondIndices.append((atom1, child2))
1332
1333
1334
    for atom in data.atoms:
        for child in data.excludeAtomWith[atom.index]:
            bondIndices.append((child, atom.index))
1335
1336
    return bondIndices

1337
1338
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
1339
1340
    counts = {}
    for element in elements:
1341
        if element in counts:
1342
1343
1344
            counts[element] += 1
        else:
            counts[element] = 1
1345
1346
1347
1348
1349
1350
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
1351
1352
    sig = []
    for c in counts:
1353
1354
        if c is not None:
            sig.append((c, counts[c]))
1355
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
1356

1357
    # Convert it to a string.
1358
1359

    s = ''
1360
    for element, count in sig:
1361
1362
1363
1364
        s += element.symbol+str(count)
    return s


1365
def _applyPatchesToMatchResidues(forcefield, data, residues, bondedToAtom, ignoreExternalBonds):
1366
1367
1368
1369
    """Try to apply patches to find matches for residues."""
    # Start by creating all templates than can be created by applying a combination of one-residue patches
    # to a single template.  The number of these is usually not too large, and they often cover a large fraction
    # of residues.
1370

1371
1372
1373
1374
1375
1376
1377
1378
    patchedTemplateSignatures = {}
    patchedTemplates = {}
    for name, template in forcefield._templates.items():
        if name in forcefield._templatePatches:
            patches = [forcefield._patches[patchName] for patchName, patchResidueIndex in forcefield._templatePatches[name] if forcefield._patches[patchName].numResidues == 1]
            if len(patches) > 0:
                newTemplates = []
                patchedTemplates[name] = newTemplates
peastman's avatar
peastman committed
1379
                _generatePatchedSingleResidueTemplates(template, patches, 0, newTemplates, set())
1380
1381
1382
1383
1384
1385
                for patchedTemplate in newTemplates:
                    signature = _createResidueSignature([atom.element for atom in patchedTemplate.atoms])
                    if signature in patchedTemplateSignatures:
                        patchedTemplateSignatures[signature].append(patchedTemplate)
                    else:
                        patchedTemplateSignatures[signature] = [patchedTemplate]
1386

1387
    # Now see if any of those templates matches any of the residues.
1388

1389
1390
    unmatchedResidues = []
    for res in residues:
1391
        [template, matches] = forcefield._getResidueTemplateMatches(res, bondedToAtom, patchedTemplateSignatures, ignoreExternalBonds)
1392
1393
1394
1395
        if matches is None:
            unmatchedResidues.append(res)
        else:
            data.recordMatchedAtomParameters(res, template, matches)
1396
1397
    if len(unmatchedResidues) == 0:
        return []
1398

1399
1400
1401
    # We need to consider multi-residue patches.  This can easily lead to a combinatorial explosion, so we make a simplifying
    # assumption: that no residue is affected by more than one multi-residue patch (in addition to any number of single-residue
    # patches).  Record all multi-residue patches, and the templates they can be applied to.
1402

1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
    patches = {}
    maxPatchSize = 0
    for patch in forcefield._patches.values():
        if patch.numResidues > 1:
            patches[patch.name] = [[] for i in range(patch.numResidues)]
            maxPatchSize = max(maxPatchSize, patch.numResidues)
    if maxPatchSize == 0:
        return unmatchedResidues # There aren't any multi-residue patches
    for templateName in forcefield._templatePatches:
        for patchName, patchResidueIndex in forcefield._templatePatches[templateName]:
            if patchName in patches:
                # The patch should accept this template, *and* all patched versions of it generated above.
                patches[patchName][patchResidueIndex].append(forcefield._templates[templateName])
                if templateName in patchedTemplates:
                    patches[patchName][patchResidueIndex] += patchedTemplates[templateName]
1418

1419
    # Record which unmatched residues are bonded to each other.
1420

1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
    bonds = set()
    topology = residues[0].chain.topology
    for atom1, atom2 in topology.bonds():
        if atom1.residue != atom2.residue:
            res1 = atom1.residue
            res2 = atom2.residue
            if res1 in unmatchedResidues and res2 in unmatchedResidues:
                bond = tuple(sorted((res1, res2), key=lambda x: x.index))
                if bond not in bonds:
                    bonds.add(bond)
1431

1432
1433
    # Identify clusters of unmatched residues that are all bonded to each other.  These are the ones we'll
    # try to apply multi-residue patches to.
1434

1435
1436
1437
1438
    clusterSize = 2
    clusters = bonds
    while clusterSize <= maxPatchSize:
        # Try to apply patches to clusters of this size.
1439

1440
1441
1442
        for patchName in patches:
            patch = forcefield._patches[patchName]
            if patch.numResidues == clusterSize:
1443
                matchedClusters = _matchToMultiResiduePatchedTemplates(data, clusters, patch, patches[patchName], bondedToAtom, ignoreExternalBonds)
1444
1445
1446
1447
1448
1449
                for cluster in matchedClusters:
                    for residue in cluster:
                        unmatchedResidues.remove(residue)
                bonds = set(bond for bond in bonds if bond[0] in unmatchedResidues and bond[1] in unmatchedResidues)

        # Now extend the clusters to find ones of the next size up.
1450

1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
        largerClusters = set()
        for cluster in clusters:
            for bond in bonds:
                if bond[0] in cluster and bond[1] not in cluster:
                    newCluster = tuple(sorted(cluster+(bond[1],), key=lambda x: x.index))
                    largerClusters.add(newCluster)
                elif bond[1] in cluster and bond[0] not in cluster:
                    newCluster = tuple(sorted(cluster+(bond[0],), key=lambda x: x.index))
                    largerClusters.add(newCluster)
        if len(largerClusters) == 0:
            # There are no clusters of this size or larger
            break
        clusters = largerClusters
        clusterSize += 1

1466
1467
1468
    return unmatchedResidues


peastman's avatar
peastman committed
1469
def _generatePatchedSingleResidueTemplates(template, patches, index, newTemplates, alteredAtoms):
1470
1471
    """Apply all possible combinations of a set of single-residue patches to a template."""
    try:
peastman's avatar
peastman committed
1472
1473
1474
1475
1476
1477
1478
        if len(alteredAtoms.intersection(patches[index].allAtomNames)) > 0:
            # This patch would alter an atom that another patch has already altered,
            # so don't apply it.
            patchedTemplate = None
        else:
            patchedTemplate = patches[index].createPatchedTemplates([template])[0]
            newTemplates.append(patchedTemplate)
1479
1480
1481
1482
    except:
        # This probably means the patch is inconsistent with another one that has already been applied,
        # so just ignore it.
        patchedTemplate = None
1483

1484
    # Call this function recursively to generate combinations of patches.
1485

1486
    if index+1 < len(patches):
peastman's avatar
peastman committed
1487
        _generatePatchedSingleResidueTemplates(template, patches, index+1, newTemplates, alteredAtoms)
1488
        if patchedTemplate is not None:
peastman's avatar
peastman committed
1489
1490
            newAlteredAtoms = alteredAtoms.union(patches[index].allAtomNames)
            _generatePatchedSingleResidueTemplates(patchedTemplate, patches, index+1, newTemplates, newAlteredAtoms)
1491
1492


1493
def _matchToMultiResiduePatchedTemplates(data, clusters, patch, residueTemplates, bondedToAtom, ignoreExternalBonds):
1494
1495
1496
    """Apply a multi-residue patch to templates, then try to match them against clusters of residues."""
    matchedClusters = []
    selectedTemplates = [None]*patch.numResidues
1497
    _applyMultiResiduePatch(data, clusters, patch, residueTemplates, selectedTemplates, 0, matchedClusters, bondedToAtom, ignoreExternalBonds)
1498
1499
1500
    return matchedClusters


1501
def _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index, matchedClusters, bondedToAtom, ignoreExternalBonds):
1502
1503
1504
1505
1506
    """This is called recursively to apply a multi-residue patch to all possible combinations of templates."""

    if index < patch.numResidues:
        for template in candidateTemplates[index]:
            selectedTemplates[index] = template
1507
            _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index+1, matchedClusters, bondedToAtom, ignoreExternalBonds)
1508
1509
1510
    else:
        # We're at the deepest level of the recursion.  We've selected a template for each residue, so apply the patch,
        # then try to match it against clusters.
1511

1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
        try:
            patchedTemplates = patch.createPatchedTemplates(selectedTemplates)
        except:
            # This probably means the patch is inconsistent with another one that has already been applied,
            # so just ignore it.
            return
        newlyMatchedClusters = []
        for cluster in clusters:
            for residues in itertools.permutations(cluster):
                residueMatches = []
                for residue, template in zip(residues, patchedTemplates):
peastman's avatar
peastman committed
1523
                    matches = compiled.matchResidueToTemplate(residue, template, bondedToAtom, ignoreExternalBonds)
1524
1525
1526
1527
1528
1529
                    if matches is None:
                        residueMatches = None
                        break
                    else:
                        residueMatches.append(matches)
                if residueMatches is not None:
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
                    # Each residue individually matches.  Now make sure they're bonded in the correct way.

                    bondsMatch = True
                    for a1, a2 in patch.addedBonds:
                        res1 = a1.residue
                        res2 = a2.residue
                        if res1 != res2:
                            # The patch adds a bond between residues.  Make sure that bond exists.

                            atoms1 = patchedTemplates[res1].atoms
                            atoms2 = patchedTemplates[res2].atoms
                            index1 = next(i for i in range(len(atoms1)) if atoms1[residueMatches[res1][i]].name == a1.name)
                            index2 = next(i for i in range(len(atoms2)) if atoms2[residueMatches[res2][i]].name == a2.name)
                            atom1 = list(residues[res1].atoms())[index1]
                            atom2 = list(residues[res2].atoms())[index2]
                            bondsMatch &= atom2.index in bondedToAtom[atom1.index]
                    if bondsMatch:
                        # We successfully matched the template to the residues.  Record the parameters.
1548

1549
1550
1551
1552
                        for i in range(patch.numResidues):
                            data.recordMatchedAtomParameters(residues[i], patchedTemplates[i], residueMatches[i])
                        newlyMatchedClusters.append(cluster)
                        break
1553

1554
        # Record which clusters were successfully matched.
1555

1556
1557
1558
        matchedClusters += newlyMatchedClusters
        for cluster in newlyMatchedClusters:
            clusters.remove(cluster)
1559

1560

1561
1562
1563
def _findMatchErrors(forcefield, res):
    """Try to guess why a residue failed to match any template and return an error message."""
    residueCounts = _countResidueAtoms([atom.element for atom in res.atoms()])
1564
    numResidueAtoms = sum(residueCounts.values())
1565
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
1566

1567
    # Loop over templates and see how closely each one might match.
1568

1569
1570
1571
1572
1573
1574
    bestMatchName = None
    numBestMatchAtoms = 3*numResidueAtoms
    numBestMatchHeavyAtoms = 2*numResidueHeavyAtoms
    for templateName in forcefield._templates:
        template = forcefield._templates[templateName]
        templateCounts = _countResidueAtoms([atom.element for atom in template.atoms])
1575

1576
        # Does the residue have any atoms that clearly aren't in the template?
1577

1578
1579
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
1580

1581
        # If there are too many missing atoms, discard this template.
1582

1583
        numTemplateAtoms = sum(templateCounts.values())
1584
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
1585
1586
1587
1588
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
1589

1590
1591
        # If this template has the same number of missing atoms as our previous best one, look at the name
        # to decide which one to use.
1592

1593
1594
1595
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
1596

1597
        # Accept this as our new best match.
1598

1599
1600
1601
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
1602
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
1603

1604
    # Return an appropriate error message.
1605

1606
    if numBestMatchAtoms == numResidueAtoms:
1607
1608
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
1609
1610
1611
1612
            return 'The set of atoms matches %s, but the bonds are different.  Perhaps the chain is missing a terminal group?' % bestMatchName
        return 'The set of atoms matches %s, but the bonds are different.' % bestMatchName
    if bestMatchName is not None:
        if numBestMatchHeavyAtoms == numResidueHeavyAtoms:
1613
1614
1615
1616
1617
            numResidueExtraParticles = len([atom for atom in res.atoms() if atom.element is None])
            if numResidueExtraParticles == 0 and numBestMatchExtraParticles == 0:
                return 'The set of atoms is similar to %s, but it is missing %d hydrogen atoms.' % (bestMatchName, numBestMatchAtoms-numResidueAtoms)
            if numBestMatchExtraParticles-numResidueExtraParticles == numBestMatchAtoms-numResidueAtoms:
                return 'The set of atoms is similar to %s, but it is missing %d extra particles.  You can add them with Modeller.addExtraParticles().' % (bestMatchName, numBestMatchAtoms-numResidueAtoms)
1618
1619
1620
        return 'The set of atoms is similar to %s, but it is missing %d atoms.' % (bestMatchName, numBestMatchAtoms-numResidueAtoms)
    return 'This might mean your input topology is missing some atoms or bonds, or possibly that you are using the wrong force field.'

1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
def _createResidueTemplate(residue):
    """Create a _TemplateData template from a Residue object.

    Parameters
    ----------
    residue : Residue
        The Residue from which the template is to be constructed.

    Returns
    -------
    template : _TemplateData
        The residue template, with atom types set to None.

    This method may be useful in creating new residue templates for residues without templates defined by the ForceField.

    """
    template = ForceField._TemplateData(residue.name)
    for atom in residue.atoms():
1639
        template.addAtom(ForceField._TemplateAtomData(atom.name, None, atom.element))
1640
1641
1642
1643
1644
1645
1646
1647
1648
    for (atom1,atom2) in residue.internal_bonds():
        template.addBondByName(atom1.name, atom2.name)
    residue_atoms = [ atom for atom in residue.atoms() ]
    for (atom1,atom2) in residue.external_bonds():
        if atom1 in residue_atoms:
            template.addExternalBondByName(atom1.name)
        elif atom2 in residue_atoms:
            template.addExternalBondByName(atom2.name)
    return template
1649

1650
1651
1652
1653
1654
def _matchImproper(data, torsion, generator):
    type1 = data.atomType[data.atoms[torsion[0]]]
    type2 = data.atomType[data.atoms[torsion[1]]]
    type3 = data.atomType[data.atoms[torsion[2]]]
    type4 = data.atomType[data.atoms[torsion[3]]]
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1655
    wildcard = generator.ff._atomClasses['']
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
    match = None
    for tordef in generator.improper:
        types1 = tordef.types1
        types2 = tordef.types2
        types3 = tordef.types3
        types4 = tordef.types4
        hasWildcard = (wildcard in (types1, types2, types3, types4))
        if match is not None and hasWildcard:
            # Prefer specific definitions over ones with wildcards
            continue
        if type1 in types1:
            for (t2, t3, t4) in itertools.permutations(((type2, 1), (type3, 2), (type4, 3))):
                if t2[0] in types2 and t3[0] in types3 and t4[0] in types4:
                    if tordef.ordering == 'default':
                        # Workaround to be more consistent with AMBER.  It uses wildcards to define most of its
                        # impropers, which leaves the ordering ambiguous.  It then follows some bizarre rules
                        # to pick the order.
                        a1 = torsion[t2[1]]
                        a2 = torsion[t3[1]]
                        e1 = data.atoms[a1].element
                        e2 = data.atoms[a2].element
                        if e1 == e2 and a1 > a2:
                            (a1, a2) = (a2, a1)
                        elif e1 != elem.carbon and (e2 == elem.carbon or e1.mass < e2.mass):
                            (a1, a2) = (a2, a1)
                        match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
                        break
                    elif tordef.ordering == 'charmm':
                        if hasWildcard:
                            # Workaround to be more consistent with AMBER.  It uses wildcards to define most of its
                            # impropers, which leaves the ordering ambiguous.  It then follows some bizarre rules
                            # to pick the order.
                            a1 = torsion[t2[1]]
                            a2 = torsion[t3[1]]
                            e1 = data.atoms[a1].element
                            e2 = data.atoms[a2].element
                            if e1 == e2 and a1 > a2:
                                (a1, a2) = (a2, a1)
                            elif e1 != elem.carbon and (e2 == elem.carbon or e1.mass < e2.mass):
                                (a1, a2) = (a2, a1)
                            match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
                        else:
                            # There are no wildcards, so the order is unambiguous.
                            match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
                        break
                    elif tordef.ordering == 'amber':
                        # topology atom indexes
                        a2 = torsion[t2[1]]
                        a3 = torsion[t3[1]]
                        a4 = torsion[t4[1]]
                        # residue indexes
                        r2 = data.atoms[a2].residue.index
                        r3 = data.atoms[a3].residue.index
                        r4 = data.atoms[a4].residue.index
                        # template atom indexes
                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                        # elements
                        e2 = data.atoms[a2].element
                        e3 = data.atoms[a3].element
                        e4 = data.atoms[a4].element
                        if not hasWildcard:
                            if t2[0] == t4[0] and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                (a2, a4) = (a4, a2)
                                r2 = data.atoms[a2].residue.index
                                r4 = data.atoms[a4].residue.index
                                ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                            if t3[0] == t4[0] and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
                                (a3, a4) = (a4, a3)
                                r3 = data.atoms[a3].residue.index
                                r4 = data.atoms[a4].residue.index
                                ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                            if t2[0] == t3[0] and (r2 > r3 or (r2 == r3 and ta2 > ta3)):
                                (a2, a3) = (a3, a2)
                        else:
                            if e2 == e4 and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                (a2, a4) = (a4, a2)
                                r2 = data.atoms[a2].residue.index
                                r4 = data.atoms[a4].residue.index
                                ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                            if e3 == e4 and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
                                (a3, a4) = (a4, a3)
                                r3 = data.atoms[a3].residue.index
                                r4 = data.atoms[a4].residue.index
                                ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                            if r2 > r3 or (r2 == r3 and ta2 > ta3):
                                (a2, a3) = (a3, a2)
                        match = (a2, a3, torsion[0], a4, tordef)
                        break
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1750
    return match
1751

1752
1753
1754
1755
1756
# The following classes are generators that know how to create Force subclasses and add them to a System that is being
# created.  Each generator class must define two methods: 1) a static method that takes an etree Element and a ForceField,
# and returns the corresponding generator object; 2) a createForce() method that constructs the Force object and adds it
# to the System.  The static method should be added to the parsers map.

1757
## @private
1758
class HarmonicBondGenerator(object):
1759
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1760

1761
1762
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
1763
        self.bondsForAtomType = defaultdict(set)
1764
1765
1766
1767
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
1768

1769
1770
1771
    def registerBond(self, parameters):
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
peastman's avatar
peastman committed
1772
            index = len(self.types1)
1773
1774
            self.types1.append(types[0])
            self.types2.append(types[1])
peastman's avatar
peastman committed
1775
1776
1777
1778
            for t in types[0]:
                self.bondsForAtomType[t].add(index)
            for t in types[1]:
                self.bondsForAtomType[t].add(index)
1779
1780
            self.length.append(_convertParameterToNumber(parameters['length']))
            self.k.append(_convertParameterToNumber(parameters['k']))
Justin MacCallum's avatar
Justin MacCallum committed
1781

1782
1783
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1784
1785
1786
1787
1788
1789
        existing = [f for f in ff._forces if isinstance(f, HarmonicBondGenerator)]
        if len(existing) == 0:
            generator = HarmonicBondGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
1790
        for bond in element.findall('Bond'):
1791
            generator.registerBond(bond.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1792

1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
        if len(existing) == 0:
            force = mm.HarmonicBondForce()
            sys.addForce(force)
        else:
            force = existing[0]
        for bond in data.bonds:
            type1 = data.atomType[data.atoms[bond.atom1]]
            type2 = data.atomType[data.atoms[bond.atom2]]
peastman's avatar
peastman committed
1804
            for i in self.bondsForAtomType[type1]:
1805
1806
1807
1808
1809
                types1 = self.types1[i]
                types2 = self.types2[i]
                if (type1 in types1 and type2 in types2) or (type1 in types2 and type2 in types1):
                    bond.length = self.length[i]
                    if bond.isConstrained:
1810
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
1811
1812
1813
1814
1815
                    if self.k[i] != 0:
                        # flexibleConstraints allows us to add parameters even if the DOF is
                        # constrained
                        if not bond.isConstrained or args.get('flexibleConstraints', False):
                            force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
1816
1817
1818
1819
1820
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


1821
## @private
1822
class HarmonicAngleGenerator(object):
1823
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1824

1825
1826
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
1827
        self.anglesForAtom2Type = defaultdict(list)
1828
1829
1830
1831
1832
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
1833

1834
1835
1836
    def registerAngle(self, parameters):
        types = self.ff._findAtomTypes(parameters, 3)
        if None not in types:
peastman's avatar
peastman committed
1837
            index = len(self.types1)
1838
1839
1840
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.types3.append(types[2])
peastman's avatar
peastman committed
1841
1842
            for t in types[1]:
                self.anglesForAtom2Type[t].append(index)
1843
1844
1845
            self.angle.append(_convertParameterToNumber(parameters['angle']))
            self.k.append(_convertParameterToNumber(parameters['k']))

1846
1847
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1848
1849
1850
1851
1852
1853
        existing = [f for f in ff._forces if isinstance(f, HarmonicAngleGenerator)]
        if len(existing) == 0:
            generator = HarmonicAngleGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
1854
        for angle in element.findall('Angle'):
1855
            generator.registerAngle(angle.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1856

1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.HarmonicAngleForce]
        if len(existing) == 0:
            force = mm.HarmonicAngleForce()
            sys.addForce(force)
        else:
            force = existing[0]
        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
peastman's avatar
peastman committed
1869
            for i in self.anglesForAtom2Type[type2]:
1870
1871
1872
1873
1874
1875
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
                if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                    if isConstrained:
                        # Find the two bonds that make this angle.
Justin MacCallum's avatar
Justin MacCallum committed
1876

1877
1878
1879
1880
1881
1882
1883
1884
1885
                        bond1 = None
                        bond2 = None
                        for bond in data.atomBonds[angle[1]]:
                            atom1 = data.bonds[bond].atom1
                            atom2 = data.bonds[bond].atom2
                            if atom1 == angle[0] or atom2 == angle[0]:
                                bond1 = bond
                            elif atom1 == angle[2] or atom2 == angle[2]:
                                bond2 = bond
Justin MacCallum's avatar
Justin MacCallum committed
1886

1887
                        # Compute the distance between atoms and add a constraint
Justin MacCallum's avatar
Justin MacCallum committed
1888

1889
1890
1891
1892
1893
                        if bond1 is not None and bond2 is not None:
                            l1 = data.bonds[bond1].length
                            l2 = data.bonds[bond2].length
                            if l1 is not None and l2 is not None:
                                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(self.angle[i]))
1894
                                data.addConstraint(sys, angle[0], angle[2], length)
1895
1896
1897
                    if self.k[i] != 0:
                        if not isConstrained or args.get('flexibleConstraints', False):
                            force.addAngle(angle[0], angle[1], angle[2], self.angle[i], self.k[i])
1898
1899
1900
1901
1902
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


1903
## @private
1904
class PeriodicTorsion(object):
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
    """A PeriodicTorsion records the information for a periodic torsion definition."""

    def __init__(self, types):
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.periodicity = []
        self.phase = []
        self.k = []
1915
        self.ordering = 'default'
1916

1917
## @private
1918
class PeriodicTorsionGenerator(object):
1919
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1920

1921
1922
    def __init__(self, forcefield):
        self.ff = forcefield
1923
1924
        self.proper = []
        self.improper = []
1925
        self.propersForAtomType = defaultdict(set)
Justin MacCallum's avatar
Justin MacCallum committed
1926

1927
1928
1929
    def registerProperTorsion(self, parameters):
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
1930
            index = len(self.proper)
1931
            self.proper.append(torsion)
1932
1933
1934
1935
            for t in torsion.types2:
                self.propersForAtomType[t].add(index)
            for t in torsion.types3:
                self.propersForAtomType[t].add(index)
1936

1937
    def registerImproperTorsion(self, parameters, ordering='default'):
1938
1939
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
1940
            if ordering in ['default', 'charmm', 'amber']:
1941
1942
                torsion.ordering = ordering
            else:
1943
                raise ValueError('Illegal ordering type %s for improper torsion %s' % (ordering, torsion))
1944
1945
            self.improper.append(torsion)

1946
1947
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1948
1949
1950
1951
1952
1953
        existing = [f for f in ff._forces if isinstance(f, PeriodicTorsionGenerator)]
        if len(existing) == 0:
            generator = PeriodicTorsionGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
1954
        for torsion in element.findall('Proper'):
1955
            generator.registerProperTorsion(torsion.attrib)
1956
        for torsion in element.findall('Improper'):
1957
1958
1959
1960
            if 'ordering' in element.attrib:
                generator.registerImproperTorsion(torsion.attrib, element.attrib['ordering'])
            else:
                generator.registerImproperTorsion(torsion.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1961

1962
1963
1964
1965
1966
1967
1968
1969
1970
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
        if len(existing) == 0:
            force = mm.PeriodicTorsionForce()
            sys.addForce(force)
        else:
            force = existing[0]
        wildcard = self.ff._atomClasses['']
tic20's avatar
tic20 committed
1971
        proper_cache = {}
1972
        for torsion in data.propers:
tic20's avatar
tic20 committed
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
            type1, type2, type3, type4 = [data.atomType[data.atoms[torsion[i]]] for i in range(4)]
            sig = (type1, type2, type3, type4)
            sig = frozenset((sig, sig[:-1]))
            match = proper_cache.get(sig, None)
            if match == -1:
                continue
            if match is None:
                for index in self.propersForAtomType[type2]:
                    tordef = self.proper[index]
                    types1 = tordef.types1
                    types2 = tordef.types2
                    types3 = tordef.types3
                    types4 = tordef.types4
                    if (type2 in types2 and type3 in types3 and type4 in types4 and type1 in types1) or (type2 in types3 and type3 in types2 and type4 in types1 and type1 in types4):
                        hasWildcard = (wildcard in (types1, types2, types3, types4))
                        if match is None or not hasWildcard: # Prefer specific definitions over ones with wildcards
                            match = tordef
                        if not hasWildcard:
                            break
                if match is None:
                    proper_cache[sig] = -1
                else:
                    proper_cache[sig] = match
1996
1997
1998
1999
            if match is not None:
                for i in range(len(match.phase)):
                    if match.k[i] != 0:
                        force.addTorsion(torsion[0], torsion[1], torsion[2], torsion[3], match.periodicity[i], match.phase[i], match.k[i])
2000
        impr_cache = {}
2001
        for torsion in data.impropers:
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
            t1, t2, t3, t4 = [data.atomType[data.atoms[torsion[i]]] for i in range(4)]
            sig = (t1, frozenset((t2,t3,t4)))
            match = impr_cache.get(sig, None)
            if match == -1:
                # Previously checked, and doesn't appear in the database
                continue
            if match is None:
                match = _matchImproper(data, torsion, self)
                if match is not None:
                    impr_cache[sig] = match
                else:
                    impr_cache[sig] = -1
2014
2015
2016
2017
2018
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                for i in range(len(tordef.phase)):
                    if tordef.k[i] != 0:
                        force.addTorsion(a1, a2, a3, a4, tordef.periodicity[i], tordef.phase[i], tordef.k[i])
2019
2020
2021
parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


2022
## @private
2023
class RBTorsion(object):
2024
2025
    """An RBTorsion records the information for a Ryckaert-Bellemans torsion definition."""

2026
    def __init__(self, types, c, ordering='charmm'):
2027
2028
2029
2030
2031
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.c = c
2032
        if ordering in ['default', 'charmm', 'amber']:
2033
2034
            self.ordering = ordering
        else:
2035
            raise ValueError('Illegal ordering type %s for RBTorsion (%s,%s,%s,%s)' % (ordering, types[0], types[1], types[2], types[3]))
2036

2037
## @private
2038
class RBTorsionGenerator(object):
2039
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2040

2041
2042
    def __init__(self, forcefield):
        self.ff = forcefield
2043
2044
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
2045

2046
2047
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2048
2049
2050
2051
2052
2053
        existing = [f for f in ff._forces if isinstance(f, RBTorsionGenerator)]
        if len(existing) == 0:
            generator = RBTorsionGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
2054
        for torsion in element.findall('Proper'):
2055
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2056
            if None not in types:
2057
2058
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
2059
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2060
            if None not in types:
2061
2062
2063
2064
                if 'ordering' in element.attrib:
                    generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)], element.attrib['ordering']))
                else:
                    generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
Justin MacCallum's avatar
Justin MacCallum committed
2065

2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.RBTorsionForce]
        if len(existing) == 0:
            force = mm.RBTorsionForce()
            sys.addForce(force)
        else:
            force = existing[0]
        wildcard = self.ff._atomClasses['']
        for torsion in data.propers:
            type1 = data.atomType[data.atoms[torsion[0]]]
            type2 = data.atomType[data.atoms[torsion[1]]]
            type3 = data.atomType[data.atoms[torsion[2]]]
            type4 = data.atomType[data.atoms[torsion[3]]]
            match = None
            for tordef in self.proper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
                if (type2 in types2 and type3 in types3 and type4 in types4 and type1 in types1) or (type2 in types3 and type3 in types2 and type4 in types1 and type1 in types4):
                    hasWildcard = (wildcard in (types1, types2, types3, types4))
                    if match is None or not hasWildcard: # Prefer specific definitions over ones with wildcards
                        match = tordef
                    if not hasWildcard:
                        break
            if match is not None:
                force.addTorsion(torsion[0], torsion[1], torsion[2], torsion[3], match.c[0], match.c[1], match.c[2], match.c[3], match.c[4], match.c[5])
        for torsion in data.impropers:
2095
            match = _matchImproper(data, torsion, self)
2096
2097
2098
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.c[0], tordef.c[1], tordef.c[2], tordef.c[3], tordef.c[4], tordef.c[5])
2099
2100
2101
2102

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


2103
## @private
2104
class CMAPTorsion(object):
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
    """A CMAPTorsion records the information for a CMAP torsion definition."""

    def __init__(self, types, map):
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.types5 = types[4]
        self.map = map

2115
## @private
2116
class CMAPTorsionGenerator(object):
2117
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2118

2119
2120
    def __init__(self, forcefield):
        self.ff = forcefield
2121
2122
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
2123

2124
2125
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2126
2127
2128
2129
2130
2131
        existing = [f for f in ff._forces if isinstance(f, CMAPTorsionGenerator)]
        if len(existing) == 0:
            generator = CMAPTorsionGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
2132
2133
2134
2135
2136
2137
2138
        for map in element.findall('Map'):
            values = [float(x) for x in map.text.split()]
            size = sqrt(len(values))
            if size*size != len(values):
                raise ValueError('CMAP must have the same number of elements along each dimension')
            generator.maps.append(values)
        for torsion in element.findall('Torsion'):
2139
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
2140
            if None not in types:
2141
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
2142

2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.CMAPTorsionForce]
        if len(existing) == 0:
            force = mm.CMAPTorsionForce()
            sys.addForce(force)
        else:
            force = existing[0]
        for map in self.maps:
            force.addMap(int(sqrt(len(map))), map)
Justin MacCallum's avatar
Justin MacCallum committed
2153

2154
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
2155

2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
        uniqueTorsions = set()
        for torsion in data.propers:
            for bond in (data.bonds[x] for x in data.atomBonds[torsion[0]]):
                if bond.atom1 == torsion[0]:
                    atom = bond.atom2
                else:
                    atom = bond.atom1
                if atom != torsion[1]:
                    uniqueTorsions.add((atom, torsion[0], torsion[1], torsion[2], torsion[3]))
            for bond in (data.bonds[x] for x in data.atomBonds[torsion[3]]):
                if bond.atom1 == torsion[3]:
                    atom = bond.atom2
                else:
                    atom = bond.atom1
                if atom != torsion[2]:
                    uniqueTorsions.add((torsion[0], torsion[1], torsion[2], torsion[3], atom))
        torsions = sorted(list(uniqueTorsions))
        wildcard = self.ff._atomClasses['']
        for torsion in torsions:
            type1 = data.atomType[data.atoms[torsion[0]]]
            type2 = data.atomType[data.atoms[torsion[1]]]
            type3 = data.atomType[data.atoms[torsion[2]]]
            type4 = data.atomType[data.atoms[torsion[3]]]
            type5 = data.atomType[data.atoms[torsion[4]]]
            match = None
            for tordef in self.torsions:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
                types5 = tordef.types5
                if (type1 in types1 and type2 in types2 and type3 in types3 and type4 in types4 and type5 in types5) or (type1 in types5 and type2 in types4 and type3 in types3 and type4 in types2 and type5 in types1):
                    hasWildcard = (wildcard in (types1, types2, types3, types4, types5))
                    if match is None or not hasWildcard: # Prefer specific definitions over ones with wildcards
                        match = tordef
                    if not hasWildcard:
                        break
            if match is not None:
                force.addTorsion(match.map, torsion[0], torsion[1], torsion[2], torsion[3], torsion[1], torsion[2], torsion[3], torsion[4])

parsers["CMAPTorsionForce"] = CMAPTorsionGenerator.parseElement


2199
## @private
2200
class NonbondedGenerator(object):
2201
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2202

2203
2204
    SCALETOL = 1e-5

2205
2206
    def __init__(self, forcefield, coulomb14scale, lj14scale):
        self.ff = forcefield
2207
2208
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
2209
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
2210

2211
    def registerAtom(self, parameters):
2212
        self.params.registerAtom(parameters)
2213

2214
2215
2216
2217
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
        if len(existing) == 0:
2218
2219
            generator = NonbondedGenerator(ff, float(element.attrib['coulomb14scale']), float(element.attrib['lj14scale']))
            ff.registerGenerator(generator)
2220
2221
2222
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2223
2224
            if abs(generator.coulomb14scale - float(element.attrib['coulomb14scale'])) > NonbondedGenerator.SCALETOL or \
                    abs(generator.lj14scale - float(element.attrib['lj14scale'])) > NonbondedGenerator.SCALETOL:
Justin MacCallum's avatar
Justin MacCallum committed
2225
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
2226
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2227

2228
2229
2230
2231
2232
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.NonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.NonbondedForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.NonbondedForce.CutoffPeriodic,
                     Ewald:mm.NonbondedForce.Ewald,
Peter Eastman's avatar
Peter Eastman committed
2233
2234
                     PME:mm.NonbondedForce.PME,
                     LJPME:mm.NonbondedForce.LJPME}
2235
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2236
            raise ValueError('Illegal nonbonded method for NonbondedForce')
2237
2238
        force = mm.NonbondedForce()
        for atom in data.atoms:
2239
2240
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
2241
2242
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2243
2244
2245
        if args['switchDistance'] is not None:
            force.setUseSwitchingFunction(True)
            force.setSwitchingDistance(args['switchDistance'])
peastman's avatar
peastman committed
2246
2247
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
2248
2249
        if 'useDispersionCorrection' in args:
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
peastman's avatar
peastman committed
2250
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
2251

peastman's avatar
peastman committed
2252
    def postprocessSystem(self, sys, data, args):
2253
2254
        # Create the exceptions.

2255
        bondIndices = _findBondsForExclusions(data, sys)
peastman's avatar
peastman committed
2256
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
2257
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
2258
2259
2260

parsers["NonbondedForce"] = NonbondedGenerator.parseElement

2261
2262

## @private
ChayaSt's avatar
ChayaSt committed
2263
class LennardJonesGenerator(object):
ChayaSt's avatar
ChayaSt committed
2264
2265
    """A NBFix generator to construct the L-J force with NBFIX implemented as a lookup table"""

ChayaSt's avatar
ChayaSt committed
2266
    def __init__(self, forcefield, lj14scale):
ChayaSt's avatar
ChayaSt committed
2267
        self.ff = forcefield
2268
        self.nbfixTypes = {}
ChayaSt's avatar
ChayaSt committed
2269
        self.lj14scale = lj14scale
2270
        self.ljTypes = ForceField._AtomTypeParameters(forcefield, 'LennardJonesForce', 'Atom', ('sigma', 'epsilon'))
ChayaSt's avatar
ChayaSt committed
2271

ChayaSt's avatar
ChayaSt committed
2272
    def registerNBFIX(self, parameters):
ChayaSt's avatar
ChayaSt committed
2273
2274
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
2275
2276
            type1 = types[0][0]
            type2 = types[1][0]
ChayaSt's avatar
ChayaSt committed
2277
2278
            epsilon = _convertParameterToNumber(parameters['epsilon'])
            sigma = _convertParameterToNumber(parameters['sigma'])
2279
2280
            self.nbfixTypes[(type1, type2)] = [sigma, epsilon]
            self.nbfixTypes[(type2, type1)] = [sigma, epsilon]
2281

ChayaSt's avatar
ChayaSt committed
2282
    def registerLennardJones(self, parameters):
2283
        self.ljTypes.registerAtom(parameters)
ChayaSt's avatar
ChayaSt committed
2284
2285
2286

    @staticmethod
    def parseElement(element, ff):
ChayaSt's avatar
ChayaSt committed
2287
        existing = [f for f in ff._forces if isinstance(f, LennardJonesGenerator)]
ChayaSt's avatar
ChayaSt committed
2288
        if len(existing) == 0:
ChayaSt's avatar
ChayaSt committed
2289
            generator = LennardJonesGenerator(ff, float(element.attrib['lj14scale']))
ChayaSt's avatar
ChayaSt committed
2290
2291
            ff.registerGenerator(generator)
        else:
ChayaSt's avatar
ChayaSt committed
2292
            # Multiple <LennardJonesForce> tags were found, probably in different files
ChayaSt's avatar
ChayaSt committed
2293
            generator = existing[0]
2294
2295
            if abs(generator.lj14scale - float(element.attrib['lj14scale'])) > NonbondedGenerator.SCALETOL:
                raise ValueError('Found multiple LennardJonesForce tags with different 1-4 scales')
ChayaSt's avatar
ChayaSt committed
2296
2297
        for LJ in element.findall('Atom'):
            generator.registerLennardJones(LJ.attrib)
ChayaSt's avatar
ChayaSt committed
2298
        for Nbfix in element.findall('NBFixPair'):
ChayaSt's avatar
ChayaSt committed
2299
            generator.registerNBFIX(Nbfix.attrib)
ChayaSt's avatar
ChayaSt committed
2300

ChayaSt's avatar
ChayaSt committed
2301
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
peastman's avatar
peastman committed
2302
2303
        # First derive the lookup tables.  We need to include entries for every type
        # that a) appears in the system and b) has unique parameters.
2304
2305

        nbfixTypeSet = set().union(*self.nbfixTypes)
peastman's avatar
peastman committed
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
        allTypes = set(data.atomType[atom] for atom in data.atoms)
        mergedTypes = []
        mergedTypeParams = []
        paramsToMergedType = {}
        typeToMergedType = {}
        for t in allTypes:
            typeParams = self.ljTypes.paramsForType[t]
            params = (typeParams['sigma'], typeParams['epsilon'])
            if t in nbfixTypeSet:
                # NBFIX types cannot be merged.
                typeToMergedType[t] = len(mergedTypes)
                mergedTypes.append(t)
                mergedTypeParams.append(params)
            elif params in paramsToMergedType:
                # We can merge this with another type.
                typeToMergedType[t] = paramsToMergedType[params]
2322
            else:
peastman's avatar
peastman committed
2323
2324
2325
2326
2327
                # This is a new type.
                typeToMergedType[t] = len(mergedTypes)
                paramsToMergedType[params] = len(mergedTypes)
                mergedTypes.append(t)
                mergedTypeParams.append(params)
2328

ChayaSt's avatar
ChayaSt committed
2329
        # Now everything is assigned. Create the A- and B-coefficient arrays
2330

peastman's avatar
peastman committed
2331
        numLjTypes = len(mergedTypes)
2332
        acoef = [0]*(numLjTypes*numLjTypes)
ChayaSt's avatar
ChayaSt committed
2333
        bcoef = acoef[:]
2334
2335
        for m in range(numLjTypes):
            for n in range(numLjTypes):
peastman's avatar
peastman committed
2336
                pair = (mergedTypes[m], mergedTypes[n])
2337
2338
2339
2340
2341
2342
                if pair in self.nbfixTypes:
                    epsilon = self.nbfixTypes[pair][1]
                    sigma = self.nbfixTypes[pair][0]
                    sigma6 = sigma**6
                    acoef[m+numLjTypes*n] = 4*epsilon*sigma6*sigma6
                    bcoef[m+numLjTypes*n] = 4*epsilon*sigma6
ChayaSt's avatar
cleanup  
ChayaSt committed
2343
                    continue
ChayaSt's avatar
ChayaSt committed
2344
                else:
peastman's avatar
peastman committed
2345
                    sigma = 0.5*(mergedTypeParams[m][0]+mergedTypeParams[n][0])
2346
                    sigma6 = sigma**6
peastman's avatar
peastman committed
2347
                    epsilon = math.sqrt(mergedTypeParams[m][1]*mergedTypeParams[n][1])
2348
2349
2350
2351
2352
2353
                    acoef[m+numLjTypes*n] = 4*epsilon*sigma6*sigma6
                    bcoef[m+numLjTypes*n] = 4*epsilon*sigma6

        self.force = mm.CustomNonbondedForce('acoef(type1, type2)/r^12 - bcoef(type1, type2)/r^6;')
        self.force.addTabulatedFunction('acoef', mm.Discrete2DFunction(numLjTypes, numLjTypes, acoef))
        self.force.addTabulatedFunction('bcoef', mm.Discrete2DFunction(numLjTypes, numLjTypes, bcoef))
ChayaSt's avatar
ChayaSt committed
2354
        self.force.addPerParticleParameter('type')
Peter Eastman's avatar
Peter Eastman committed
2355
        if nonbondedMethod in [CutoffPeriodic, Ewald, PME, LJPME]:
ChayaSt's avatar
ChayaSt committed
2356
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
ChayaSt's avatar
ChayaSt committed
2357
        elif nonbondedMethod is NoCutoff:
ChayaSt's avatar
ChayaSt committed
2358
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
ChayaSt's avatar
ChayaSt committed
2359
        elif nonbondedMethod is CutoffNonPeriodic:
ChayaSt's avatar
ChayaSt committed
2360
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic)
ChayaSt's avatar
ChayaSt committed
2361
        else:
2362
            raise AssertionError('Unrecognized nonbonded method [%s]' % nonbondedMethod)
2363
        if args['switchDistance'] is not None:
2364
2365
            self.force.setUseSwitchingFunction(True)
            self.force.setSwitchingDistance(args['switchDistance'])
2366

ChayaSt's avatar
ChayaSt committed
2367
        # Add the particles
2368

peastman's avatar
peastman committed
2369
2370
        for atom in data.atoms:
            self.force.addParticle((typeToMergedType[data.atomType[atom]],))
ChayaSt's avatar
ChayaSt committed
2371
2372
        self.force.setUseLongRangeCorrection(True)
        self.force.setCutoffDistance(nonbondedCutoff)
ChayaSt's avatar
ChayaSt committed
2373
2374
2375
        sys.addForce(self.force)

    def postprocessSystem(self, sys, data, args):
ChayaSt's avatar
ChayaSt committed
2376
        # Create the exceptions.
2377

2378
        bondIndices = _findBondsForExclusions(data, sys)
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
        forceCopy = deepcopy(self.force)
        forceCopy.createExclusionsFromBonds(bondIndices, 2)
        self.force.createExclusionsFromBonds(bondIndices, 3)
        if self.force.getNumExclusions() > forceCopy.getNumExclusions() and self.lj14scale != 0:
            # We need to create a CustomBondForce and use it to implement the scaled 1-4 interactions.

            bonded = mm.CustomBondForce('%g*epsilon*((sigma/r)^12-(sigma/r)^6)' % (4*self.lj14scale))
            bonded.addPerBondParameter('sigma')
            bonded.addPerBondParameter('epsilon')
            sys.addForce(bonded)
            skip = set(tuple(forceCopy.getExclusionParticles(i)) for i in range(forceCopy.getNumExclusions()))
            for i in range(self.force.getNumExclusions()):
                p1,p2 = self.force.getExclusionParticles(i)
                a1 = data.atoms[p1]
                a2 = data.atoms[p2]
                if (p1,p2) not in skip and (p2,p1) not in skip:
                    type1 = data.atomType[a1]
                    type2 = data.atomType[a2]
                    if (type1, type2) in self.nbfixTypes:
                        sigma, epsilon = self.nbfixTypes[(type1, type2)]
                    else:
                        values1 = self.ljTypes.getAtomParameters(a1, data)
                        values2 = self.ljTypes.getAtomParameters(a2, data)
2402
2403
2404
2405
2406
2407
2408
2409
                        extra1 = self.ljTypes.getExtraParameters(a1, data)
                        extra2 = self.ljTypes.getExtraParameters(a2, data)
                        sigma1 = float(extra1['sigma14']) if 'sigma14' in extra1 else values1[0]
                        sigma2 = float(extra2['sigma14']) if 'sigma14' in extra2 else values2[0]
                        epsilon1 = float(extra1['epsilon14']) if 'epsilon14' in extra1 else values1[1]
                        epsilon2 = float(extra2['epsilon14']) if 'epsilon14' in extra2 else values2[1]
                        sigma = 0.5*(sigma1+sigma2)
                        epsilon = sqrt(epsilon1*epsilon2)
2410
                    bonded.addBond(p1, p2, (sigma, epsilon))
ChayaSt's avatar
ChayaSt committed
2411

ChayaSt's avatar
ChayaSt committed
2412
parsers["LennardJonesForce"] = LennardJonesGenerator.parseElement
2413

2414

2415
## @private
2416
class GBSAOBCGenerator(object):
2417
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2418

2419
2420
    def __init__(self, forcefield):
        self.ff = forcefield
2421
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
2422

2423
    def registerAtom(self, parameters):
2424
        self.params.registerAtom(parameters)
2425

2426
2427
    @staticmethod
    def parseElement(element, ff):
2428
2429
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
2430
2431
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
2432
2433
2434
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2435
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2436

2437
2438
2439
2440
2441
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.NonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.NonbondedForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.NonbondedForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2442
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
2443
2444
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
2445
2446
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
2447
2448
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2449
2450
2451
2452
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
2453
2454
        sys.addForce(force)

2455
2456
    def postprocessSystem(self, sys, data, args):
        # Disable the reaction field approximation, since it produces bad results when combined with GB.
Justin MacCallum's avatar
Justin MacCallum committed
2457

2458
2459
2460
2461
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

2462
2463
2464
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


2465
## @private
2466
class CustomBondGenerator(object):
2467
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2468

2469
2470
    def __init__(self, forcefield):
        self.ff = forcefield
2471
2472
2473
2474
2475
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2476

2477
2478
    @staticmethod
    def parseElement(element, ff):
2479
2480
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
2481
2482
2483
2484
2485
2486
        generator.energy = element.attrib['energy']
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerBondParameter'):
            generator.perBondParams.append(param.attrib['name'])
        for bond in element.findall('Bond'):
2487
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2488
            if None not in types:
2489
2490
2491
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.paramValues.append([float(bond.attrib[param]) for param in generator.perBondParams])
Justin MacCallum's avatar
Justin MacCallum committed
2492

2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        force = mm.CustomBondForce(self.energy)
        sys.addForce(force)
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perBondParams:
            force.addPerBondParameter(param)
        for bond in data.bonds:
            type1 = data.atomType[data.atoms[bond.atom1]]
            type2 = data.atomType[data.atoms[bond.atom2]]
            for i in range(len(self.types1)):
                types1 = self.types1[i]
                types2 = self.types2[i]
                if (type1 in types1 and type2 in types2) or (type1 in types2 and type2 in types1):
                    force.addBond(bond.atom1, bond.atom2, self.paramValues[i])
                    break

parsers["CustomBondForce"] = CustomBondGenerator.parseElement


2513
## @private
2514
class CustomAngleGenerator(object):
2515
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2516

2517
2518
    def __init__(self, forcefield):
        self.ff = forcefield
2519
2520
2521
2522
2523
2524
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2525

2526
2527
    @staticmethod
    def parseElement(element, ff):
2528
2529
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
2530
2531
2532
2533
2534
2535
        generator.energy = element.attrib['energy']
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerAngleParameter'):
            generator.perAngleParams.append(param.attrib['name'])
        for angle in element.findall('Angle'):
2536
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2537
            if None not in types:
2538
2539
2540
2541
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.paramValues.append([float(angle.attrib[param]) for param in generator.perAngleParams])
Justin MacCallum's avatar
Justin MacCallum committed
2542

2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        force = mm.CustomAngleForce(self.energy)
        sys.addForce(force)
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perAngleParams:
            force.addPerAngleParameter(param)
        for angle in data.angles:
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
            for i in range(len(self.types1)):
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
                if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                    force.addAngle(angle[0], angle[1], angle[2], self.paramValues[i])
                    break

parsers["CustomAngleForce"] = CustomAngleGenerator.parseElement


2565
## @private
2566
class CustomTorsion(object):
2567
2568
    """A CustomTorsion records the information for a custom torsion definition."""

2569
    def __init__(self, types, paramValues, ordering='charmm'):
2570
2571
2572
2573
2574
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.paramValues = paramValues
2575
        if ordering in ['default', 'charmm', 'amber']:
2576
2577
            self.ordering = ordering
        else:
2578
            raise ValueError('Illegal ordering type %s for CustomTorsion (%s,%s,%s,%s)' % (ordering, types[0], types[1], types[2], types[3]))
2579

2580
## @private
2581
class CustomTorsionGenerator(object):
2582
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2583

2584
2585
    def __init__(self, forcefield):
        self.ff = forcefield
2586
2587
2588
2589
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
2590

2591
2592
    @staticmethod
    def parseElement(element, ff):
2593
2594
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
2595
2596
2597
2598
2599
2600
        generator.energy = element.attrib['energy']
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerTorsionParameter'):
            generator.perTorsionParams.append(param.attrib['name'])
        for torsion in element.findall('Proper'):
2601
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2602
            if None not in types:
2603
2604
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
2605
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2606
            if None not in types:
2607
2608
2609
2610
                if 'ordering' in element.attrib:
                    generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams], element.attrib['ordering']))
                else:
                    generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
Justin MacCallum's avatar
Justin MacCallum committed
2611

2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        force = mm.CustomTorsionForce(self.energy)
        sys.addForce(force)
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perTorsionParams:
            force.addPerTorsionParameter(param)
        wildcard = self.ff._atomClasses['']
        for torsion in data.propers:
            type1 = data.atomType[data.atoms[torsion[0]]]
            type2 = data.atomType[data.atoms[torsion[1]]]
            type3 = data.atomType[data.atoms[torsion[2]]]
            type4 = data.atomType[data.atoms[torsion[3]]]
            match = None
            for tordef in self.proper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
                if (type2 in types2 and type3 in types3 and type4 in types4 and type1 in types1) or (type2 in types3 and type3 in types2 and type4 in types1 and type1 in types4):
                    hasWildcard = (wildcard in (types1, types2, types3, types4))
                    if match is None or not hasWildcard: # Prefer specific definitions over ones with wildcards
                        match = tordef
                    if not hasWildcard:
                        break
            if match is not None:
                force.addTorsion(torsion[0], torsion[1], torsion[2], torsion[3], match.paramValues)
        for torsion in data.impropers:
2640
            match = _matchImproper(data, torsion, self)
2641
2642
2643
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
2644
2645
2646
2647

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


2648
## @private
2649
class CustomNonbondedGenerator(object):
2650
2651
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

2652
2653
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
2654
2655
2656
2657
2658
2659
2660
2661
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2662
2663
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
2664
2665
2666
2667
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerParticleParameter'):
            generator.perParticleParams.append(param.attrib['name'])
2668
2669
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2670
        generator.functions += _parseFunctions(element)
2671
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomNonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomNonbondedForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.CustomNonbondedForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
            raise ValueError('Illegal nonbonded method for CustomNonbondedForce')
        force = mm.CustomNonbondedForce(self.energy)
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perParticleParams:
            force.addPerParticleParameter(param)
2683
        _createFunctions(force, self.functions)
2684
        for atom in data.atoms:
2685
2686
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2687
2688
2689
2690
2691
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

    def postprocessSystem(self, sys, data, args):
2692
        # Create the exclusions.
2693

2694
        bondIndices = _findBondsForExclusions(data, sys)
2695
2696
2697
2698
2699
2700
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomNonbondedForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


2701
## @private
2702
class CustomGBGenerator(object):
2703
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2704

2705
2706
    def __init__(self, forcefield):
        self.ff = forcefield
2707
2708
2709
2710
2711
2712
2713
2714
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2715
2716
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
2717
2718
2719
2720
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerParticleParameter'):
            generator.perParticleParams.append(param.attrib['name'])
2721
2722
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2723
2724
2725
2726
2727
2728
2729
        computationMap = {"SingleParticle" : mm.CustomGBForce.SingleParticle,
                          "ParticlePair" : mm.CustomGBForce.ParticlePair,
                          "ParticlePairNoExclusions" : mm.CustomGBForce.ParticlePairNoExclusions}
        for value in element.findall('ComputedValue'):
            generator.computedValues.append((value.attrib['name'], value.text, computationMap[value.attrib['type']]))
        for term in element.findall('EnergyTerm'):
            generator.energyTerms.append((term.text, computationMap[term.attrib['type']]))
2730
        generator.functions += _parseFunctions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2731

2732
2733
2734
2735
2736
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomGBForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomGBForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.CustomGBForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2737
            raise ValueError('Illegal nonbonded method for CustomGBForce')
2738
2739
2740
2741
2742
2743
2744
2745
2746
        force = mm.CustomGBForce()
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perParticleParams:
            force.addPerParticleParameter(param)
        for value in self.computedValues:
            force.addComputedValue(value[0], value[1], value[2])
        for term in self.energyTerms:
            force.addEnergyTerm(term[0], term[1])
2747
        _createFunctions(force, self.functions)
2748
        for atom in data.atoms:
2749
2750
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2751
2752
2753
2754
2755
2756
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

2757

2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
## @private
class CustomHbondGenerator(object):
    """A CustomHbondGenerator constructs a CustomHbondForce."""

    def __init__(self, forcefield):
        self.ff = forcefield
        self.donorTypes1 = []
        self.donorTypes2 = []
        self.donorTypes3 = []
        self.acceptorTypes1 = []
        self.acceptorTypes2 = []
        self.acceptorTypes3 = []
        self.globalParams = {}
        self.perDonorParams = []
        self.perAcceptorParams = []
        self.donorParamValues = []
        self.acceptorParamValues = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
        generator = CustomHbondGenerator(ff)
        ff.registerGenerator(generator)
        generator.energy = element.attrib['energy']
        generator.bondCutoff = int(element.attrib['bondCutoff'])
        generator.particlesPerDonor = int(element.attrib['particlesPerDonor'])
        generator.particlesPerAcceptor = int(element.attrib['particlesPerAcceptor'])
        if generator.particlesPerDonor < 1 or generator.particlesPerDonor > 3:
            raise ValueError('Illegal value for particlesPerDonor for CustomHbondForce')
        if generator.particlesPerAcceptor < 1 or generator.particlesPerAcceptor > 3:
            raise ValueError('Illegal value for particlesPerAcceptor for CustomHbondForce')
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerDonorParameter'):
            generator.perDonorParams.append(param.attrib['name'])
        for param in element.findall('PerAcceptorParameter'):
            generator.perAcceptorParams.append(param.attrib['name'])
        for donor in element.findall('Donor'):
2796
            types = ff._findAtomTypes(donor.attrib, 3)[:generator.particlesPerDonor]
2797
2798
2799
2800
2801
2802
2803
2804
            if None not in types:
                generator.donorTypes1.append(types[0])
                if len(types) > 1:
                    generator.donorTypes2.append(types[1])
                if len(types) > 2:
                    generator.donorTypes3.append(types[2])
                generator.donorParamValues.append([float(donor.attrib[param]) for param in generator.perDonorParams])
        for acceptor in element.findall('Acceptor'):
2805
            types = ff._findAtomTypes(acceptor.attrib, 3)[:generator.particlesPerAcceptor]
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
            if None not in types:
                generator.acceptorTypes1.append(types[0])
                if len(types) > 1:
                    generator.acceptorTypes2.append(types[1])
                if len(types) > 2:
                    generator.acceptorTypes3.append(types[2])
                generator.acceptorParamValues.append([float(acceptor.attrib[param]) for param in generator.perAcceptorParams])
        generator.functions += _parseFunctions(element)

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomHbondForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomHbondForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.CustomHbondForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
            raise ValueError('Illegal nonbonded method for CustomNonbondedForce')
        force = mm.CustomHbondForce(self.energy)
        sys.addForce(force)
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perDonorParams:
            force.addPerDonorParameter(param)
        for param in self.perAcceptorParams:
            force.addPerAcceptorParameter(param)
        _createFunctions(force, self.functions)
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)

        # Add donors.

        if self.particlesPerDonor == 1:
            for atom in data.atoms:
                type1 = data.atomType[atom]
                for i in range(len(self.donorTypes1)):
                    types1 = self.donorTypes1[i]
                    if type1 in self.donorTypes1[i]:
                        force.addDonor(atom.index, -1, -1, self.donorParamValues[i])
        elif self.particlesPerDonor == 2:
            for bond in data.bonds:
                type1 = data.atomType[data.atoms[bond.atom1]]
                type2 = data.atomType[data.atoms[bond.atom2]]
                for i in range(len(self.donorTypes1)):
                    types1 = self.donorTypes1[i]
                    types2 = self.donorTypes2[i]
2849
2850
2851
2852
                    if type1 in types1 and type2 in types2:
                        force.addDonor(bond.atom1, bond.atom2, -1, self.donorParamValues[i])
                    elif type1 in types2 and type2 in types1:
                        force.addDonor(bond.atom2, bond.atom1, -1, self.donorParamValues[i])
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
        else:
            for angle in data.angles:
                type1 = data.atomType[data.atoms[angle[0]]]
                type2 = data.atomType[data.atoms[angle[1]]]
                type3 = data.atomType[data.atoms[angle[2]]]
                for i in range(len(self.donorTypes1)):
                    types1 = self.donorTypes1[i]
                    types2 = self.donorTypes2[i]
                    types3 = self.donorTypes3[i]
                    if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                        force.addDonor(angle[0], angle[1], angle[2], self.donorParamValues[i])

        # Add acceptors.

        if self.particlesPerAcceptor == 1:
            for atom in data.atoms:
                type1 = data.atomType[atom]
                for i in range(len(self.acceptorTypes1)):
                    types1 = self.acceptorTypes1[i]
                    if type1 in self.acceptorTypes1[i]:
                        force.addAcceptor(atom.index, -1, -1, self.acceptorParamValues[i])
        elif self.particlesPerAcceptor == 2:
            for bond in data.bonds:
                type1 = data.atomType[data.atoms[bond.atom1]]
                type2 = data.atomType[data.atoms[bond.atom2]]
                for i in range(len(self.acceptorTypes1)):
                    types1 = self.acceptorTypes1[i]
                    types2 = self.acceptorTypes2[i]
2881
                    if type1 in types1 and type2 in types2:
2882
                        force.addAcceptor(bond.atom1, bond.atom2, -1, self.acceptorParamValues[i])
2883
2884
                    elif type1 in types2 and type2 in types1:
                        force.addAcceptor(bond.atom2, bond.atom1, -1, self.acceptorParamValues[i])
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920
2921
2922
        else:
            for angle in data.angles:
                type1 = data.atomType[data.atoms[angle[0]]]
                type2 = data.atomType[data.atoms[angle[1]]]
                type3 = data.atomType[data.atoms[angle[2]]]
                for i in range(len(self.acceptorTypes1)):
                    types1 = self.acceptorTypes1[i]
                    types2 = self.acceptorTypes2[i]
                    types3 = self.acceptorTypes3[i]
                    if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                        force.addAcceptor(angle[0], angle[1], angle[2], self.acceptorParamValues[i])

        # Add exclusions.

        for donor in range(force.getNumDonors()):
            (d1, d2, d3, params) = force.getDonorParameters(donor)
            outerAtoms = set((d1, d2, d3))
            if -1 in outerAtoms:
                outerAtoms.remove(-1)
            excludedAtoms = set(outerAtoms)
            for i in range(self.bondCutoff):
                newOuterAtoms = set()
                for atom in outerAtoms:
                    for bond in data.atomBonds[atom]:
                        b = data.bonds[bond]
                        bondedAtom = (b.atom2 if b.atom1 == atom else b.atom1)
                        if bondedAtom not in excludedAtoms:
                            newOuterAtoms.add(bondedAtom)
                            excludedAtoms.add(bondedAtom)
                outerAtoms = newOuterAtoms
            for acceptor in range(force.getNumAcceptors()):
                (a1, a2, a3, params) = force.getAcceptorParameters(acceptor)
                if a1 in excludedAtoms or a2 in excludedAtoms or a3 in excludedAtoms:
                    force.addExclusion(donor, acceptor)

parsers["CustomHbondForce"] = CustomHbondGenerator.parseElement


2923
## @private
2924
class CustomManyParticleGenerator(object):
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
    """A CustomManyParticleGenerator constructs a CustomManyParticleForce."""

    def __init__(self, forcefield, particlesPerSet, energy, permutationMode, bondCutoff):
        self.ff = forcefield
        self.particlesPerSet = particlesPerSet
        self.energy = energy
        self.permutationMode = permutationMode
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
        self.functions = []
        self.typeFilters = []
2937

2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
    @staticmethod
    def parseElement(element, ff):
        permutationMap = {"SinglePermutation" : mm.CustomManyParticleForce.SinglePermutation,
                          "UniqueCentralParticle" : mm.CustomManyParticleForce.UniqueCentralParticle}
        generator = CustomManyParticleGenerator(ff, int(element.attrib['particlesPerSet']), element.attrib['energy'], permutationMap[element.attrib['permutationMode']], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerParticleParameter'):
            generator.perParticleParams.append(param.attrib['name'])
        for param in element.findall('TypeFilter'):
            generator.typeFilters.append((int(param.attrib['index']), [int(x) for x in param.attrib['types'].split(',')]))
2950
2951
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomManyParticleForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomManyParticleForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.CustomManyParticleForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
            raise ValueError('Illegal nonbonded method for CustomManyParticleForce')
        force = mm.CustomManyParticleForce(self.particlesPerSet, self.energy)
        force.setPermutationMode(self.permutationMode)
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perParticleParams:
            force.addPerParticleParameter(param)
        for index, types in self.typeFilters:
            force.setTypeFilter(index, types)
        for (name, type, values, params) in self.functions:
            if type == 'Continuous1D':
                force.addTabulatedFunction(name, mm.Continuous1DFunction(values, params['min'], params['max']))
            elif type == 'Continuous2D':
                force.addTabulatedFunction(name, mm.Continuous2DFunction(params['xsize'], params['ysize'], values, params['xmin'], params['xmax'], params['ymin'], params['ymax']))
            elif type == 'Continuous3D':
                force.addTabulatedFunction(name, mm.Continuous2DFunction(params['xsize'], params['ysize'], params['zsize'], values, params['xmin'], params['xmax'], params['ymin'], params['ymax'], params['zmin'], params['zmax']))
            elif type == 'Discrete1D':
                force.addTabulatedFunction(name, mm.Discrete1DFunction(values))
            elif type == 'Discrete2D':
                force.addTabulatedFunction(name, mm.Discrete2DFunction(params['xsize'], params['ysize'], values))
            elif type == 'Discrete3D':
                force.addTabulatedFunction(name, mm.Discrete2DFunction(params['xsize'], params['ysize'], params['zsize'], values))
        for atom in data.atoms:
2981
2982
2983
            values = self.params.getAtomParameters(atom, data)
            type = int(self.params.getExtraParameters(atom, data)['filterType'])
            force.addParticle(values, type)
2984
2985
2986
2987
2988
2989
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

    def postprocessSystem(self, sys, data, args):
        # Create exclusions based on bonds.
2990

2991
2992
2993
2994
2995
2996
2997
2998
2999
3000
3001
        bondIndices = []
        for bond in data.bonds:
            bondIndices.append((bond.atom1, bond.atom2))

        # If a virtual site does *not* share exclusions with another atom, add a bond between it and its first parent atom.

        for i in range(sys.getNumParticles()):
            if sys.isVirtualSite(i):
                (site, atoms, excludeWith) = data.virtualSites[data.atoms[i]]
                if excludeWith is None:
                    bondIndices.append((i, site.getParticle(0)))
3002

3003
3004
        # Certain particles, such as lone pairs and Drude particles, share exclusions with a parent atom.
        # If the parent atom does not interact with an atom, the child particle does not either.
3005

3006
3007
3008
3009
3010
3011
3012
3013
3014
        for atom1, atom2 in bondIndices:
            for child1 in data.excludeAtomWith[atom1]:
                bondIndices.append((child1, atom2))
                for child2 in data.excludeAtomWith[atom2]:
                    bondIndices.append((child1, child2))
            for child2 in data.excludeAtomWith[atom2]:
                bondIndices.append((atom1, child2))

        # Create the exclusions.
3015

3016
3017
3018
3019
3020
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomManyParticleForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomManyParticleForce"] = CustomManyParticleGenerator.parseElement

Peter Eastman's avatar
Peter Eastman committed
3021
def getAtomPrint(data, atomIndex):
3022

Peter Eastman's avatar
Peter Eastman committed
3023
3024
3025
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
3026
    else:
Peter Eastman's avatar
Peter Eastman committed
3027
        returnString = "NA"
3028
3029
3030
3031
3032

    return returnString

#=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3033
def countConstraint(data):
3034

Peter Eastman's avatar
Peter Eastman committed
3035
    bondCount = 0
3036
3037
3038
3039
3040
3041
3042
    angleCount = 0
    for bond in data.bonds:
        if bond.isConstrained:
            bondCount += 1

    angleCount = 0
    for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
Peter Eastman's avatar
Peter Eastman committed
3043
        if (isConstrained):
3044
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
3045

3046
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
3047

3048
## @private
3049
class AmoebaBondGenerator(object):
3050
3051
3052

    #=============================================================================================

3053
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
3054
3055

    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3056

3057
3058
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
3059
3060
3061
3062
3063
3064
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3065

3066
3067
3068
3069
3070
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

3071
        # <AmoebaBondForce bond-cubic="-25.5" bond-quartic="379.3125">
3072
        # <Bond class1="1" class2="2" length="0.1437" k="156900.0"/>
Justin MacCallum's avatar
Justin MacCallum committed
3073

3074
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
3075
3076
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
3077
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
3078
            if None not in types:
3079
3080
3081
3082
3083
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.length.append(float(bond.attrib['length']))
                generator.k.append(float(bond.attrib['k']))
            else:
3084
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
3085
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
3086
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3087
3088
                raise ValueError(outputString)

3089
3090
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3091
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3092

Mark Friedrichs's avatar
Mark Friedrichs committed
3093
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3094

Peter Eastman's avatar
Peter Eastman committed
3095
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3096
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
3097
        if len(existing) == 0:
3098
            force = mm.AmoebaBondForce()
3099
3100
3101
3102
            sys.addForce(force)
        else:
            force = existing[0]

3103
3104
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
3105
3106
3107
3108
3109
3110
3111
3112
3113
3114

        for bond in data.bonds:
            type1 = data.atomType[data.atoms[bond.atom1]]
            type2 = data.atomType[data.atoms[bond.atom2]]
            for i in range(len(self.types1)):
                types1 = self.types1[i]
                types2 = self.types2[i]
                if (type1 in types1 and type2 in types2) or (type1 in types2 and type2 in types1):
                    bond.length = self.length[i]
                    if bond.isConstrained:
3115
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
3116
3117
3118
                    if self.k[i] != 0:
                        if not bond.isConstrained or args.get('flexibleConstraints', False):
                            force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
3119
3120
                    break

3121
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
3122
3123
3124
3125

#=============================================================================================
# Add angle constraint
#=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3126

Peter Eastman's avatar
Peter Eastman committed
3127
def addAngleConstraint(angle, idealAngle, data, sys):
3128
3129

    # Find the two bonds that make this angle.
Justin MacCallum's avatar
Justin MacCallum committed
3130

Peter Eastman's avatar
Peter Eastman committed
3131
3132
    bond1 = None
    bond2 = None
3133
3134
3135
3136
3137
3138
3139
    for bond in data.atomBonds[angle[1]]:
        atom1 = data.bonds[bond].atom1
        atom2 = data.bonds[bond].atom2
        if atom1 == angle[0] or atom2 == angle[0]:
            bond1 = bond
        elif atom1 == angle[2] or atom2 == angle[2]:
            bond2 = bond
Justin MacCallum's avatar
Justin MacCallum committed
3140

3141
        # Compute the distance between atoms and add a constraint
Justin MacCallum's avatar
Justin MacCallum committed
3142

3143
3144
3145
3146
        if bond1 is not None and bond2 is not None:
            l1 = data.bonds[bond1].length
            l2 = data.bonds[bond2].length
            if l1 is not None and l2 is not None:
Peter Eastman's avatar
Peter Eastman committed
3147
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
3148
                data.addConstraint(sys, angle[0], angle[2], length)
3149
3150
3151
                return

#=============================================================================================
3152
## @private
3153
class AmoebaAngleGenerator(object):
3154
3155

    #=============================================================================================
3156
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
3157
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3158

3159
3160
    def __init__(self, forceField, cubic, quartic, pentic, sextic):

Peter Eastman's avatar
Peter Eastman committed
3161
3162
3163
3164
3165
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3166

Peter Eastman's avatar
Peter Eastman committed
3167
3168
3169
        self.types1 = []
        self.types2 = []
        self.types3 = []
3170

Peter Eastman's avatar
Peter Eastman committed
3171
3172
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3173

3174
3175
3176
3177
3178
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

3179
        # <AmoebaAngleForce angle-cubic="-0.014" angle-quartic="5.6e-05" angle-pentic="-7e-07" angle-sextic="2.2e-08">
3180
3181
        #   <Angle class1="2" class2="1" class3="3" k="0.0637259642196" angle1="122.00"  />

3182
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
3183
3184
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
3185
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
3186
            if None not in types:
3187
3188
3189
3190
3191
3192

                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
3193
                angleList.append(float(angle.attrib['angle1']))
3194
3195

                try:
Peter Eastman's avatar
Peter Eastman committed
3196
                    angleList.append(float(angle.attrib['angle2']))
3197
                    try:
Peter Eastman's avatar
Peter Eastman committed
3198
                        angleList.append(float(angle.attrib['angle3']))
3199
3200
3201
3202
3203
3204
3205
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
3206
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
3207
3208
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
3209
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
3210
3211
                raise ValueError(outputString)

3212
3213
3214
3215
    #=============================================================================================
    # createForce is bypassed here since the AmoebaOutOfPlaneBendForce generator must first execute
    # and partition angles into in-plane and non-in-plane angles
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3216

Peter Eastman's avatar
Peter Eastman committed
3217
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3218
3219
3220
3221
3222
3223
        pass

    #=============================================================================================
    # createForcePostOpBendAngle is called by AmoebaOutOfPlaneBendForce with the list of
    # non-in-plane angles
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3224

Peter Eastman's avatar
Peter Eastman committed
3225
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3226
3227
3228
3229

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3230
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
3231
3232

        if len(existing) == 0:
3233
            force = mm.AmoebaAngleForce()
3234
3235
3236
3237
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3238
        # set scalars
3239

3240
3241
3242
3243
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
3244

3245
3246
        DEG_TO_RAD = math.pi / 180

3247
        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
3248
3249
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3250

Peter Eastman's avatar
Peter Eastman committed
3251
3252
3253
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3254
3255
3256
3257
3258
3259
3260
            for i in range(len(self.types1)):
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
                if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                    if isConstrained and self.k[i] != 0.0:
                        angleDict['idealAngle'] = self.angle[i][0]
3261
3262
                        addAngleConstraint(angle, self.angle[i][0]*DEG_TO_RAD, data, sys)
                    if self.k[i] != 0 and (not isConstrained or args.get('flexibleConstraints', False)):
Peter Eastman's avatar
Peter Eastman committed
3263
3264
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
3265
3266
3267
3268
3269
3270
                            # get k-index by counting number of non-angle hydrogens on the central atom
                            # based on kangle.f
                            numberOfHydrogens = 0
                            for bond in data.atomBonds[angle[1]]:
                                atom1 = data.bonds[bond].atom1
                                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3271
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
3272
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3273
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
3274
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3275
                            if (numberOfHydrogens < lenAngle):
3276
3277
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
3278
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
3279
                                raise ValueError(outputString)
3280
3281
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
3282

3283
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
3284
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
3285
3286
3287
3288
3289
3290
                    break

    #=============================================================================================
    # createForcePostOpBendInPlaneAngle is called by AmoebaOutOfPlaneBendForce with the list of
    # in-plane angles
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3291

Peter Eastman's avatar
Peter Eastman committed
3292
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3293
3294
3295
3296

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3297
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
3298
3299

        if len(existing) == 0:
3300
            force = mm.AmoebaInPlaneAngleForce()
3301
3302
3303
3304
3305
3306
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

3307
3308
3309
3310
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
3311
3312

        for angleDict in angleList:
Justin MacCallum's avatar
Justin MacCallum committed
3313

Peter Eastman's avatar
Peter Eastman committed
3314
3315
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3316

Peter Eastman's avatar
Peter Eastman committed
3317
3318
3319
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3320
3321
3322
3323
3324
3325
3326
3327
3328

            for i in range(len(self.types1)):

                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]

                if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                    angleDict['idealAngle'] = self.angle[i][0]
Peter Eastman's avatar
Peter Eastman committed
3329
                    if (isConstrained and self.k[i] != 0.0):
3330
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3331
                    if self.k[i] != 0.0 and (not isConstrained or args.get('flexibleConstraints', False)):
3332
3333
3334
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

3335
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
3336
3337
3338

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
3339
3340
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
3341
3342
#=============================================================================================

3343
## @private
3344
class AmoebaOutOfPlaneBendGenerator(object):
3345
3346
3347
3348

    #=============================================================================================

    """An AmoebaOutOfPlaneBendGenerator constructs a AmoebaOutOfPlaneBendForce."""
Justin MacCallum's avatar
Justin MacCallum committed
3349

3350
3351
3352
3353
    #=============================================================================================

    def __init__(self, forceField, type, cubic, quartic, pentic, sextic):

Peter Eastman's avatar
Peter Eastman committed
3354
3355
3356
3357
3358
3359
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3360

Peter Eastman's avatar
Peter Eastman committed
3361
3362
3363
3364
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3365

Peter Eastman's avatar
Peter Eastman committed
3366
        self.ks = []
3367
3368
3369
3370
3371

    #=============================================================================================
    # Local version of findAtomTypes needed since class indices are 0 (i.e., not recognized)
    # for types3 and 4
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3372

3373
3374
3375
3376
3377
3378
3379
3380
3381
3382
3383
3384
3385
3386
3387
3388
3389
3390
3391
3392
3393
3394
3395
3396
3397
    def findAtomTypes(self, forceField, node, num):
        """Parse the attributes on an XML tag to find the set of atom types for each atom it involves."""
        types = []
        attrib = node.attrib
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
            if classAttrib in attrib:
                if attrib[classAttrib] in forceField._atomClasses:
                    types.append(forceField._atomClasses[attrib[classAttrib]])
                else:
                    types.append(set())
        return types

    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

        #  <AmoebaOutOfPlaneBendForce type="ALLINGER" opbend-cubic="-0.014" opbend-quartic="5.6e-05" opbend-pentic="-7e-07" opbend-sextic="2.2e-08">
        #   <Angle class1="2" class2="1" class3="0" class4="0" k="0.0531474541591"/>
        #   <Angle class1="3" class2="1" class3="0" class4="0" k="0.0898536095496"/>
Justin MacCallum's avatar
Justin MacCallum committed
3398

3399
3400
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
3401
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
3402
3403
3404
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
3405
                                                   float(element.attrib['opbend-sextic']))
3406
3407
3408
3409

        forceField._forces.append(generator)

        for angle in element.findall('Angle'):
Peter Eastman's avatar
Peter Eastman committed
3410
            types = generator.findAtomTypes(forceField, angle, 4)
3411
3412
            if types is not None:

Peter Eastman's avatar
Peter Eastman committed
3413
3414
3415
3416
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
3417

Peter Eastman's avatar
Peter Eastman committed
3418
                generator.ks.append(float(angle.attrib['k']))
3419
3420

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3421
3422
                outputString = "AmoebaOutOfPlaneBendGenerator error getting types: %s %s %s %s." % (
                               angle.attrib['class1'], angle.attrib['class2'], angle.attrib['class3'], angle.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
3423
3424
                raise ValueError(outputString)

3425
3426
3427
3428
3429
3430
    #=============================================================================================
    # Get middle atom in a angle
    # return index of middle atom or -1 if no middle is found
    # This method appears not to be needed since the angle[1] entry appears to always
    # be the middle atom. However, was unsure if this is guaranteed
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3431

Peter Eastman's avatar
Peter Eastman committed
3432
    def getMiddleAtom(self, angle, data):
3433
3434
3435

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
3436
        middleAtom = -1
Justin MacCallum's avatar
Justin MacCallum committed
3437
        for atomIndex in angle:
Peter Eastman's avatar
Peter Eastman committed
3438
            isMiddle = 0
3439
3440
3441
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3442
                if (atom1 != atomIndex):
3443
3444
3445
                    partner = atom1
                else:
                    partner = atom2
Justin MacCallum's avatar
Justin MacCallum committed
3446
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]):
3447
3448
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
3449
            if (isMiddle == 2):
3450
3451
3452
3453
3454
                return atomIndex
        return -1

    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3455
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3456
3457
3458
3459
3460
3461
3462
3463
3464
3465
3466
3467
3468

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.AmoebaOutOfPlaneBendForce]
        if len(existing) == 0:
            force = mm.AmoebaOutOfPlaneBendForce()
            sys.addForce(force)
        else:
            force = existing[0]

        # set scalars

Peter Eastman's avatar
Peter Eastman committed
3469
3470
3471
3472
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
3473
3474
3475
3476

        # this hash is used to insure the out-of-plane-bend bonds
        # are only added once

Peter Eastman's avatar
Peter Eastman committed
3477
        skipAtoms = dict()
3478
3479
3480
3481

        # these lists are used in the partitioning of the angles into
        # angle and inPlane angles

Peter Eastman's avatar
Peter Eastman committed
3482
3483
        inPlaneAngles = []
        nonInPlaneAngles = []
3484
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
3485
        idealAngles = []*len(data.angles)
3486
3487
3488

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):

Peter Eastman's avatar
Peter Eastman committed
3489
3490
3491
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
3492
3493
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
3494
                middleType = -1
3495
3496
                middleCovalency = -1

Justin MacCallum's avatar
Justin MacCallum committed
3497
            # if middle atom has covalency of 3 and
3498
            # the types of the middle atom and the partner atom (atom bonded to
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3499
            # middle atom, but not in angle) match types1 and types2, then
Justin MacCallum's avatar
Justin MacCallum committed
3500
            # three out-of-plane bend angles are generated. Three in-plane angle
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3501
            # are also generated. If the conditions are not satisfied, the angle is marked as 'generic' angle (not a in-plane angle)
3502

Peter Eastman's avatar
Peter Eastman committed
3503
            if (middleAtom > -1 and middleCovalency == 3 and middleAtom not in skipAtoms):
3504

Peter Eastman's avatar
Peter Eastman committed
3505
3506
3507
3508
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
3509
3510
3511
3512

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3513
                    if (atom1 != middleAtom):
3514
3515
3516
3517
3518
3519
3520
3521
                        partner = atom1
                    else:
                        partner = atom2

                    partnerType = data.atomType[data.atoms[partner]]
                    for i in range(len(self.types1)):
                        types1 = self.types1[i]
                        types2 = self.types2[i]
Peter Eastman's avatar
Peter Eastman committed
3522
3523
3524
3525
3526
                        if (middleType in types2 and partnerType in types1):
                            partners.append(partner)
                            partnerSet.add(partner)
                            partnerTypes.append(partnerType)
                            partnerK.append(self.ks[i])
Justin MacCallum's avatar
Justin MacCallum committed
3527

Peter Eastman's avatar
Peter Eastman committed
3528
                if (len(partners) == 3):
3529

Peter Eastman's avatar
Peter Eastman committed
3530
3531
3532
                    force.addOutOfPlaneBend(partners[0], middleAtom, partners[1], partners[2], partnerK[2])
                    force.addOutOfPlaneBend(partners[0], middleAtom, partners[2], partners[1], partnerK[1])
                    force.addOutOfPlaneBend(partners[1], middleAtom, partners[2], partners[0], partnerK[0])
3533

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3534
3535
                    # skipAtoms is used to insure angles are only included once

3536
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
3537
3538
3539
3540
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3541
3542
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
3543
3544
3545
3546
3547
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
3548
3549
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
3550
                    angleDict['isConstrained'] = 0
3551

Peter Eastman's avatar
Peter Eastman committed
3552
3553
3554
3555
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3556
3557

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3558
3559
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3560

Peter Eastman's avatar
Peter Eastman committed
3561
                    inPlaneAngles.append(angleDict)
3562
3563

                else:
Peter Eastman's avatar
Peter Eastman committed
3564
3565
3566
3567
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
3568
            else:
Peter Eastman's avatar
Peter Eastman committed
3569
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
3570

Peter Eastman's avatar
Peter Eastman committed
3571
                    partnerSet = skipAtoms[middleAtom]
Justin MacCallum's avatar
Justin MacCallum committed
3572

Peter Eastman's avatar
Peter Eastman committed
3573
                    angleDict = {}
3574

Peter Eastman's avatar
Peter Eastman committed
3575
3576
3577
3578
3579
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
3580

Peter Eastman's avatar
Peter Eastman committed
3581
                    angleDict['isConstrained'] = isConstrained
3582

Peter Eastman's avatar
Peter Eastman committed
3583
3584
3585
3586
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3587
3588

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3589
3590
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3591

Peter Eastman's avatar
Peter Eastman committed
3592
                    inPlaneAngles.append(angleDict)
3593
3594

                else:
Peter Eastman's avatar
Peter Eastman committed
3595
3596
                    angleDict = {}
                    angleDict['angle'] = angle
3597
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
3598
                    nonInPlaneAngles.append(angleDict)
3599

3600
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
3601
3602

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3603
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
3604
3605
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
3606
3607

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3608
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
3609
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
3610
                    nonInPlaneAngles.append(angleDict)
3611
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
3612
3613
3614
3615
3616

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

#=============================================================================================

3617
## @private
3618
class AmoebaTorsionGenerator(object):
3619
3620
3621
3622
3623
3624
3625

    #=============================================================================================
    """An AmoebaTorsionGenerator constructs a AmoebaTorsionForce."""
    #=============================================================================================

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
3626
        self.torsionUnit = torsionUnit
3627

Peter Eastman's avatar
Peter Eastman committed
3628
3629
3630
3631
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3632

Peter Eastman's avatar
Peter Eastman committed
3633
3634
3635
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
3636

3637
3638
3639
3640
3641
3642
3643
3644
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

        #  <AmoebaTorsionForce torsionUnit="0.5">
        #   <Torsion class1="3" class2="1" class3="2" class4="3"   amp1="0.0" angle1="0.0"   amp2="0.0" angle2="3.14159265359"   amp3="0.0" angle3="0.0" />
        #   <Torsion class1="3" class2="1" class3="2" class4="6"   amp1="0.0" angle1="0.0"   amp2="0.0" angle2="3.14159265359"   amp3="-0.263592" angle3="0.0" />
Justin MacCallum's avatar
Justin MacCallum committed
3645

Peter Eastman's avatar
Peter Eastman committed
3646
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
3647
3648
3649
3650
3651
3652
        forceField._forces.append(generator)

        # collect particle classes and t1,t2,t3,
        # where ti=[amplitude_i,angle_i]

        for torsion in element.findall('Torsion'):
3653
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
3654
            if None not in types:
3655
3656
3657
3658
3659
3660
3661

                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])

                for ii in range(1,4):
Peter Eastman's avatar
Peter Eastman committed
3662
3663
                    tInfo = []
                    suffix = str(ii)
3664
3665
3666
3667
3668
3669
                    ampName = 'amp' + suffix
                    tInfo.append(float(torsion.attrib[ampName]))

                    angName = 'angle' + suffix
                    tInfo.append(float(torsion.attrib[angName]))

Peter Eastman's avatar
Peter Eastman committed
3670
3671
3672
3673
3674
3675
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
3676
3677
3678

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
3679
3680
3681
3682
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
3683
3684
                raise ValueError(outputString)

3685
3686
3687
3688
3689
    #=============================================================================================

    def createForce(self, sys, data, nontorsionedMethod, nontorsionedCutoff, args):

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3690
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
3691
        if len(existing) == 0:
3692
            force = mm.PeriodicTorsionForce()
3693
3694
3695
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3696

3697
3698
3699
3700
3701
3702
3703
3704
3705
3706
3707
3708
3709
3710
3711
3712
        for torsion in data.propers:

            type1 = data.atomType[data.atoms[torsion[0]]]
            type2 = data.atomType[data.atoms[torsion[1]]]
            type3 = data.atomType[data.atoms[torsion[2]]]
            type4 = data.atomType[data.atoms[torsion[3]]]

            for i in range(len(self.types1)):

                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
                types4 = self.types4[i]

                # match types in forward or reverse direction

Peter Eastman's avatar
Peter Eastman committed
3713
                if (type1 in types1 and type2 in types2 and type3 in types3 and type4 in types4) or (type4 in types1 and type3 in types2 and type2 in types3 and type1 in types4):
3714
3715
3716
3717
3718
3719
                    if self.t1[i][0] != 0:
                        force.addTorsion(torsion[0], torsion[1], torsion[2], torsion[3], 1, self.t1[i][1], self.t1[i][0])
                    if self.t2[i][0] != 0:
                        force.addTorsion(torsion[0], torsion[1], torsion[2], torsion[3], 2, self.t2[i][1], self.t2[i][0])
                    if self.t3[i][0] != 0:
                        force.addTorsion(torsion[0], torsion[1], torsion[2], torsion[3], 3, self.t3[i][1], self.t3[i][0])
3720
3721
3722
3723
3724
3725
                    break

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

#=============================================================================================

3726
## @private
3727
class AmoebaPiTorsionGenerator(object):
3728
3729
3730
3731
3732
3733

    #=============================================================================================

    """An AmoebaPiTorsionGenerator constructs a AmoebaPiTorsionForce."""

    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3734

3735
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
3736
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
3737
3738
3739
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3740

3741
3742
3743
3744
3745
3746
3747
3748
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

        #  <AmoebaPiTorsionForce piTorsionUnit="1.0">
        #   <PiTorsion class1="1" class2="3" k="28.6604" />

Peter Eastman's avatar
Peter Eastman committed
3749
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
3750
3751
3752
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
3753
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
3754
            if None not in types:
3755
3756
3757
3758
3759
3760
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.k.append(float(piTorsion.attrib['k']))
            else:
                outputString = "AmoebaPiTorsionGenerator: error getting types: %s %s " % (
                                    piTorsion.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
3761
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3762
3763
                raise ValueError(outputString)

3764
3765
3766
3767
3768
3769
3770
3771
3772
3773
3774
3775
3776
3777
3778
3779
    #=============================================================================================

    def createForce(self, sys, data, nonpiTorsionedMethod, nonpiTorsionedCutoff, args):

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.AmoebaPiTorsionForce]

        if len(existing) == 0:
            force = mm.AmoebaPiTorsionForce()
            sys.addForce(force)
        else:
            force = existing[0]

        for bond in data.bonds:

            # search for bonds with both atoms in bond having covalency == 3
Justin MacCallum's avatar
Justin MacCallum committed
3780

3781
3782
            atom1 = bond.atom1
            atom2 = bond.atom2
Justin MacCallum's avatar
Justin MacCallum committed
3783

3784
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
3785
3786
3787
3788
3789
3790
3791
3792
3793
3794
3795

                type1 = data.atomType[data.atoms[atom1]]
                type2 = data.atomType[data.atoms[atom2]]

                for i in range(len(self.types1)):

                   types1 = self.types1[i]
                   types2 = self.types2[i]

                   if (type1 in types1 and type2 in types2) or (type1 in types2 and type2 in types1):

Justin MacCallum's avatar
Justin MacCallum committed
3796
3797
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
3798
3799
3800
3801
3802
3803
3804
3805
3806
3807
3808
3809

                       piTorsionAtom1 = -1
                       piTorsionAtom2 = -1
                       piTorsionAtom3 = atom1

                       piTorsionAtom4 = atom2
                       piTorsionAtom5 = -1
                       piTorsionAtom6 = -1

                       for bond in data.atomBonds[atom1]:
                           bondedAtom1 = data.bonds[bond].atom1
                           bondedAtom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3810
                           if (bondedAtom1 != atom1):
3811
3812
3813
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
3814
3815
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3816
                                   piTorsionAtom1 = b1
3817
3818
3819
3820
3821
3822
                               else:
                                   piTorsionAtom2 = b1

                       for bond in data.atomBonds[atom2]:
                           bondedAtom1 = data.bonds[bond].atom1
                           bondedAtom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3823
                           if (bondedAtom1 != atom2):
3824
3825
3826
3827
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3828
3829
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3830
                                   piTorsionAtom5 = b1
3831
3832
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
3833

Peter Eastman's avatar
Peter Eastman committed
3834
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
3835
3836
3837
3838
3839

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

#=============================================================================================

3840
## @private
3841
class AmoebaTorsionTorsionGenerator(object):
3842
3843
3844
3845
3846

    #=============================================================================================
    """An AmoebaTorsionTorsionGenerator constructs a AmoebaTorsionTorsionForce."""
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3847
    def __init__(self):
3848

Peter Eastman's avatar
Peter Eastman committed
3849
3850
3851
3852
3853
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
3854

Peter Eastman's avatar
Peter Eastman committed
3855
        self.gridIndex = []
3856

Peter Eastman's avatar
Peter Eastman committed
3857
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
3858

3859
3860
3861
3862
3863
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
3864
        generator = AmoebaTorsionTorsionGenerator()
3865
3866
3867
3868
3869
3870
3871
        forceField._forces.append(generator)
        maxGridIndex = -1

        # <AmoebaTorsionTorsionForce >
        # <TorsionTorsion class1="3" class2="1" class3="2" class4="3" class5="1" grid="0" nx="25" ny="25" />

        for torsionTorsion in element.findall('TorsionTorsion'):
3872
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
3873
            if None not in types:
3874
3875
3876
3877
3878
3879
3880

                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
                generator.types5.append(types[4])

Peter Eastman's avatar
Peter Eastman committed
3881
3882
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
3883
3884
3885
3886
3887
3888
3889
3890
3891
                    maxGridIndex = gridIndex

                generator.gridIndex.append(gridIndex)
            else:
                outputString = "AmoebaTorsionTorsionGenerator: error getting types: %s %s %s %s %s" % (
                                    torsionTorsion.attrib['class1'],
                                    torsionTorsion.attrib['class2'],
                                    torsionTorsion.attrib['class3'],
                                    torsionTorsion.attrib['class4'],
Peter Eastman's avatar
Peter Eastman committed
3892
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
3893
3894
                raise ValueError(outputString)

3895
3896
3897
3898
3899
3900
3901
3902
3903
3904
        # load grid

        # xml source

        # <TorsionTorsionGrid grid="0" nx="25" ny="25" >
        # <Grid angle1="-180.00" angle2="-180.00" f="0.0" fx="2.31064374824e-05" fy="0.0" fxy="-0.0052801799672" />
        # <Grid angle1="-165.00" angle2="-180.00" f="-0.66600912" fx="-0.06983370052" fy="-0.075058725744" fxy="-0.0044462732032" />

        # output grid:

Justin MacCallum's avatar
Justin MacCallum committed
3905
3906
3907
3908
3909
3910
        #     grid[x][y][0] = x value
        #     grid[x][y][1] = y value
        #     grid[x][y][2] = function value
        #     grid[x][y][3] = dfdx value
        #     grid[x][y][4] = dfdy value
        #     grid[x][y][5] = dfd(xy) value
3911
3912

        maxGridIndex    += 1
Peter Eastman's avatar
Peter Eastman committed
3913
        generator.grids = maxGridIndex*[]
3914
3915
        for torsionTorsionGrid in element.findall('TorsionTorsionGrid'):

Peter Eastman's avatar
Peter Eastman committed
3916
3917
3918
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
3919

Peter Eastman's avatar
Peter Eastman committed
3920
3921
            grid = []
            gridCol = []
3922
3923
3924
3925
3926

            gridColIndex = 0

            for gridEntry in torsionTorsionGrid.findall('Grid'):

Peter Eastman's avatar
Peter Eastman committed
3927
3928
3929
3930
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
3931
                if 'fx' in gridEntry.attrib:
3932
3933
3934
                    gridRow.append(float(gridEntry.attrib['fx']))
                    gridRow.append(float(gridEntry.attrib['fy']))
                    gridRow.append(float(gridEntry.attrib['fxy']))
Peter Eastman's avatar
Peter Eastman committed
3935
                gridCol.append(gridRow)
3936
3937

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
3938
3939
3940
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
3941
3942
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
3943

Peter Eastman's avatar
Peter Eastman committed
3944
3945
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
3946
            else:
Peter Eastman's avatar
Peter Eastman committed
3947
3948
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
3949
3950
3951
3952
                generator.grids[gridIndex] = grid

    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3953
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
3954
3955
3956
3957
3958
3959
3960
3961

        chiralAtomIndex = -1

        # if atomC has four bonds, find the
        # two bonds that do not include atomB and atomD
        # set chiralAtomIndex to one of these, if they are
        # not the same atom(type/mass)

Peter Eastman's avatar
Peter Eastman committed
3962
        if (len(data.atomBonds[atomC]) == 4):
3963
3964
3965
3966
3967
            atomE = -1
            atomF = -1
            for bond in data.atomBonds[atomC]:
                bondedAtom1 = data.bonds[bond].atom1
                bondedAtom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3968
3969
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
3970
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
3971
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
3972
3973
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
3974
3975
                if (hit > -1):
                    if (atomE == -1):
3976
3977
3978
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
3979

3980
3981
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
3982
            if (atomE == -1 or atomF == -1):
3983
                outputString = "getChiralAtomIndex: error getting bonded partners of atomC=%s %d %s" % (atomC.name, atomC.residue.index, atomC.residue.name,)
Justin MacCallum's avatar
Justin MacCallum committed
3984
                raise ValueError(outputString)
3985
3986
3987
3988
3989

            # check for different type/mass between atoms E & F

            typeE = int(data.atomType[data.atoms[atomE]])
            typeF = int(data.atomType[data.atoms[atomF]])
Peter Eastman's avatar
Peter Eastman committed
3990
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
3991
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
3992
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
3993
                chiralAtomIndex = atomF
3994

Peter Eastman's avatar
Peter Eastman committed
3995
3996
3997
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
3998
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
3999
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
4000
                chiralAtomIndex = massF
4001
4002
4003
4004

        return chiralAtomIndex

    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
4005

4006
4007
4008
4009
4010
4011
4012
4013
4014
4015
4016
4017
4018
    def createForce(self, sys, data, nonpiTorsionedMethod, nonpiTorsionedCutoff, args):

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.AmoebaTorsionTorsionForce]

        if len(existing) == 0:
            force = mm.AmoebaTorsionTorsionForce()
            sys.addForce(force)
        else:
            force = existing[0]

        for angle in data.angles:

Justin MacCallum's avatar
Justin MacCallum committed
4019
4020
            # search for bitorsions; based on TINKER subroutine bitors()

4021
4022
4023
4024
4025
4026
4027
            ib = angle[0]
            ic = angle[1]
            id = angle[2]

            for bondIndex in data.atomBonds[ib]:
                bondedAtom1 = data.bonds[bondIndex].atom1
                bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
4028
                if (bondedAtom1 != ib):
4029
4030
4031
4032
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4033
                if (ia != ic and ia != id):
4034
4035
4036
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
4037
                        if (bondedAtom1 != id):
4038
4039
4040
4041
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4042
                        if (ie != ic and ie != ib and ie != ia):
4043
4044
4045
4046
4047
4048
4049
4050
4051
4052
4053
4054
4055
4056
4057
4058
4059
4060
4061
4062

                            # found candidate set of atoms
                            # check if types match in order or reverse order

                            type1 = data.atomType[data.atoms[ia]]
                            type2 = data.atomType[data.atoms[ib]]
                            type3 = data.atomType[data.atoms[ic]]
                            type4 = data.atomType[data.atoms[id]]
                            type5 = data.atomType[data.atoms[ie]]

                            for i in range(len(self.types1)):

                                types1 = self.types1[i]
                                types2 = self.types2[i]
                                types3 = self.types3[i]
                                types4 = self.types4[i]
                                types5 = self.types5[i]

                                # match in order

Peter Eastman's avatar
Peter Eastman committed
4063
4064
4065
                                if (type1 in types1 and type2 in types2 and type3 in types3 and type4 in types4 and type5 in types5):
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ia, ib, ic, id, ie, chiralAtomIndex, self.gridIndex[i])
4066
4067
4068

                                # match in reverse order

4069
                                elif (type5 in types1 and type4 in types2 and type3 in types3 and type2 in types4 and type1 in types5):
Peter Eastman's avatar
Peter Eastman committed
4070
4071
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
4072
4073
4074
4075

        # set grids

        for (index, grid) in enumerate(self.grids):
Peter Eastman's avatar
Peter Eastman committed
4076
            force.setTorsionTorsionGrid(index, grid)
Justin MacCallum's avatar
Justin MacCallum committed
4077

4078
4079
4080
4081
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

#=============================================================================================

4082
## @private
4083
class AmoebaStretchBendGenerator(object):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4084
4085

    #=============================================================================================
4086
4087
4088
4089
4090
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
4091
4092
4093
        self.types1 = []
        self.types2 = []
        self.types3 = []
4094

Peter Eastman's avatar
Peter Eastman committed
4095
4096
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
4097

4098
4099
4100
4101
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
4102
        generator = AmoebaStretchBendGenerator()
4103
4104
4105
4106
4107
4108
4109
        forceField._forces.append(generator)

        # <AmoebaStretchBendForce stretchBendUnit="1.0">
        # <StretchBend class1="2" class2="1" class3="3" k1="5.25776946506" k2="5.25776946506" />
        # <StretchBend class1="2" class2="1" class3="4" k1="3.14005676385" k2="3.14005676385" />

        for stretchBend in element.findall('StretchBend'):
4110
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
4111
            if None not in types:
4112
4113
4114
4115
4116
4117
4118
4119
4120
4121
4122
4123

                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])

                generator.k1.append(float(stretchBend.attrib['k1']))
                generator.k2.append(float(stretchBend.attrib['k2']))

            else:
                outputString = "AmoebaStretchBendGenerator : error getting types: %s %s %s" % (
                                    stretchBend.attrib['class1'],
                                    stretchBend.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
4124
                                    stretchBend.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
4125
4126
                raise ValueError(outputString)

4127
4128
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
4129
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
4130
4131
    # having been called since the ideal bond lengths and angle are needed here.
    # As a conseqeunce, createForce() is not implemented since it is not guaranteed that the generator for
Justin MacCallum's avatar
Justin MacCallum committed
4132
4133
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
4134
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
4135
4136
4137

    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4138
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4139
4140
4141
4142
4143
4144
4145
4146
        pass

    #=============================================================================================

    # Note: request for constrained bonds is ignored.

    #=============================================================================================

4147
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
4148
4149
4150
4151
4152
4153
4154
4155
4156
4157
4158
4159

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.AmoebaStretchBendForce]
        if len(existing) == 0:
            force = mm.AmoebaStretchBendForce()
            sys.addForce(force)
        else:
            force = existing[0]

        for angleDict in angleList:

            angle = angleDict['angle']
Peter Eastman's avatar
Peter Eastman committed
4160
            if ('isConstrained' in angleDict):
4161
4162
4163
4164
4165
4166
4167
4168
                isConstrained = angleDict['isConstrained']
            else:
                isConstrained = 0

            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]

Peter Eastman's avatar
Peter Eastman committed
4169
            radian = 57.2957795130
4170
4171
4172
4173
4174
4175
4176
4177
4178
4179
            for i in range(len(self.types1)):

                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]

                # match types
                # get ideal bond lengths, bondAB, bondCB
                # get ideal angle

Justin MacCallum's avatar
Justin MacCallum committed
4180
                if (type2 in types2 and ((type1 in types1 and type3 in types3) or (type3 in types1 and type1 in types3))):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4181
4182
4183
4184
4185
4186
4187
4188
4189
4190
4191
4192
4193
4194
4195
                    bondAB = -1.0
                    bondCB = -1.0
                    swap = 0
                    for bond in data.atomBonds[angle[1]]:
                        atom1 = data.bonds[bond].atom1
                        atom2 = data.bonds[bond].atom2
                        length = data.bonds[bond].length
                        if (atom1 == angle[0]):
                            bondAB = length
                        if (atom1 == angle[2]):
                            bondCB = length
                        if (atom2 == angle[2]):
                            bondCB = length
                        if (atom2 == angle[0]):
                            bondAB = length
Justin MacCallum's avatar
Justin MacCallum committed
4196

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4197
                    # check that ideal angle and bonds are set
4198

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4199
                    if ('idealAngle' not in angleDict):
4200

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4201
4202
4203
4204
4205
                       outputString = "AmoebaStretchBendGenerator: ideal angle is not set for following entry:\n"
                       outputString += "   types: %5s %5s %5s atoms: " % (type1, type2, type3)
                       outputString += getAtomPrint( data, angle[0] ) + ' '
                       outputString += getAtomPrint( data, angle[1] ) + ' '
                       outputString += getAtomPrint( data, angle[2] )
Justin MacCallum's avatar
Justin MacCallum committed
4206
                       raise ValueError(outputString)
4207

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4208
4209
4210
4211
4212
4213
4214
                    elif (bondAB < 0 or bondCB < 0):

                       outputString = "AmoebaStretchBendGenerator: bonds not set: %15.7e %15.7e. for following entry:" % (bondAB, bondCB)
                       outputString += "     types: [%5s %5s %5s] atoms: " % (type1, type2, type3)
                       outputString += getAtomPrint( data, angle[0] ) + ' '
                       outputString += getAtomPrint( data, angle[1] ) + ' '
                       outputString += getAtomPrint( data, angle[2] )
Justin MacCallum's avatar
Justin MacCallum committed
4215
                       raise ValueError(outputString)
4216

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4217
                    else:
4218
                        force.addStretchBend(angle[0], angle[1], angle[2], bondAB, bondCB, angleDict['idealAngle']/radian, self.k1[i], self.k2[i])
4219

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4220
                    break
4221
4222
4223
4224
4225

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

#=============================================================================================

4226
## @private
4227
class AmoebaVdwGenerator(object):
4228
4229

    """A AmoebaVdwGenerator constructs a AmoebaVdwForce."""
Justin MacCallum's avatar
Justin MacCallum committed
4230

4231
4232
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4233
    def __init__(self, type, radiusrule, radiustype, radiussize, epsilonrule, vdw13Scale, vdw14Scale, vdw15Scale):
4234

Justin MacCallum's avatar
Justin MacCallum committed
4235
        self.type = type
4236

Peter Eastman's avatar
Peter Eastman committed
4237
4238
4239
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
4240

Peter Eastman's avatar
Peter Eastman committed
4241
        self.epsilonrule = epsilonrule
4242

Peter Eastman's avatar
Peter Eastman committed
4243
4244
4245
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
4246
4247
4248
4249
4250
4251
4252

    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

        # <AmoebaVdwForce type="BUFFERED-14-7" radiusrule="CUBIC-MEAN" radiustype="R-MIN" radiussize="DIAMETER" epsilonrule="HHG" vdw-13-scale="0.0" vdw-14-scale="1.0" vdw-15-scale="1.0" >
Justin MacCallum's avatar
Justin MacCallum committed
4253
4254
4255
        #   <Vdw class="1" sigma="0.371" epsilon="0.46024" reduction="1.0" />
        #   <Vdw class="2" sigma="0.382" epsilon="0.422584" reduction="1.0" />

4256
4257
4258
        existing = [f for f in forceField._forces if isinstance(f, AmoebaVdwGenerator)]
        if len(existing) == 0:
            generator = AmoebaVdwGenerator(element.attrib['type'], element.attrib['radiusrule'], element.attrib['radiustype'], element.attrib['radiussize'], element.attrib['epsilonrule'],
Justin MacCallum's avatar
Justin MacCallum committed
4259
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
4260
4261
4262
4263
4264
4265
4266
4267
4268
4269
4270
4271
            forceField.registerGenerator(generator)
            generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaVdwForce', 'Vdw', ('sigma', 'epsilon', 'reduction'))
        else:
            # Multiple <AmoebaVdwForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
            if abs(generator.vdw13Scale - float(element.attrib['vdw-13-scale'])) > NonbondedGenerator.SCALETOL or \
                    abs(generator.vdw14Scale - float(element.attrib['vdw-14-scale'])) > NonbondedGenerator.SCALETOL or \
                    abs(generator.vdw15Scale - float(element.attrib['vdw-15-scale'])) > NonbondedGenerator.SCALETOL:
                raise ValueError('Found multiple AmoebaVdwForce tags with different scale factors')
            if generator.radiusrule != element.attrib['radiusrule'] or generator.epsilonrule != element.attrib['epsilonrule'] or \
                    generator.radiustype != element.attrib['radiustype'] or generator.radiussize != element.attrib['radiussize']:
                raise ValueError('Found multiple AmoebaVdwForce tags with different combining rules')
4272
        generator.params.parseDefinitions(element)
4273
4274
4275
4276
4277
        two_six = 1.122462048309372

    #=============================================================================================

    @staticmethod
4278
    def getBondedParticleSets(sys, data):
4279

4280
4281
4282
4283
4284
4285
        bondedParticleSets = [set() for i in range(len(data.atoms))]
        bondIndices = _findBondsForExclusions(data, sys)
        for atom1, atom2 in bondIndices:
            bondedParticleSets[atom1].add(atom2)
            bondedParticleSets[atom2].add(atom1)
        return bondedParticleSets
Justin MacCallum's avatar
Justin MacCallum committed
4286

4287
4288
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4289
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4290

Peter Eastman's avatar
Peter Eastman committed
4291
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
4292
4293
        epsilonMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'HARMONIC':1, 'HHG':1}

4294
4295
        force = mm.AmoebaVdwForce()
        sys.addForce(force)
4296

4297
        # sigma and epsilon combining rules
4298

4299
4300
4301
4302
        if ('sigmaCombiningRule' in args):
            sigmaRule = args['sigmaCombiningRule'].upper()
            if (sigmaRule.upper() in sigmaMap):
                force.setSigmaCombiningRule(sigmaRule.upper())
4303
            else:
4304
4305
4306
4307
                stringList = ' ' . join(str(x) for x in sigmaMap.keys())
                raise ValueError( "AmoebaVdwGenerator: sigma combining rule %s not recognized; valid values are %s; using default." % (sigmaRule, stringList) )
        else:
            force.setSigmaCombiningRule(self.radiusrule)
4308

4309
4310
4311
4312
        if ('epsilonCombiningRule' in args):
            epsilonRule = args['epsilonCombiningRule'].upper()
            if (epsilonRule.upper() in epsilonMap):
                force.setEpsilonCombiningRule(epsilonRule.upper())
4313
            else:
4314
4315
4316
4317
                stringList = ' ' . join(str(x) for x in epsilonMap.keys())
                raise ValueError( "AmoebaVdwGenerator: epsilon combining rule %s not recognized; valid values are %s; using default." % (epsilonRule, stringList) )
        else:
            force.setEpsilonCombiningRule(self.epsilonrule)
Justin MacCallum's avatar
Justin MacCallum committed
4318

4319
        # cutoff
4320

4321
4322
4323
4324
        if ('vdwCutoff' in args):
            force.setCutoff(args['vdwCutoff'])
        else:
            force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4325

4326
        # dispersion correction
4327

4328
4329
        if ('useDispersionCorrection' in args):
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
4330

4331
4332
        if (nonbondedMethod == PME):
            force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
4333
4334
4335

        # add particles to force

4336
4337
4338
4339
4340
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
4341
        for (i, atom) in enumerate(data.atoms):
4342
4343
4344
4345
4346
4347
4348
4349
4350
4351
            values = self.params.getAtomParameters(atom, data)
            # ivIndex = index of bonded partner for hydrogens; otherwise ivIndex = particle index

            ivIndex = i
            if atom.element == elem.hydrogen and len(data.atomBonds[i]) == 1:
                bondIndex = data.atomBonds[i][0]
                if (data.bonds[bondIndex].atom1 == i):
                    ivIndex = data.bonds[bondIndex].atom2
                else:
                    ivIndex = data.bonds[bondIndex].atom1
4352

4353
            force.addParticle(ivIndex, values[0]*sigmaScale, values[1], values[2])
4354
4355
4356
4357
4358
4359
4360

        # set combining rules

        # set particle exclusions: self, 1-2 and 1-3 bonds
        # (1) collect in bondedParticleSets[i], 1-2 indices for all bonded partners of particle i
        # (2) add 1-2,1-3 and self to exclusion set

4361
        bondedParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4362
4363

        for (i,atom) in enumerate(data.atoms):
Justin MacCallum's avatar
Justin MacCallum committed
4364

4365
4366
4367
4368
4369
4370
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
4371
            if (self.vdw13Scale == 0.0):
4372
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
4373
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
4374
4375
4376
4377
4378

            # self

            exclusionSet.add(i)

4379
            force.setParticleExclusions(i, tuple(exclusionSet))
4380
4381
4382
4383
4384

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

#=============================================================================================

4385
## @private
4386
class AmoebaMultipoleGenerator(object):
4387
4388
4389
4390

    #=============================================================================================

    """A AmoebaMultipoleGenerator constructs a AmoebaMultipoleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
4391

4392
4393
    #=============================================================================================

4394
    def __init__(self, forceField):
Peter Eastman's avatar
Peter Eastman committed
4395
4396
        self.forceField = forceField
        self.typeMap = {}
4397
4398
4399
4400
4401
4402

    #=============================================================================================
    # Set axis type
    #=============================================================================================

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
4403
    def setAxisType(kIndices):
4404
4405
4406

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
4407
4408
4409
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
4410
                else:
Peter Eastman's avatar
Peter Eastman committed
4411
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
4412

Peter Eastman's avatar
Peter Eastman committed
4413
4414
                if (kIndicesLen > 2):
                    kx = kIndices[2]
4415
                else:
Peter Eastman's avatar
Peter Eastman committed
4416
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
4417

Peter Eastman's avatar
Peter Eastman committed
4418
4419
                if (kIndicesLen > 1):
                    kz = kIndices[1]
4420
                else:
Peter Eastman's avatar
Peter Eastman committed
4421
                    kz = 0
4422

Peter Eastman's avatar
Peter Eastman committed
4423
4424
                while(len(kIndices) < 4):
                    kIndices.append(0)
4425
4426

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
4427
                if (kz == 0):
4428
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
4429
                if (kz != 0 and kx == 0):
4430
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
4431
                if (kz < 0 or kx < 0):
4432
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
4433
                if (kx < 0 and ky < 0):
4434
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
4435
                if (kz < 0 and kx < 0 and ky  < 0):
4436
4437
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
4438
4439
4440
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
4441
4442
4443
4444
4445
4446
4447
4448

                return axisType

    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
4449
        #   <AmoebaMultipoleForce  direct11Scale="0.0"  direct12Scale="1.0"  direct13Scale="1.0"  direct14Scale="1.0"  mpole12Scale="0.0"  mpole13Scale="0.0"  mpole14Scale="0.4"  mpole15Scale="0.8"  mutual11Scale="1.0"  mutual12Scale="1.0"  mutual13Scale="1.0"  mutual14Scale="1.0"  polar12Scale="0.0"  polar13Scale="0.0"  polar14Intra="0.5"  polar14Scale="1.0"  polar15Scale="1.0"  >
4450
4451
4452
        # <Multipole class="1"    kz="2"    kx="4"    c0="-0.22620" d1="0.08214" d2="0.00000" d3="0.34883" q11="0.11775" q21="0.00000" q22="-1.02185" q31="-0.17555" q32="0.00000" q33="0.90410"  />
        # <Multipole class="2"    kz="1"    kx="3"    c0="-0.15245" d1="0.19517" d2="0.00000" d3="0.19687" q11="-0.20677" q21="0.00000" q22="-0.48084" q31="-0.01672" q32="0.00000" q33="0.68761"  />

4453
4454
4455
4456
4457
4458
4459
        existing = [f for f in forceField._forces if isinstance(f, AmoebaMultipoleGenerator)]
        if len(existing) == 0:
            generator = AmoebaMultipoleGenerator(forceField)
            forceField.registerGenerator(generator)
        else:
            # Multiple <AmoebaMultipoleForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
4460
4461
4462
4463

        # set type map: [ kIndices, multipoles, AMOEBA/OpenMM axis type]

        for atom in element.findall('Multipole'):
4464
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4465
            if None not in types:
4466
4467
4468
4469
4470
4471
4472
4473

                # k-indices not provided default to 0

                kIndices = [int(atom.attrib['type'])]

                kStrings = [ 'kz', 'kx', 'ky' ]
                for kString in kStrings:
                    try:
Peter Eastman's avatar
Peter Eastman committed
4474
4475
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
4476
                    except:
4477
4478
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
4479
                # set axis type based on k-Indices
4480

Peter Eastman's avatar
Peter Eastman committed
4481
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
4482
4483
4484

                # set multipole

Peter Eastman's avatar
Peter Eastman committed
4485
                charge = float(atom.attrib['c0'])
Justin MacCallum's avatar
Justin MacCallum committed
4486

Peter Eastman's avatar
Peter Eastman committed
4487
4488
4489
4490
4491
4492
4493
4494
4495
4496
4497
4498
4499
                conversion = 1.0
                dipole = [ conversion*float(atom.attrib['d1']), conversion*float(atom.attrib['d2']), conversion*float(atom.attrib['d3'])]

                quadrupole = []
                quadrupole.append(conversion*float(atom.attrib['q11']))
                quadrupole.append(conversion*float(atom.attrib['q21']))
                quadrupole.append(conversion*float(atom.attrib['q31']))
                quadrupole.append(conversion*float(atom.attrib['q21']))
                quadrupole.append(conversion*float(atom.attrib['q22']))
                quadrupole.append(conversion*float(atom.attrib['q32']))
                quadrupole.append(conversion*float(atom.attrib['q31']))
                quadrupole.append(conversion*float(atom.attrib['q32']))
                quadrupole.append(conversion*float(atom.attrib['q33']))
4500
4501

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4502
                    if (t not in generator.typeMap):
4503
4504
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
4505
4506
4507
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
4508
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
4509
4510
4511
4512
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
4513

4514
            else:
Peter Eastman's avatar
Peter Eastman committed
4515
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4516
4517
                raise ValueError(outputString)

4518
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
4519

4520
        for atom in element.findall('Polarize'):
4521
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4522
            if None not in types:
4523

Peter Eastman's avatar
Peter Eastman committed
4524
4525
4526
4527
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
4528
4529
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
4530
                    pdamp = pow(polarizability, 1.0/6.0)
4531

Peter Eastman's avatar
Peter Eastman committed
4532
4533
                pgrpMap = dict()
                for index in range(1, 7):
4534
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
4535
                    if (pgrp in atom.attrib):
4536
4537
4538
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4539
4540
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
4541
                        raise ValueError(outputString)
4542
4543
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
4544
4545
4546
4547
4548
4549
4550
                        hit = 0
                        for (ii, typeMap) in enumerate(typeMapList):

                            if (typeMap['classIndex'] == classIndex):
                                typeMap['polarizability'] = polarizability
                                typeMap['thole'] = thole
                                typeMap['pdamp'] = pdamp
Justin MacCallum's avatar
Justin MacCallum committed
4551
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
4552
4553
4554
4555
4556
                                typeMapList[ii] = typeMap
                                hit = 1

                        if (hit == 0):
                            outputString = "AmoebaMultipoleGenerator: error getting type for polarize: class index=%s not in multipole list?" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4557
4558
                            raise ValueError(outputString)

4559
            else:
Peter Eastman's avatar
Peter Eastman committed
4560
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4561
4562
                raise ValueError(outputString)

4563
4564
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4565
    def setPolarGroups(self, data, bonded12ParticleSets, force):
4566
4567
4568
4569
4570

        for (atomIndex, atom) in enumerate(data.atoms):

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
4571
4572
4573
4574
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
4575
4576
4577
4578
            atom.polarizationGroupSet = list()
            atom.polarizationGroups[atomIndex] = 1
            for bondedAtomIndex in bondedAtomIndices:
                bondedAtomType = int(data.atomType[data.atoms[bondedAtomIndex]])
Peter Eastman's avatar
Peter Eastman committed
4579
4580
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
4581
4582
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
4583

4584
4585
4586
4587
        # pgrp11

        for (atomIndex, atom) in enumerate(data.atoms):

Peter Eastman's avatar
Peter Eastman committed
4588
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
4589
4590
                continue

Peter Eastman's avatar
Peter Eastman committed
4591
4592
            group = set()
            visited = set()
4593
4594
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4595
4596
4597
4598
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
4599
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
4600
4601
                if (nextAtom not in visited):
                   visited.add(nextAtom)
4602
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4603
4604
4605
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
4606
4607
4608

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
4609
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
4610
4611

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4612
4613
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
4614
4615
4616
4617
4618

        # pgrp12

        for (atomIndex, atom) in enumerate(data.atoms):

Peter Eastman's avatar
Peter Eastman committed
4619
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
4620
4621
                continue

Peter Eastman's avatar
Peter Eastman committed
4622
            pgrp11 = set(atom.polarizationGroupSet[0])
4623
4624
4625
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4626
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
4627
4628
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4629
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
4630

4631
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4632
4633
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
4634
4635
4636
4637
4638

        # pgrp13

        for (atomIndex, atom) in enumerate(data.atoms):

Peter Eastman's avatar
Peter Eastman committed
4639
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
4640
4641
                continue

Peter Eastman's avatar
Peter Eastman committed
4642
4643
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
4644
4645
4646
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4647
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
4648
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4649
            pgrp13 = pgrp13 - set(pgrp11)
4650
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4651
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
4652

4653
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4654
4655
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
4656
4657
4658
4659
4660

        # pgrp14

        for (atomIndex, atom) in enumerate(data.atoms):

Peter Eastman's avatar
Peter Eastman committed
4661
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
4662
4663
                continue

Peter Eastman's avatar
Peter Eastman committed
4664
4665
4666
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
4667
4668
4669
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4670
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
4671
4672
4673

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4674
            pgrp14 = pgrp14 - set(pgrp11)
4675
4676

            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4677
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp14)
Justin MacCallum's avatar
Justin MacCallum committed
4678

4679
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4680
4681
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
4682
4683
4684

    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4685
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4686
4687
4688
4689

        methodMap = {NoCutoff:mm.AmoebaMultipoleForce.NoCutoff,
                     PME:mm.AmoebaMultipoleForce.PME}

4690
4691
4692
4693
4694
4695
4696
        force = mm.AmoebaMultipoleForce()
        sys.addForce(force)
        if (nonbondedMethod not in methodMap):
            raise ValueError( "AmoebaMultipoleForce: input cutoff method not available." )
        else:
            force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
4697

4698
4699
        if ('ewaldErrorTolerance' in args):
            force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))
4700

4701
4702
4703
4704
4705
4706
4707
4708
        if ('polarization' in args):
            polarizationType = args['polarization']
            if (polarizationType.lower() == 'direct'):
                force.setPolarizationType(mm.AmoebaMultipoleForce.Direct)
            elif (polarizationType.lower() == 'extrapolated'):
                force.setPolarizationType(mm.AmoebaMultipoleForce.Extrapolated)
            else:
                force.setPolarizationType(mm.AmoebaMultipoleForce.Mutual)
4709

4710
4711
        if ('aEwald' in args):
            force.setAEwald(float(args['aEwald']))
4712

4713
4714
        if ('pmeGridDimensions' in args):
            force.setPmeGridDimensions(args['pmeGridDimensions'])
4715

4716
4717
        if ('mutualInducedMaxIterations' in args):
            force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
4718

4719
4720
        if ('mutualInducedTargetEpsilon' in args):
            force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))
4721
4722

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
4723
        # throw error if particle type not available
4724
4725
4726
4727
4728

        # get 1-2, 1-3, 1-4, 1-5 bonded sets

        # 1-2

4729
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4730
4731
4732
4733
4734

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4735
            bonded13Set = set()
4736
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4737
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4738
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
4739
4740
4741
4742

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4743
            selfSet = set()
4744
4745
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4746
4747
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
4748
4749
4750
4751
4752

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4753
4754
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4755
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4756
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
4757

4758
4759
4760
4761
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4762
            selfSet = set()
4763
4764
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4765
4766
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
4767
4768
4769
4770
4771

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4772
4773
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4774
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4775
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
4776
4777
4778
4779
4780
4781

            # remove 1-4, 1-3, 1-2 and self from set

            bonded15Set = bonded15Set - bonded12ParticleSets[i]
            bonded15Set = bonded15Set - bonded13ParticleSets[i]
            bonded15Set = bonded15Set - bonded14ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4782
            selfSet = set()
4783
4784
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4785
4786
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
4787
4788
4789
4790
4791

        for (atomIndex, atom) in enumerate(data.atoms):
            t = data.atomType[atom]
            if t in self.typeMap:

Peter Eastman's avatar
Peter Eastman committed
4792
4793
                multipoleList = self.typeMap[t]
                hit = 0
4794
4795
4796
4797
4798
4799
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4800
                    if (hit != 0):
4801
4802
                        break

Peter Eastman's avatar
Peter Eastman committed
4803
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4804
4805

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4806
4807
                    kx = kIndices[2]
                    ky = kIndices[3]
4808
4809
4810
4811

                    # assign multipole parameters
                    #    (1) get bonded partners
                    #    (2) match parameter types
Justin MacCallum's avatar
Justin MacCallum committed
4812

4813
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
4814
4815
4816
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4817
4818
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
4819
                       if (hit != 0):
4820
4821
4822
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
4823
4824
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
4825
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4826
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4827
4828
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4829
4830
4831
4832
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4833
4834
4835
4836
4837
4838
4839
4840
4841
4842
                                      if( bondedAtomXType == bondedAtomZType and xaxis < zaxis ):
                                          swapI = zaxis
                                          zaxis = xaxis
                                          xaxis = swapI
                                      else:
                                          for bondedAtomXIndex in bondedAtomIndices:
                                              bondedAtomX1Type = int(data.atomType[data.atoms[bondedAtomXIndex]])
                                              if( bondedAtomX1Type == kx and bondedAtomXIndex != bondedAtomZIndex and bondedAtomXIndex < xaxis ):
                                                  xaxis = bondedAtomXIndex

4843
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4844
                                      hit = 1
4845
4846
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4847
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4848
4849
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4850
4851
4852
4853
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4854
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4855
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
4856

4857
4858
4859
4860
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4861
                    if (hit != 0):
4862
4863
                        break

Peter Eastman's avatar
Peter Eastman committed
4864
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4865
4866

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4867
4868
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
4869

4870
4871
4872
                    # assign multipole parameters
                    #    (1) get bonded partners
                    #    (2) match parameter types
Justin MacCallum's avatar
Justin MacCallum committed
4873

4874
4875
4876
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
4877
4878
4879
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4880
4881
4882

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4883
                       if (hit != 0):
4884
4885
4886
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
4887
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
4888

Peter Eastman's avatar
Peter Eastman committed
4889
                       if (bondedAtomZType == kz):
4890
4891
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
4892
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4893
4894
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4895
4896
4897
4898
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4899
4900
4901
4902
4903
4904
4905
4906

                                      # select xaxis w/ smallest index

                                      for bondedAtomXIndex in bondedAtom13Indices:
                                          bondedAtomX1Type = int(data.atomType[data.atoms[bondedAtomXIndex]])
                                          if( bondedAtomX1Type == kx and bondedAtomXIndex != bondedAtomZIndex and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex] and bondedAtomXIndex < xaxis ):
                                              xaxis = bondedAtomXIndex

4907
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4908
                                      hit = 3
4909
4910
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
4911
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4912
4913
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4914
4915
4916
4917
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4918
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4919
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
4920

4921
4922
4923
4924
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4925
                    if (hit != 0):
4926
4927
                        break

Peter Eastman's avatar
Peter Eastman committed
4928
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4929
4930
4931
4932

                    kz = kIndices[1]
                    kx = kIndices[2]

Peter Eastman's avatar
Peter Eastman committed
4933
4934
4935
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4936
4937
4938

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4939
                        if (hit != 0):
4940
4941
4942
                            break

                        bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
4943
                        bondedAtomZ = data.atoms[bondedAtomZIndex]
4944

Peter Eastman's avatar
Peter Eastman committed
4945
                        if (kx == 0 and kz == bondedAtomZType):
4946
                            zaxis = bondedAtomZIndex
4947
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4948
                            hit = 5
4949
4950
4951
4952
4953

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4954
                    if (hit != 0):
4955
4956
                        break

Peter Eastman's avatar
Peter Eastman committed
4957
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4958
4959
4960

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
4961
4962
4963
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4964

Peter Eastman's avatar
Peter Eastman committed
4965
                    if (kz == 0):
4966
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4967
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
4968

4969
4970
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
4971
                if (hit != 0):
4972

Peter Eastman's avatar
Peter Eastman committed
4973
                    atom.multipoleDict = savedMultipoleDict
4974
                    atom.polarizationGroups = dict()
4975
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
4976
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
4977
                    if (atomIndex == newIndex):
4978
4979
4980
4981
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent12, tuple(bonded12ParticleSets[atomIndex]))
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent13, tuple(bonded13ParticleSets[atomIndex]))
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent14, tuple(bonded14ParticleSets[atomIndex]))
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent15, tuple(bonded15ParticleSets[atomIndex]))
4982
                    else:
4983
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
4984
                else:
Peter Eastman's avatar
Peter Eastman committed
4985
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
4986
            else:
Peter Eastman's avatar
Peter Eastman committed
4987
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
4988
4989
4990

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
4991
        self.setPolarGroups(data, bonded12ParticleSets, force)
4992
4993
4994
4995
4996

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

#=============================================================================================

4997
## @private
4998
class AmoebaWcaDispersionGenerator(object):
4999
5000

    """A AmoebaWcaDispersionGenerator constructs a AmoebaWcaDispersionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
5001

5002
5003
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5004
    def __init__(self, epso, epsh, rmino, rminh, awater, slevy, dispoff, shctd):
5005

Justin MacCallum's avatar
Justin MacCallum committed
5006
5007
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
5008
5009
5010
5011
5012
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
5013
        self.shctd = shctd
5014
5015
5016
5017
5018
5019
5020
5021
5022

    #=========================================================================================

    @staticmethod
    def parseElement(element, forceField):

        #  <AmoebaWcaDispersionForce epso="0.46024" epsh="0.056484" rmino="0.17025" rminh="0.13275" awater="33.428" slevy="1.0"  dispoff="0.026" shctd="0.81" >
        #   <WcaDispersion class="1" radius="0.1855" epsilon="0.46024" />
        #   <WcaDispersion class="2" radius="0.191" epsilon="0.422584" />
Justin MacCallum's avatar
Justin MacCallum committed
5023

Peter Eastman's avatar
Peter Eastman committed
5024
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
5025
5026
5027
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
Justin MacCallum's avatar
Justin MacCallum committed
5028
                                                  element.attrib['awater'],
5029
5030
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Justin MacCallum's avatar
Justin MacCallum committed
5031
                                                  element.attrib['shctd'])
5032
        forceField._forces.append(generator)
5033
5034
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaWcaDispersionForce', 'WcaDispersion', ('radius', 'epsilon'))
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
5035

5036
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
5037

Peter Eastman's avatar
Peter Eastman committed
5038
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
5039
5040
5041
5042
5043
5044
5045
5046
5047
5048
5049
5050

        # get or create force depending on whether it has already been added to the system

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.AmoebaWcaDispersionForce]
        if len(existing) == 0:
            force = mm.AmoebaWcaDispersionForce()
            sys.addForce(force)
        else:
            force = existing[0]

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
5051
        # throw error if particle type not available
5052

Peter Eastman's avatar
Peter Eastman committed
5053
5054
5055
5056
5057
5058
5059
5060
        force.setEpso(   float(self.epso   ))
        force.setEpsh(   float(self.epsh   ))
        force.setRmino(  float(self.rmino  ))
        force.setRminh(  float(self.rminh  ))
        force.setDispoff(float(self.dispoff))
        force.setSlevy(  float(self.slevy  ))
        force.setAwater( float(self.awater ))
        force.setShctd(  float(self.shctd  ))
5061

5062
5063
5064
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
5065
5066
5067
5068
5069

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

#=============================================================================================

5070
## @private
5071
class AmoebaGeneralizedKirkwoodGenerator(object):
5072
5073

    """A AmoebaGeneralizedKirkwoodGenerator constructs a AmoebaGeneralizedKirkwoodForce."""
Justin MacCallum's avatar
Justin MacCallum committed
5074

5075
5076
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5077
5078
5079
5080
5081
5082
5083
5084
5085
5086
5087
    def __init__(self, forceField, solventDielectric, soluteDielectric, includeCavityTerm, probeRadius, surfaceAreaFactor):

        self.forceField = forceField
        self.solventDielectric = solventDielectric
        self.soluteDielectric = soluteDielectric
        self.includeCavityTerm = includeCavityTerm
        self.probeRadius = probeRadius
        self.surfaceAreaFactor = surfaceAreaFactor

        self.radiusTypeMap = {}
        self.radiusTypeMap['Bondi'] = {}
Justin MacCallum's avatar
Justin MacCallum committed
5088
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5089
5090
5091
5092
5093
5094
5095
5096
5097
5098
5099
5100
5101
5102
5103
5104
5105
5106
5107
5108
5109
5110
5111
5112
5113
        rscale = 1.03

        bondiMap[0] = 0.00
        bondiMap[1] = 0.12*rscale
        bondiMap[2] = 0.14*rscale
        bondiMap[5] = 0.18*rscale

        bondiMap[6] = 0.170*rscale
        bondiMap[7] = 0.155*rscale
        bondiMap[8] = 0.152*rscale
        bondiMap[9] = 0.147*rscale

        bondiMap[10] = 0.154*rscale
        bondiMap[14] = 0.210*rscale
        bondiMap[15] = 0.180*rscale
        bondiMap[16] = 0.180*rscale

        bondiMap[17] = 0.175 *rscale
        bondiMap[18] = 0.188*rscale
        bondiMap[34] = 0.190*rscale
        bondiMap[35] = 0.185*rscale

        bondiMap[36] = 0.202*rscale
        bondiMap[53] = 0.198*rscale
        bondiMap[54] = 0.216*rscale
5114
5115
5116

    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5117
    def getObcShct(self, data, atomIndex):
5118

Peter Eastman's avatar
Peter Eastman committed
5119
        atom = data.atoms[atomIndex]
5120
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5121
        shct = -1.0
5122
5123

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
5124

Peter Eastman's avatar
Peter Eastman committed
5125
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
5126
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5127
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5128
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
5129
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
5130
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
5131
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
5132
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5133
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
5134
5135
5136
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
5137
        elif (atomicNumber == 16):              # S(16)
5138
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
5139
        elif (atomicNumber == 26):              # Fe(26)
5140
5141
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
5142
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5143
            raise ValueError( "getObcShct: no GK overlap scale factor for atom %s of %s %d" % (atom.name, atom.residue.name, atom.residue.index) )
Justin MacCallum's avatar
Justin MacCallum committed
5144
5145

        return shct
5146
5147
5148

    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5149
    def getAmoebaTypeRadius(self, data, bondedAtomIndices, atomIndex):
5150

Peter Eastman's avatar
Peter Eastman committed
5151
        atom = data.atoms[atomIndex]
5152
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5153
        radius = -1.0
5154

Peter Eastman's avatar
Peter Eastman committed
5155
        if (atomicNumber == 1):                  # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
5156

Peter Eastman's avatar
Peter Eastman committed
5157
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
5158

Peter Eastman's avatar
Peter Eastman committed
5159
            if (len(bondedAtomIndices) < 1):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5160
                 raise ValueError( "AmoebaGeneralizedKirkwoodGenerator: error getting atom bonded to %s of %s %d " % (atom.name, atom.residue.name, atom.residue.index) )
Justin MacCallum's avatar
Justin MacCallum committed
5161

5162
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5163
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
5164

Peter Eastman's avatar
Peter Eastman committed
5165
            if (bondedAtomAtomicNumber == 7):
5166
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
5167
            if (bondedAtomAtomicNumber == 8):
5168
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
5169

Peter Eastman's avatar
Peter Eastman committed
5170
        elif (atomicNumber == 3):               # Li(3)
5171
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
5172
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5173

5174
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
5175
            if (len(bondedAtomIndices) == 3):
5176
5177
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
5178
            elif (len(bondedAtomIndices) == 4):
5179
5180
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5181
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
5182
5183
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
5184
        elif (atomicNumber == 7):               # N(7)
5185
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
5186
        elif (atomicNumber == 8):               # O(8)
5187
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
5188
            if (len(bondedAtomIndices) == 2):
5189
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
5190
        elif (atomicNumber == 9):               # F(9)
5191
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
5192
        elif (atomicNumber == 10):
5193
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
5194
        elif (atomicNumber == 11):
5195
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
5196
        elif (atomicNumber == 12):
5197
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5198
        elif (atomicNumber == 14):
5199
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
5200
        elif (atomicNumber == 15):              # P(15)
5201
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
5202
        elif (atomicNumber == 16):              # S(16)
5203
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
5204
        elif (atomicNumber == 17):
5205
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
5206
        elif (atomicNumber == 18):
5207
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5208
        elif (atomicNumber == 19):
5209
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
5210
        elif (atomicNumber == 20):
5211
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
5212
        elif (atomicNumber == 35):
5213
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
5214
        elif (atomicNumber == 36):
5215
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
5216
        elif (atomicNumber == 37):
5217
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
5218
        elif (atomicNumber == 53):
5219
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
5220
        elif (atomicNumber == 54):
5221
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
5222
        elif (atomicNumber == 55):
5223
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
5224
        elif (atomicNumber == 56):
5225
5226
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
5227
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5228
5229
            outputString = "No GK radius for atom %s of %s %d" % (atom.name, atom.residue.name, atom.residue.index)
            raise ValueError( outputString )
Justin MacCallum's avatar
Justin MacCallum committed
5230

5231
5232
5233
5234
        return radius

    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5235
    def getBondiTypeRadius(self, data, bondedAtomIndices, atomIndex):
5236

Justin MacCallum's avatar
Justin MacCallum committed
5237
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5238
        atom = data.atoms[atomIndex]
5239
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
5240
        if (atomicNumber in bondiMap):
5241
5242
            radius = bondiMap[atomicNumber]
        else:
peastman's avatar
peastman committed
5243
            outputString = "Warning no Bondi radius for atom %s of %s %d using default value" % (atom.name, atom.residue.name, atom.residue.index)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5244
            raise ValueError( outputString )
Justin MacCallum's avatar
Justin MacCallum committed
5245

5246
5247
5248
5249
5250
5251
5252
5253
5254
5255
        return radius

    #=========================================================================================

    @staticmethod
    def parseElement(element, forceField):

        #  <AmoebaGeneralizedKirkwoodForce solventDielectric="78.3" soluteDielectric="1.0" includeCavityTerm="1" probeRadius="0.14" surfaceAreaFactor="-170.351730663">
        #   <GeneralizedKirkwood type="1" charge="-0.22620" shct="0.79"  />
        #   <GeneralizedKirkwood type="2" charge="-0.15245" shct="0.72"  />
Justin MacCallum's avatar
Justin MacCallum committed
5256

Peter Eastman's avatar
Peter Eastman committed
5257
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
Justin MacCallum's avatar
Justin MacCallum committed
5258
5259
                                                        element.attrib['includeCavityTerm'],
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor'])
5260
5261
5262
        forceField._forces.append(generator)

    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
5263

Peter Eastman's avatar
Peter Eastman committed
5264
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
5265

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5266
5267
5268
        if( nonbondedMethod != NoCutoff ):
            raise ValueError( "Only the nonbondedMethod=NoCutoff option is available for implicit solvent simulations." )

5269
5270
5271
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
5272
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5273
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
5274
        if (len(amoebaMultipoleForceList) > 0):
5275
5276
5277
5278
5279
            amoebaMultipoleForce = amoebaMultipoleForceList[0]
        else:
            # call AmoebaMultipoleForceGenerator.createForce() to ensure charges have been set

            for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
5280
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
5281
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
5282

5283
5284
5285
5286
5287
5288
5289
        # get or create force depending on whether it has already been added to the system

        existing = [f for f in existing if type(f) == mm.AmoebaGeneralizedKirkwoodForce]
        if len(existing) == 0:

            force = mm.AmoebaGeneralizedKirkwoodForce()
            sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
5290

Peter Eastman's avatar
Peter Eastman committed
5291
5292
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
5293
            else:
Peter Eastman's avatar
Peter Eastman committed
5294
                force.setSolventDielectric(   float(self.solventDielectric))
5295

Peter Eastman's avatar
Peter Eastman committed
5296
5297
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
5298
            else:
Peter Eastman's avatar
Peter Eastman committed
5299
                force.setSoluteDielectric(    float(self.soluteDielectric))
5300

Peter Eastman's avatar
Peter Eastman committed
5301
5302
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
5303
            else:
Peter Eastman's avatar
Peter Eastman committed
5304
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
5305
5306
5307
5308
5309

        else:
            force = existing[0]

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
5310
        # throw error if particle type not available
5311

Peter Eastman's avatar
Peter Eastman committed
5312
5313
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
5314
5315
5316

        # 1-2

5317
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5318
5319

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
5320
5321
5322
5323
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5324
            else:
Peter Eastman's avatar
Peter Eastman committed
5325
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5326
5327
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
5328
            force.addParticle(multipoleParameters[0], radius, shct)
5329
5330
5331
5332
5333

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

#=============================================================================================

5334
## @private
5335
class AmoebaUreyBradleyGenerator(object):
5336
5337
5338
5339

    #=============================================================================================
    """An AmoebaUreyBradleyGenerator constructs a AmoebaUreyBradleyForce."""
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
5340

5341
    def __init__(self):
5342

peastman's avatar
peastman committed
5343
        self.anglesForAtom2Type = defaultdict(list)
Peter Eastman's avatar
Peter Eastman committed
5344
5345
5346
        self.types1 = []
        self.types2 = []
        self.types3 = []
5347

Peter Eastman's avatar
Peter Eastman committed
5348
5349
        self.length = []
        self.k = []
5350
5351
5352
5353
5354
5355

    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

5356
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
5357
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
5358

5359
        generator = AmoebaUreyBradleyGenerator()
5360
5361
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
5362
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
5363
            if None not in types:
peastman's avatar
peastman committed
5364
                index = len(generator.types1)
5365
5366
5367
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
peastman's avatar
peastman committed
5368
5369
                for t in types[1]:
                    generator.anglesForAtom2Type[t].append(index)
5370
5371
5372
5373
5374
                generator.length.append(float(bond.attrib['d']))
                generator.k.append(float(bond.attrib['k']))

            else:
                outputString = "AmoebaUreyBradleyGenerator: error getting types: %s %s %s" % (
Peter Eastman's avatar
Peter Eastman committed
5375
                                    bond.attrib['class1'], bond.attrib['class2'], bond.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
5376
5377
                raise ValueError(outputString)

5378
5379
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
5380
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
5381

Peter Eastman's avatar
Peter Eastman committed
5382
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5383
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
5384
5385

        if len(existing) == 0:
5386
            force = mm.HarmonicBondForce()
5387
5388
5389
5390
5391
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
5392
            if (isConstrained and not args.get('flexibleConstraints', False)):
5393
5394
5395
5396
                continue
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
peastman's avatar
peastman committed
5397
            for i in self.anglesForAtom2Type[type2]:
5398
5399
5400
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
Peter Eastman's avatar
Peter Eastman committed
5401
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
5402
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
5403
5404
5405
5406
5407
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
5408
5409
5410


## @private
5411
class DrudeGenerator(object):
peastman's avatar
peastman committed
5412
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
5413

5414
5415
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
5416
5417
5418
5419
5420
5421
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
5422
5423
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
5424
5425
5426
5427
        else:
            # Multiple <DrudeForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
        for particle in element.findall('Particle'):
5428
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
5429
5430
5431
5432
5433
5434
5435
5436
5437
5438
            if None not in types[:2]:
                aniso12 = 0.0
                aniso34 = 0.0
                if 'aniso12' in particle.attrib:
                    aniso12 = float(particle.attrib['aniso12'])
                if 'aniso34' in particle.attrib:
                    aniso34 = float(particle.attrib['aniso34'])
                values = (types[1], types[2], types[3], types[4], float(particle.attrib['charge']), float(particle.attrib['polarizability']), aniso12, aniso34, float(particle.attrib['thole']))
                for t in types[0]:
                    generator.typeMap[t] = values
Justin MacCallum's avatar
Justin MacCallum committed
5439

peastman's avatar
peastman committed
5440
5441
5442
5443
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        force = mm.DrudeForce()
        if not any(isinstance(f, mm.NonbondedForce) for f in sys.getForces()):
            raise ValueError('<DrudeForce> must come after <NonbondedForce> in XML file')
Justin MacCallum's avatar
Justin MacCallum committed
5444

peastman's avatar
peastman committed
5445
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
5446

peastman's avatar
peastman committed
5447
5448
5449
5450
5451
5452
5453
5454
5455
5456
5457
5458
5459
5460
5461
5462
        for atom in data.atoms:
            t = data.atomType[atom]
            if t in self.typeMap:
                # Find other atoms in the residue that affect the Drude particle.
                p = [-1, -1, -1, -1]
                values = self.typeMap[t]
                for atom2 in atom.residue.atoms():
                    type2 = data.atomType[atom2]
                    if type2 in values[0]:
                        p[0] = atom2.index
                    elif values[1] is not None and type2 in values[1]:
                        p[1] = atom2.index
                    elif values[2] is not None and type2 in values[2]:
                        p[2] = atom2.index
                    elif values[3] is not None and type2 in values[3]:
                        p[3] = atom2.index
5463
5464
                force.addParticle(atom.index, p[0], p[1], p[2], p[3], values[4], values[5], values[6], values[7])
                data.excludeAtomWith[p[0]].append(atom.index)
peastman's avatar
peastman committed
5465
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
5466

peastman's avatar
peastman committed
5467
5468
    def postprocessSystem(self, sys, data, args):
        # For every nonbonded exclusion between Drude particles, add a screened pair.
Justin MacCallum's avatar
Justin MacCallum committed
5469

peastman's avatar
peastman committed
5470
5471
5472
5473
5474
5475
5476
        drude = [f for f in sys.getForces() if isinstance(f, mm.DrudeForce)][0]
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
        particleMap = {}
        for i in range(drude.getNumParticles()):
            particleMap[drude.getParticleParameters(i)[0]] = i
        for i in range(nonbonded.getNumExceptions()):
            (particle1, particle2, charge, sigma, epsilon) = nonbonded.getExceptionParameters(i)
5477
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
5478
5479
5480
5481
5482
                # This is an exclusion.
                if particle1 in particleMap and particle2 in particleMap:
                    # It connects two Drude particles, so add a screened pair.
                    drude1 = particleMap[particle1]
                    drude2 = particleMap[particle2]
5483
5484
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
5485
5486
5487
5488
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

Justin MacCallum's avatar
Justin MacCallum committed
5489
parsers["DrudeForce"] = DrudeGenerator.parseElement
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
5490
5491

#=============================================================================================