forcefield.py 195 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
from __future__ 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
            else:
                if typeAttrib not in attrib or attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
325
326
327
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
328
329
        return types

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

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

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

379
380
381
382
383
384
385
        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
386
387
            msg =  "Atom name '%s' not found in residue template '%s'.\n" % (atom_name, self.name)
            msg += "Possible atom names are: %s" % str(atomIndices.keys())
388
389
            raise ValueError(msg)

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

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

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

    class _BondData:
        """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
429

430
431
    class _VirtualSiteData:
        """Inner class used to encapsulate data about a virtual site."""
432
        def __init__(self, node, atomIndices):
433
434
435
            attrib = node.attrib
            self.type = attrib['type']
            if self.type == 'average2':
436
                numAtoms = 2
437
438
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
439
                numAtoms = 3
440
441
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
442
                numAtoms = 3
443
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
444
            elif self.type == 'localCoords':
445
                numAtoms = 3
446
447
448
449
                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'])]
450
451
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
452
453
454
455
456
457
            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)]
458
459
460
461
            if 'excludeWith' in attrib:
                self.excludeWith = int(attrib['excludeWith'])
            else:
                self.excludeWith = self.atoms[0]
462

463
464
465
466
467
468
469
470
    class _AtomType:
        """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

471
472
473
474
475
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
    class _AtomTypeParameters:
        """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))


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

        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]

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

Robert McGibbon's avatar
Robert McGibbon committed
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
        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
606
607
        """
        data = ForceField._SystemData()
608
        data.atoms = list(topology.atoms())
609
610
        for atom in data.atoms:
            data.excludeAtomWith.append([])
611
612

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

614
        for bond in topology.bonds():
615
            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
616
617

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

619
620
621
622
623
624
625
626
627
628
629
630
        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
631
        # 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
632

633
634
        for chain in topology.chains():
            for res in chain.residues():
635
                # Attempt to match one of the existing templates.
636
                [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
637
638
                if matches is None:
                    # No existing templates match.  Try any registered residue template generators.
639
                    for generator in self._templateGenerators:
640
                        if generator(self, res):
641
                            # This generator has registered a new residue template that should match.
642
                            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
643
644
645
646
647
648
649
650
                            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.
651
                if matches is None:
652
                    raise ValueError('No template found for residue %d (%s).  %s' % (res.index+1, res.name, _findMatchErrors(self, res)))
653
654

                # Store parameters for the matched residue template.
655
                matchAtoms = dict(zip(matches, res.atoms()))
656
657
                for atom, match in zip(res.atoms(), matches):
                    data.atomType[atom] = template.atoms[match].type
658
                    data.atomParameters[atom] = template.atoms[match].parameters
659
660
                    for site in template.virtualSites:
                        if match == site.index:
661
                            data.virtualSites[atom] = (site, [matchAtoms[i].index for i in site.atoms], matchAtoms[site.excludeWith].index)
662
663

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

665
666
        sys = mm.System()
        for atom in topology.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
667
            # Look up the atom type name, returning a helpful error message if it cannot be found.
668
669
670
671
            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
672
            # Look up the type name in the list of registered atom types, returning a helpful error message if it cannot be found.
673
674
675
676
            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
677
678

            # Add the particle to the OpenMM system.
679
            mass = self._atomTypes[typename].mass
680
            sys.addParticle(mass)
681

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
682
        # Adjust hydroten masses if requested.
683

684
685
686
687
688
689
690
691
        if hydrogenMass is not None:
            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
692

693
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
694

695
696
697
        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
698
        elif nonbondedMethod not in [NoCutoff, CutoffNonPeriodic]:
699
700
701
            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
702

703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
        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
718

719
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
720

721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
        uniquePropers = set()
        for angle in data.angles:
            for atom in bondedToAtom[angle[0]]:
                if atom != angle[1]:
                    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]]:
                if atom != angle[1]:
                    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
736

737
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
738

739
740
741
742
743
        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
744

745
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
746

747
748
749
750
751
752
753
754
755
756
757
758
759
760
        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
761

762
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
763

764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
        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
785

786
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
787

788
        for atom in data.virtualSites:
789
            (site, atoms, excludeWith) = data.virtualSites[atom]
790
            index = atom.index
791
            data.excludeAtomWith[excludeWith].append(index)
792
            if site.type == 'average2':
793
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(atoms[0], atoms[1], site.weights[0], site.weights[1]))
794
            elif site.type == 'average3':
795
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
796
            elif site.type == 'outOfPlane':
797
798
799
800
801
802
803
                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
804

805
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
806

807
808
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
809
810
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
811

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

peastman's avatar
peastman committed
814
815
816
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
817

818
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
819

820
        for script in self._scripts:
821
            exec(script, locals())
822
823
824
        return sys


825
826
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
827
828
    counts = {}
    for element in elements:
829
        if element in counts:
830
831
832
            counts[element] += 1
        else:
            counts[element] = 1
833
834
835
836
837
838
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
839
840
    sig = []
    for c in counts:
841
842
        if c is not None:
            sig.append((c, counts[c]))
843
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
844

845
    # Convert it to a string.
846
847

    s = ''
848
    for element, count in sig:
849
850
851
852
        s += element.symbol+str(count)
    return s


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

Robert McGibbon's avatar
Robert McGibbon committed
856
857
858
859
860
861
862
863
864
865
866
    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
    -------
867
868
869
    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
870
871
    """
    atoms = list(res.atoms())
872
873
    if len(atoms) != len(template.atoms):
        return None
874
875
    matches = len(atoms)*[0]
    hasMatch = len(atoms)*[False]
Justin MacCallum's avatar
Justin MacCallum committed
876

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

879
880
    renumberAtoms = {}
    for i in range(len(atoms)):
881
        renumberAtoms[atoms[i].index] = i
882
883
884
    bondedTo = []
    externalBonds = []
    for atom in atoms:
885
        bonds = [renumberAtoms[x] for x in bondedToAtom[atom.index] if x in renumberAtoms]
886
        bondedTo.append(bonds)
887
        externalBonds.append(len([x for x in bondedToAtom[atom.index] if x not in renumberAtoms]))
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908

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

909
910
911
912
913
914
915
916
917
918
    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
919
    name = atoms[position].name
920
921
    for i in range(len(atoms)):
        atom = template.atoms[i]
922
        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]:
923
            # See if the bonds for this identification are consistent
Justin MacCallum's avatar
Justin MacCallum committed
924

925
926
927
            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
928

929
930
931
932
933
934
935
936
                matches[position] = i
                hasMatch[i] = True
                if _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch, position+1):
                    return True
                hasMatch[i] = False
    return False


937
938
939
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()])
940
    numResidueAtoms = sum(residueCounts.values())
941
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
942

943
    # Loop over templates and see how closely each one might match.
944

945
946
947
948
949
950
    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])
951

952
        # Does the residue have any atoms that clearly aren't in the template?
953

954
955
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
956

957
        # If there are too many missing atoms, discard this template.
958

959
        numTemplateAtoms = sum(templateCounts.values())
960
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
961
962
963
964
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
965

966
967
        # 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.
968

969
970
971
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
972

973
        # Accept this as our new best match.
974

975
976
977
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
978
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
979

980
    # Return an appropriate error message.
981

982
    if numBestMatchAtoms == numResidueAtoms:
983
984
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
985
986
987
988
            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:
989
990
991
992
993
            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)
994
995
996
997
        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.'


998
999
1000
1001
1002
# 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.

1003
## @private
1004
1005
class HarmonicBondGenerator:
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1006

1007
1008
    def __init__(self, forcefield):
        self.ff = forcefield
1009
1010
1011
1012
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
1013

1014
1015
1016
1017
1018
1019
1020
    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
1021

1022
1023
    @staticmethod
    def parseElement(element, ff):
1024
1025
        generator = HarmonicBondGenerator(ff)
        ff.registerGenerator(generator)
1026
        for bond in element.findall('Bond'):
1027
            generator.registerBond(bond.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1028

1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
    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:
1046
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
1047
1048
1049
1050
1051
1052
1053
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


1054
## @private
1055
1056
class HarmonicAngleGenerator:
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1057

1058
1059
    def __init__(self, forcefield):
        self.ff = forcefield
1060
1061
1062
1063
1064
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
1065

1066
1067
1068
1069
1070
1071
1072
1073
1074
    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']))

1075
1076
    @staticmethod
    def parseElement(element, ff):
1077
1078
        generator = HarmonicAngleGenerator(ff)
        ff.registerGenerator(generator)
1079
        for angle in element.findall('Angle'):
1080
            generator.registerAngle(angle.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1081

1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
    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
1101

1102
1103
1104
1105
1106
1107
1108
1109
1110
                        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
1111

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

1114
1115
1116
1117
1118
                        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]))
1119
                                data.addConstraint(sys, angle[0], angle[2], length)
1120
1121
1122
1123
1124
1125
1126
                    elif self.k[i] != 0:
                        force.addAngle(angle[0], angle[1], angle[2], self.angle[i], self.k[i])
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


1127
## @private
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
class PeriodicTorsion:
    """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 = []

