"docs-source/vscode:/vscode.git/clone" did not exist on "1a3411a5edc2e91cdc9dac625e2d4b6267426601"
forcefield.py 204 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-2015 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
41
42
from math import sqrt, cos
import simtk.openmm as mm
import simtk.unit as unit
43
from . import element as elem
44
45
from simtk.openmm.app import Topology

46
47
def _convertParameterToNumber(param):
    if unit.is_quantity(param):
48
49
50
        if param.unit.is_compatible(unit.bar):
            return param / unit.bar
        return param.value_in_unit_system(unit.md_unit_system)
51
52
    return float(param)

53
54
# Enumerated values for nonbonded method

55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
class NoCutoff(object):
    def __repr__(self):
        return 'NoCutoff'
NoCutoff = NoCutoff()

class CutoffNonPeriodic(object):
    def __repr__(self):
        return 'CutoffNonPeriodic'
CutoffNonPeriodic = CutoffNonPeriodic()

class CutoffPeriodic(object):
    def __repr__(self):
        return 'CutoffPeriodic'
CutoffPeriodic = CutoffPeriodic()

class Ewald(object):
    def __repr__(self):
        return 'Ewald'
Ewald = Ewald()

class PME(object):
    def __repr__(self):
        return 'PME'
PME = PME()
79
80
81

# Enumerated values for constraint type

82
83
84
85
86
87
88
89
90
91
92
93
94
95
class HBonds(object):
    def __repr__(self):
        return 'HBonds'
HBonds = HBonds()

class AllBonds(object):
    def __repr__(self):
        return 'AllBonds'
AllBonds = AllBonds()

class HAngles(object):
    def __repr__(self):
        return 'HAngles'
HAngles = HAngles()
96
97
98
99
100
101
102
103
104
105

# 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
106

Robert McGibbon's avatar
Robert McGibbon committed
107
108
109
110
111
112
113
114
        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.
115
116
117
        """
        self._atomTypes = {}
        self._templates = {}
118
        self._templateSignatures = {None:[]}
119
        self._atomClasses = {'':set()}
120
        self._forces = []
121
        self._scripts = []
122
        self._templateGenerators = []
123
        for file in files:
peastman's avatar
peastman committed
124
            self.loadFile(file)
125

peastman's avatar
peastman committed
126
    def loadFile(self, file):
127
        """Load an XML file and add the definitions from it to this ForceField.
128

Robert McGibbon's avatar
Robert McGibbon committed
129
130
131
132
133
134
135
136
        Parameters
        ----------
        file : string or file
            An XML file containing force field definitions.  It may be either 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.
peastman's avatar
peastman committed
137
138
139
140
141
142
        """
        try:
            # this handles either filenames or open file-like objects
            tree = etree.parse(file)
        except IOError:
            tree = etree.parse(os.path.join(os.path.dirname(__file__), 'data', file))
143
        except Exception as e:
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
144
            # Fail with an error message about which file could not be read.
145
146
147
            # 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'
148
            if hasattr(file, 'name'):
149
150
                filename = file.name
            else:
151
152
                filename = str(file)
            msg += "ForceField.loadFile() encountered an error reading file '%s'\n" % filename
153
154
            raise Exception(msg)

peastman's avatar
peastman committed
155
156
157
158
159
160
161
162
163
164
165
166
167
168
        root = tree.getroot()

        # Load the atom types.

        if tree.getroot().find('AtomTypes') is not None:
            for type in tree.getroot().find('AtomTypes').findall('Type'):
                self.registerAtomType(type.attrib)

        # Load the residue templates.

        if tree.getroot().find('Residues') is not None:
            for residue in root.find('Residues').findall('Residue'):
                resName = residue.attrib['name']
                template = ForceField._TemplateData(resName)
169
                atomIndices = {}
peastman's avatar
peastman committed
170
                for atom in residue.findall('Atom'):
171
172
173
174
                    params = {}
                    for key in atom.attrib:
                        if key not in ('name', 'type'):
                            params[key] = _convertParameterToNumber(atom.attrib[key])
175
176
177
178
                    atomName = atom.attrib['name']
                    if atomName in atomIndices:
                        raise ValueError('Residue '+resName+' contains multiple atoms named '+atomName)
                    atomIndices[atomName] = len(template.atoms)
179
180
                    typeName = atom.attrib['type']
                    template.atoms.append(ForceField._TemplateAtomData(atomName, typeName, self._atomTypes[typeName].element, params))
peastman's avatar
peastman committed
181
                for site in residue.findall('VirtualSite'):
182
                    template.virtualSites.append(ForceField._VirtualSiteData(site, atomIndices))
peastman's avatar
peastman committed
183
                for bond in residue.findall('Bond'):
184
                    if 'atomName1' in bond.attrib:
185
                        template.addBondByName(bond.attrib['atomName1'], bond.attrib['atomName2'])
186
187
                    else:
                        template.addBond(int(bond.attrib['from']), int(bond.attrib['to']))
peastman's avatar
peastman committed
188
                for bond in residue.findall('ExternalBond'):
189
                    if 'atomName' in bond.attrib:
190
                        template.addExternalBondByName(bond.attrib['atomName'])
191
                    else:
192
                        template.addExternalBond(int(bond.attrib['from']))
peastman's avatar
peastman committed
193
194
195
196
197
198
199
200
201
202
203
204
                self.registerResidueTemplate(template)

        # Load force definitions

        for child in root:
            if child.tag in parsers:
                parsers[child.tag](child, self)

        # Load scripts

        for node in tree.getroot().findall('Script'):
            self.registerScript(node.text)
205

206
207
208
    def getGenerators(self):
        """Get the list of all registered generators."""
        return self._forces
209

210
211
212
    def registerGenerator(self, generator):
        """Register a new generator."""
        self._forces.append(generator)
213

214
215
216
217
218
219
220
221
222
223
224
225
    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)
226
        self._atomTypes[name] = ForceField._AtomType(name, atomClass, mass, element)
227
228
229
230
231
232
233
        if atomClass in self._atomClasses:
            typeSet = self._atomClasses[atomClass]
        else:
            typeSet = set()
            self._atomClasses[atomClass] = typeSet
        typeSet.add(name)
        self._atomClasses[''].add(name)
234

235
236
237
238
239
240
241
242
    def registerResidueTemplate(self, template):
        """Register a new residue template."""
        self._templates[template.name] = template
        signature = _createResidueSignature([atom.element for atom in template.atoms])
        if signature in self._templateSignatures:
            self._templateSignatures[signature].append(template)
        else:
            self._templateSignatures[signature] = [template]
243

244
245
246
    def registerScript(self, script):
        """Register a new script to be executed after building the System."""
        self._scripts.append(script)
247

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
248
    def registerTemplateGenerator(self, generator):
249
250
251
252
        """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
253
254
        .. CAUTION:: This method is experimental, and its API is subject to change.

255
256
        Parameters
        ----------
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
257
        generator : function
258
259
            A function that will be called when a residue is encountered that does not match an existing forcefield template.

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
260
        When a residue without a template is encountered, the `generator` function is called with:
261

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
262
263
        ::
           success = generator(forcefield, residue)
264
265
        ```

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
266
267
268
269
270
271
272
273
274
275
        where `forcefield` is the calling `ForceField` object and `residue` is a simtk.openmm.app.topology.Residue object.

        `generator` must conform to the following API:
        ::
          Parameters
           ----------
           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.
276

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
277
278
279
280
281
           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`.
282

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
283
284
285
286
287
           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.
288
289

        """
290
        self._templateGenerators.append(generator)
291

292
    def _findAtomTypes(self, attrib, num):
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
        """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.

        """
308
309
310
311
312
313
314
        types = []
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
315
            typeAttrib = 'type'+suffix
316
            if classAttrib in attrib:
317
                if typeAttrib in attrib:
318
                    raise ValueError('Specified both a type and a class for the same atom: '+str(attrib))
319
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
320
321
322
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
323
324
            elif typeAttrib in attrib:
                if attrib[typeAttrib] == '':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
325
                    types.append(self._atomClasses[''])
326
                elif attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
327
328
329
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
330
331
            else:
                types.append(None) # Unknown atom type
332
333
        return types

334
    def _parseTorsion(self, attrib):
335
        """Parse the node defining a torsion."""
336
        types = self._findAtomTypes(attrib, 4)
peastman's avatar
peastman committed
337
        if None in types:
338
339
340
341
342
            return None
        torsion = PeriodicTorsion(types)
        index = 1
        while 'phase%d'%index in attrib:
            torsion.periodicity.append(int(attrib['periodicity%d'%index]))
343
344
            torsion.phase.append(_convertParameterToNumber(attrib['phase%d'%index]))
            torsion.k.append(_convertParameterToNumber(attrib['k%d'%index]))
345
            index += 1
Justin MacCallum's avatar
Justin MacCallum committed
346
347
        return torsion

348
    class _SystemData(object):
349
350
351
        """Inner class used to encapsulate data about the system being created."""
        def __init__(self):
            self.atomType = {}
352
            self.atomParameters = {}
353
            self.atoms = []
354
            self.excludeAtomWith = []
355
            self.virtualSites = {}
356
357
358
359
360
361
            self.bonds = []
            self.angles = []
            self.propers = []
            self.impropers = []
            self.atomBonds = []
            self.isAngleConstrained = []
362
363
364
365
366
367
368
369
370
371
372
            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:
                if self.constraints(key) != distance:
                    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)
373

374
    class _TemplateData(object):
375
376
377
378
        """Inner class used to encapsulate data about a residue template definition."""
        def __init__(self, name):
            self.name = name
            self.atoms = []
379
            self.virtualSites = []
380
381
            self.bonds = []
            self.externalBonds = []
382

383
384
385
386
387
388
389
        def getAtomIndexByName(self, atom_name):
            """Look up an atom index by atom name, providing a helpful error message if not found."""
            for (index, atom) in enumerate(self.atoms):
                if atom.name == atom_name:
                    return index

            # Provide a helpful error message if atom name not found.
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
390
391
            msg =  "Atom name '%s' not found in residue template '%s'.\n" % (atom_name, self.name)
            msg += "Possible atom names are: %s" % str(atomIndices.keys())
392
393
            raise ValueError(msg)

394
        def addBond(self, atom1, atom2):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
395
            """Add a bond between two atoms in a template given their indices in the template."""
396
397
398
            self.bonds.append((atom1, atom2))
            self.atoms[atom1].bondedTo.append(atom2)
            self.atoms[atom2].bondedTo.append(atom1)
399

400
        def addBondByName(self, atom1_name, atom2_name):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
401
            """Add a bond between two atoms in a template given their atom names."""
402
403
404
405
406
            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
407
            """Designate that an atom in a residue template has an external bond, using atom index within template."""
408
409
410
411
            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
412
            """Designate that an atom in a residue template has an external bond, using atom name within template."""
413
414
415
            atom = self.getAtomIndexByName(atom_name)
            self.addExternalBond(atom)

416
    class _TemplateAtomData(object):
417
        """Inner class used to encapsulate data about an atom in a residue template definition."""
418
        def __init__(self, name, type, element, parameters={}):
419
420
421
            self.name = name
            self.type = type
            self.element = element
422
            self.parameters = parameters
423
424
425
            self.bondedTo = []
            self.externalBonds = 0

426
    class _BondData(object):
427
428
429
430
431
432
        """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