1140
## @private
1141
1142
class PeriodicTorsionGenerator:
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1143

1144
1145
    def __init__(self, forcefield):
        self.ff = forcefield
1146
1147
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
1148

1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
    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)

1159
1160
    @staticmethod
    def parseElement(element, ff):
1161
1162
        generator = PeriodicTorsionGenerator(ff)
        ff.registerGenerator(generator)
1163
        for torsion in element.findall('Proper'):
1164
            generator.registerProperTorsion(torsion.attrib)
1165
        for torsion in element.findall('Improper'):
1166
            generator.registerImproperTorsion(torsion.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1167

1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
    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]]]
            done = False
            for tordef in self.improper:
                if done:
                    break
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
                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:
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
                            # 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)
                            for i in range(len(tordef.phase)):
                                if tordef.k[i] != 0:
                                    force.addTorsion(a1, a2, torsion[0], torsion[t4[1]], tordef.periodicity[i], tordef.phase[i], tordef.k[i])
1228
1229
1230
1231
1232
1233
                            done = True
                            break

parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


1234
## @private
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
class RBTorsion:
    """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

1245
## @private
1246
1247
class RBTorsionGenerator:
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1248

1249
1250
    def __init__(self, forcefield):
        self.ff = forcefield
1251
1252
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
1253

1254
1255
    @staticmethod
    def parseElement(element, ff):
1256
1257
        generator = RBTorsionGenerator(ff)
        ff.registerGenerator(generator)
1258
        for torsion in element.findall('Proper'):
1259
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1260
            if None not in types:
1261
1262
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
1263
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1264
            if None not in types:
1265
                generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
Justin MacCallum's avatar
Justin MacCallum committed
1266

1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
    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]]]
            done = False
            for tordef in self.improper:
                if done:
                    break
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
                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:
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
                            if wildcard in (types1, types2, types3, types4):
                                # 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)
                                force.addTorsion(a1, a2, torsion[0], torsion[t4[1]], tordef.c[0], tordef.c[1], tordef.c[2], tordef.c[3], tordef.c[4], tordef.c[5])
                            else:
                                # There are no wildcards, so the order is unambiguous.
                                force.addTorsion(torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef.c[0], tordef.c[1], tordef.c[2], tordef.c[3], tordef.c[4], tordef.c[5])
1327
1328
1329
1330
1331
1332
                            done = True
                            break

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


1333
## @private
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
class CMAPTorsion:
    """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

1345
## @private
1346
1347
class CMAPTorsionGenerator:
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1348

1349
1350
    def __init__(self, forcefield):
        self.ff = forcefield
1351
1352
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
1353

1354
1355
    @staticmethod
    def parseElement(element, ff):
1356
1357
        generator = CMAPTorsionGenerator(ff)
        ff.registerGenerator(generator)
1358
1359
1360
1361
1362
1363
1364
        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'):
1365
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
1366
            if None not in types:
1367
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
1368

1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
    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
1379

1380
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
1381

1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
        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


1425
## @private
1426
1427
class NonbondedGenerator:
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1428

1429
1430
    def __init__(self, forcefield, coulomb14scale, lj14scale):
        self.ff = forcefield
1431
1432
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
1433
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
1434

1435
    def registerAtom(self, parameters):
1436
        self.params.registerAtom(parameters)
1437

1438
1439
1440
1441
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
        if len(existing) == 0:
1442
1443
            generator = NonbondedGenerator(ff, float(element.attrib['coulomb14scale']), float(element.attrib['lj14scale']))
            ff.registerGenerator(generator)
1444
1445
1446
1447
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
            if generator.coulomb14scale != float(element.attrib['coulomb14scale']) or generator.lj14scale != float(element.attrib['lj14scale']):
Justin MacCallum's avatar
Justin MacCallum committed
1448
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
1449
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
1450

1451
1452
1453
1454
1455
1456
1457
    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
1458
            raise ValueError('Illegal nonbonded method for NonbondedForce')
1459
1460
        force = mm.NonbondedForce()
        for atom in data.atoms:
1461
1462
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
1463
1464
1465
1466
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
1467
1468
        if 'useDispersionCorrection' in args:
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
peastman's avatar
peastman committed
1469
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
1470

peastman's avatar
peastman committed
1471
    def postprocessSystem(self, sys, data, args):
1472
        # Create exceptions based on bonds.
1473

1474
1475
1476
        bondIndices = []
        for bond in data.bonds:
            bondIndices.append((bond.atom1, bond.atom2))
1477
1478
1479

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

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

1486
1487
        # 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.
1488

1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
        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
1499
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
1500
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
1501
1502
1503
1504

parsers["NonbondedForce"] = NonbondedGenerator.parseElement


1505
## @private
1506
1507
class GBSAOBCGenerator:
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1508

1509
1510
    def __init__(self, forcefield):
        self.ff = forcefield
1511
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
1512

1513
    def registerAtom(self, parameters):
1514
        self.params.registerAtom(parameters)
1515

1516
1517
    @staticmethod
    def parseElement(element, ff):
1518
1519
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
1520
1521
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
1522
1523
1524
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
1525
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
1526

1527
1528
1529
1530
1531
    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
1532
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
1533
1534
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
1535
1536
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
1537
1538
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
1539
1540
1541
1542
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
1543
1544
        sys.addForce(force)

1545
1546
    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
1547

1548
1549
1550
1551
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

1552
1553
1554
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


1555
## @private
1556
1557
class CustomBondGenerator:
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1558

1559
1560
    def __init__(self, forcefield):
        self.ff = forcefield
1561
1562
1563
1564
1565
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
1566

1567
1568
    @staticmethod
    def parseElement(element, ff):
1569
1570
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
1571
1572
1573
1574
1575
1576
        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'):
1577
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
1578
            if None not in types:
1579
1580
1581
                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
1582

1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
    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


1603
## @private
1604
1605
class CustomAngleGenerator:
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1606

1607
1608
    def __init__(self, forcefield):
        self.ff = forcefield
1609
1610
1611
1612
1613
1614
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
1615

1616
1617
    @staticmethod
    def parseElement(element, ff):
1618
1619
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
1620
1621
1622
1623
1624
1625
        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'):
1626
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
1627
            if None not in types:
1628
1629
1630
1631
                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
1632

1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
    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


1655
## @private
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
class CustomTorsion:
    """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

1666
## @private
1667
1668
class CustomTorsionGenerator:
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1669

1670
1671
    def __init__(self, forcefield):
        self.ff = forcefield
1672
1673
1674
1675
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
1676

1677
1678
    @staticmethod
    def parseElement(element, ff):
1679
1680
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
1681
1682
1683
1684
1685
1686
        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'):
1687
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1688
            if None not in types:
1689
1690
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
1691
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1692
            if None not in types:
1693
                generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
Justin MacCallum's avatar
Justin MacCallum committed
1694

1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
    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]]]
            done = False
            for tordef in self.improper:
                if done:
                    break
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
                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:
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
                            if wildcard in (types1, types2, types3, types4):
                                # 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)
                                force.addTorsion(a1, a2, torsion[0], torsion[t4[1]], tordef.paramValues)
                            else:
                                # There are no wildcards, so the order is unambiguous.
                                force.addTorsion(torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef.paramValues)
1754
1755
1756
1757
1758
1759
                            done = True
                            break

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


1760
1761
1762
1763
## @private
class CustomNonbondedGenerator:
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

1764
1765
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
1766
1767
1768
1769
1770
1771
1772
1773
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
1774
1775
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
1776
1777
1778
1779
        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'])
1780
1781
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807

    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:
1808
1809
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
1810
1811
1812
1813
1814
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

1817
1818
1819
        bondIndices = []
        for bond in data.bonds:
            bondIndices.append((bond.atom1, bond.atom2))
1820
1821
1822

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

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

1829
1830
        # 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.
1831

1832
1833
1834
1835
1836
1837
1838
1839
1840
        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.
1841

1842
1843
1844
1845
1846
1847
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomNonbondedForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


1848
## @private
1849
1850
class CustomGBGenerator:
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1851

1852
1853
    def __init__(self, forcefield):
        self.ff = forcefield
1854
1855
1856
1857
1858
1859
1860
1861
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
1862
1863
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
1864
1865
1866
1867
        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'])
1868
1869
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
1870
1871
1872
1873
1874
1875
1876
1877
1878
        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()]
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
            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
1890

1891
1892
1893
1894
1895
    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
1896
            raise ValueError('Illegal nonbonded method for CustomGBForce')
1897
1898
1899
1900
1901
1902
1903
1904
1905
        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])
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
        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))
1919
        for atom in data.atoms:
1920
1921
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
1922
1923
1924
1925
1926
1927
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942

## @private
class CustomManyParticleGenerator:
    """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 = []
1943

1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
    @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(',')]))
1956
1957
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
1958
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
1985
1986

    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:
1987
1988
1989
            values = self.params.getAtomParameters(atom, data)
            type = int(self.params.getExtraParameters(atom, data)['filterType'])
            force.addParticle(values, type)
1990
1991
1992
1993
1994
1995
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
        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)))
2008

2009
2010
        # 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.
2011

2012
2013
2014
2015
2016
2017
2018
2019
2020
        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.
2021

2022
2023
2024
2025
2026
        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
2027
def getAtomPrint(data, atomIndex):
2028

Peter Eastman's avatar
Peter Eastman committed
2029
2030
2031
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
2032
    else:
Peter Eastman's avatar
Peter Eastman committed
2033
        returnString = "NA"
2034
2035
2036
2037
2038

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
2039
def countConstraint(data):
2040

Peter Eastman's avatar
Peter Eastman committed
2041
    bondCount = 0
2042
2043
2044
2045
2046
2047
2048
    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
2049
        if (isConstrained):
2050
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
2051

2052
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
2053

2054
## @private
2055
class AmoebaBondGenerator:
2056
2057
2058

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

2059
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
2060
2061

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

2063
2064
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
2065
2066
2067
2068
2069
2070
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2071

2072
2073
2074
2075
2076
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

2080
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
2081
2082
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
2083
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2084
            if None not in types:
2085
2086
2087
2088
2089
                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:
2090
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
2091
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
2092
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
2093
2094
                raise ValueError(outputString)

2095
2096
    #=============================================================================================

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

Mark Friedrichs's avatar
Mark Friedrichs committed
2099
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2100

Peter Eastman's avatar
Peter Eastman committed
2101
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2102
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
2103
        if len(existing) == 0:
2104
            force = mm.AmoebaBondForce()
2105
2106
2107
2108
            sys.addForce(force)
        else:
            force = existing[0]

2109
2110
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120

        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:
2121
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
2122
2123
2124
2125
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

2126
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
2127
2128
2129
2130

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

Peter Eastman's avatar
Peter Eastman committed
2132
def addAngleConstraint(angle, idealAngle, data, sys):
2133
2134

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

Peter Eastman's avatar
Peter Eastman committed
2136
2137
    bond1 = None
    bond2 = None
2138
2139
2140
2141
2142
2143
2144
    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
2145

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

2148
2149
2150
2151
        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
2152
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
2153
                data.addConstraint(sys, angle[0], angle[2], length)
2154
2155
2156
                return

#=============================================================================================
2157
## @private
2158
class AmoebaAngleGenerator:
2159
2160

    #=============================================================================================
2161
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
2162
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
2163

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

Peter Eastman's avatar
Peter Eastman committed
2166
2167
2168
2169
2170
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
2171

Peter Eastman's avatar
Peter Eastman committed
2172
2173
2174
        self.types1 = []
        self.types2 = []
        self.types3 = []
2175

Peter Eastman's avatar
Peter Eastman committed
2176
2177
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2178

2179
2180
2181
2182
2183
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

2187
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
2188
2189
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
2190
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2191
            if None not in types:
2192
2193
2194
2195
2196
2197

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
2198
                angleList.append(float(angle.attrib['angle1']))
2199
2200

                try:
Peter Eastman's avatar
Peter Eastman committed
2201
                    angleList.append(float(angle.attrib['angle2']))
2202
                    try:
Peter Eastman's avatar
Peter Eastman committed
2203
                        angleList.append(float(angle.attrib['angle3']))
2204
2205
2206
2207
2208
2209
2210
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
2211
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
2212
2213
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
2214
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
2215
2216
                raise ValueError(outputString)

2217
2218
2219
2220
    #=============================================================================================
    # 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
2221

Peter Eastman's avatar
Peter Eastman committed
2222
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2223
2224
2225
2226
2227
2228
        pass

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

Peter Eastman's avatar
Peter Eastman committed
2230
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
2231
2232
2233
2234

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2235
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
2236
2237

        if len(existing) == 0:
2238
            force = mm.AmoebaAngleForce()
2239
2240
2241
2242
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2243
        # set scalars
2244

2245
2246
2247
2248
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
2249
2250

        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
2251
2252
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
2253

Peter Eastman's avatar
Peter Eastman committed
2254
2255
2256
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
2257
2258
2259
2260
2261
2262
2263
            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]
2264
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
2265
                    elif self.k[i] != 0:
Peter Eastman's avatar
Peter Eastman committed
2266
2267
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
2268
2269
2270
2271
2272
2273
                            # 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
2274
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
2275
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
2276
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
2277
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
2278
                            if (numberOfHydrogens < lenAngle):
2279
2280
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
2281
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
2282
                                raise ValueError(outputString)
2283
2284
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
2285

2286
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
2287
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
2288
2289
2290
2291
2292
2293
                    break

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

Peter Eastman's avatar
Peter Eastman committed
2295
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
2296
2297
2298
2299

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2300
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
2301
2302

        if len(existing) == 0:
2303
            force = mm.AmoebaInPlaneAngleForce()
2304
2305
2306
2307
2308
2309
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

2310
2311
2312
2313
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
2314
2315

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