433

434
    class _VirtualSiteData(object):
435
        """Inner class used to encapsulate data about a virtual site."""
436
        def __init__(self, node, atomIndices):
437
438
439
            attrib = node.attrib
            self.type = attrib['type']
            if self.type == 'average2':
440
                numAtoms = 2
441
442
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
443
                numAtoms = 3
444
445
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
446
                numAtoms = 3
447
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
448
            elif self.type == 'localCoords':
449
                numAtoms = 3
450
451
452
453
                self.originWeights = [float(attrib['wo1']), float(attrib['wo2']), float(attrib['wo3'])]
                self.xWeights = [float(attrib['wx1']), float(attrib['wx2']), float(attrib['wx3'])]
                self.yWeights = [float(attrib['wy1']), float(attrib['wy2']), float(attrib['wy3'])]
                self.localPos = [float(attrib['p1']), float(attrib['p2']), float(attrib['p3'])]
454
455
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
456
457
458
459
460
461
            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)]
462
463
464
465
            if 'excludeWith' in attrib:
                self.excludeWith = int(attrib['excludeWith'])
            else:
                self.excludeWith = self.atoms[0]
466

467
    class _AtomType(object):
468
469
470
471
472
473
474
        """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

475
    class _AtomTypeParameters(object):
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
        """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))


546
    def _getResidueTemplateMatches(self, res, bondedToAtom):
547
548
549
550
551
552
        """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.
553
554
        bondedToAtom : list of set of int
            bondedToAtom[i] is the set of atoms bonded to atom index i
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575

        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
        signature = _createResidueSignature([atom.element for atom in res.atoms()])
        if signature in self._templateSignatures:
            for t in self._templateSignatures[signature]:
                matches = _matchResidue(res, t, bondedToAtom)
                if matches is not None:
                    template = t
                    break
        return [template, matches]

576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
    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())
593
594
595
        for (atom1, atom2) in topology.bonds():
            bondedToAtom[atom1.index].add(atom2.index)
            bondedToAtom[atom2.index].add(atom1.index)
596
        return bondedToAtom
597

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

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

603
604
605
        Parameters
        ----------
        topology : Topology
606
            The Topology whose residues are to be checked against the forcefield residue templates.
607
608
609
610

        Returns
        -------
        unmatched_residues : list of Residue
611
            List of Residue objects from `topology` for which no forcefield residue templates are available.
612
613
614
615
616
            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.
617
        bondedToAtom = self._buildBondedToAtomList(topology)
618
        unmatched_residues = list() # list of unmatched residues
619
620
621
622
623
624
        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)
625
626
627

        return unmatched_residues

628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
    def getMatchingTemplates(self, topology):
        """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.

        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.
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
            # 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

660
661
    def generateTemplatesForUnmatchedResidues(self, topology):
        """Generate forcefield residue templates for residues in specified topology for which no forcefield templates are available.
662

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

665
666
667
        Parameters
        ----------
        topology : Topology
668
            The Topology whose residues are to be checked against the forcefield residue templates.
669
670
671

        Returns
        -------
672
673
674
675
676
677
        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]`
678
679
680
681
682

        """
        # Get a non-unique list of unmatched residues.
        unmatched_residues = self.getUnmatchedResidues(topology)
        # Generate a unique list of unmatched residues by comparing fingerprints.
683
        bondedToAtom = self._buildBondedToAtomList(topology)
684
685
        unique_unmatched_residues = list() # list of unique unmatched Residue objects from topology
        templates = list() # corresponding _TemplateData templates
686
687
688
        signatures = set()
        for residue in unmatched_residues:
            signature = _createResidueSignature([ atom.element for atom in residue.atoms() ])
689
            template = _createResidueTemplate(residue)
690
691
692
693
694
695
696
697
698
699
700
            is_unique = True
            if signature in signatures:
                # Signature is the same as an existing residue; check connectivity.
                for check_residue in unique_unmatched_residues:
                    matches = _matchResidue(check_residue, template, bondedToAtom)
                    if matches is not None:
                        is_unique = False
            if is_unique:
                # Residue is unique.
                unique_unmatched_residues.append(residue)
                signatures.add(signature)
701
                templates.append(template)
702

703
        return [templates, unique_unmatched_residues]
704

705
    def createSystem(self, topology, nonbondedMethod=NoCutoff, nonbondedCutoff=1.0*unit.nanometer,
706
                     constraints=None, rigidWater=True, removeCMMotion=True, hydrogenMass=None, **args):
707
        """Construct an OpenMM System representing a Topology with this force field.
Justin MacCallum's avatar
Justin MacCallum committed
708

Robert McGibbon's avatar
Robert McGibbon committed
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
        Parameters
        ----------
        topology : Topology
            The Topology for which to create a System
        nonbondedMethod : object=NoCutoff
            The method to use for nonbonded interactions.  Allowed values are
            NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, or PME.
        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.
        args
             Arbitrary additional keyword arguments may also be specified.
             This allows extra parameters to be specified that are specific to
             particular force fields.

        Returns
        -------
        system
            the newly created System
739
740
        """
        data = ForceField._SystemData()
741
        data.atoms = list(topology.atoms())
742
743
        for atom in data.atoms:
            data.excludeAtomWith.append([])
744
745

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

747
        for bond in topology.bonds():
748
            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
749
750

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

752
753
754
755
756
757
758
759
760
761
762
763
        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.
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
764
        # If no templates are found, attempt to use residue template generators to create new templates (and potentially atom types/parameters).
Justin MacCallum's avatar
Justin MacCallum committed
765

766
767
        for chain in topology.chains():
            for res in chain.residues():
768
                # Attempt to match one of the existing templates.
769
                [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
770
771
                if matches is None:
                    # No existing templates match.  Try any registered residue template generators.
772
                    for generator in self._templateGenerators:
773
                        if generator(self, res):
774
                            # This generator has registered a new residue template that should match.
775
                            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
776
777
778
779
780
781
782
783
                            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

                # Raise an exception if we have found no templates that match.
784
                if matches is None:
785
                    raise ValueError('No template found for residue %d (%s).  %s' % (res.index+1, res.name, _findMatchErrors(self, res)))
786
787

                # Store parameters for the matched residue template.
788
                matchAtoms = dict(zip(matches, res.atoms()))
789
790
                for atom, match in zip(res.atoms(), matches):
                    data.atomType[atom] = template.atoms[match].type
791
                    data.atomParameters[atom] = template.atoms[match].parameters
792
793
                    for site in template.virtualSites:
                        if match == site.index:
794
                            data.virtualSites[atom] = (site, [matchAtoms[i].index for i in site.atoms], matchAtoms[site.excludeWith].index)
795
796

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

798
799
        sys = mm.System()
        for atom in topology.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
800
            # Look up the atom type name, returning a helpful error message if it cannot be found.
801
802
803
804
            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
805
            # Look up the type name in the list of registered atom types, returning a helpful error message if it cannot be found.
806
807
808
809
            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
810
811

            # Add the particle to the OpenMM system.
812
            mass = self._atomTypes[typename].mass
813
            sys.addParticle(mass)
814

815
        # Adjust hydrogen masses if requested.
816

817
        if hydrogenMass is not None:
818
819
            if not unit.is_quantity(hydrogenMass):
                hydrogenMass *= unit.dalton
820
821
822
823
824
825
826
            for atom1, atom2 in topology.bonds():
                if atom1.element == elem.hydrogen:
                    (atom1, atom2) = (atom2, atom1)
                if atom2.element == elem.hydrogen and atom1.element not in (elem.hydrogen, None):
                    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
827

828
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
829

830
831
832
        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
833
        elif nonbondedMethod not in [NoCutoff, CutoffNonPeriodic]:
834
835
836
            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
837

838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
        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
853

854
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
855

856
857
858
        uniquePropers = set()
        for angle in data.angles:
            for atom in bondedToAtom[angle[0]]:
pgrinaway's avatar
pgrinaway committed
859
                if atom != angle[1] and atom != angle[2]:
860
861
862
863
864
                    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
865
                if atom != angle[1] and atom != angle[0]:
866
867
868
869
870
                    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
871

872
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
873

874
875
876
877
878
        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
879

880
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
881

882
883
884
885
886
887
888
889
890
891
892
893
894
895
        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]
                bond.isConstrained = atom1.name.startswith('H') or atom2.name.startswith('H')
        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
896

897
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
898

899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
        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
                if atom1.name.startswith('H'):
                    numH += 1
                if atom3.name.startswith('H'):
                    numH += 1
                data.isAngleConstrained.append(numH == 2 or (numH == 1 and atom2.name.startswith('O')))
        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
920

921
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
922

923
        for atom in data.virtualSites:
924
            (site, atoms, excludeWith) = data.virtualSites[atom]
925
            index = atom.index
926
            data.excludeAtomWith[excludeWith].append(index)
927
            if site.type == 'average2':
928
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(atoms[0], atoms[1], site.weights[0], site.weights[1]))
929
            elif site.type == 'average3':
930
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
931
            elif site.type == 'outOfPlane':