Peter Eastman's avatar
Peter Eastman committed
2317
2318
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
2319

Peter Eastman's avatar
Peter Eastman committed
2320
2321
2322
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
2323
2324
2325
2326
2327
2328
2329
2330
2331

            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
2332
                    if (isConstrained and self.k[i] != 0.0):
2333
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
2334
2335
2336
2337
                    else:
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

2338
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
2339
2340
2341

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
2342
2343
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
2344
2345
#=============================================================================================

2346
## @private
2347
2348
2349
2350
2351
class AmoebaOutOfPlaneBendGenerator:

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

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

2353
2354
2355
2356
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
2357
2358
2359
2360
2361
2362
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
2363

Peter Eastman's avatar
Peter Eastman committed
2364
2365
2366
2367
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
2368

Peter Eastman's avatar
Peter Eastman committed
2369
        self.ks = []
2370
2371
2372
2373
2374

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

2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
    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
2401

2402
2403
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
2404
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
2405
2406
2407
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
2408
                                                   float(element.attrib['opbend-sextic']))
2409
2410
2411
2412

        forceField._forces.append(generator)

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

Peter Eastman's avatar
Peter Eastman committed
2416
2417
2418
2419
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
2420

Peter Eastman's avatar
Peter Eastman committed
2421
                generator.ks.append(float(angle.attrib['k']))
2422
2423

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2424
2425
                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
2426
2427
                raise ValueError(outputString)

2428
2429
2430
2431
2432
2433
    #=============================================================================================
    # 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
2434

Peter Eastman's avatar
Peter Eastman committed
2435
    def getMiddleAtom(self, angle, data):
2436
2437
2438

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
2439
        middleAtom = -1
Justin MacCallum's avatar
Justin MacCallum committed
2440
        for atomIndex in angle:
Peter Eastman's avatar
Peter Eastman committed
2441
            isMiddle = 0
2442
2443
2444
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
2445
                if (atom1 != atomIndex):
2446
2447
2448
                    partner = atom1
                else:
                    partner = atom2
Justin MacCallum's avatar
Justin MacCallum committed
2449
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]):
2450
2451
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
2452
            if (isMiddle == 2):
2453
2454
2455
2456
2457
                return atomIndex
        return -1

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

Peter Eastman's avatar
Peter Eastman committed
2458
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471

        # 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
2472
2473
2474
2475
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
2476
2477
2478
2479

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

Peter Eastman's avatar
Peter Eastman committed
2480
        skipAtoms = dict()
2481
2482
2483
2484

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

Peter Eastman's avatar
Peter Eastman committed
2485
2486
        inPlaneAngles = []
        nonInPlaneAngles = []
2487
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
2488
        idealAngles = []*len(data.angles)
2489
2490
2491

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

Peter Eastman's avatar
Peter Eastman committed
2492
2493
2494
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
2495
2496
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
2497
                middleType = -1
2498
2499
                middleCovalency = -1

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

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

Peter Eastman's avatar
Peter Eastman committed
2508
2509
2510
2511
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
2512
2513
2514
2515

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
2516
                    if (atom1 != middleAtom):
2517
2518
2519
2520
2521
2522
2523
2524
                        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
2525
2526
2527
2528
2529
                        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
2530

Peter Eastman's avatar
Peter Eastman committed
2531
                if (len(partners) == 3):
2532

Peter Eastman's avatar
Peter Eastman committed
2533
2534
2535
                    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])
2536

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

2539
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
2540
2541
2542
2543
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2544
2545
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
2546
2547
2548
2549
2550
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
2551
2552
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
2553
                    angleDict['isConstrained'] = 0
2554

Peter Eastman's avatar
Peter Eastman committed
2555
2556
2557
2558
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
2559
2560

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
2561
2562
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
2563

Peter Eastman's avatar
Peter Eastman committed
2564
                    inPlaneAngles.append(angleDict)
2565
2566

                else:
Peter Eastman's avatar
Peter Eastman committed
2567
2568
2569
2570
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
2571
            else:
Peter Eastman's avatar
Peter Eastman committed
2572
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
2573

Peter Eastman's avatar
Peter Eastman committed
2574
                    partnerSet = skipAtoms[middleAtom]
Justin MacCallum's avatar
Justin MacCallum committed
2575

Peter Eastman's avatar
Peter Eastman committed
2576
                    angleDict = {}
2577

Peter Eastman's avatar
Peter Eastman committed
2578
2579
2580
2581
2582
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
2583

Peter Eastman's avatar
Peter Eastman committed
2584
                    angleDict['isConstrained'] = isConstrained
2585

Peter Eastman's avatar
Peter Eastman committed
2586
2587
2588
2589
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
2590
2591

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
2592
2593
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
2594

Peter Eastman's avatar
Peter Eastman committed
2595
                    inPlaneAngles.append(angleDict)
2596
2597

                else:
Peter Eastman's avatar
Peter Eastman committed
2598
2599
                    angleDict = {}
                    angleDict['angle'] = angle
2600
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
2601
                    nonInPlaneAngles.append(angleDict)
2602

2603
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
2604
2605

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
2606
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
2607
2608
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
2609
2610

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
2611
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
2612
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
2613
                    nonInPlaneAngles.append(angleDict)
2614
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
2615
2616
2617
2618
2619

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

2620
## @private
2621
2622
2623
2624
2625
2626
2627
2628
class AmoebaTorsionGenerator:

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
2629
        self.torsionUnit = torsionUnit
2630

Peter Eastman's avatar
Peter Eastman committed
2631
2632
2633
2634
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
2635

Peter Eastman's avatar
Peter Eastman committed
2636
2637
2638
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
2639

2640
2641
2642
2643
2644
2645
2646
2647
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
2649
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
2650
2651
2652
2653
2654
2655
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
2656
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2657
            if None not in types:
2658
2659
2660
2661
2662
2663
2664

                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
2665
2666
                    tInfo = []
                    suffix = str(ii)
2667
2668
2669
2670
2671
2672
                    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
2673
2674
2675
2676
2677
2678
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
2679
2680
2681

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
2682
2683
2684
2685
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
2686
2687
                raise ValueError(outputString)

2688
2689
2690
2691
2692
    #=============================================================================================

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2693
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
2694
        if len(existing) == 0:
2695
            force = mm.PeriodicTorsionForce()
2696
2697
2698
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2699

2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
        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
2716
                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):
2717
2718
2719
2720
2721
2722
                    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])
2723
2724
2725
2726
2727
2728
                    break

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

2729
## @private
2730
2731
2732
2733
2734
2735
2736
class AmoebaPiTorsionGenerator:

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

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

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

2738
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
2739
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
2740
2741
2742
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2743

2744
2745
2746
2747
2748
2749
2750
2751
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
2752
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
2753
2754
2755
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
2756
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
2757
            if None not in types:
2758
2759
2760
2761
2762
2763
                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
2764
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
2765
2766
                raise ValueError(outputString)

2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
    #=============================================================================================

    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
2783

2784
2785
            atom1 = bond.atom1
            atom2 = bond.atom2
Justin MacCallum's avatar
Justin MacCallum committed
2786

2787
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798

                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
2799
2800
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812

                       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
2813
                           if (bondedAtom1 != atom1):
2814
2815
2816
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
2817
2818
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
2819
                                   piTorsionAtom1 = b1
2820
2821
2822
2823
2824
2825
                               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
2826
                           if (bondedAtom1 != atom2):
2827
2828
2829
2830
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
2831
2832
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
2833
                                   piTorsionAtom5 = b1
2834
2835
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
2836

Peter Eastman's avatar
Peter Eastman committed
2837
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
2838
2839
2840
2841
2842

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

2843
## @private
2844
2845
2846
2847
2848
2849
class AmoebaTorsionTorsionGenerator:

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

Peter Eastman's avatar
Peter Eastman committed
2850
    def __init__(self):
2851

Peter Eastman's avatar
Peter Eastman committed
2852
2853
2854
2855
2856
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
2857

Peter Eastman's avatar
Peter Eastman committed
2858
        self.gridIndex = []
2859

Peter Eastman's avatar
Peter Eastman committed
2860
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
2861

2862
2863
2864
2865
2866
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
2867
        generator = AmoebaTorsionTorsionGenerator()
2868
2869
2870
2871
2872
2873
2874
        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'):
2875
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
2876
            if None not in types:
2877
2878
2879
2880
2881
2882
2883

                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
2884
2885
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
2886
2887
2888
2889
2890
2891
2892
2893
2894
                    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
2895
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
2896
2897
                raise ValueError(outputString)

2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
        # 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
2908
2909
2910
2911
2912
2913
        #     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
2914
2915

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

Peter Eastman's avatar
Peter Eastman committed
2919
2920
2921
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
2922

Peter Eastman's avatar
Peter Eastman committed
2923
2924
            grid = []
            gridCol = []
2925
2926
2927
2928
2929

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
2930
2931
2932
2933
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
2934
                if 'fx' in gridEntry.attrib:
2935
2936
2937
                    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
2938
                gridCol.append(gridRow)
2939
2940

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
2941
2942
2943
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
2944
2945
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
2946

Peter Eastman's avatar
Peter Eastman committed
2947
2948
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
2949
            else:
Peter Eastman's avatar
Peter Eastman committed
2950
2951
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
2952
2953
2954
2955
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
2956
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
2957
2958
2959
2960
2961
2962
2963
2964

        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
2965
        if (len(data.atomBonds[atomC]) == 4):
2966
2967
2968
2969
2970
            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
2971
2972
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
2973
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
2974
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
2975
2976
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
2977
2978
                if (hit > -1):
                    if (atomE == -1):
2979
2980
2981
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
2982

2983
2984
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
2985
2986
            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
2987
                raise ValueError(outputString)
2988
2989
2990
2991
2992

            # 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
2993
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
2994
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
2995
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
2996
                chiralAtomIndex = atomF
2997

Peter Eastman's avatar
Peter Eastman committed
2998
2999
3000
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
3001
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
3002
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
3003
                chiralAtomIndex = massF
3004
3005
3006
3007

        return chiralAtomIndex

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

3009
3010
3011
3012
3013
3014
3015
3016
3017
3018
3019
3020
3021
    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
3022
3023
            # search for bitorsions; based on TINKER subroutine bitors()

3024
3025
3026
3027
3028
3029
3030
            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
3031
                if (bondedAtom1 != ib):
3032
3033
3034
3035
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3036
                if (ia != ic and ia != id):
3037
3038
3039
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
3040
                        if (bondedAtom1 != id):
3041
3042
3043
3044
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3045
                        if (ie != ic and ie != ib and ie != ia):
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
3062
3063
3064
3065

                            # 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
3066
3067
3068
                                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])
3069
3070
3071

                                # match in reverse order

3072
                                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
3073
3074
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
3075
3076
3077
3078

        # set grids

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

3081
3082
3083
3084
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

3085
## @private
3086
class AmoebaStretchBendGenerator:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3087
3088

    #=============================================================================================
3089
3090
3091
3092
3093
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
3094
3095
3096
        self.types1 = []
        self.types2 = []
        self.types3 = []
3097

Peter Eastman's avatar
Peter Eastman committed
3098
3099
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
3100

3101
3102
3103
3104
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
3105
        generator = AmoebaStretchBendGenerator()
3106
3107
3108
3109
3110
3111
3112
        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'):
3113
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
3114
            if None not in types:
3115
3116
3117
3118
3119
3120
3121
3122
3123
3124
3125
3126

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

3130
3131
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
3132
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
3133
3134
    # 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
3135
3136
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
3137
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
3138
3139
3140

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

Peter Eastman's avatar
Peter Eastman committed
3141
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3142
3143
3144
3145
3146
3147
3148
3149
        pass

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

    # Note: request for constrained bonds is ignored.

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

3150
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162

        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
3163
            if ('isConstrained' in angleDict):
3164
3165
3166
3167
3168
3169
3170
3171
                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
3172
            radian = 57.2957795130
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
            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
3183
                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
3184
3185
3186
3187
3188
3189
3190
3191
3192
3193
3194
3195
3196
3197
3198
                    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
3199

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3202
                    if ('idealAngle' not in angleDict):
3203

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3204
3205
3206
3207
3208
                       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
3209
                       raise ValueError(outputString)
3210

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3211
3212
3213
3214
3215
3216
3217
                    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
3218
                       raise ValueError(outputString)
3219

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3223
                    break
3224
3225
3226
3227
3228

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

3229
## @private
3230
3231
3232
class AmoebaVdwGenerator:

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

3234
3235
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
3238
        self.type = type
3239

Peter Eastman's avatar
Peter Eastman committed
3240
3241
3242
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
3243

Peter Eastman's avatar
Peter Eastman committed
3244
        self.epsilonrule = epsilonrule
3245

Peter Eastman's avatar
Peter Eastman committed
3246
3247
3248
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
3249
3250
3251
3252
3253
3254
3255

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

    @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
3256
3257
3258
3259
3260
        #   <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']))
3261
        forceField._forces.append(generator)
3262
3263
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaVdwForce', 'Vdw', ('sigma', 'epsilon', 'reduction'))
        generator.params.parseDefinitions(element)
3264
3265
3266
3267
3268
3269
3270
3271
3272
        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
3273
    def getBondedParticleSet(particleIndex, data):
3274
3275
3276
3277
3278
3279

        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
3280
3281
            if (atom1 != particleIndex):
                bondedParticleSet.add(atom1)
3282
            else:
Peter Eastman's avatar
Peter Eastman committed
3283
                bondedParticleSet.add(atom2)
3284
3285

        return bondedParticleSet
Justin MacCallum's avatar
Justin MacCallum committed
3286

3287
3288
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3291
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
3292
3293
3294
3295
3296
3297
3298
3299
3300
3301
3302
3303
        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
3304
            if ('sigmaCombiningRule' in args):
3305
                sigmaRule = args['sigmaCombiningRule'].upper()
Peter Eastman's avatar
Peter Eastman committed
3306
3307
                if (sigmaRule.upper() in sigmaMap):
                    force.setSigmaCombiningRule(sigmaRule.upper())
3308
                else:
Peter Eastman's avatar
Peter Eastman committed
3309
                    stringList = ' ' . join(str(x) for x in sigmaMap.keys())
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3310
                    raise ValueError( "AmoebaVdwGenerator: sigma combining rule %s not recognized; valid values are %s; using default." % (sigmaRule, stringList) )
3311
            else:
Peter Eastman's avatar
Peter Eastman committed
3312
                force.setSigmaCombiningRule(self.radiusrule)
3313

Peter Eastman's avatar
Peter Eastman committed
3314
            if ('epsilonCombiningRule' in args):
3315
                epsilonRule = args['epsilonCombiningRule'].upper()
Peter Eastman's avatar
Peter Eastman committed
3316
3317
                if (epsilonRule.upper() in epsilonMap):
                    force.setEpsilonCombiningRule(epsilonRule.upper())
3318
                else:
Peter Eastman's avatar
Peter Eastman committed
3319
                    stringList = ' ' . join(str(x) for x in epsilonMap.keys())
Justin MacCallum's avatar
Justin MacCallum committed
3320
                    raise ValueError( "AmoebaVdwGenerator: epsilon combining rule %s not recognized; valid values are %s; using default." % (epsilonRule, stringList) )