932
933
934
935
936
937
938
                sys.setVirtualSite(index, mm.OutOfPlaneSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
            elif site.type == 'localCoords':
                sys.setVirtualSite(index, mm.LocalCoordinatesSite(atoms[0], atoms[1], atoms[2],
                                                                  mm.Vec3(site.originWeights[0], site.originWeights[1], site.originWeights[2]),
                                                                  mm.Vec3(site.xWeights[0], site.xWeights[1], site.xWeights[2]),
                                                                  mm.Vec3(site.yWeights[0], site.yWeights[1], site.yWeights[2]),
                                                                  mm.Vec3(site.localPos[0], site.localPos[1], site.localPos[2])))
Justin MacCallum's avatar
Justin MacCallum committed
939

940
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
941

942
943
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
944
945
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
946

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

peastman's avatar
peastman committed
949
950
951
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
952

953
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
954

955
        for script in self._scripts:
956
            exec(script, locals())
957
958
959
        return sys


960
961
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
962
963
    counts = {}
    for element in elements:
964
        if element in counts:
965
966
967
            counts[element] += 1
        else:
            counts[element] = 1
968
969
970
971
972
973
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
974
975
    sig = []
    for c in counts:
976
977
        if c is not None:
            sig.append((c, counts[c]))
978
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
979

980
    # Convert it to a string.
981
982

    s = ''
983
    for element, count in sig:
984
985
986
        s += element.symbol+str(count)
    return s

987
def _matchResidue(res, template, bondedToAtom):
988
    """Determine whether a residue matches a template and return a list of corresponding atoms.
Justin MacCallum's avatar
Justin MacCallum committed
989

Robert McGibbon's avatar
Robert McGibbon committed
990
991
992
993
994
995
996
997
998
999
1000
    Parameters
    ----------
    res : Residue
        The residue to check
    template : _TemplateData
        The template to compare it to
    bondedToAtom : list
        Enumerates which other atoms each atom is bonded to

    Returns
    -------
1001
1002
1003
    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
1004
1005
    """
    atoms = list(res.atoms())
1006
1007
    if len(atoms) != len(template.atoms):
        return None
1008
1009
    matches = len(atoms)*[0]
    hasMatch = len(atoms)*[False]
Justin MacCallum's avatar
Justin MacCallum committed
1010

1011
    # Translate from global to local atom indices, and record the bonds for each atom.
Justin MacCallum's avatar
Justin MacCallum committed
1012

1013
1014
    renumberAtoms = {}
    for i in range(len(atoms)):
1015
        renumberAtoms[atoms[i].index] = i
1016
1017
1018
    bondedTo = []
    externalBonds = []
    for atom in atoms:
1019
        bonds = [renumberAtoms[x] for x in bondedToAtom[atom.index] if x in renumberAtoms]
1020
        bondedTo.append(bonds)
1021
        externalBonds.append(len([x for x in bondedToAtom[atom.index] if x not in renumberAtoms]))
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042

    # For each unique combination of element and number of bonds, make sure the residue and
    # template have the same number of atoms.

    residueTypeCount = {}
    for i, atom in enumerate(atoms):
        key = (atom.element, len(bondedTo[i]), externalBonds[i])
        if key not in residueTypeCount:
            residueTypeCount[key] = 1
        residueTypeCount[key] += 1
    templateTypeCount = {}
    for i, atom in enumerate(template.atoms):
        key = (atom.element, len(atom.bondedTo), atom.externalBonds)
        if key not in templateTypeCount:
            templateTypeCount[key] = 1
        templateTypeCount[key] += 1
    if residueTypeCount != templateTypeCount:
        return None

    # Recursively match atoms.

1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
    if _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch, 0):
        return matches
    return None


def _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch, position):
    """This is called recursively from inside _matchResidue() to identify matching atoms."""
    if position == len(atoms):
        return True
    elem = atoms[position].element
peastman's avatar
peastman committed
1053
    name = atoms[position].name
1054
1055
    for i in range(len(atoms)):
        atom = template.atoms[i]
1056
        if ((atom.element is not None and atom.element == elem) or (atom.element is None and atom.name == name)) and not hasMatch[i] and len(atom.bondedTo) == len(bondedTo[position]) and atom.externalBonds == externalBonds[position]:
1057
            # See if the bonds for this identification are consistent
Justin MacCallum's avatar
Justin MacCallum committed
1058

1059
1060
1061
            allBondsMatch = all((bonded > position or matches[bonded] in atom.bondedTo for bonded in bondedTo[position]))
            if allBondsMatch:
                # This is a possible match, so trying matching the rest of the residue.
Justin MacCallum's avatar
Justin MacCallum committed
1062

1063
1064
1065
1066
1067
1068
1069
1070
                matches[position] = i
                hasMatch[i] = True
                if _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch, position+1):
                    return True
                hasMatch[i] = False
    return False


1071
1072
1073
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()])
1074
    numResidueAtoms = sum(residueCounts.values())
1075
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
1076

1077
    # Loop over templates and see how closely each one might match.
1078

1079
1080
1081
1082
1083
1084
    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])
1085

1086
        # Does the residue have any atoms that clearly aren't in the template?
1087

1088
1089
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
1090

1091
        # If there are too many missing atoms, discard this template.
1092

1093
        numTemplateAtoms = sum(templateCounts.values())
1094
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
1095
1096
1097
1098
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
1099

1100
1101
        # 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.
1102

1103
1104
1105
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
1106

1107
        # Accept this as our new best match.
1108

1109
1110
1111
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
1112
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
1113

1114
    # Return an appropriate error message.
1115

1116
    if numBestMatchAtoms == numResidueAtoms:
1117
1118
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
1119
1120
1121
1122
            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:
1123
1124
1125
1126
1127
            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)
1128
1129
1130
        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.'

1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
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():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1149
        template.atoms.append(ForceField._TemplateAtomData(atom.name, None, atom.element))
1150
1151
1152
1153
1154
1155
1156
1157
1158
    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
1159

1160
1161
1162
1163
1164
# 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.

1165
## @private
1166
class HarmonicBondGenerator(object):
1167
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1168

1169
1170
    def __init__(self, forcefield):
        self.ff = forcefield
1171
1172
1173
1174
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
1175

1176
1177
1178
1179
1180
1181
1182
    def registerBond(self, parameters):
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.length.append(_convertParameterToNumber(parameters['length']))
            self.k.append(_convertParameterToNumber(parameters['k']))
Justin MacCallum's avatar
Justin MacCallum committed
1183

1184
1185
    @staticmethod
    def parseElement(element, ff):
1186
1187
        generator = HarmonicBondGenerator(ff)
        ff.registerGenerator(generator)
1188
        for bond in element.findall('Bond'):
1189
            generator.registerBond(bond.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1190

1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
    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]]
            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:
1208
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
1209
1210
1211
1212
1213
1214
1215
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


1216
## @private
1217
class HarmonicAngleGenerator(object):
1218
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1219

1220
1221
    def __init__(self, forcefield):
        self.ff = forcefield
1222
1223
1224
1225
1226
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
1227

1228
1229
1230
1231
1232
1233
1234
1235
1236
    def registerAngle(self, parameters):
        types = self.ff._findAtomTypes(parameters, 3)
        if None not in types:
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.types3.append(types[2])
            self.angle.append(_convertParameterToNumber(parameters['angle']))
            self.k.append(_convertParameterToNumber(parameters['k']))

1237
1238
    @staticmethod
    def parseElement(element, ff):
1239
1240
        generator = HarmonicAngleGenerator(ff)
        ff.registerGenerator(generator)
1241
        for angle in element.findall('Angle'):
1242
            generator.registerAngle(angle.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1243

1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
    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]]]
            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:
                        # Find the two bonds that make this angle.
Justin MacCallum's avatar
Justin MacCallum committed
1263

1264
1265
1266
1267
1268
1269
1270
1271
1272
                        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
1273

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

1276
1277
1278
1279
1280
                        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]))
1281
                                data.addConstraint(sys, angle[0], angle[2], length)
1282
1283
1284
1285
1286
1287
1288
                    elif self.k[i] != 0:
                        force.addAngle(angle[0], angle[1], angle[2], self.angle[i], self.k[i])
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


1289
## @private
1290
class PeriodicTorsion(object):
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
    """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 = []

1302
## @private
1303
class PeriodicTorsionGenerator(object):
1304
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1305

1306
1307
    def __init__(self, forcefield):
        self.ff = forcefield
1308
1309
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
1310

1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
    def registerProperTorsion(self, parameters):
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
            self.proper.append(torsion)

    def registerImproperTorsion(self, parameters):
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
            self.improper.append(torsion)

1321
1322
    @staticmethod
    def parseElement(element, ff):
1323
1324
        generator = PeriodicTorsionGenerator(ff)
        ff.registerGenerator(generator)
1325
        for torsion in element.findall('Proper'):
1326
            generator.registerProperTorsion(torsion.attrib)
1327
        for torsion in element.findall('Improper'):
1328
            generator.registerImproperTorsion(torsion.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1329

1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
    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['']
        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:
                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])
        for torsion in data.impropers:
            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]]]
1365
            match = None
1366
1367
1368
1369
1370
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
1371
1372
1373
1374
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
1375
1376
1377
                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:
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
                            # 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)
1389
                            match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
1390
                            break
1391
1392
1393
1394
1395
            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])
1396
1397
1398
1399

parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


1400
## @private
1401
class RBTorsion(object):
1402
1403
1404
1405
1406
1407
1408
1409
1410
    """An RBTorsion records the information for a Ryckaert-Bellemans torsion definition."""

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

1411
## @private
1412
class RBTorsionGenerator(object):
1413
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1414

1415
1416
    def __init__(self, forcefield):
        self.ff = forcefield
1417
1418
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
1419

1420
1421
    @staticmethod
    def parseElement(element, ff):
1422
1423
        generator = RBTorsionGenerator(ff)
        ff.registerGenerator(generator)
1424
        for torsion in element.findall('Proper'):
1425
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1426
            if None not in types:
1427
1428
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
1429
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1430
            if None not in types:
1431
                generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
Justin MacCallum's avatar
Justin MacCallum committed
1432

1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
    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:
            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]]]
1466
            match = None
1467
1468
1469
1470
1471
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
1472
1473
1474
1475
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
1476
1477
1478
                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:
1479
                            if hasWildcard:
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
                                # 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)
1491
                                match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
1492
1493
                            else:
                                # There are no wildcards, so the order is unambiguous.
1494
                                match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
1495
                            break
1496
1497
1498
            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])
1499
1500
1501
1502

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


1503
## @private
1504
class CMAPTorsion(object):
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
    """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

1515
## @private
1516
class CMAPTorsionGenerator(object):
1517
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1518

1519
1520
    def __init__(self, forcefield):
        self.ff = forcefield
1521
1522
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
1523

1524
1525
    @staticmethod
    def parseElement(element, ff):
1526
1527
        generator = CMAPTorsionGenerator(ff)
        ff.registerGenerator(generator)
1528
1529
1530
1531
1532
1533
1534
        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'):
1535
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
1536
            if None not in types:
1537
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
1538

1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
    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
1549

1550
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
1551

1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
        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


1595
## @private
1596
class NonbondedGenerator(object):
1597
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1598

1599
1600
    SCALETOL = 1e-5

1601
1602
    def __init__(self, forcefield, coulomb14scale, lj14scale):
        self.ff = forcefield
1603
1604
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
1605
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
1606

1607
    def registerAtom(self, parameters):
1608
        self.params.registerAtom(parameters)
1609