3321
            else:
Peter Eastman's avatar
Peter Eastman committed
3322
                force.setEpsilonCombiningRule(self.epsilonrule)
Justin MacCallum's avatar
Justin MacCallum committed
3323

3324
3325
            # cutoff

Peter Eastman's avatar
Peter Eastman committed
3326
            if ('vdwCutoff' in args):
3327
                force.setCutoff(args['vdwCutoff'])
3328
            else:
Peter Eastman's avatar
Peter Eastman committed
3329
                force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3330

3331
3332
3333
            # dispersion correction

            if ('useDispersionCorrection' in args):
3334
                force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
3335

Peter Eastman's avatar
Peter Eastman committed
3336
            if (nonbondedMethod == PME):
3337
                force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
Justin MacCallum's avatar
Justin MacCallum committed
3338

3339
3340
3341
3342
3343
        else:
            force = existing[0]

        # add particles to force

3344
3345
3346
3347
3348
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
3349
        for (i, atom) in enumerate(data.atoms):
3350
3351
3352
3353
3354
3355
3356
3357
3358
3359
            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
3360

3361
            force.addParticle(ivIndex, values[0]*sigmaScale, values[1], values[2])
3362
3363
3364
3365
3366
3367
3368
3369
3370

        # 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
3371
            bondedParticleSets.append(AmoebaVdwGenerator.getBondedParticleSet(i, data))
3372
3373

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

3375
3376
3377
3378
3379
3380
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
3381
            if (self.vdw13Scale == 0.0):
3382
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
3383
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
3384
3385
3386
3387
3388

            # self

            exclusionSet.add(i)

3389
            force.setParticleExclusions(i, tuple(exclusionSet))
3390
3391
3392
3393
3394

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

3395
## @private
3396
3397
3398
3399
3400
class AmoebaMultipoleGenerator:

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

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

3402
3403
3404
3405
3406
3407
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3410
        self.forceField = forceField
3411

Justin MacCallum's avatar
Justin MacCallum committed
3412
3413
3414
3415
        self.direct11Scale = direct11Scale
        self.direct12Scale = direct12Scale
        self.direct13Scale = direct13Scale
        self.direct14Scale = direct14Scale
3416

Justin MacCallum's avatar
Justin MacCallum committed
3417
3418
3419
3420
        self.mpole12Scale = mpole12Scale
        self.mpole13Scale = mpole13Scale
        self.mpole14Scale = mpole14Scale
        self.mpole15Scale = mpole15Scale
3421

Justin MacCallum's avatar
Justin MacCallum committed
3422
3423
3424
3425
        self.mutual11Scale = mutual11Scale
        self.mutual12Scale = mutual12Scale
        self.mutual13Scale = mutual13Scale
        self.mutual14Scale = mutual14Scale
3426

Justin MacCallum's avatar
Justin MacCallum committed
3427
3428
3429
3430
        self.polar12Scale = polar12Scale
        self.polar13Scale = polar13Scale
        self.polar14Scale = polar14Scale
        self.polar15Scale = polar15Scale
3431

Peter Eastman's avatar
Peter Eastman committed
3432
        self.typeMap = {}
3433
3434
3435
3436
3437
3438

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
3439
    def setAxisType(kIndices):
3440
3441
3442

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
3443
3444
3445
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
3446
                else:
Peter Eastman's avatar
Peter Eastman committed
3447
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
3448

Peter Eastman's avatar
Peter Eastman committed
3449
3450
                if (kIndicesLen > 2):
                    kx = kIndices[2]
3451
                else:
Peter Eastman's avatar
Peter Eastman committed
3452
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
3453

Peter Eastman's avatar
Peter Eastman committed
3454
3455
                if (kIndicesLen > 1):
                    kz = kIndices[1]
3456
                else:
Peter Eastman's avatar
Peter Eastman committed
3457
                    kz = 0
3458

Peter Eastman's avatar
Peter Eastman committed
3459
3460
                while(len(kIndices) < 4):
                    kIndices.append(0)
3461
3462

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
3463
                if (kz == 0):
3464
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
3465
                if (kz != 0 and kx == 0):
3466
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
3467
                if (kz < 0 or kx < 0):
3468
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
3469
                if (kx < 0 and ky < 0):
3470
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
3471
                if (kz < 0 and kx < 0 and ky  < 0):
3472
3473
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
3474
3475
3476
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
3477
3478
3479
3480
3481
3482
3483
3484

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
3485
        #   <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"  >
3486
3487
3488
        # <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
3489
        generator = AmoebaMultipoleGenerator(forceField,
3490
3491
3492
3493
3494
3495
3496
3497
3498
3499
3500
3501
3502
3503
3504
3505
3506
3507
                                              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
3508
                                              element.attrib['polar15Scale'])
3509
3510
3511
3512
3513
3514
3515
3516



        forceField._forces.append(generator)

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

        for atom in element.findall('Multipole'):
3517
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
3518
            if None not in types:
3519
3520
3521
3522
3523
3524
3525
3526

                # 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
3527
3528
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
3529
                    except:
3530
3531
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
3532
                # set axis type based on k-Indices
3533

Peter Eastman's avatar
Peter Eastman committed
3534
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
3535
3536
3537

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
3540
3541
3542
3543
3544
3545
3546
3547
3548
3549
3550
3551
3552
                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']))
3553
3554

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
3555
                    if (t not in generator.typeMap):
3556
3557
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
3558
3559
3560
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
3561
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
3562
3563
3564
3565
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
3566

3567
            else:
Peter Eastman's avatar
Peter Eastman committed
3568
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
3569
3570
                raise ValueError(outputString)

3571
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
3572

3573
        for atom in element.findall('Polarize'):
3574
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
3575
            if None not in types:
3576

Peter Eastman's avatar
Peter Eastman committed
3577
3578
3579
3580
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
3581
3582
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
3583
                    pdamp = pow(polarizability, 1.0/6.0)
3584

Peter Eastman's avatar
Peter Eastman committed
3585
3586
                pgrpMap = dict()
                for index in range(1, 7):
3587
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
3588
                    if (pgrp in atom.attrib):
3589
3590
3591
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
3592
3593
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
3594
                        raise ValueError(outputString)
3595
3596
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
3597
3598
3599
3600
3601
3602
3603
                        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
3604
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
3605
3606
3607
3608
3609
                                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
3610
3611
                            raise ValueError(outputString)

3612
            else:
Peter Eastman's avatar
Peter Eastman committed
3613
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
3614
3615
                raise ValueError(outputString)

3616
3617
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3618
    def setPolarGroups(self, data, bonded12ParticleSets, force):
3619
3620
3621
3622
3623

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
3624
3625
3626
3627
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
3628
3629
3630
3631
            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
3632
3633
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
3634
3635
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
3636

3637
3638
3639
3640
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
3641
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
3642
3643
                continue

Peter Eastman's avatar
Peter Eastman committed
3644
3645
            group = set()
            visited = set()
3646
3647
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
3648
3649
3650
3651
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
3652
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
3653
3654
                if (nextAtom not in visited):
                   visited.add(nextAtom)
3655
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
3656
3657
3658
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
3659
3660
3661

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
3662
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
3663
3664

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3665
3666
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
3667
3668
3669
3670
3671

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
3672
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
3673
3674
                continue

Peter Eastman's avatar
Peter Eastman committed
3675
            pgrp11 = set(atom.polarizationGroupSet[0])
3676
3677
3678
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3679
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
3680
3681
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3682
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
3683

3684
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3685
3686
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
3687
3688
3689
3690
3691

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
3692
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
3693
3694
                continue

Peter Eastman's avatar
Peter Eastman committed
3695
3696
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
3697
3698
3699
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3700
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
3701
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
3702
            pgrp13 = pgrp13 - set(pgrp11)
3703
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3704
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
3705

3706
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3707
3708
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
3709
3710
3711
3712
3713

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
3714
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
3715
3716
                continue

Peter Eastman's avatar
Peter Eastman committed
3717
3718
3719
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
3720
3721
3722
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3723
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
3724
3725
3726

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
3727
            pgrp14 = pgrp14 - set(pgrp11)
3728
3729

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

3732
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3733
3734
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
3735
3736
3737

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

Peter Eastman's avatar
Peter Eastman committed
3738
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3739
3740
3741
3742
3743
3744
3745
3746
3747
3748
3749

        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
3750
            if (nonbondedMethod not in methodMap):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3751
                raise ValueError( "AmoebaMultipoleForce: input cutoff method not available." )
3752
            else:
Peter Eastman's avatar
Peter Eastman committed
3753
                force.setNonbondedMethod(methodMap[nonbondedMethod])
3754
3755
            force.setCutoffDistance(nonbondedCutoff)

Peter Eastman's avatar
Peter Eastman committed
3756
            if ('ewaldErrorTolerance' in args):
3757
3758
                force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))

Peter Eastman's avatar
Peter Eastman committed
3759
            if ('polarization' in args):
3760
                polarizationType = args['polarization']
Peter Eastman's avatar
Peter Eastman committed
3761
                if (polarizationType.lower() == 'direct'):
3762
3763
3764
3765
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Direct)
                else:
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Mutual)

Peter Eastman's avatar
Peter Eastman committed
3766
3767
            if ('aEwald' in args):
                force.setAEwald(float(args['aEwald']))
3768

Peter Eastman's avatar
Peter Eastman committed
3769
            if ('pmeGridDimensions' in args):
3770
3771
                force.setPmeGridDimensions(args['pmeGridDimensions'])

Peter Eastman's avatar
Peter Eastman committed
3772
3773
            if ('mutualInducedMaxIterations' in args):
                force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
3774

Peter Eastman's avatar
Peter Eastman committed
3775
            if ('mutualInducedTargetEpsilon' in args):
3776
3777
3778
3779
3780
3781
                force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))

        else:
            force = existing[0]

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
3782
        # throw error if particle type not available
3783
3784
3785
3786
3787
3788
3789

        # 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
3790
3791
3792
            bonded12ParticleSet = AmoebaVdwGenerator.getBondedParticleSet(i, data)
            bonded12ParticleSet = set(sorted(bonded12ParticleSet))
            bonded12ParticleSets.append(bonded12ParticleSet)
3793
3794
3795
3796
3797

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3798
            bonded13Set = set()
3799
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
3800
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
3801
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
3802
3803
3804
3805

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
3806
            selfSet = set()
3807
3808
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3809
3810
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
3811
3812
3813
3814
3815

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3816
3817
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
3818
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
3819
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
3820

3821
3822
3823
3824
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
3825
            selfSet = set()
3826
3827
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3828
3829
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
3830
3831
3832
3833
3834

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3835
3836
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
3837
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
3838
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
3839
3840
3841
3842
3843
3844

            # 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
3845
            selfSet = set()
3846
3847
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3848
3849
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
3850
3851
3852
3853
3854

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

Peter Eastman's avatar
Peter Eastman committed
3855
3856
                multipoleList = self.typeMap[t]
                hit = 0
3857
3858
3859
3860
3861
3862
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3863
                    if (hit != 0):
3864
3865
                        break

Peter Eastman's avatar
Peter Eastman committed
3866
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
3867
3868

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
3869
3870
                    kx = kIndices[2]
                    ky = kIndices[3]
3871
3872
3873
3874

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

3876
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
3877
3878
3879
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3880
3881
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
3882
                       if (hit != 0):
3883
3884
3885
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
3886
3887
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
3888
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
3889
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
3890
3891
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
3892
3893
3894
3895
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
3896
3897
3898
3899
3900
3901
3902
3903
3904
3905
                                      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

3906
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3907
                                      hit = 1
3908
3909
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
3910
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
3911
3912
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
3913
3914
3915
3916
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
3917
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3918
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
3919

3920
3921
3922
3923
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3924
                    if (hit != 0):
3925
3926
                        break

Peter Eastman's avatar
Peter Eastman committed
3927
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
3928
3929

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
3930
3931
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
3932

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

3937
3938
3939
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
3940
3941
3942
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3943
3944
3945

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
3946
                       if (hit != 0):
3947
3948
3949
                           break

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

Peter Eastman's avatar
Peter Eastman committed
3952
                       if (bondedAtomZType == kz):
3953
3954
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
3955
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
3956
3957
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
3958
3959
3960
3961
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
3962
3963
3964
3965
3966
3967
3968
3969

                                      # 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

3970
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3971
                                      hit = 3
3972
3973
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
3974
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
3975
3976
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
3977
3978
3979
3980
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
3981
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3982
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
3983

3984
3985
3986
3987
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3988
                    if (hit != 0):
3989
3990
                        break

Peter Eastman's avatar
Peter Eastman committed
3991
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
3992
3993
3994
3995

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

Peter Eastman's avatar
Peter Eastman committed
3996
3997
3998
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3999
4000
4001

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4002
                        if (hit != 0):
4003
4004
4005
                            break

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

Peter Eastman's avatar
Peter Eastman committed
4008
4009
                        if (kx == 0 and kz == bondedAtomZType):
                            kz = bondedAtomZIndex
4010
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4011
                            hit = 5
4012
4013
4014
4015
4016

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4017
                    if (hit != 0):
4018
4019
                        break

Peter Eastman's avatar
Peter Eastman committed
4020
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4021
4022
4023

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
4024
4025
4026
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4027

Peter Eastman's avatar
Peter Eastman committed
4028
                    if (kz == 0):
4029
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4030
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
4031

4032
4033
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
4034
                if (hit != 0):
4035

Peter Eastman's avatar
Peter Eastman committed
4036
                    atom.multipoleDict = savedMultipoleDict
4037
                    atom.polarizationGroups = dict()
4038
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
4039
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
4040
                    if (atomIndex == newIndex):
4041
4042
4043
4044
                        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]))
4045
                    else:
4046
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
4047
                else:
Peter Eastman's avatar
Peter Eastman committed
4048
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
4049
            else:
Peter Eastman's avatar
Peter Eastman committed
4050
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
4051
4052
4053

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
4054
        self.setPolarGroups(data, bonded12ParticleSets, force)
4055
4056
4057
4058
4059

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

4060
## @private
4061
4062
4063
class AmoebaWcaDispersionGenerator:

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

4065
4066
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4069
4070
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
4071
4072
4073
4074
4075
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
4076
        self.shctd = shctd
4077
4078
4079
4080
4081
4082
4083
4084
4085

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

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

Peter Eastman's avatar
Peter Eastman committed
4087
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
4088
4089
4090
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
Justin MacCallum's avatar
Justin MacCallum committed
4091
                                                  element.attrib['awater'],
4092
4093
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Justin MacCallum's avatar
Justin MacCallum committed
4094
                                                  element.attrib['shctd'])
4095
        forceField._forces.append(generator)
4096
4097
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaWcaDispersionForce', 'WcaDispersion', ('radius', 'epsilon'))
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
4098

4099
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
4100

Peter Eastman's avatar
Peter Eastman committed
4101
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4102
4103
4104
4105
4106
4107
4108
4109
4110
4111
4112
4113

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

Peter Eastman's avatar
Peter Eastman committed
4116
4117
4118
4119
4120
4121
4122
4123
        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  ))
4124