1610
1611
1612
1613
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
        if len(existing) == 0:
1614
1615
            generator = NonbondedGenerator(ff, float(element.attrib['coulomb14scale']), float(element.attrib['lj14scale']))
            ff.registerGenerator(generator)
1616
1617
1618
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
1619
1620
            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
1621
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
1622
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
1623

1624
1625
1626
1627
1628
1629
1630
    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,
                     PME:mm.NonbondedForce.PME}
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
1631
            raise ValueError('Illegal nonbonded method for NonbondedForce')
1632
1633
        force = mm.NonbondedForce()
        for atom in data.atoms:
1634
1635
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
1636
1637
1638
1639
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
1640
1641
        if 'useDispersionCorrection' in args:
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
peastman's avatar
peastman committed
1642
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
1643

peastman's avatar
peastman committed
1644
    def postprocessSystem(self, sys, data, args):
1645
        # Create exceptions based on bonds.
1646

1647
1648
1649
        bondIndices = []
        for bond in data.bonds:
            bondIndices.append((bond.atom1, bond.atom2))
1650
1651
1652

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

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

1659
1660
        # 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.
1661

1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
        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 exceptions.

peastman's avatar
peastman committed
1672
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
1673
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
1674
1675
1676
1677

parsers["NonbondedForce"] = NonbondedGenerator.parseElement


1678
## @private
1679
class GBSAOBCGenerator(object):
1680
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1681

1682
1683
    def __init__(self, forcefield):
        self.ff = forcefield
1684
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
1685

1686
    def registerAtom(self, parameters):
1687
        self.params.registerAtom(parameters)
1688

1689
1690
    @staticmethod
    def parseElement(element, ff):
1691
1692
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
1693
1694
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
1695
1696
1697
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
1698
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
1699

1700
1701
1702
1703
1704
    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
1705
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
1706
1707
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
1708
1709
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
1710
1711
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
1712
1713
1714
1715
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
1716
1717
        sys.addForce(force)

1718
1719
    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
1720

1721
1722
1723
1724
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

1725
1726
1727
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


1728
## @private
1729
class CustomBondGenerator(object):
1730
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1731

1732
1733
    def __init__(self, forcefield):
        self.ff = forcefield
1734
1735
1736
1737
1738
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
1739

1740
1741
    @staticmethod
    def parseElement(element, ff):
1742
1743
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
1744
1745
1746
1747
1748
1749
        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'):
1750
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
1751
            if None not in types:
1752
1753
1754
                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
1755

1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
    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


1776
## @private
1777
class CustomAngleGenerator(object):
1778
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1779

1780
1781
    def __init__(self, forcefield):
        self.ff = forcefield
1782
1783
1784
1785
1786
1787
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
1788

1789
1790
    @staticmethod
    def parseElement(element, ff):
1791
1792
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
1793
1794
1795
1796
1797
1798
        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'):
1799
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
1800
            if None not in types:
1801
1802
1803
1804
                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
1805

1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
    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


1828
## @private
1829
class CustomTorsion(object):
1830
1831
1832
1833
1834
1835
1836
1837
1838
    """A CustomTorsion records the information for a custom torsion definition."""

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

1839
## @private
1840
class CustomTorsionGenerator(object):
1841
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1842

1843
1844
    def __init__(self, forcefield):
        self.ff = forcefield
1845
1846
1847
1848
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
1849

1850
1851
    @staticmethod
    def parseElement(element, ff):
1852
1853
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
1854
1855
1856
1857
1858
1859
        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'):
1860
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1861
            if None not in types:
1862
1863
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
1864
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1865
            if None not in types:
1866
                generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
Justin MacCallum's avatar
Justin MacCallum committed
1867

1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
    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:
            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]]]
1900
            match = None
1901
1902
1903
1904
1905
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
1906
1907
1908
1909
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
1910
1911
1912
                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:
1913
                            if hasWildcard:
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
                                # 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)
1925
                                match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
1926
1927
                            else:
                                # There are no wildcards, so the order is unambiguous.
1928
                                match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
1929
                            break
1930
1931
1932
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
1933
1934
1935
1936

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


1937
## @private
1938
class CustomNonbondedGenerator(object):
1939
1940
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

1941
1942
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
1943
1944
1945
1946
1947
1948
1949
1950
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
1951
1952
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
1953
1954
1955
1956
        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'])
1957
1958
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984

    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)
        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:
1985
1986
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
1987
1988
1989
1990
1991
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

1994
1995
1996
        bondIndices = []
        for bond in data.bonds:
            bondIndices.append((bond.atom1, bond.atom2))
1997
1998
1999

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

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

2006
2007
        # 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.
2008

2009
2010
2011
2012
2013
2014
2015
2016
2017
        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.
2018

2019
2020
2021
2022
2023
2024
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomNonbondedForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


2025
## @private
2026
class CustomGBGenerator(object):
2027
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2028

2029
2030
    def __init__(self, forcefield):
        self.ff = forcefield
2031
2032
2033
2034
2035
2036
2037
2038
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2039
2040
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
2041
2042
2043
2044
        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'])
2045
2046
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2047
2048
2049
2050
2051
2052
2053
2054
2055
        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']]))
        for function in element.findall("Function"):
            values = [float(x) for x in function.text.split()]
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
            if 'type' in function.attrib:
                type = function.attrib['type']
            else:
                type = '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])
            generator.functions.append((function.attrib['name'], type, values, params))
Justin MacCallum's avatar
Justin MacCallum committed
2067

2068
2069
2070
2071
2072
    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
2073
            raise ValueError('Illegal nonbonded method for CustomGBForce')
2074
2075
2076
2077
2078
2079
2080
2081
2082
        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])
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
        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))
2096
        for atom in data.atoms:
2097
2098
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2099
2100
2101
2102
2103
2104
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

2105
2106

## @private
2107
class CustomManyParticleGenerator(object):
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
    """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 = []
2120

2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
    @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(',')]))
2133
2134
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163

    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:
2164
2165
2166
            values = self.params.getAtomParameters(atom, data)
            type = int(self.params.getExtraParameters(atom, data)['filterType'])
            force.addParticle(values, type)
2167
2168
2169
2170
2171
2172
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
        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)))
2185

2186
2187
        # 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.
2188

2189
2190
2191
2192
2193
2194
2195
2196
2197
        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.
2198

2199
2200
2201
2202
2203
        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
2204
def getAtomPrint(data, atomIndex):
2205

Peter Eastman's avatar
Peter Eastman committed
2206
2207
2208
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
2209
    else:
Peter Eastman's avatar
Peter Eastman committed
2210
        returnString = "NA"
2211
2212
2213
2214
2215

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
2216
def countConstraint(data):
2217

Peter Eastman's avatar
Peter Eastman committed
2218
    bondCount = 0
2219
2220
2221
2222
2223
2224
2225
    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
2226
        if (isConstrained):
2227
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
2228

2229
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
2230

2231
## @private
2232
class AmoebaBondGenerator(object):
2233
2234
2235

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

2236
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
2237
2238

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

2240
2241
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
2242
2243
2244
2245
2246
2247
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2248

2249
2250
2251
2252
2253
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

2257
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
2258
2259
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
2260
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2261
            if None not in types:
2262
2263
2264
2265
2266
                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:
2267
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
2268
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
2269
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
2270
2271
                raise ValueError(outputString)

2272
2273
    #=============================================================================================

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

Mark Friedrichs's avatar
Mark Friedrichs committed
2276
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2277

Peter Eastman's avatar
Peter Eastman committed
2278
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2279
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
2280
        if len(existing) == 0:
2281
            force = mm.AmoebaBondForce()
2282
2283
2284
2285
            sys.addForce(force)
        else:
            force = existing[0]

2286
2287
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297

        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:
2298
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
2299
2300
2301
2302
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

2303
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
2304
2305
2306
2307

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

Peter Eastman's avatar
Peter Eastman committed
2309
def addAngleConstraint(angle, idealAngle, data, sys):
2310
2311

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

Peter Eastman's avatar
Peter Eastman committed
2313
2314
    bond1 = None
    bond2 = None
2315
2316
2317
2318
2319
2320
2321
    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
2322

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

2325
2326
2327
2328
        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
2329
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
2330
                data.addConstraint(sys, angle[0], angle[2], length)
2331
2332
2333
                return

#=============================================================================================
2334
## @private
2335
class AmoebaAngleGenerator(object):
2336
2337

    #=============================================================================================
2338
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
2339
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
2340

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

Peter Eastman's avatar
Peter Eastman committed
2343
2344
2345
2346
2347
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
2348

Peter Eastman's avatar
Peter Eastman committed
2349
2350
2351
        self.types1 = []
        self.types2 = []
        self.types3 = []
2352

Peter Eastman's avatar
Peter Eastman committed
2353
2354
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2355

2356
2357
2358
2359
2360
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

2364
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
2365
2366
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
2367
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2368
            if None not in types:
2369
2370
2371
2372
2373
2374

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
2375
                angleList.append(float(angle.attrib['angle1']))
2376
2377

                try:
Peter Eastman's avatar
Peter Eastman committed
2378
                    angleList.append(float(angle.attrib['angle2']))
2379
                    try:
Peter Eastman's avatar
Peter Eastman committed
2380
                        angleList.append(float(angle.attrib['angle3']))
2381
2382
2383
2384
2385
2386
2387
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
2388
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
2389
2390
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
2391
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
2392
2393
                raise ValueError(outputString)

2394
2395
2396
2397
    #=============================================================================================
    # 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
2398

Peter Eastman's avatar
Peter Eastman committed
2399
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2400
2401
2402
2403
2404
2405
        pass

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

Peter Eastman's avatar
Peter Eastman committed
2407
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
2408
2409
2410
2411

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2412
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
2413
2414

        if len(existing) == 0:
2415
            force = mm.AmoebaAngleForce()
2416
2417
2418
2419
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2420
        # set scalars
2421

2422
2423
2424
2425
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
2426
2427

        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
2428
2429
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
2430

Peter Eastman's avatar
Peter Eastman committed
2431
2432
2433
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
2434
2435
2436
2437
2438
2439
2440
            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]
2441
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
2442
                    elif self.k[i] != 0:
Peter Eastman's avatar
Peter Eastman committed
2443
2444
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
2445
2446
2447
2448
2449
2450
                            # 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
2451
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
2452
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
2453
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
2454
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
2455
                            if (numberOfHydrogens < lenAngle):
2456
2457
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
2458
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
2459
                                raise ValueError(outputString)
2460
2461
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
2462

2463
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
2464
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
2465
2466
2467
2468
2469
2470
                    break

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

Peter Eastman's avatar
Peter Eastman committed
2472
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
2473
2474
2475
2476

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2477
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
2478
2479

        if len(existing) == 0:
2480
            force = mm.AmoebaInPlaneAngleForce()
2481
2482
2483
2484
2485
2486
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

2487
2488
2489
2490
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
2491
2492

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

Peter Eastman's avatar
Peter Eastman committed
2494
2495
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
2496

Peter Eastman's avatar
Peter Eastman committed
2497
2498
2499
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
2500
2501
2502
2503
2504
2505
2506
2507
2508

            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
2509
                    if (isConstrained and self.k[i] != 0.0):
2510
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
2511
2512
2513
2514
                    else:
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

2515
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
2516
2517
2518

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
2519
2520
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
2521
2522
#=============================================================================================

2523
## @private
2524
class AmoebaOutOfPlaneBendGenerator(object):
2525
2526
2527
2528

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

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

2530
2531
2532
2533
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
2534
2535
2536
2537
2538
2539
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
2540

Peter Eastman's avatar
Peter Eastman committed
2541
2542
2543
2544
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
2545

Peter Eastman's avatar
Peter Eastman committed
2546
        self.ks = []
2547
2548
2549
2550
2551

    #=============================================================================================
    # 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
2552

2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
    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
2578

2579
2580
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
2581
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
2582
2583
2584
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
2585
                                                   float(element.attrib['opbend-sextic']))
2586
2587
2588
2589

        forceField._forces.append(generator)

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

Peter Eastman's avatar
Peter Eastman committed
2593
2594
2595
2596
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
2597

Peter Eastman's avatar
Peter Eastman committed
2598
                generator.ks.append(float(angle.attrib['k']))
2599
2600

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2601
2602
                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
2603
2604
                raise ValueError(outputString)

2605
2606
2607
2608
2609
2610
    #=============================================================================================
    # 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
2611

Peter Eastman's avatar
Peter Eastman committed
2612
    def getMiddleAtom(self, angle, data):
2613
2614
2615

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
2616
        middleAtom = -1
Justin MacCallum's avatar
Justin MacCallum committed
2617
        for atomIndex in angle:
Peter Eastman's avatar
Peter Eastman committed
2618
            isMiddle = 0
2619
2620
2621
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
2622
                if (atom1 != atomIndex):
2623
2624
2625
                    partner = atom1
                else:
                    partner = atom2
Justin MacCallum's avatar
Justin MacCallum committed
2626
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]):
2627
2628
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
2629
            if (isMiddle == 2):
2630
2631
2632
2633
2634
                return atomIndex
        return -1

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

Peter Eastman's avatar
Peter Eastman committed
2635
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648

        # 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
2649
2650
2651
2652
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
2653
2654
2655
2656

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

Peter Eastman's avatar
Peter Eastman committed
2657
        skipAtoms = dict()
2658
2659
2660
2661

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

Peter Eastman's avatar
Peter Eastman committed
2662
2663
        inPlaneAngles = []
        nonInPlaneAngles = []
2664
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
2665
        idealAngles = []*len(data.angles)
2666
2667
2668

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

Peter Eastman's avatar
Peter Eastman committed
2669
2670
2671
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
2672
2673
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
2674
                middleType = -1
2675
2676
                middleCovalency = -1

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

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

Peter Eastman's avatar
Peter Eastman committed
2685
2686
2687
2688
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
2689
2690
2691
2692

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
2693
                    if (atom1 != middleAtom):
2694
2695
2696
2697
2698
2699
2700
2701
                        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
2702
2703
2704
2705
2706
                        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
2707

Peter Eastman's avatar
Peter Eastman committed
2708
                if (len(partners) == 3):
2709

Peter Eastman's avatar
Peter Eastman committed
2710
2711
2712
                    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])
2713

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

2716
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
2717
2718
2719
2720
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2721
2722
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
2723
2724
2725
2726
2727
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
2728
2729
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
2730
                    angleDict['isConstrained'] = 0
2731

Peter Eastman's avatar
Peter Eastman committed
2732
2733
2734
2735
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
2736
2737

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
2738
2739
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
2740

Peter Eastman's avatar
Peter Eastman committed
2741
                    inPlaneAngles.append(angleDict)
2742
2743

                else:
Peter Eastman's avatar
Peter Eastman committed
2744
2745
2746
2747
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
2748
            else:
Peter Eastman's avatar
Peter Eastman committed
2749
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
2750

Peter Eastman's avatar
Peter Eastman committed
2751
                    partnerSet = skipAtoms[middleAtom]
Justin MacCallum's avatar
Justin MacCallum committed
2752

Peter Eastman's avatar
Peter Eastman committed
2753
                    angleDict = {}
2754

Peter Eastman's avatar
Peter Eastman committed
2755
2756
2757
2758
2759
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
2760

Peter Eastman's avatar
Peter Eastman committed
2761
                    angleDict['isConstrained'] = isConstrained
2762

Peter Eastman's avatar
Peter Eastman committed
2763
2764
2765
2766
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
2767
2768

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
2769
2770
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
2771

Peter Eastman's avatar
Peter Eastman committed
2772
                    inPlaneAngles.append(angleDict)
2773
2774

                else:
Peter Eastman's avatar
Peter Eastman committed
2775
2776
                    angleDict = {}
                    angleDict['angle'] = angle
2777
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
2778
                    nonInPlaneAngles.append(angleDict)
2779

2780
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
2781
2782

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
2783
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
2784
2785
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
2786
2787

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
2788
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
2789
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
2790
                    nonInPlaneAngles.append(angleDict)
2791
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
2792
2793
2794
2795
2796

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

2797
## @private
2798
class AmoebaTorsionGenerator(object):
2799
2800
2801
2802
2803
2804
2805

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
2806
        self.torsionUnit = torsionUnit
2807

Peter Eastman's avatar
Peter Eastman committed
2808
2809
2810
2811
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
2812

Peter Eastman's avatar
Peter Eastman committed
2813
2814
2815
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
2816

2817
2818
2819
2820
2821
2822
2823
2824
    #=============================================================================================

    @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
2825

Peter Eastman's avatar
Peter Eastman committed
2826
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
2827
2828
2829
2830
2831
2832
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
2833
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2834
            if None not in types:
2835
2836
2837
2838
2839
2840
2841

                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
2842
2843
                    tInfo = []
                    suffix = str(ii)
2844
2845
2846
2847
2848
2849
                    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
2850
2851
2852
2853
2854
2855
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
2856
2857
2858

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
2859
2860
2861
2862
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
2863
2864
                raise ValueError(outputString)

2865
2866
2867
2868
2869
    #=============================================================================================

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2870
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
2871
        if len(existing) == 0:
2872
            force = mm.PeriodicTorsionForce()
2873
2874
2875
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2876

2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
        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
2893
                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):
2894
2895
2896
2897
2898
2899
                    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])
2900
2901
2902
2903
2904
2905
                    break

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

2906
## @private
2907
class AmoebaPiTorsionGenerator(object):
2908
2909
2910
2911
2912
2913

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

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

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

2915
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
2916
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
2917
2918
2919
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2920

2921
2922
2923
2924
2925
2926
2927
2928
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
2929
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
2930
2931
2932
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
2933
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
2934
            if None not in types:
2935
2936
2937
2938
2939
2940
                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
2941
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
2942
2943
                raise ValueError(outputString)

2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
    #=============================================================================================

    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
2960

2961
2962
            atom1 = bond.atom1
            atom2 = bond.atom2
Justin MacCallum's avatar
Justin MacCallum committed
2963

2964
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975

                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
2976
2977
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989

                       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
2990
                           if (bondedAtom1 != atom1):
2991
2992
2993
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
2994
2995
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
2996
                                   piTorsionAtom1 = b1
2997
2998
2999
3000
3001
3002
                               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
3003
                           if (bondedAtom1 != atom2):
3004
3005
3006
3007
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3008
3009
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3010
                                   piTorsionAtom5 = b1
3011
3012
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
3013

Peter Eastman's avatar
Peter Eastman committed
3014
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
3015
3016
3017
3018
3019

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

3020
## @private
3021
class AmoebaTorsionTorsionGenerator(object):
3022
3023
3024
3025
3026

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

Peter Eastman's avatar
Peter Eastman committed
3027
    def __init__(self):
3028

Peter Eastman's avatar
Peter Eastman committed
3029
3030
3031
3032
3033
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
3034

Peter Eastman's avatar
Peter Eastman committed
3035
        self.gridIndex = []
3036

Peter Eastman's avatar
Peter Eastman committed
3037
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
3038

3039
3040
3041
3042
3043
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
3044
        generator = AmoebaTorsionTorsionGenerator()
3045
3046
3047
3048
3049
3050
3051
        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'):
3052
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
3053
            if None not in types:
3054
3055
3056
3057
3058
3059
3060

                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
3061
3062
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
3063
3064
3065
3066
3067
3068
3069
3070
3071
                    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
3072
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
3073
3074
                raise ValueError(outputString)

3075
3076
3077
3078
3079
3080
3081
3082
3083
3084
        # 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
3085
3086
3087
3088
3089
3090
        #     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
3091
3092

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

Peter Eastman's avatar
Peter Eastman committed
3096
3097
3098
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
3099

Peter Eastman's avatar
Peter Eastman committed
3100
3101
            grid = []
            gridCol = []
3102
3103
3104
3105
3106

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
3107
3108
3109
3110
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
3111
                if 'fx' in gridEntry.attrib:
3112
3113
3114
                    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
3115
                gridCol.append(gridRow)
3116
3117

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
3118
3119
3120
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
3121
3122
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
3123

Peter Eastman's avatar
Peter Eastman committed
3124
3125
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
3126
            else:
Peter Eastman's avatar
Peter Eastman committed
3127
3128
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
3129
3130
3131
3132
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
3133
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
3134
3135
3136
3137
3138
3139
3140
3141

        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
3142
        if (len(data.atomBonds[atomC]) == 4):
3143
3144
3145
3146
3147
            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
3148
3149
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
3150
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
3151
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
3152
3153
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
3154
3155
                if (hit > -1):
                    if (atomE == -1):
3156
3157
3158
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
3159

3160
3161
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
3162
3163
            if (atomE == -1 or atomF == -1):
                outputString = "getChiralAtomIndex: error getting bonded partners of atomC=%s %d %s" % (atomC.name, atomC.resiude.index, atomC.resiude.name,)
Justin MacCallum's avatar
Justin MacCallum committed
3164
                raise ValueError(outputString)
3165
3166
3167
3168
3169

            # 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
3170
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
3171
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
3172
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
3173
                chiralAtomIndex = atomF
3174

Peter Eastman's avatar
Peter Eastman committed
3175
3176
3177
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
3178
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
3179
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
3180
                chiralAtomIndex = massF
3181
3182
3183
3184

        return chiralAtomIndex

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

3186
3187
3188
3189
3190
3191
3192
3193
3194
3195
3196
3197
3198
    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
3199
3200
            # search for bitorsions; based on TINKER subroutine bitors()

3201
3202
3203
3204
3205
3206
3207
            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
3208
                if (bondedAtom1 != ib):
3209
3210
3211
3212
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3213
                if (ia != ic and ia != id):
3214
3215
3216
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
3217
                        if (bondedAtom1 != id):
3218
3219
3220
3221
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3222
                        if (ie != ic and ie != ib and ie != ia):
3223
3224
3225
3226
3227
3228
3229
3230
3231
3232
3233
3234
3235
3236
3237
3238
3239
3240
3241
3242

                            # 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
3243
3244
3245
                                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])
3246
3247
3248

                                # match in reverse order

3249
                                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
3250
3251
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
3252
3253
3254
3255

        # set grids

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

3258
3259
3260
3261
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

3262
## @private
3263
class AmoebaStretchBendGenerator(object):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3264
3265

    #=============================================================================================
3266
3267
3268
3269
3270
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
3271
3272
3273
        self.types1 = []
        self.types2 = []
        self.types3 = []
3274

Peter Eastman's avatar
Peter Eastman committed
3275
3276
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
3277

3278
3279
3280
3281
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
3282
        generator = AmoebaStretchBendGenerator()
3283
3284
3285
3286
3287
3288
3289
        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'):
3290
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
3291
            if None not in types:
3292
3293
3294
3295
3296
3297
3298
3299
3300
3301
3302
3303

                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
3304
                                    stretchBend.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
3305
3306
                raise ValueError(outputString)

3307
3308
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
3309
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
3310
3311
    # 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
3312
3313
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
3314
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
3315
3316
3317

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

Peter Eastman's avatar
Peter Eastman committed
3318
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3319
3320
3321
3322
3323
3324
3325
3326
        pass

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

    # Note: request for constrained bonds is ignored.

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

3327
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3328
3329
3330
3331
3332
3333
3334
3335
3336
3337
3338
3339

        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
3340
            if ('isConstrained' in angleDict):
3341
3342
3343
3344
3345
3346
3347
3348
                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
3349
            radian = 57.2957795130
3350
3351
3352
3353
3354
3355
3356
3357
3358
3359
            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
3360
                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
3361
3362
3363
3364
3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
                    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
3376

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3379
                    if ('idealAngle' not in angleDict):
3380

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3381
3382
3383
3384
3385
                       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
3386
                       raise ValueError(outputString)
3387

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3388
3389
3390
3391
3392
3393
3394
                    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
3395
                       raise ValueError(outputString)
3396

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3400
                    break
3401
3402
3403
3404
3405

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

3406
## @private
3407
class AmoebaVdwGenerator(object):
3408
3409

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

3411
3412
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
3415
        self.type = type
3416

Peter Eastman's avatar
Peter Eastman committed
3417
3418
3419
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
3420

Peter Eastman's avatar
Peter Eastman committed
3421
        self.epsilonrule = epsilonrule
3422

Peter Eastman's avatar
Peter Eastman committed
3423
3424
3425
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
3426
3427
3428
3429
3430
3431
3432

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

    @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
3433
3434
3435
3436
3437
        #   <Vdw class="1" sigma="0.371" epsilon="0.46024" reduction="1.0" />
        #   <Vdw class="2" sigma="0.382" epsilon="0.422584" reduction="1.0" />

        generator = AmoebaVdwGenerator(element.attrib['type'], element.attrib['radiusrule'], element.attrib['radiustype'], element.attrib['radiussize'], element.attrib['epsilonrule'],
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
3438
        forceField._forces.append(generator)
3439
3440
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaVdwForce', 'Vdw', ('sigma', 'epsilon', 'reduction'))
        generator.params.parseDefinitions(element)
3441
3442
3443
3444
3445
3446
3447
3448
3449
        two_six = 1.122462048309372

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

    # Return a set containing the indices of particles bonded to particle with index=particleIndex

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
3450
    def getBondedParticleSet(particleIndex, data):
3451
3452
3453
3454
3455
3456

        bondedParticleSet = set()

        for bond in data.atomBonds[particleIndex]:
            atom1 = data.bonds[bond].atom1
            atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3457
3458
            if (atom1 != particleIndex):
                bondedParticleSet.add(atom1)
3459
            else:
Peter Eastman's avatar
Peter Eastman committed
3460
                bondedParticleSet.add(atom2)
3461
3462

        return bondedParticleSet
Justin MacCallum's avatar
Justin MacCallum committed
3463

3464
3465
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3468
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
3469
3470
3471
3472
3473
3474
3475
3476
3477
3478
3479
3480
        epsilonMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'HARMONIC':1, 'HHG':1}

        # 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.AmoebaVdwForce]
        if len(existing) == 0:
            force = mm.AmoebaVdwForce()
            sys.addForce(force)

            # sigma and epsilon combining rules

Peter Eastman's avatar
Peter Eastman committed
3481
            if ('sigmaCombiningRule' in args):
3482
                sigmaRule = args['sigmaCombiningRule'].upper()
Peter Eastman's avatar
Peter Eastman committed
3483
3484
                if (sigmaRule.upper() in sigmaMap):
                    force.setSigmaCombiningRule(sigmaRule.upper())
3485
                else:
Peter Eastman's avatar
Peter Eastman committed
3486
                    stringList = ' ' . join(str(x) for x in sigmaMap.keys())
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3487
                    raise ValueError( "AmoebaVdwGenerator: sigma combining rule %s not recognized; valid values are %s; using default." % (sigmaRule, stringList) )
3488
            else:
Peter Eastman's avatar
Peter Eastman committed
3489
                force.setSigmaCombiningRule(self.radiusrule)
3490

Peter Eastman's avatar
Peter Eastman committed
3491
            if ('epsilonCombiningRule' in args):
3492
                epsilonRule = args['epsilonCombiningRule'].upper()
Peter Eastman's avatar
Peter Eastman committed
3493
3494
                if (epsilonRule.upper() in epsilonMap):
                    force.setEpsilonCombiningRule(epsilonRule.upper())
3495
                else:
Peter Eastman's avatar
Peter Eastman committed
3496
                    stringList = ' ' . join(str(x) for x in epsilonMap.keys())
Justin MacCallum's avatar
Justin MacCallum committed
3497
                    raise ValueError( "AmoebaVdwGenerator: epsilon combining rule %s not recognized; valid values are %s; using default." % (epsilonRule, stringList) )
3498
            else:
Peter Eastman's avatar
Peter Eastman committed
3499
                force.setEpsilonCombiningRule(self.epsilonrule)
Justin MacCallum's avatar
Justin MacCallum committed
3500

3501
3502
            # cutoff

Peter Eastman's avatar
Peter Eastman committed
3503
            if ('vdwCutoff' in args):
3504
                force.setCutoff(args['vdwCutoff'])
3505
            else:
Peter Eastman's avatar
Peter Eastman committed
3506
                force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3507

3508
3509
3510
            # dispersion correction

            if ('useDispersionCorrection' in args):
3511
                force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
3512

Peter Eastman's avatar
Peter Eastman committed
3513
            if (nonbondedMethod == PME):
3514
                force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
Justin MacCallum's avatar
Justin MacCallum committed
3515

3516
3517
3518
3519
3520
        else:
            force = existing[0]

        # add particles to force

3521
3522
3523
3524
3525
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
3526
        for (i, atom) in enumerate(data.atoms):
3527
3528
3529
3530
3531
3532
3533
3534
3535
3536
            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
3537

3538
            force.addParticle(ivIndex, values[0]*sigmaScale, values[1], values[2])
3539
3540
3541
3542
3543
3544
3545
3546
3547

        # 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

        bondedParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3548
            bondedParticleSets.append(AmoebaVdwGenerator.getBondedParticleSet(i, data))
3549
3550

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

3552
3553
3554
3555
3556
3557
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
3558
            if (self.vdw13Scale == 0.0):
3559
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
3560
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
3561
3562
3563
3564
3565

            # self

            exclusionSet.add(i)

3566
            force.setParticleExclusions(i, tuple(exclusionSet))
3567
3568
3569
3570
3571

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

3572
## @private
3573
class AmoebaMultipoleGenerator(object):
3574
3575
3576
3577

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

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

3579
3580
3581
3582
3583
3584
    #=============================================================================================

    def __init__(self, forceField,
                       direct11Scale, direct12Scale, direct13Scale, direct14Scale,
                       mpole12Scale,  mpole13Scale,  mpole14Scale,  mpole15Scale,
                       mutual11Scale, mutual12Scale, mutual13Scale, mutual14Scale,
Peter Eastman's avatar
Peter Eastman committed
3585
                       polar12Scale,  polar13Scale,  polar14Scale,  polar15Scale):
3586

Peter Eastman's avatar
Peter Eastman committed
3587
        self.forceField = forceField
3588

Justin MacCallum's avatar
Justin MacCallum committed
3589
3590
3591
3592
        self.direct11Scale = direct11Scale
        self.direct12Scale = direct12Scale
        self.direct13Scale = direct13Scale
        self.direct14Scale = direct14Scale
3593

Justin MacCallum's avatar
Justin MacCallum committed
3594
3595
3596
3597
        self.mpole12Scale = mpole12Scale
        self.mpole13Scale = mpole13Scale
        self.mpole14Scale = mpole14Scale
        self.mpole15Scale = mpole15Scale
3598

Justin MacCallum's avatar
Justin MacCallum committed
3599
3600
3601
3602
        self.mutual11Scale = mutual11Scale
        self.mutual12Scale = mutual12Scale
        self.mutual13Scale = mutual13Scale
        self.mutual14Scale = mutual14Scale
3603

Justin MacCallum's avatar
Justin MacCallum committed
3604
3605
3606
3607
        self.polar12Scale = polar12Scale
        self.polar13Scale = polar13Scale
        self.polar14Scale = polar14Scale
        self.polar15Scale = polar15Scale
3608

Peter Eastman's avatar
Peter Eastman committed
3609
        self.typeMap = {}
3610
3611
3612
3613
3614
3615

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
3616
    def setAxisType(kIndices):
3617
3618
3619

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
3620
3621
3622
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
3623
                else:
Peter Eastman's avatar
Peter Eastman committed
3624
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
3625

Peter Eastman's avatar
Peter Eastman committed
3626
3627
                if (kIndicesLen > 2):
                    kx = kIndices[2]
3628
                else:
Peter Eastman's avatar
Peter Eastman committed
3629
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
3630

Peter Eastman's avatar
Peter Eastman committed
3631
3632
                if (kIndicesLen > 1):
                    kz = kIndices[1]
3633
                else:
Peter Eastman's avatar
Peter Eastman committed
3634
                    kz = 0
3635

Peter Eastman's avatar
Peter Eastman committed
3636
3637
                while(len(kIndices) < 4):
                    kIndices.append(0)
3638
3639

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
3640
                if (kz == 0):
3641
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
3642
                if (kz != 0 and kx == 0):
3643
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
3644
                if (kz < 0 or kx < 0):
3645
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
3646
                if (kx < 0 and ky < 0):
3647
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
3648
                if (kz < 0 and kx < 0 and ky  < 0):
3649
3650
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
3651
3652
3653
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
3654
3655
3656
3657
3658
3659
3660
3661

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
3662
        #   <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"  >
3663
3664
3665
        # <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"  />

Peter Eastman's avatar
Peter Eastman committed
3666
        generator = AmoebaMultipoleGenerator(forceField,
3667
3668
3669
3670
3671
3672
3673
3674
3675
3676
3677
3678
3679
3680
3681
3682
3683
3684
                                              element.attrib['direct11Scale'],
                                              element.attrib['direct12Scale'],
                                              element.attrib['direct13Scale'],
                                              element.attrib['direct14Scale'],

                                              element.attrib['mpole12Scale'],
                                              element.attrib['mpole13Scale'],
                                              element.attrib['mpole14Scale'],
                                              element.attrib['mpole15Scale'],

                                              element.attrib['mutual11Scale'],
                                              element.attrib['mutual12Scale'],
                                              element.attrib['mutual13Scale'],
                                              element.attrib['mutual14Scale'],

                                              element.attrib['polar12Scale'],
                                              element.attrib['polar13Scale'],
                                              element.attrib['polar14Scale'],
Peter Eastman's avatar
Peter Eastman committed
3685
                                              element.attrib['polar15Scale'])
3686
3687
3688
3689
3690
3691
3692
3693



        forceField._forces.append(generator)

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

        for atom in element.findall('Multipole'):
3694
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
3695
            if None not in types:
3696
3697
3698
3699
3700
3701
3702
3703

                # 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
3704
3705
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
3706
                    except:
3707
3708
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
3709
                # set axis type based on k-Indices
3710

Peter Eastman's avatar
Peter Eastman committed
3711
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
3712
3713
3714

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
3717
3718
3719
3720
3721
3722
3723
3724
3725
3726
3727
3728
3729
                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']))
3730
3731

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
3732
                    if (t not in generator.typeMap):
3733
3734
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
3735
3736
3737
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
3738
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
3739
3740
3741
3742
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
3743

3744
            else:
Peter Eastman's avatar
Peter Eastman committed
3745
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
3746
3747
                raise ValueError(outputString)

3748
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
3749

3750
        for atom in element.findall('Polarize'):
3751
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
3752
            if None not in types:
3753

Peter Eastman's avatar
Peter Eastman committed
3754
3755
3756
3757
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
3758
3759
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
3760
                    pdamp = pow(polarizability, 1.0/6.0)
3761

Peter Eastman's avatar
Peter Eastman committed
3762
3763
                pgrpMap = dict()
                for index in range(1, 7):
3764
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
3765
                    if (pgrp in atom.attrib):
3766
3767
3768
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
3769
3770
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
3771
                        raise ValueError(outputString)
3772
3773
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
3774
3775
3776
3777
3778
3779
3780
                        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
3781
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
3782
3783
3784
3785
3786
                                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
3787
3788
                            raise ValueError(outputString)

3789
            else:
Peter Eastman's avatar
Peter Eastman committed
3790
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
3791
3792
                raise ValueError(outputString)

3793
3794
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3795
    def setPolarGroups(self, data, bonded12ParticleSets, force):
3796
3797
3798
3799
3800

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
3801
3802
3803
3804
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
3805
3806
3807
3808
            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
3809
3810
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
3811
3812
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
3813

3814
3815
3816
3817
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
3818
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
3819
3820
                continue

Peter Eastman's avatar
Peter Eastman committed
3821
3822
            group = set()
            visited = set()
3823
3824
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
3825
3826
3827
3828
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
3829
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
3830
3831
                if (nextAtom not in visited):
                   visited.add(nextAtom)
3832
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
3833
3834
3835
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
3836
3837
3838

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
3839
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
3840
3841

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3842
3843
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
3844
3845
3846
3847
3848

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
3849
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
3850
3851
                continue

Peter Eastman's avatar
Peter Eastman committed
3852
            pgrp11 = set(atom.polarizationGroupSet[0])
3853
3854
3855
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3856
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
3857
3858
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3859
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
3860

3861
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3862
3863
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
3864
3865
3866
3867
3868

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
3869
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
3870
3871
                continue

Peter Eastman's avatar
Peter Eastman committed
3872
3873
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
3874
3875
3876
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3877
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
3878
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
3879
            pgrp13 = pgrp13 - set(pgrp11)
3880
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3881
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
3882

3883
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3884
3885
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
3886
3887
3888
3889
3890

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
3891
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
3892
3893
                continue

Peter Eastman's avatar
Peter Eastman committed
3894
3895
3896
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
3897
3898
3899
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3900
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
3901
3902
3903

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
3904
            pgrp14 = pgrp14 - set(pgrp11)
3905
3906

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

3909
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3910
3911
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
3912
3913
3914

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

Peter Eastman's avatar
Peter Eastman committed
3915
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3916
3917
3918
3919
3920
3921
3922
3923
3924
3925
3926

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

        # 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.AmoebaMultipoleForce]
        if len(existing) == 0:
            force = mm.AmoebaMultipoleForce()
            sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
3927
            if (nonbondedMethod not in methodMap):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3928
                raise ValueError( "AmoebaMultipoleForce: input cutoff method not available." )
3929
            else:
Peter Eastman's avatar
Peter Eastman committed
3930
                force.setNonbondedMethod(methodMap[nonbondedMethod])
3931
3932
            force.setCutoffDistance(nonbondedCutoff)

Peter Eastman's avatar
Peter Eastman committed
3933
            if ('ewaldErrorTolerance' in args):
3934
3935
                force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))

Peter Eastman's avatar
Peter Eastman committed
3936
            if ('polarization' in args):
3937
                polarizationType = args['polarization']
Peter Eastman's avatar
Peter Eastman committed
3938
                if (polarizationType.lower() == 'direct'):
3939
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Direct)
3940
3941
                elif (polarizationType.lower() == 'extrapolated'):
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Extrapolated)
3942
3943
3944
                else:
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Mutual)

Peter Eastman's avatar
Peter Eastman committed
3945
3946
            if ('aEwald' in args):
                force.setAEwald(float(args['aEwald']))
3947

Peter Eastman's avatar
Peter Eastman committed
3948
            if ('pmeGridDimensions' in args):
3949
3950
                force.setPmeGridDimensions(args['pmeGridDimensions'])

Peter Eastman's avatar
Peter Eastman committed
3951
3952
            if ('mutualInducedMaxIterations' in args):
                force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
3953

Peter Eastman's avatar
Peter Eastman committed
3954
            if ('mutualInducedTargetEpsilon' in args):
3955
3956
3957
3958
3959
3960
                force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))

        else:
            force = existing[0]

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
3961
        # throw error if particle type not available
3962
3963
3964
3965
3966
3967
3968

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

        # 1-2

        bonded12ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3969
3970
3971
            bonded12ParticleSet = AmoebaVdwGenerator.getBondedParticleSet(i, data)
            bonded12ParticleSet = set(sorted(bonded12ParticleSet))
            bonded12ParticleSets.append(bonded12ParticleSet)
3972
3973
3974
3975
3976

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3977
            bonded13Set = set()
3978
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
3979
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
3980
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
3981
3982
3983
3984

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
3985
            selfSet = set()
3986
3987
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3988
3989
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
3990
3991
3992
3993
3994

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3995
3996
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
3997
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
3998
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
3999

4000
4001
4002
4003
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4004
            selfSet = set()
4005
4006
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4007
4008
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
4009
4010
4011
4012
4013

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4014
4015
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4016
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4017
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
4018
4019
4020
4021
4022
4023

            # 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
4024
            selfSet = set()
4025
4026
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4027
4028
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
4029
4030
4031
4032
4033

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

Peter Eastman's avatar
Peter Eastman committed
4034
4035
                multipoleList = self.typeMap[t]
                hit = 0
4036
4037
4038
4039
4040
4041
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4042
                    if (hit != 0):
4043
4044
                        break

Peter Eastman's avatar
Peter Eastman committed
4045
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4046
4047

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4048
4049
                    kx = kIndices[2]
                    ky = kIndices[3]
4050
4051
4052
4053

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

4055
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
4056
4057
4058
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4059
4060
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
4061
                       if (hit != 0):
4062
4063
4064
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
4065
4066
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
4067
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4068
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4069
4070
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4071
4072
4073
4074
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4075
4076
4077
4078
4079
4080
4081
4082
4083
4084
                                      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

4085
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4086
                                      hit = 1
4087
4088
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4089
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4090
4091
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4092
4093
4094
4095
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4096
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4097
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
4098

4099
4100
4101
4102
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4103
                    if (hit != 0):
4104
4105
                        break

Peter Eastman's avatar
Peter Eastman committed
4106
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4107
4108

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4109
4110
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
4111

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

4116
4117
4118
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
4119
4120
4121
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4122
4123
4124

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4125
                       if (hit != 0):
4126
4127
4128
                           break

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

Peter Eastman's avatar
Peter Eastman committed
4131
                       if (bondedAtomZType == kz):
4132
4133
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
4134
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4135
4136
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4137
4138
4139
4140
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4141
4142
4143
4144
4145
4146
4147
4148

                                      # 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

4149
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4150
                                      hit = 3
4151
4152
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
4153
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4154
4155
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4156
4157
4158
4159
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4160
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4161
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
4162

4163
4164
4165
4166
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4167
                    if (hit != 0):
4168
4169
                        break

Peter Eastman's avatar
Peter Eastman committed
4170
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4171
4172
4173
4174

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

Peter Eastman's avatar
Peter Eastman committed
4175
4176
4177
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4178
4179
4180

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4181
                        if (hit != 0):
4182
4183
4184
                            break

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

Peter Eastman's avatar
Peter Eastman committed
4187
4188
                        if (kx == 0 and kz == bondedAtomZType):
                            kz = bondedAtomZIndex
4189
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4190
                            hit = 5
4191
4192
4193
4194
4195

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4196
                    if (hit != 0):
4197
4198
                        break

Peter Eastman's avatar
Peter Eastman committed
4199
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4200
4201
4202

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
4203
4204
4205
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4206

Peter Eastman's avatar
Peter Eastman committed
4207
                    if (kz == 0):
4208
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4209
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
4210

4211
4212
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
4213
                if (hit != 0):
4214

Peter Eastman's avatar
Peter Eastman committed
4215
                    atom.multipoleDict = savedMultipoleDict
4216
                    atom.polarizationGroups = dict()
4217
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
4218
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
4219
                    if (atomIndex == newIndex):
4220
4221
4222
4223
                        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]))
4224
                    else:
4225
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
4226
                else:
Peter Eastman's avatar
Peter Eastman committed
4227
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
4228
            else:
Peter Eastman's avatar
Peter Eastman committed
4229
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
4230
4231
4232

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
4233
        self.setPolarGroups(data, bonded12ParticleSets, force)
4234
4235
4236
4237
4238

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

4239
## @private
4240
class AmoebaWcaDispersionGenerator(object):
4241
4242

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

4244
4245
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4248
4249
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
4250
4251
4252
4253
4254
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
4255
        self.shctd = shctd
4256
4257
4258
4259
4260
4261
4262
4263
4264

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

    @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
4265

Peter Eastman's avatar
Peter Eastman committed
4266
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
4267
4268
4269
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
Justin MacCallum's avatar
Justin MacCallum committed
4270
                                                  element.attrib['awater'],
4271
4272
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Justin MacCallum's avatar
Justin MacCallum committed
4273
                                                  element.attrib['shctd'])
4274
        forceField._forces.append(generator)
4275
4276
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaWcaDispersionForce', 'WcaDispersion', ('radius', 'epsilon'))
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
4277

4278
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
4279

Peter Eastman's avatar
Peter Eastman committed
4280
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4281
4282
4283
4284
4285
4286
4287
4288
4289
4290
4291
4292

        # 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
4293
        # throw error if particle type not available
4294

Peter Eastman's avatar
Peter Eastman committed
4295
4296
4297
4298
4299
4300
4301
4302
        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  ))
4303

4304
4305
4306
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
4307
4308
4309
4310
4311

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

4312
## @private
4313
class AmoebaGeneralizedKirkwoodGenerator(object):
4314
4315

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

4317
4318
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
4319
4320
4321
4322
4323
4324
4325
4326
4327
4328
4329
    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
4330
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
4331
4332
4333
4334
4335
4336
4337
4338
4339
4340
4341
4342
4343
4344
4345
4346
4347
4348
4349
4350
4351
4352
4353
4354
4355
        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
4356
4357
4358

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

Peter Eastman's avatar
Peter Eastman committed
4359
    def getObcShct(self, data, atomIndex):
4360

Peter Eastman's avatar
Peter Eastman committed
4361
        atom = data.atoms[atomIndex]
4362
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4363
        shct = -1.0
4364
4365

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
4366

Peter Eastman's avatar
Peter Eastman committed
4367
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
4368
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
4369
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
4370
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
4371
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
4372
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
4373
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
4374
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
4375
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
4376
4377
4378
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
4379
        elif (atomicNumber == 16):              # S(16)
4380
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
4381
        elif (atomicNumber == 26):              # Fe(26)
4382
4383
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
4384
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4385
            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
4386
4387

        return shct
4388
4389
4390

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

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

Peter Eastman's avatar
Peter Eastman committed
4393
        atom = data.atoms[atomIndex]
4394
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4395
        radius = -1.0
4396

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

Peter Eastman's avatar
Peter Eastman committed
4399
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
4400

Peter Eastman's avatar
Peter Eastman committed
4401
            if (len(bondedAtomIndices) < 1):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4402
                 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
4403

4404
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4405
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
4406

Peter Eastman's avatar
Peter Eastman committed
4407
            if (bondedAtomAtomicNumber == 7):
4408
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
4409
            if (bondedAtomAtomicNumber == 8):
4410
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
4411

Peter Eastman's avatar
Peter Eastman committed
4412
        elif (atomicNumber == 3):               # Li(3)
4413
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
4414
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
4415

4416
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
4417
            if (len(bondedAtomIndices) == 3):
4418
4419
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
4420
            elif (len(bondedAtomIndices) == 4):
4421
4422
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4423
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
4424
4425
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
4426
        elif (atomicNumber == 7):               # N(7)
4427
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
4428
        elif (atomicNumber == 8):               # O(8)
4429
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
4430
            if (len(bondedAtomIndices) == 2):
4431
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
4432
        elif (atomicNumber == 9):               # F(9)
4433
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
4434
        elif (atomicNumber == 10):
4435
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
4436
        elif (atomicNumber == 11):
4437
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
4438
        elif (atomicNumber == 12):
4439
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
4440
        elif (atomicNumber == 14):
4441
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
4442
        elif (atomicNumber == 15):              # P(15)
4443
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
4444
        elif (atomicNumber == 16):              # S(16)
4445
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
4446
        elif (atomicNumber == 17):
4447
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
4448
        elif (atomicNumber == 18):
4449
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
4450
        elif (atomicNumber == 19):
4451
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
4452
        elif (atomicNumber == 20):
4453
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
4454
        elif (atomicNumber == 35):
4455
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
4456
        elif (atomicNumber == 36):
4457
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
4458
        elif (atomicNumber == 37):
4459
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
4460
        elif (atomicNumber == 53):
4461
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
4462
        elif (atomicNumber == 54):
4463
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
4464
        elif (atomicNumber == 55):
4465
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
4466
        elif (atomicNumber == 56):
4467
4468
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
4469
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4470
4471
            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
4472

4473
4474
4475
4476
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
4479
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
4480
        atom = data.atoms[atomIndex]
4481
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
4482
        if (atomicNumber in bondiMap):
4483
4484
            radius = bondiMap[atomicNumber]
        else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4485
4486
            outputString = "Warning no Bondi radius for atom %s of %s %d using default value=%f" % (atom.name, atom.residue.name, atom.residue.index, radius)
            raise ValueError( outputString )
Justin MacCallum's avatar
Justin MacCallum committed
4487

4488
4489
4490
4491
4492
4493
4494
4495
4496
4497
        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
4498

Peter Eastman's avatar
Peter Eastman committed
4499
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
Justin MacCallum's avatar
Justin MacCallum committed
4500
4501
                                                        element.attrib['includeCavityTerm'],
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor'])
4502
4503
4504
        forceField._forces.append(generator)

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

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

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

4511
4512
4513
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
4514
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
4515
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
4516
        if (len(amoebaMultipoleForceList) > 0):
4517
4518
4519
4520
4521
            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
4522
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
4523
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
4524

4525
4526
4527
4528
4529
4530
4531
        # 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
4532

Peter Eastman's avatar
Peter Eastman committed
4533
4534
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
4535
            else:
Peter Eastman's avatar
Peter Eastman committed
4536
                force.setSolventDielectric(   float(self.solventDielectric))
4537

Peter Eastman's avatar
Peter Eastman committed
4538
4539
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
4540
            else:
Peter Eastman's avatar
Peter Eastman committed
4541
                force.setSoluteDielectric(    float(self.soluteDielectric))
4542

Peter Eastman's avatar
Peter Eastman committed
4543
4544
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
4545
            else:
Peter Eastman's avatar
Peter Eastman committed
4546
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
4547
4548
4549
4550
4551

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
4554
4555
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
4556
4557
4558
4559
4560

        # 1-2

        bonded12ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4561
4562
4563
            bonded12ParticleSet = AmoebaVdwGenerator.getBondedParticleSet(i, data)
            bonded12ParticleSet = set(sorted(bonded12ParticleSet))
            bonded12ParticleSets.append(bonded12ParticleSet)
4564
4565

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
4566
4567
4568
4569
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
4570
            else:
Peter Eastman's avatar
Peter Eastman committed
4571
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
4572
4573
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
4574
            force.addParticle(multipoleParameters[0], radius, shct)
4575
4576
4577
4578
4579

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

4580
## @private
4581
class AmoebaUreyBradleyGenerator(object):
4582
4583
4584
4585

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

4587
    def __init__(self):
4588

Peter Eastman's avatar
Peter Eastman committed
4589
4590
4591
        self.types1 = []
        self.types2 = []
        self.types3 = []
4592

Peter Eastman's avatar
Peter Eastman committed
4593
4594
        self.length = []
        self.k = []
4595
4596
4597
4598
4599
4600

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

    @staticmethod
    def parseElement(element, forceField):

4601
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
4602
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
4603

4604
        generator = AmoebaUreyBradleyGenerator()
4605
4606
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
4607
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
4608
            if None not in types:
4609
4610
4611
4612
4613
4614
4615
4616
4617
4618

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

                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
4619
                                    bond.attrib['class1'], bond.attrib['class2'], bond.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
4620
4621
                raise ValueError(outputString)

4622
4623
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
4626
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
4627
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
4628
4629

        if len(existing) == 0:
4630
            force = mm.HarmonicBondForce()
4631
4632
4633
4634
4635
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
Peter Eastman's avatar
Peter Eastman committed
4636
            if (isConstrained):
4637
4638
4639
4640
4641
4642
4643
4644
                continue
            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]
Peter Eastman's avatar
Peter Eastman committed
4645
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
4646
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
4647
4648
4649
4650
4651
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
4652
4653
4654


## @private
4655
class DrudeGenerator(object):
peastman's avatar
peastman committed
4656
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
4657

4658
4659
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
4660
4661
4662
4663
4664
4665
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
4666
4667
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
4668
4669
4670
4671
        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'):
4672
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
4673
4674
4675
4676
4677
4678
4679
4680
4681
4682
            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
4683

peastman's avatar
peastman committed
4684
4685
4686
4687
    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
4688

peastman's avatar
peastman committed
4689
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
4690

peastman's avatar
peastman committed
4691
4692
4693
4694
4695
4696
4697
4698
4699
4700
4701
4702
4703
4704
4705
4706
        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
4707
4708
                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
4709
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
4710

peastman's avatar
peastman committed
4711
4712
    def postprocessSystem(self, sys, data, args):
        # For every nonbonded exclusion between Drude particles, add a screened pair.
Justin MacCallum's avatar
Justin MacCallum committed
4713

peastman's avatar
peastman committed
4714
4715
4716
4717
4718
4719
4720
        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)
4721
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
4722
4723
4724
4725
4726
                # 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]
4727
4728
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
4729
4730
4731
4732
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

Justin MacCallum's avatar
Justin MacCallum committed
4733
parsers["DrudeForce"] = DrudeGenerator.parseElement
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
4734
4735

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