4125
4126
4127
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
4128
4129
4130
4131
4132

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

4133
## @private
4134
4135
4136
class AmoebaGeneralizedKirkwoodGenerator:

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

4138
4139
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
4140
4141
4142
4143
4144
4145
4146
4147
4148
4149
4150
    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
4151
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
4152
4153
4154
4155
4156
4157
4158
4159
4160
4161
4162
4163
4164
4165
4166
4167
4168
4169
4170
4171
4172
4173
4174
4175
4176
        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
4177
4178
4179

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

Peter Eastman's avatar
Peter Eastman committed
4180
    def getObcShct(self, data, atomIndex):
4181

Peter Eastman's avatar
Peter Eastman committed
4182
        atom = data.atoms[atomIndex]
4183
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4184
        shct = -1.0
4185
4186

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
4187

Peter Eastman's avatar
Peter Eastman committed
4188
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
4189
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
4190
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
4191
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
4192
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
4193
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
4194
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
4195
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
4196
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
4197
4198
4199
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
4200
        elif (atomicNumber == 16):              # S(16)
4201
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
4202
        elif (atomicNumber == 26):              # Fe(26)
4203
4204
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
4205
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4206
            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
4207
4208

        return shct
4209
4210
4211

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

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

Peter Eastman's avatar
Peter Eastman committed
4214
        atom = data.atoms[atomIndex]
4215
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4216
        radius = -1.0
4217

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

Peter Eastman's avatar
Peter Eastman committed
4220
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
4221

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

4225
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4226
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
4227

Peter Eastman's avatar
Peter Eastman committed
4228
            if (bondedAtomAtomicNumber == 7):
4229
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
4230
            if (bondedAtomAtomicNumber == 8):
4231
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
4232

Peter Eastman's avatar
Peter Eastman committed
4233
        elif (atomicNumber == 3):               # Li(3)
4234
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
4235
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
4236

4237
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
4238
            if (len(bondedAtomIndices) == 3):
4239
4240
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
4241
            elif (len(bondedAtomIndices) == 4):
4242
4243
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4244
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
4245
4246
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
4247
        elif (atomicNumber == 7):               # N(7)
4248
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
4249
        elif (atomicNumber == 8):               # O(8)
4250
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
4251
            if (len(bondedAtomIndices) == 2):
4252
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
4253
        elif (atomicNumber == 9):               # F(9)
4254
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
4255
        elif (atomicNumber == 10):
4256
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
4257
        elif (atomicNumber == 11):
4258
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
4259
        elif (atomicNumber == 12):
4260
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
4261
        elif (atomicNumber == 14):
4262
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
4263
        elif (atomicNumber == 15):              # P(15)
4264
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
4265
        elif (atomicNumber == 16):              # S(16)
4266
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
4267
        elif (atomicNumber == 17):
4268
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
4269
        elif (atomicNumber == 18):
4270
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
4271
        elif (atomicNumber == 19):
4272
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
4273
        elif (atomicNumber == 20):
4274
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
4275
        elif (atomicNumber == 35):
4276
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
4277
        elif (atomicNumber == 36):
4278
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
4279
        elif (atomicNumber == 37):
4280
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
4281
        elif (atomicNumber == 53):
4282
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
4283
        elif (atomicNumber == 54):
4284
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
4285
        elif (atomicNumber == 55):
4286
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
4287
        elif (atomicNumber == 56):
4288
4289
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
4290
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4291
4292
            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
4293

4294
4295
4296
4297
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
4300
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
4301
        atom = data.atoms[atomIndex]
4302
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
4303
        if (atomicNumber in bondiMap):
4304
4305
            radius = bondiMap[atomicNumber]
        else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4306
4307
            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
4308

4309
4310
4311
4312
4313
4314
4315
4316
4317
4318
        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
4319

Peter Eastman's avatar
Peter Eastman committed
4320
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
Justin MacCallum's avatar
Justin MacCallum committed
4321
4322
                                                        element.attrib['includeCavityTerm'],
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor'])
4323
4324
4325
        forceField._forces.append(generator)

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

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

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

4332
4333
4334
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
4335
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
4336
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
4337
        if (len(amoebaMultipoleForceList) > 0):
4338
4339
4340
4341
4342
            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
4343
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
4344
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
4345

4346
4347
4348
4349
4350
4351
4352
        # 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
4353

Peter Eastman's avatar
Peter Eastman committed
4354
4355
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
4356
            else:
Peter Eastman's avatar
Peter Eastman committed
4357
                force.setSolventDielectric(   float(self.solventDielectric))
4358

Peter Eastman's avatar
Peter Eastman committed
4359
4360
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
4361
            else:
Peter Eastman's avatar
Peter Eastman committed
4362
                force.setSoluteDielectric(    float(self.soluteDielectric))
4363

Peter Eastman's avatar
Peter Eastman committed
4364
4365
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
4366
            else:
Peter Eastman's avatar
Peter Eastman committed
4367
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
4368
4369
4370
4371
4372

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
4375
4376
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
4377
4378
4379
4380
4381

        # 1-2

        bonded12ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4382
4383
4384
            bonded12ParticleSet = AmoebaVdwGenerator.getBondedParticleSet(i, data)
            bonded12ParticleSet = set(sorted(bonded12ParticleSet))
            bonded12ParticleSets.append(bonded12ParticleSet)
4385
4386

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
4387
4388
4389
4390
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
4391
            else:
Peter Eastman's avatar
Peter Eastman committed
4392
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
4393
4394
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
4395
            force.addParticle(multipoleParameters[0], radius, shct)
4396
4397
4398
4399
4400

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

4401
## @private
4402
4403
4404
4405
4406
class AmoebaUreyBradleyGenerator:

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

4408
    def __init__(self):
4409

Peter Eastman's avatar
Peter Eastman committed
4410
4411
4412
        self.types1 = []
        self.types2 = []
        self.types3 = []
4413

Peter Eastman's avatar
Peter Eastman committed
4414
4415
        self.length = []
        self.k = []
4416
4417
4418
4419
4420
4421

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

    @staticmethod
    def parseElement(element, forceField):

4422
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
4423
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
4424

4425
        generator = AmoebaUreyBradleyGenerator()
4426
4427
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
4428
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
4429
            if None not in types:
4430
4431
4432
4433
4434
4435
4436
4437
4438
4439

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

4443
4444
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
4447
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
4448
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
4449
4450

        if len(existing) == 0:
4451
            force = mm.HarmonicBondForce()
4452
4453
4454
4455
4456
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
Peter Eastman's avatar
Peter Eastman committed
4457
            if (isConstrained):
4458
4459
4460
4461
4462
4463
4464
4465
                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
4466
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
4467
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
4468
4469
4470
4471
4472
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
4473
4474
4475
4476
4477


## @private
class DrudeGenerator:
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
4478

4479
4480
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
4481
4482
4483
4484
4485
4486
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
4487
4488
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
4489
4490
4491
4492
        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'):
4493
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
4494
4495
4496
4497
4498
4499
4500
4501
4502
4503
            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
4504

peastman's avatar
peastman committed
4505
4506
4507
4508
    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
4509

peastman's avatar
peastman committed
4510
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
4511

peastman's avatar
peastman committed
4512
4513
4514
4515
4516
4517
4518
4519
4520
4521
4522
4523
4524
4525
4526
4527
        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
4528
4529
                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
4530
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
4531

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

peastman's avatar
peastman committed
4535
4536
4537
4538
4539
4540
4541
        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)
4542
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
4543
4544
4545
4546
4547
                # 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]
4548
4549
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
4550
4551
4552
4553
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

Justin MacCallum's avatar
Justin MacCallum committed
4554
parsers["DrudeForce"] = DrudeGenerator.parseElement