"platforms/opencl/include/OpenCLKernels.h" did not exist on "9b6ddf423ffe1cd6bd4f3e85ce65ac5641f4ecaf"
forcefield.py 170 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-2013 Stanford University and the Authors.
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
Authors: Peter Eastman, Mark Friedrichs
Contributors:

Permission is hereby granted, free of charge, to any person obtaining a 
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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
"""
__author__ = "Peter Eastman"
__version__ = "1.0"

import os
import itertools
import xml.etree.ElementTree as etree
from math import sqrt, cos
import simtk.openmm as mm
import simtk.unit as unit
import element as elem
from simtk.openmm.app import Topology

# Enumerated values for nonbonded method

NoCutoff = object()
CutoffNonPeriodic = object()
CutoffPeriodic = object()
Ewald = object()
PME = object()

# Enumerated values for constraint type

HBonds = object()
AllBonds = object()
HAngles = object()

# 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.
        
        Parameters:
68
         - files A list of XML files defining the force field.  Each entry may be an absolute file path, a path relative to the
69
70
71
72
           current working directory, or a path relative to this module's data subdirectory (for built in force fields).
        """
        self._atomTypes = {}
        self._templates = {}
73
        self._templateSignatures = {None:[]}
74
        self._atomClasses = {'':set()}
75
        self._forces = []
76
        self._scripts = []
77
78
79
80
81
82
83
84
85
86
87
        for file in files:
            try:
                tree = etree.parse(file)
            except IOError:
                tree = etree.parse(os.path.join(os.path.dirname(__file__), 'data', file))
            root = tree.getroot()
            
            # Load the atom types.
            
            if tree.getroot().find('AtomTypes') is not None:
                for type in tree.getroot().find('AtomTypes').findall('Type'):
88
89
90
                    element = None
                    if 'element' in type.attrib:
                        element = elem.get_by_symbol(type.attrib['element'])
91
92
93
94
                    name = type.attrib['name']
                    if name in self._atomTypes:
                        raise ValueError('Found multiple definitions for atom type: '+name)
                    self._atomTypes[name] = (type.attrib['class'], float(type.attrib['mass']), element)
95
96
97
98
99
100
101
102
103
104
            
            # 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)
                    self._templates[resName] = template
                    for atom in residue.findall('Atom'):
                        template.atoms.append(ForceField._TemplateAtomData(atom.attrib['name'], atom.attrib['type'], self._atomTypes[atom.attrib['type']][2]))
105
106
                    for site in residue.findall('VirtualSite'):
                        template.virtualSites.append(ForceField._VirtualSiteData(site))
107
108
109
110
111
112
113
114
115
116
                    for bond in residue.findall('Bond'):
                        b = (int(bond.attrib['from']), int(bond.attrib['to']))
                        template.bonds.append(b)
                        template.atoms[b[0]].bondedTo.append(b[1])
                        template.atoms[b[1]].bondedTo.append(b[0])
                    for bond in residue.findall('ExternalBond'):
                        b = int(bond.attrib['from'])
                        template.externalBonds.append(b)
                        template.atoms[b].externalBonds += 1
            for template in self._templates.values():
117
118
119
                signature = _createResidueSignature([atom.element for atom in template.atoms])
                if signature in self._templateSignatures:
                    self._templateSignatures[signature].append(template)
120
                else:
121
                    self._templateSignatures[signature] = [template]
122
123
124
125
126
127
128
129
130
131
132
            
            # Build sets of every atom type belonging to each class
            
            for type in self._atomTypes:
                atomClass = self._atomTypes[type][0]
                if atomClass in self._atomClasses:
                    typeSet = self._atomClasses[atomClass]
                else:
                    typeSet = set()
                    self._atomClasses[atomClass] = typeSet
                typeSet.add(type)
133
                self._atomClasses[''].add(type)
134
135
136
137
138
139
            
            # Load force definitions
            
            for child in root:
                if child.tag in parsers:
                    parsers[child.tag](child, self)
140
141
142
143
144
            
            # Load scripts
            
            for node in tree.getroot().findall('Script'):
                self._scripts.append(node.text)
145
146
147
148
149
150
151
152
153
154
155

    def _findAtomTypes(self, 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
156
            typeAttrib = 'type'+suffix
157
            if classAttrib in attrib:
158
159
                if typeAttrib in attrib:
                    raise ValueError('Tag specifies both a type and a class for the same atom: '+etree.tostring(node))
160
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
161
162
163
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
164
165
            else:
                if typeAttrib not in attrib or attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
166
167
168
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
169
170
171
172
173
        return types

    def _parseTorsion(self, node):
        """Parse the node defining a torsion."""
        types = self._findAtomTypes(node, 4)
peastman's avatar
peastman committed
174
        if None in types:
175
176
177
178
179
180
181
182
183
            return None
        torsion = PeriodicTorsion(types)
        attrib = node.attrib
        index = 1
        while 'phase%d'%index in attrib:
            torsion.periodicity.append(int(attrib['periodicity%d'%index]))
            torsion.phase.append(float(attrib['phase%d'%index]))
            torsion.k.append(float(attrib['k%d'%index]))
            index += 1
184
        return torsion            
185
186
187
188
189
190
        
    class _SystemData:
        """Inner class used to encapsulate data about the system being created."""
        def __init__(self):
            self.atomType = {}
            self.atoms = []
191
            self.virtualSites = {}
192
193
194
195
196
197
198
199
200
201
202
203
            self.bonds = []
            self.angles = []
            self.propers = []
            self.impropers = []
            self.atomBonds = []
            self.isAngleConstrained = []

    class _TemplateData:
        """Inner class used to encapsulate data about a residue template definition."""
        def __init__(self, name):
            self.name = name
            self.atoms = []
204
            self.virtualSites = []
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
            self.bonds = []
            self.externalBonds = []

    class _TemplateAtomData:
        """Inner class used to encapsulate data about an atom in a residue template definition."""
        def __init__(self, name, type, element):
            self.name = name
            self.type = type
            self.element = element
            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
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
    
    class _VirtualSiteData:
        """Inner class used to encapsulate data about a virtual site."""
        def __init__(self, node):
            attrib = node.attrib
            self.index = int(attrib['index'])
            self.type = attrib['type']
            if self.type == 'average2':
                self.atoms = [int(attrib['atom1']), int(attrib['atom2'])]
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
                self.atoms = [int(attrib['atom1']), int(attrib['atom2']), int(attrib['atom3'])]
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
                self.atoms = [int(attrib['atom1']), int(attrib['atom2']), int(attrib['atom3'])]
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
            self.atoms = [x-self.index for x in self.atoms]
243
244

    def createSystem(self, topology, nonbondedMethod=NoCutoff, nonbondedCutoff=1.0*unit.nanometer,
245
                     constraints=None, rigidWater=True, removeCMMotion=True, **args):
246
247
248
249
250
251
252
        """Construct an OpenMM System representing a Topology with this force field.
        
        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
Peter Eastman's avatar
Peter Eastman committed
253
         - constraints (object=None) Specifies which bonds and angles should be implemented with constraints.
254
255
           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
256
         - removeCMMotion (boolean=True) If true, a CMMotionRemover will be added to the System
257
         - args Arbitrary additional keyword arguments may also be specified.  This allows extra parameters to be specified that are specific to
258
259
260
261
           particular force fields.
        Returns: the newly created System
        """
        data = ForceField._SystemData()
262
        data.atoms = list(topology.atoms())
263
264
265
266

        # Make a list of all bonds
        
        for bond in topology.bonds():
267
            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287

        # Record which atoms are bonded to each other atom
        
        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.
        
        for chain in topology.chains():
            for res in chain.residues():
                template = None
                matches = None
288
289
290
291
                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)
292
293
294
295
                        if matches is not None:
                            template = t
                            break
                if matches is None:
296
                    raise ValueError('No template found for residue %d (%s).  This might mean your input topology is missing some atoms or bonds, or possibly that you are using the wrong force field.' % (res.index+1, res.name))
297
298
                for atom, match in zip(res.atoms(), matches):
                    data.atomType[atom] = template.atoms[match].type
299
300
301
                    for site in template.virtualSites:
                        if match == site.index:
                            data.virtualSites[atom] = site
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400

        # Create the System and add atoms
        
        sys = mm.System()
        for atom in topology.atoms():
            sys.addParticle(self._atomTypes[data.atomType[atom]][1])
        
        # Set periodic boundary conditions.
        
        boxSize = topology.getUnitCellDimensions()
        if boxSize is not None:
            sys.setDefaultPeriodicBoxVectors((boxSize[0], 0, 0), (0, boxSize[1], 0), (0, 0, boxSize[2]))
        elif nonbondedMethod is not NoCutoff and nonbondedMethod is not CutoffNonPeriodic:
            raise ValueError('Requested periodic boundary conditions for a Topology that does not specify periodic box dimensions')

        # Make a list of all unique angles
        
        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))
        
        # Make a list of all unique proper torsions
        
        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))
        
        # Make a list of all unique improper torsions
        
        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]))
        
        # Identify bonds that should be implemented with constraints
        
        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
        
        # Identify angles that should be implemented with constraints
        
        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
401
402
403
404
405
        
        # Add virtual sites
        
        for atom in data.virtualSites:
            site = data.virtualSites[atom]
406
            index = atom.index
407
408
409
410
411
412
413
            if site.type == 'average2':
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(index+site.atoms[0], index+site.atoms[1], site.weights[0], site.weights[1]))
            elif site.type == 'average3':
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(index+site.atoms[0], index+site.atoms[1], index+site.atoms[2], site.weights[0], site.weights[1], site.weights[2]))
            elif site.type == 'outOfPlane':
                sys.setVirtualSite(index, mm.OutOfPlaneSite(index+site.atoms[0], index+site.atoms[1], index+site.atoms[2], site.weights[0], site.weights[1], site.weights[2]))
        
414
415
416
417
        # Add forces to the System
        
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
418
419
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
420
        
peastman's avatar
peastman committed
421
422
423
424
425
426
        # Let generators do postprocessing
        
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
        
427
428
429
430
        # Execute scripts found in the XML files.
        
        for script in self._scripts:
            exec script
431
432
433
434
435
436
437
        return sys


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = {}
    for element in elements:
438
        if element is None:
439
440
            pass # This residue contains "atoms" (probably virtual sites) that should match any element
        elif element in counts:
441
442
443
444
445
446
447
            counts[element] += 1
        else:
            counts[element] = 1
    sig = []
    for c in counts:
        sig.append((c, counts[c]))
    sig.sort(key=lambda x: -x[0].mass)
448
449
    
    # Convert it to a string.
450
451

    s = ''
452
    for element, count in sig:
453
454
455
456
        s += element.symbol+str(count)
    return s


457
def _matchResidue(res, template, bondedToAtom):
458
459
460
461
462
463
464
465
466
467
    """Determine whether a residue matches a template and return a list of corresponding atoms.
    
    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: a list specifying which atom of the template each atom of the residue corresponds to,
    or None if it does not match the template
    """
    atoms = list(res.atoms())
468
469
    if len(atoms) != len(template.atoms):
        return None
470
471
472
473
474
475
476
    matches = len(atoms)*[0]
    hasMatch = len(atoms)*[False]
    
    # Translate from global to local atom indices, and record the bonds for each atom.
    
    renumberAtoms = {}
    for i in range(len(atoms)):
477
        renumberAtoms[atoms[i].index] = i
478
479
480
    bondedTo = []
    externalBonds = []
    for atom in atoms:
481
        bonds = [renumberAtoms[x] for x in bondedToAtom[atom.index] if x in renumberAtoms]
482
        bondedTo.append(bonds)
483
        externalBonds.append(len([x for x in bondedToAtom[atom.index] if x not in renumberAtoms]))
484
485
486
487
488
489
490
491
492
493
    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
494
    name = atoms[position].name
495
496
    for i in range(len(atoms)):
        atom = template.atoms[i]
peastman's avatar
peastman committed
497
        if (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]:
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
            # See if the bonds for this identification are consistent
            
            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.
                
                matches[position] = i
                hasMatch[i] = True
                if _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch, position+1):
                    return True
                hasMatch[i] = False
    return False


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

517
## @private
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
class HarmonicBondGenerator:
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
    
    def __init__(self):
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
    
    @staticmethod
    def parseElement(element, ff):
        generator = HarmonicBondGenerator()
        ff._forces.append(generator)
        for bond in element.findall('Bond'):
            types = ff._findAtomTypes(bond, 2)
peastman's avatar
peastman committed
533
            if None not in types:
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.length.append(float(bond.attrib['length']))
                generator.k.append(float(bond.attrib['k']))
    
    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:
                        sys.addConstraint(bond.atom1, bond.atom2, self.length[i])
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


564
## @private
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
class HarmonicAngleGenerator:
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
    
    def __init__(self):
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
    
    @staticmethod
    def parseElement(element, ff):
        generator = HarmonicAngleGenerator()
        ff._forces.append(generator)
        for angle in element.findall('Angle'):
            types = ff._findAtomTypes(angle, 3)
peastman's avatar
peastman committed
581
            if None not in types:
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.angle.append(float(angle.attrib['angle']))
                generator.k.append(float(angle.attrib['k']))
    
    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.
                        
                        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
                        
                        # Compute the distance between atoms and add a constraint
                        
                        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]))
                                sys.addConstraint(angle[0], angle[2], length)
                    elif self.k[i] != 0:
                        force.addAngle(angle[0], angle[1], angle[2], self.angle[i], self.k[i])
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


633
## @private
634
635
636
637
638
639
640
641
642
643
644
645
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 = []

646
## @private
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
class PeriodicTorsionGenerator:
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
    
    def __init__(self):
        self.proper = []
        self.improper = []
    
    @staticmethod
    def parseElement(element, ff):
        generator = PeriodicTorsionGenerator()
        generator.ff = ff
        ff._forces.append(generator)
        for torsion in element.findall('Proper'):
            torsion = ff._parseTorsion(torsion)
            if torsion is not None:
                generator.proper.append(torsion)
        for torsion in element.findall('Improper'):
            torsion = ff._parseTorsion(torsion)
            if torsion is not None:
                generator.improper.append(torsion)
    
    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:
                            # Workaround to be more consistent with AMBER.  It uses wildcards to define most of its
                            # impropers, which leaves the ordering ambigous.  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])
                            done = True
                            break

parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


734
## @private
735
736
737
738
739
740
741
742
743
744
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

745
## @private
746
747
748
749
750
751
752
753
754
755
756
757
758
759
class RBTorsionGenerator:
    """An RBTorsionGenerator constructs an RBTorsionForce."""
    
    def __init__(self):
        self.proper = []
        self.improper = []
    
    @staticmethod
    def parseElement(element, ff):
        generator = RBTorsionGenerator()
        generator.ff = ff
        ff._forces.append(generator)
        for torsion in element.findall('Proper'):
            types = ff._findAtomTypes(torsion, 4)
peastman's avatar
peastman committed
760
            if None not in types:
761
762
763
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
            types = ff._findAtomTypes(torsion, 4)
peastman's avatar
peastman committed
764
            if None not in types:
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
                generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
    
    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:
                            # Workaround to be more consistent with AMBER.  It uses wildcards to define most of its
                            # impropers, which leaves the ordering ambigous.  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])
                            done = True
                            break

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


829
## @private
830
831
832
833
834
835
836
837
838
839
840
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

841
## @private
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
class CMAPTorsionGenerator:
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
    
    def __init__(self):
        self.torsions = []
        self.maps = []
    
    @staticmethod
    def parseElement(element, ff):
        generator = CMAPTorsionGenerator()
        generator.ff = ff
        ff._forces.append(generator)
        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'):
            types = ff._findAtomTypes(torsion, 5)
peastman's avatar
peastman committed
862
            if None not in types:
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
    
    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)
        
        # Find all chains of length 5
        
        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


921
## @private
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
class NonbondedGenerator:
    """A NonbondedGenerator constructs a NonbondedForce."""
    
    def __init__(self, coulomb14scale, lj14scale):
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
        if len(existing) == 0:
            generator = NonbondedGenerator(float(element.attrib['coulomb14scale']), float(element.attrib['lj14scale']))
            ff._forces.append(generator)
        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']):
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales') 
        for atom in element.findall('Atom'):
            types = ff._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
943
            if None not in types:
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
                values = (float(atom.attrib['charge']), float(atom.attrib['sigma']), float(atom.attrib['epsilon']))
                for t in types[0]:
                    generator.typeMap[t] = values
    
    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:
            raise ValueError('Illegal nonbonded method for NonbondedForce') 
        force = mm.NonbondedForce()
        for atom in data.atoms:
            t = data.atomType[atom]
            if t in self.typeMap:
                values = self.typeMap[t]
                force.addParticle(values[0], values[1], values[2])
            else:
963
                raise ValueError('No nonbonded parameters defined for atom type '+t)
peastman's avatar
peastman committed
964
965
966
967
968
969
970
971
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
        sys.addForce(force)
    
    def postprocessSystem(self, sys, data, args):
        # Create exceptions based on bonds, virtual sites, and Drude particles.
972
973
974
        bondIndices = []
        for bond in data.bonds:
            bondIndices.append((bond.atom1, bond.atom2))
975
976
977
978
979
        for i in range(sys.getNumParticles()):
            if sys.isVirtualSite(i):
                site = sys.getVirtualSite(i)
                for j in range(site.getNumParticles()):
                    bondIndices.append((i, site.getParticle(j)))
peastman's avatar
peastman committed
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
        drude = [f for f in sys.getForces() if isinstance(f, mm.DrudeForce)]
        if len(drude) > 0:
            drude = drude[0]
            # For purposes of creating exceptions, a Drude particle is "bonded" to anything
            # its parent atom is bonded to.
            drudeMap = {}
            for i in range(drude.getNumParticles()):
                params = drude.getParticleParameters(i)
                drudeMap[params[1]] = params[0]
            for atom1, atom2 in bondIndices:
                drude1 = drudeMap[atom1] if atom1 in drudeMap else None
                drude2 = drudeMap[atom2] if atom2 in drudeMap else None
                if drude1 is not None:
                    bondIndices.append((drude1, atom2))
                    if drude2 is not None:
                        bondIndices.append((drude1, drude2))
                if drude2 is not None:
                    bondIndices.append((atom1, drude2))
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)        
1000
1001
1002
1003

parsers["NonbondedForce"] = NonbondedGenerator.parseElement


1004
## @private
1005
1006
1007
1008
1009
1010
1011
1012
class GBSAOBCGenerator:
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
    
    def __init__(self):
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
1013
1014
1015
1016
1017
1018
1019
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
            generator = GBSAOBCGenerator()
            ff._forces.append(generator)
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
1020
1021
        for atom in element.findall('Atom'):
            types = ff._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
1022
            if None not in types:
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
                values = (float(atom.attrib['charge']), float(atom.attrib['radius']), float(atom.attrib['scale']))
                for t in types[0]:
                    generator.typeMap[t] = values
    
    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:
            raise ValueError('Illegal nonbonded method for GBSAOBCForce') 
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
            t = data.atomType[atom]
            if t in self.typeMap:
                values = self.typeMap[t]
                force.addParticle(values[0], values[1], values[2])
            else:
                raise ValueError('No GBSAOBC parameters defined for atom type '+t) 
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
1043
1044
1045
1046
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
1047
1048
1049
1050
1051
        sys.addForce(force)

parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


1052
## @private
1053
1054
1055
1056
1057
1058
class GBVIGenerator:

    """A GBVIGenerator constructs a GBVIForce."""
    
    def __init__(self,ff):

Peter Eastman's avatar
Peter Eastman committed
1059
1060
1061
1062
1063
        self.ff = ff
        self.fixedParameters = {}
        self.fixedParameters['soluteDielectric'] = 1.0
        self.fixedParameters['solventDielectric'] = 78.3
        self.fixedParameters['scalingMethod'] = 1
1064
        self.fixedParameters['quinticUpperBornRadiusLimit'] = 5.0
Peter Eastman's avatar
Peter Eastman committed
1065
        self.fixedParameters['quinticLowerLimitFactor'] = 0.8
1066

Peter Eastman's avatar
Peter Eastman committed
1067
        self.typeMap = {}
1068
1069
1070

    @staticmethod
    def parseElement(element, ff):
Peter Eastman's avatar
Peter Eastman committed
1071
        generator = GBVIGenerator(ff)
1072
        for key in generator.fixedParameters.iterkeys():
Peter Eastman's avatar
Peter Eastman committed
1073
1074
            if (key in element.attrib):
                generator.fixedParameters[key] = float(element.attrib[key])
1075
1076
1077
1078

        ff._forces.append(generator)
        for atom in element.findall('Atom'):
            types = ff._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
1079
            if None not in types:
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
                values = (float(atom.attrib['charge']), float(atom.attrib['radius']), float(atom.attrib['gamma']))
                for t in types[0]:
                    generator.typeMap[t] = values
    
    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:
            raise ValueError('Illegal nonbonded method for GB/VI Force') 

        # add particles

        force = mm.GBVIForce()
        for atom in data.atoms:
            t = data.atomType[atom]
            if t in self.typeMap:
                values = self.typeMap[t]
                force.addParticle(values[0], values[1], values[2])
            else:
                raise ValueError('No GB/VI parameters defined for atom type '+t) 

        # get HarmonicBond generator -- exit if not found

        hbGenerator = 0
        for generator in self.ff._forces:
Peter Eastman's avatar
Peter Eastman committed
1108
            if (generator.__class__.__name__ == 'HarmonicBondGenerator'): 
1109
1110
1111
               hbGenerator = generator
               break

Peter Eastman's avatar
Peter Eastman committed
1112
        if (hbGenerator == 0):
1113
1114
1115
1116
1117
1118
1119
            raise ValueError('HarmonicBondGenerator not found.') 

        # add bonds

        for bond in data.bonds:
            type1 = data.atomType[data.atoms[bond.atom1]]
            type2 = data.atomType[data.atoms[bond.atom2]]
Peter Eastman's avatar
Peter Eastman committed
1120
            hit = 0
1121
1122
1123
1124
1125
1126
1127
1128
1129
            for i in range(len(hbGenerator.types1)):
                types1 = hbGenerator.types1[i]
                types2 = hbGenerator.types2[i]
                if (type1 in types1 and type2 in types2) or (type1 in types2 and type2 in types1):
                    #bond.length = hbGenerator.length[i]
                    force.addBond(bond.atom1, bond.atom2, hbGenerator.length[i])

        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
Peter Eastman's avatar
Peter Eastman committed
1130
1131
1132
1133
1134
        force.setSolventDielectric(self.fixedParameters['solventDielectric'])
        force.setSoluteDielectric(self.fixedParameters['soluteDielectric'])
        force.setBornRadiusScalingMethod(self.fixedParameters['scalingMethod'])
        force.setQuinticLowerLimitFactor(self.fixedParameters['quinticLowerLimitFactor'])
        force.setQuinticUpperBornRadiusLimit(self.fixedParameters['quinticUpperBornRadiusLimit'])
1135
1136
1137
1138
1139
        
        sys.addForce(force)

parsers["GBVIForce"] = GBVIGenerator.parseElement

1140
## @private
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
class CustomBondGenerator:
    """A CustomBondGenerator constructs a CustomBondForce."""
    
    def __init__(self):
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
    
    @staticmethod
    def parseElement(element, ff):
        generator = CustomBondGenerator()
        ff._forces.append(generator)
        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'):
            types = ff._findAtomTypes(bond, 2)
peastman's avatar
peastman committed
1162
            if None not in types:
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.paramValues.append([float(bond.attrib[param]) for param in generator.perBondParams])
    
    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


1187
## @private
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
class CustomAngleGenerator:
    """A CustomAngleGenerator constructs a CustomAngleForce."""
    
    def __init__(self):
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
    
    @staticmethod
    def parseElement(element, ff):
        generator = CustomAngleGenerator()
        ff._forces.append(generator)
        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'):
            types = ff._findAtomTypes(angle, 3)
peastman's avatar
peastman committed
1210
            if None not in types:
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
                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])
    
    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


1238
## @private
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
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

1249
## @private
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
class CustomTorsionGenerator:
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
    
    def __init__(self):
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
    
    @staticmethod
    def parseElement(element, ff):
        generator = CustomTorsionGenerator()
        generator.ff = ff
        ff._forces.append(generator)
        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'):
            types = ff._findAtomTypes(torsion, 4)
peastman's avatar
peastman committed
1271
            if None not in types:
1272
1273
1274
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
            types = ff._findAtomTypes(torsion, 4)
peastman's avatar
peastman committed
1275
            if None not in types:
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
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
                generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
    
    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:
                            # Workaround to be more consistent with AMBER.  It uses wildcards to define most of its
                            # impropers, which leaves the ordering ambigous.  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)
                            done = True
                            break

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


1339
## @private
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
class CustomGBGenerator:
    """A CustomGBGenerator constructs a CustomGBForce."""
    
    def __init__(self):
        self.typeMap = {}
        self.globalParams = {}
        self.perParticleParams = []
        self.paramValues = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
        generator = CustomGBGenerator()
        ff._forces.append(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 atom in element.findall('Atom'):
            types = ff._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
1362
            if None not in types:
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
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
                values = [float(atom.attrib[param]) for param in generator.perParticleParams]
                for t in types[0]:
                    generator.typeMap[t] = values
        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()]
            generator.functions.append((function.attrib['name'], values, float(function.attrib['min']), float(function.attrib['max'])))
    
    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:
            raise ValueError('Illegal nonbonded method for CustomGBForce') 
        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])
        for function in self.functions:
            force.addFunction(function[0], function[1], function[2], function[3])
        for atom in data.atoms:
            t = data.atomType[atom]
            if t in self.typeMap:
                values = self.typeMap[t]
                force.addParticle(self.typeMap[t])
            else:
                raise ValueError('No CustomGB parameters defined for atom type '+t) 
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

Peter Eastman's avatar
Peter Eastman committed
1407
def getAtomPrint(data, atomIndex):
1408

Peter Eastman's avatar
Peter Eastman committed
1409
1410
1411
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
1412
    else:
Peter Eastman's avatar
Peter Eastman committed
1413
        returnString = "NA"
1414
1415
1416
1417
1418

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
1419
def countConstraint(data):
1420

Peter Eastman's avatar
Peter Eastman committed
1421
    bondCount = 0
1422
1423
1424
1425
1426
1427
1428
    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
1429
        if (isConstrained):
1430
1431
            angleCount += 1
 
Peter Eastman's avatar
Peter Eastman committed
1432
    print "Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount))
1433

1434
## @private
1435
class AmoebaBondGenerator:
1436
1437
1438

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

1439
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
1440
1441
1442
1443
1444

    #=============================================================================================
    
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
1445
1446
1447
1448
1449
1450
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
1451
1452
1453
1454
1455
1456
    
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

1457
        # <AmoebaBondForce bond-cubic="-25.5" bond-quartic="379.3125">
1458
1459
        # <Bond class1="1" class2="2" length="0.1437" k="156900.0"/>
    
1460
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
1461
1462
1463
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
            types = forceField._findAtomTypes(bond, 2)
peastman's avatar
peastman committed
1464
            if None not in types:
1465
1466
1467
1468
1469
                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:
1470
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
1471
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
1472
1473
                                    bond.attrib['class2'])
                raise ValueError(outputString) 
1474
1475
1476
    
    #=============================================================================================

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

Mark Friedrichs's avatar
Mark Friedrichs committed
1479
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1480

Peter Eastman's avatar
Peter Eastman committed
1481
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
1482
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
1483
        if len(existing) == 0:
1484
            force = mm.AmoebaBondForce()
1485
1486
1487
1488
            sys.addForce(force)
        else:
            force = existing[0]

1489
1490
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
1491
1492
1493
1494

        for bond in data.bonds:
            type1 = data.atomType[data.atoms[bond.atom1]]
            type2 = data.atomType[data.atoms[bond.atom2]]
Peter Eastman's avatar
Peter Eastman committed
1495
            hit = 0
1496
1497
1498
1499
1500
            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]
Peter Eastman's avatar
Peter Eastman committed
1501
                    hit = 1
1502
1503
1504
1505
1506
1507
                    if bond.isConstrained:
                        sys.addConstraint(bond.atom1, bond.atom2, self.length[i])
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1508
            if (hit == 0): 
1509
                outputString = "AmoebaBondGenerator missing: types=[%5s %5s] atoms=[%6d %6d] " % (type1, type2, bond.atom1, bond.atom2)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1510
                raise ValueError(outputString) 
1511

1512
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
1513
1514
1515
1516
1517

#=============================================================================================
# Add angle constraint
#=============================================================================================
    
Peter Eastman's avatar
Peter Eastman committed
1518
def addAngleConstraint(angle, idealAngle, data, sys):
1519
1520
1521

    # Find the two bonds that make this angle.
                    
Peter Eastman's avatar
Peter Eastman committed
1522
1523
    bond1 = None
    bond2 = None
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
    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
                    
        # Compute the distance between atoms and add a constraint
                    
        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
1538
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
1539
1540
1541
1542
                sys.addConstraint(angle[0], angle[2], length)
                return

#=============================================================================================
1543
## @private
1544
class AmoebaAngleGenerator:
1545
1546

    #=============================================================================================
1547
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
1548
1549
1550
1551
    #=============================================================================================
    
    def __init__(self, forceField, cubic, quartic, pentic, sextic):

Peter Eastman's avatar
Peter Eastman committed
1552
1553
1554
1555
1556
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
1557

Peter Eastman's avatar
Peter Eastman committed
1558
1559
1560
        self.types1 = []
        self.types2 = []
        self.types3 = []
1561

Peter Eastman's avatar
Peter Eastman committed
1562
1563
        self.angle = []
        self.k = []
1564
1565
1566
1567
1568
1569
    
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

1573
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
1574
1575
1576
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
            types = forceField._findAtomTypes(angle, 3)
peastman's avatar
peastman committed
1577
            if None not in types:
1578
1579
1580
1581
1582
1583

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
1584
                angleList.append(float(angle.attrib['angle1']))
1585
1586

                try:
Peter Eastman's avatar
Peter Eastman committed
1587
                    angleList.append(float(angle.attrib['angle2']))
1588
                    try:
Peter Eastman's avatar
Peter Eastman committed
1589
                        angleList.append(float(angle.attrib['angle3']))
1590
1591
1592
1593
1594
1595
1596
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
1597
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
1598
1599
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
1600
1601
                                    angle.attrib['class3'])
                raise ValueError(outputString) 
1602
1603
1604
1605
1606
1607
    
    #=============================================================================================
    # createForce is bypassed here since the AmoebaOutOfPlaneBendForce generator must first execute
    # and partition angles into in-plane and non-in-plane angles
    #=============================================================================================
    
Peter Eastman's avatar
Peter Eastman committed
1608
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
1609
1610
1611
1612
1613
1614
1615
        pass

    #=============================================================================================
    # createForcePostOpBendAngle is called by AmoebaOutOfPlaneBendForce with the list of
    # non-in-plane angles
    #=============================================================================================
    
Peter Eastman's avatar
Peter Eastman committed
1616
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
1617
1618
1619
1620

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
1621
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
1622
1623

        if len(existing) == 0:
1624
            force = mm.AmoebaAngleForce()
1625
1626
1627
1628
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1629
        # set scalars
1630

1631
1632
1633
1634
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
1635
1636

        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
1637
1638
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
1639

Peter Eastman's avatar
Peter Eastman committed
1640
1641
1642
1643
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
            hit = 0
1644
1645
1646
1647
1648
            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):
Peter Eastman's avatar
Peter Eastman committed
1649
                    hit = 1
1650
1651
                    if isConstrained and self.k[i] != 0.0:
                        angleDict['idealAngle'] = self.angle[i][0]
Peter Eastman's avatar
Peter Eastman committed
1652
                        addAngleConstraint(angle, self.angle[i][0], data, sys)
1653
                    elif self.k[i] != 0:
Peter Eastman's avatar
Peter Eastman committed
1654
1655
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
1656
1657
1658
1659
1660
1661
                            # 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
1662
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
1663
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
1664
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
1665
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
1666
                            if (numberOfHydrogens < lenAngle):
1667
1668
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
1669
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1670
                                raise ValueError(outputString) 
1671
1672
1673
1674
                        else:
                            angleValue =  self.angle[i][0]
               
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
1675
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
1676
                    break
Peter Eastman's avatar
Peter Eastman committed
1677
            if (hit == 0): 
1678
                outputString = "AmoebaAngleGenerator missing types: [%s %s %s] for atoms: " % (type1, type2, type3)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1679
1680
1681
1682
1683
                outputString += getAtomPrint( data, angle[0] ) + ' '
                outputString += getAtomPrint( data, angle[1] ) + ' '
                outputString += getAtomPrint( data, angle[2] )
                outputString += " indices: [%6d %6d %6d]" % (angle[0], angle[1], angle[2])
                raise ValueError(outputString) 
1684
1685
1686
1687
1688
1689

    #=============================================================================================
    # createForcePostOpBendInPlaneAngle is called by AmoebaOutOfPlaneBendForce with the list of
    # in-plane angles
    #=============================================================================================
    
Peter Eastman's avatar
Peter Eastman committed
1690
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
1691
1692
1693
1694

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
1695
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
1696
1697

        if len(existing) == 0:
1698
            force = mm.AmoebaInPlaneAngleForce()
1699
1700
1701
1702
1703
1704
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

1705
1706
1707
1708
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
1709
1710
1711

        for angleDict in angleList:
 
Peter Eastman's avatar
Peter Eastman committed
1712
1713
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
1714

Peter Eastman's avatar
Peter Eastman committed
1715
1716
1717
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
1718

Peter Eastman's avatar
Peter Eastman committed
1719
            hit = 0
1720
1721
1722
1723
1724
1725
1726
            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):
Peter Eastman's avatar
Peter Eastman committed
1727
                    hit = 1
1728
                    angleDict['idealAngle'] = self.angle[i][0]
Peter Eastman's avatar
Peter Eastman committed
1729
1730
                    if (isConstrained and self.k[i] != 0.0):
                        addAngleConstraint(angle, self.angle[i][0], data, sys)
1731
1732
1733
1734
                    else:
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

Peter Eastman's avatar
Peter Eastman committed
1735
            if (hit == 0): 
1736
                outputString = "AmoebaInPlaneAngleGenerator missing types: [%s %s %s] atoms: " % (type1, type2, type3)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1737
1738
1739
1740
1741
                outputString += getAtomPrint( data, angle[0] ) + ' '
                outputString += getAtomPrint( data, angle[1] ) + ' '
                outputString += getAtomPrint( data, angle[2] )
                outputString += " indices: [%6d %6d %6d]" % (angle[0], angle[1], angle[2])
                raise ValueError(outputString) 
1742

1743
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
1744
1745
1746

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
1747
1748
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
1749
1750
#=============================================================================================

1751
## @private
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
class AmoebaOutOfPlaneBendGenerator:

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

    """An AmoebaOutOfPlaneBendGenerator constructs a AmoebaOutOfPlaneBendForce."""
    
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
1762
1763
1764
1765
1766
1767
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
1768

Peter Eastman's avatar
Peter Eastman committed
1769
1770
1771
1772
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
1773

Peter Eastman's avatar
Peter Eastman committed
1774
        self.ks = []
1775
1776
1777
1778
1779
1780
1781
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
1808

    #=============================================================================================
    # Local version of findAtomTypes needed since class indices are 0 (i.e., not recognized)
    # for types3 and 4
    #=============================================================================================
    
    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"/>
         
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
1809
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
1810
1811
1812
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
1813
                                                   float(element.attrib['opbend-sextic']))
1814
1815
1816
1817

        forceField._forces.append(generator)

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

Peter Eastman's avatar
Peter Eastman committed
1821
1822
1823
1824
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
1825

Peter Eastman's avatar
Peter Eastman committed
1826
                generator.ks.append(float(angle.attrib['k']))
1827
1828

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1829
1830
                outputString = "AmoebaOutOfPlaneBendGenerator error getting types: %s %s %s %s." % (
                               angle.attrib['class1'], angle.attrib['class2'], angle.attrib['class3'], angle.attrib['class4'])
Peter Eastman's avatar
Peter Eastman committed
1831
                raise ValueError(outputString) 
1832
1833
1834
1835
1836
1837
1838
1839
    
    #=============================================================================================
    # 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
    #=============================================================================================
    
Peter Eastman's avatar
Peter Eastman committed
1840
    def getMiddleAtom(self, angle, data):
1841
1842
1843

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
1844
        middleAtom = -1
1845
        for atomIndex in angle: 
Peter Eastman's avatar
Peter Eastman committed
1846
            isMiddle = 0
1847
1848
1849
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
1850
                if (atom1 != atomIndex):
1851
1852
1853
                    partner = atom1
                else:
                    partner = atom2
Peter Eastman's avatar
Peter Eastman committed
1854
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]): 
1855
1856
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
1857
            if (isMiddle == 2):
1858
1859
1860
1861
1862
                return atomIndex
        return -1

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

Peter Eastman's avatar
Peter Eastman committed
1863
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876

        # 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
1877
1878
1879
1880
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
1881
1882
1883
1884

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

Peter Eastman's avatar
Peter Eastman committed
1885
        skipAtoms = dict()
1886
1887
1888
1889

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

Peter Eastman's avatar
Peter Eastman committed
1890
1891
        inPlaneAngles = []
        nonInPlaneAngles = []
1892
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
1893
        idealAngles = []*len(data.angles)
1894
1895
1896

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

Peter Eastman's avatar
Peter Eastman committed
1897
1898
1899
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
1900
1901
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
1902
                middleType = -1
1903
1904
1905
1906
                middleCovalency = -1

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

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

Peter Eastman's avatar
Peter Eastman committed
1913
1914
1915
1916
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
1917
1918
1919
1920

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
1921
                    if (atom1 != middleAtom):
1922
1923
1924
1925
1926
1927
1928
1929
                        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
1930
1931
1932
1933
1934
                        if (middleType in types2 and partnerType in types1):
                            partners.append(partner)
                            partnerSet.add(partner)
                            partnerTypes.append(partnerType)
                            partnerK.append(self.ks[i])
1935
             
Peter Eastman's avatar
Peter Eastman committed
1936
                if (len(partners) == 3):
1937

Peter Eastman's avatar
Peter Eastman committed
1938
1939
1940
                    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])
1941

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

1944
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
1945
1946
1947
1948
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1949
1950
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
1951
1952
1953
1954
1955
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
1956
1957
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
1958
                    angleDict['isConstrained'] = 0
1959

Peter Eastman's avatar
Peter Eastman committed
1960
1961
1962
1963
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
1964
1965

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
1966
1967
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
1968

Peter Eastman's avatar
Peter Eastman committed
1969
                    inPlaneAngles.append(angleDict)
1970
1971

                else:
Peter Eastman's avatar
Peter Eastman committed
1972
1973
1974
1975
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
1976
            else:
Peter Eastman's avatar
Peter Eastman committed
1977
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
1978

Peter Eastman's avatar
Peter Eastman committed
1979
                    partnerSet = skipAtoms[middleAtom]
1980
                  
Peter Eastman's avatar
Peter Eastman committed
1981
                    angleDict = {}
1982

Peter Eastman's avatar
Peter Eastman committed
1983
1984
1985
1986
1987
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
1988

Peter Eastman's avatar
Peter Eastman committed
1989
                    angleDict['isConstrained'] = isConstrained
1990

Peter Eastman's avatar
Peter Eastman committed
1991
1992
1993
1994
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
1995
1996

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
1997
1998
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
1999

Peter Eastman's avatar
Peter Eastman committed
2000
                    inPlaneAngles.append(angleDict)
2001
2002

                else:
Peter Eastman's avatar
Peter Eastman committed
2003
2004
                    angleDict = {}
                    angleDict['angle'] = angle
2005
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
2006
                    nonInPlaneAngles.append(angleDict)
2007

2008
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
2009
2010

        for force in self.forceField._forces:
2011
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'): 
Peter Eastman's avatar
Peter Eastman committed
2012
2013
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
2014
2015

        for force in self.forceField._forces:
Peter Eastman's avatar
Peter Eastman committed
2016
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'): 
2017
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
2018
                    nonInPlaneAngles.append(angleDict)
2019
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
2020
2021
2022
2023
2024

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

2025
## @private
2026
2027
2028
2029
2030
2031
2032
2033
class AmoebaTorsionGenerator:

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
2034
        self.torsionUnit = torsionUnit
2035

Peter Eastman's avatar
Peter Eastman committed
2036
2037
2038
2039
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
2040

Peter Eastman's avatar
Peter Eastman committed
2041
2042
2043
        self.t1 = []
        self.t2 = []
        self.t3 = []
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
    
    #=============================================================================================

    @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" />
         
Peter Eastman's avatar
Peter Eastman committed
2054
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
2055
2056
2057
2058
2059
2060
2061
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
            types = forceField._findAtomTypes(torsion, 4)
peastman's avatar
peastman committed
2062
            if None not in types:
2063
2064
2065
2066
2067
2068
2069

                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
2070
2071
                    tInfo = []
                    suffix = str(ii)
2072
2073
2074
2075
2076
2077
                    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
2078
2079
2080
2081
2082
2083
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
2084
2085
2086
2087
2088
2089

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
                                    stretchBend.attrib['class1'],
                                    stretchBend.attrib['class2'],
                                    stretchBend.attrib['class3'],
Peter Eastman's avatar
Peter Eastman committed
2090
2091
                                    stretchBend.attrib['class4'])
                raise ValueError(outputString) 
2092
2093
2094
2095
2096
2097
    
    #=============================================================================================

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2098
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
2099
        if len(existing) == 0:
2100
            force = mm.PeriodicTorsionForce()
2101
2102
2103
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2104

2105
2106
2107
2108
2109
2110
2111
        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]]]

Peter Eastman's avatar
Peter Eastman committed
2112
            hit = 0
2113
2114
2115
2116
2117
2118
2119
2120
2121
            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
2122
2123
                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):
                    hit = 1
2124
2125
2126
2127
2128
2129
                    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])
2130
2131
                    break

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2132
2133
2134
2135
2136
2137
2138
2139
            if (hit == 0): 
                outputString = "AmoebaTorsionGenerator missing type: [%s %s %s %s] atoms: " % (type1, type2, type3, type4)
                outputString += getAtomPrint( data, torsion[0] ) + ' '
                outputString += getAtomPrint( data, torsion[1] ) + ' '
                outputString += getAtomPrint( data, torsion[2] ) + ' '
                outputString += getAtomPrint( data, torsion[3] )
                outputString += " indices: [%6d %6d %6d %6d]" % (torsion[0], torsion[1], torsion[2], torsion[3])
                raise ValueError(outputString) 
2140
2141
2142
2143
2144

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

2145
## @private
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
class AmoebaPiTorsionGenerator:

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

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

    #=============================================================================================
    
    def __init__(self, piTorsionUnit):
        self.piTorsionUnit = piTorsionUnit 
Peter Eastman's avatar
Peter Eastman committed
2156
2157
2158
        self.types1 = []
        self.types2 = []
        self.k = []
2159
2160
2161
2162
2163
2164
2165
2166
2167
    
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
2168
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
2169
2170
2171
2172
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
            types = forceField._findAtomTypes(piTorsion, 2)
peastman's avatar
peastman committed
2173
            if None not in types:
2174
2175
2176
2177
2178
2179
                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
2180
2181
                                    piTorsion.attrib['class2'])
                raise ValueError(outputString) 
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
    
    #=============================================================================================

    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
 
            atom1 = bond.atom1
            atom2 = bond.atom2
 
Peter Eastman's avatar
Peter Eastman committed
2203
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom1]) == 3):
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228

                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):

                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2 
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1 

                       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
2229
                           if (bondedAtom1 != atom1):
2230
2231
2232
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
2233
2234
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
2235
2236
2237
2238
2239
2240
2241
                                   piTorsionAtom1 = b1 
                               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
2242
                           if (bondedAtom1 != atom2):
2243
2244
2245
2246
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
2247
2248
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
2249
2250
2251
2252
                                   piTorsionAtom5 = b1 
                               else:
                                   piTorsionAtom6 = b1
    
Peter Eastman's avatar
Peter Eastman committed
2253
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
2254
2255
2256
2257
2258

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

2259
## @private
2260
2261
2262
2263
2264
2265
class AmoebaTorsionTorsionGenerator:

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

Peter Eastman's avatar
Peter Eastman committed
2266
    def __init__(self):
2267

Peter Eastman's avatar
Peter Eastman committed
2268
2269
2270
2271
2272
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
2273

Peter Eastman's avatar
Peter Eastman committed
2274
        self.gridIndex = []
2275

Peter Eastman's avatar
Peter Eastman committed
2276
        self.grids = []
2277
2278
2279
2280
2281
2282
    
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
2283
        generator = AmoebaTorsionTorsionGenerator()
2284
2285
2286
2287
2288
2289
2290
        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'):
Peter Eastman's avatar
Peter Eastman committed
2291
            types = forceField._findAtomTypes(torsionTorsion, 5)
peastman's avatar
peastman committed
2292
            if None not in types:
2293
2294
2295
2296
2297
2298
2299

                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
2300
2301
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
2302
2303
2304
2305
2306
2307
2308
2309
2310
                    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
2311
2312
                                    torsionTorsion.attrib['class5'] )
                raise ValueError(outputString) 
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
    
        # 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:

        #     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 

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

Peter Eastman's avatar
Peter Eastman committed
2335
2336
2337
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
2338

Peter Eastman's avatar
Peter Eastman committed
2339
2340
            grid = []
            gridCol = []
2341
2342
2343
2344
2345

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
2346
2347
2348
2349
2350
2351
2352
2353
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
                gridRow.append(float(gridEntry.attrib['fx']))
                gridRow.append(float(gridEntry.attrib['fy']))
                gridRow.append(float(gridEntry.attrib['fxy']))
                gridCol.append(gridRow)
2354
2355

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
2356
2357
2358
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
2359
2360
2361
                    gridColIndex = 0

            
Peter Eastman's avatar
Peter Eastman committed
2362
2363
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
2364
            else:
Peter Eastman's avatar
Peter Eastman committed
2365
2366
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
2367
2368
2369
2370
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
2371
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
2372
2373
2374
2375
2376
2377
2378
2379

        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
2380
        if (len(data.atomBonds[atomC]) == 4):
2381
2382
2383
2384
2385
            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
2386
2387
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
2388
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
2389
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
2390
2391
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
2392
2393
                if (hit > -1):
                    if (atomE == -1):
2394
2395
2396
2397
2398
2399
                        atomE = hit
                    else:
                        atomF = hit
       
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
2400
2401
2402
            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,)
                raise ValueError(outputString) 
2403
2404
2405
2406
2407

            # 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
2408
            if (typeE > typeF):
2409
                chiralAtomIndex = atomE 
Peter Eastman's avatar
Peter Eastman committed
2410
            if (typeF > typeE):
2411
2412
                chiralAtomIndex = atomF 

Peter Eastman's avatar
Peter Eastman committed
2413
2414
2415
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
2416
                chiralAtomIndex = massE 
Peter Eastman's avatar
Peter Eastman committed
2417
            if (massF > massE):
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
                chiralAtomIndex = massF 

        return chiralAtomIndex

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

            # search for bitorsions; based on TINKER subroutine bitors() 
 
            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
2446
                if (bondedAtom1 != ib):
2447
2448
2449
2450
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
2451
                if (ia != ic and ia != id):
2452
2453
2454
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
2455
                        if (bondedAtom1 != id):
2456
2457
2458
2459
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
2460
                        if (ie != ic and ie != ib and ie != ia):
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480

                            # 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
2481
2482
2483
                                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])
2484
2485
2486

                                # match in reverse order

Peter Eastman's avatar
Peter Eastman committed
2487
2488
2489
                                if (type5 in types1 and type4 in types2 and type3 in types3 and type2 in types4 and type1 in types5):
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
2490
2491
2492
2493

        # set grids

        for (index, grid) in enumerate(self.grids):
Peter Eastman's avatar
Peter Eastman committed
2494
            force.setTorsionTorsionGrid(index, grid)
2495
2496
2497
2498
2499
 
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

2500
## @private
2501
class AmoebaStretchBendGenerator:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2502
2503

    #=============================================================================================
2504
2505
2506
2507
2508
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
2509
2510
2511
        self.types1 = []
        self.types2 = []
        self.types3 = []
2512

Peter Eastman's avatar
Peter Eastman committed
2513
2514
        self.k1 = []
        self.k2 = []
2515
2516
2517
2518
2519
    
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
2520
        generator = AmoebaStretchBendGenerator()
2521
2522
2523
2524
2525
2526
2527
2528
        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'):
            types = forceField._findAtomTypes(stretchBend, 3)
peastman's avatar
peastman committed
2529
            if None not in types:
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541

                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
2542
2543
                                    stretchBend.attrib['class3'])
                raise ValueError(outputString) 
2544
2545
2546
    
    #=============================================================================================

2547
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce 
2548
2549
    # 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
2550
2551
2552
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator(). 
    # Instead, createForcePostAmoebaBondForce() is called 
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
2553
2554
2555

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

Peter Eastman's avatar
Peter Eastman committed
2556
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2557
2558
2559
2560
2561
2562
2563
2564
        pass

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

    # Note: request for constrained bonds is ignored.

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

2565
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577

        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
2578
            if ('isConstrained' in angleDict):
2579
2580
2581
2582
2583
2584
2585
2586
                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
2587
            radian = 57.2957795130
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
            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

Peter Eastman's avatar
Peter Eastman committed
2598
                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
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
                    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
                    
                    # check that ideal angle and bonds are set
2616

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2617
                    if ('idealAngle' not in angleDict):
2618

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2619
2620
2621
2622
2623
2624
                       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] )
                       raise ValueError(outputString) 
2625

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2626
2627
2628
2629
2630
2631
2632
2633
                    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] )
                       raise ValueError(outputString) 
2634

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2638
                    break
2639
2640
2641
2642
2643

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

2644
## @private
2645
2646
2647
2648
2649
2650
class AmoebaVdwGenerator:

    """A AmoebaVdwGenerator constructs a AmoebaVdwForce."""
    
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
2653
        self.type = type 
2654

Peter Eastman's avatar
Peter Eastman committed
2655
2656
2657
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
2658

Peter Eastman's avatar
Peter Eastman committed
2659
        self.epsilonrule = epsilonrule
2660

Peter Eastman's avatar
Peter Eastman committed
2661
2662
2663
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
2664

Peter Eastman's avatar
Peter Eastman committed
2665
        self.typeMap = {}
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675

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

    @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" >
        #   <Vdw class="1" sigma="0.371" epsilon="0.46024" reduction="1.0" /> 
        #   <Vdw class="2" sigma="0.382" epsilon="0.422584" reduction="1.0" /> 
         
Peter Eastman's avatar
Peter Eastman committed
2676
2677
        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'])) 
2678
2679
2680
2681
2682
2683
2684
2685
        forceField._forces.append(generator)
        two_six = 1.122462048309372

        # types[] = [ sigma, epsilon, reductionFactor, class ]
        # sigma is modified based on radiustype and radiussize

        for atom in element.findall('Vdw'):
            types = forceField._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
2686
            if None not in types:
2687

Peter Eastman's avatar
Peter Eastman committed
2688
                values = [float(atom.attrib['sigma']), float(atom.attrib['epsilon']), float(atom.attrib['reduction'])]
2689
2690


Peter Eastman's avatar
Peter Eastman committed
2691
                if (generator.radiustype == 'SIGMA'):
2692
2693
                    values[0] *= two_six
      
Peter Eastman's avatar
Peter Eastman committed
2694
                if (generator.radiussize == 'DIAMETER'):
2695
2696
2697
2698
2699
2700
                    values[0] *= 0.5

                for t in types[0]:
                    generator.typeMap[t] = values
    
            else:
Peter Eastman's avatar
Peter Eastman committed
2701
2702
                outputString = "AmoebaVdwGenerator: error getting type: %s" % (atom.attrib['class'])
                raise ValueError(outputString) 
2703
2704
2705
2706
2707
2708
2709
2710
    
    #=============================================================================================

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

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
2711
    def getBondedParticleSet(particleIndex, data):
2712
2713
2714
2715
2716
2717

        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
2718
2719
            if (atom1 != particleIndex):
                bondedParticleSet.add(atom1)
2720
            else:
Peter Eastman's avatar
Peter Eastman committed
2721
                bondedParticleSet.add(atom2)
2722
2723
2724
2725
2726

        return bondedParticleSet
          
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
2729
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
        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
2742
            if ('sigmaCombiningRule' in args):
2743
                sigmaRule = args['sigmaCombiningRule'].upper()
Peter Eastman's avatar
Peter Eastman committed
2744
2745
                if (sigmaRule.upper() in sigmaMap):
                    force.setSigmaCombiningRule(sigmaRule.upper())
2746
                else:
Peter Eastman's avatar
Peter Eastman committed
2747
                    stringList = ' ' . join(str(x) for x in sigmaMap.keys())
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2748
                    raise ValueError( "AmoebaVdwGenerator: sigma combining rule %s not recognized; valid values are %s; using default." % (sigmaRule, stringList) )
2749
            else:
Peter Eastman's avatar
Peter Eastman committed
2750
                force.setSigmaCombiningRule(self.radiusrule)
2751

Peter Eastman's avatar
Peter Eastman committed
2752
            if ('epsilonCombiningRule' in args):
2753
                epsilonRule = args['epsilonCombiningRule'].upper()
Peter Eastman's avatar
Peter Eastman committed
2754
2755
                if (epsilonRule.upper() in epsilonMap):
                    force.setEpsilonCombiningRule(epsilonRule.upper())
2756
                else:
Peter Eastman's avatar
Peter Eastman committed
2757
                    stringList = ' ' . join(str(x) for x in epsilonMap.keys())
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2758
                    raise ValueError( "AmoebaVdwGenerator: epsilon combining rule %s not recognized; valid values are %s; using default." % (epsilonRule, stringList) )  
2759
            else:
Peter Eastman's avatar
Peter Eastman committed
2760
                force.setEpsilonCombiningRule(self.epsilonrule)
2761
2762
2763
               
            # cutoff

Peter Eastman's avatar
Peter Eastman committed
2764
            if ('vdwCutoff' in args):
2765
                force.setCutoff(args['vdwCutoff'])
2766
            else:
Peter Eastman's avatar
Peter Eastman committed
2767
                force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2768

2769
2770
2771
2772
2773
            # dispersion correction

            if ('useDispersionCorrection' in args):
                force.setUseDispersionCorrection(int(args['useDispersionCorrection']))

Peter Eastman's avatar
Peter Eastman committed
2774
            if (nonbondedMethod == PME):
2775
                force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
           
        else:
            force = existing[0]

        # add particles to force
        # throw error if particle type not available 

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

Peter Eastman's avatar
Peter Eastman committed
2787
                values = self.typeMap[t]
2788
2789
2790
   
                # ivIndex = index of bonded partner for hydrogens; otherwise ivIndex = particle index

Peter Eastman's avatar
Peter Eastman committed
2791
2792
2793
                ivIndex = i
                mass = sys.getParticleMass(i)/unit.dalton
                if (mass < 1.9 and len(data.atomBonds[i]) == 1):
2794
                    bondIndex = data.atomBonds[i][0]
Peter Eastman's avatar
Peter Eastman committed
2795
                    if (data.bonds[bondIndex].atom1 == i):
2796
2797
2798
2799
                        ivIndex = data.bonds[bondIndex].atom2
                    else:
                        ivIndex = data.bonds[bondIndex].atom1

2800
                force.addParticle(ivIndex, values[0], values[1], values[2])
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
            else:
                raise ValueError('No vdw type for atom %s' % (atom.name))

        # 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
2812
            bondedParticleSets.append(AmoebaVdwGenerator.getBondedParticleSet(i, data))
2813
2814
2815
2816
2817
2818
2819
2820
2821

        for (i,atom) in enumerate(data.atoms):
 
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
2822
            if (self.vdw13Scale == 0.0):
2823
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
2824
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
2825
2826
2827
2828
2829

            # self

            exclusionSet.add(i)

Peter Eastman's avatar
Peter Eastman committed
2830
            force.setParticleExclusions(i, exclusionSet)
2831
2832
2833
2834
2835

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

2836
## @private
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
class AmoebaMultipoleGenerator:

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

    """A AmoebaMultipoleGenerator constructs a AmoebaMultipoleForce."""
    
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
2851
        self.forceField = forceField
2852

Peter Eastman's avatar
Peter Eastman committed
2853
2854
2855
2856
        self.direct11Scale = direct11Scale 
        self.direct12Scale = direct12Scale 
        self.direct13Scale = direct13Scale 
        self.direct14Scale = direct14Scale 
2857

Peter Eastman's avatar
Peter Eastman committed
2858
2859
2860
2861
        self.mpole12Scale = mpole12Scale 
        self.mpole13Scale = mpole13Scale 
        self.mpole14Scale = mpole14Scale 
        self.mpole15Scale = mpole15Scale 
2862

Peter Eastman's avatar
Peter Eastman committed
2863
2864
2865
2866
        self.mutual11Scale = mutual11Scale 
        self.mutual12Scale = mutual12Scale 
        self.mutual13Scale = mutual13Scale 
        self.mutual14Scale = mutual14Scale 
2867

Peter Eastman's avatar
Peter Eastman committed
2868
2869
2870
2871
        self.polar12Scale = polar12Scale 
        self.polar13Scale = polar13Scale 
        self.polar14Scale = polar14Scale 
        self.polar15Scale = polar15Scale 
2872

Peter Eastman's avatar
Peter Eastman committed
2873
        self.typeMap = {}
2874
2875
2876
2877
2878
2879

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
2880
    def setAxisType(kIndices):
2881
2882
2883

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
2884
2885
2886
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
2887
                else:
Peter Eastman's avatar
Peter Eastman committed
2888
                    ky = 0
2889
   
Peter Eastman's avatar
Peter Eastman committed
2890
2891
                if (kIndicesLen > 2):
                    kx = kIndices[2]
2892
                else:
Peter Eastman's avatar
Peter Eastman committed
2893
                    kx = 0
2894
   
Peter Eastman's avatar
Peter Eastman committed
2895
2896
                if (kIndicesLen > 1):
                    kz = kIndices[1]
2897
                else:
Peter Eastman's avatar
Peter Eastman committed
2898
                    kz = 0
2899

Peter Eastman's avatar
Peter Eastman committed
2900
2901
                while(len(kIndices) < 4):
                    kIndices.append(0)
2902
2903

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
2904
                if (kz == 0):
2905
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
2906
                if (kz != 0 and kx == 0):
2907
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
2908
                if (kz < 0 or kx < 0):
2909
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
2910
                if (kx < 0 and ky < 0):
2911
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
2912
                if (kz < 0 and kx < 0 and ky  < 0):
2913
2914
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Peter Eastman's avatar
Peter Eastman committed
2915
2916
2917
                kIndices[1] = abs(kz)   
                kIndices[2] = abs(kx)   
                kIndices[3] = abs(ky)   
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

        #   <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"  > 
        # <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
2930
        generator = AmoebaMultipoleGenerator(forceField,
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
                                              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
2949
                                              element.attrib['polar15Scale'])
2950
2951
2952
2953
2954
2955
2956
2957
2958



        forceField._forces.append(generator)

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

        for atom in element.findall('Multipole'):
            types = forceField._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
2959
            if None not in types:
2960
2961
2962
2963
2964
2965
2966
2967

                # 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
2968
2969
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
2970
2971
2972
2973
2974
                    except: 
                        pass

                # set axis type based on k-Indices 

Peter Eastman's avatar
Peter Eastman committed
2975
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
2976
2977
2978

                # set multipole

Peter Eastman's avatar
Peter Eastman committed
2979
                charge = float(atom.attrib['c0'])
2980
 
Peter Eastman's avatar
Peter Eastman committed
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
                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']))
2994
2995

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
2996
                    if (t not in generator.typeMap):
2997
2998
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
2999
3000
3001
3002
3003
3004
3005
3006
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
                    valueMap['charge'] = charge 
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
3007
3008
    
            else:
Peter Eastman's avatar
Peter Eastman committed
3009
3010
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
                raise ValueError(outputString) 
3011
3012
3013
3014
3015
    
        # polarization parameters
 
        for atom in element.findall('Polarize'):
            types = forceField._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
3016
            if None not in types:
3017

Peter Eastman's avatar
Peter Eastman committed
3018
3019
3020
3021
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
3022
3023
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
3024
                    pdamp = pow(polarizability, 1.0/6.0)
3025

Peter Eastman's avatar
Peter Eastman committed
3026
3027
                pgrpMap = dict()
                for index in range(1, 7):
3028
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
3029
                    if (pgrp in atom.attrib):
3030
3031
3032
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
3033
3034
3035
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
                        raise ValueError(outputString) 
3036
3037
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
3038
3039
3040
3041
3042
3043
3044
3045
3046
3047
3048
3049
3050
3051
                        hit = 0
                        for (ii, typeMap) in enumerate(typeMapList):

                            if (typeMap['classIndex'] == classIndex):
                                typeMap['polarizability'] = polarizability
                                typeMap['thole'] = thole
                                typeMap['pdamp'] = pdamp
                                typeMap['pgrpMap'] = pgrpMap 
                                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'])
                            raise ValueError(outputString) 
3052
3053
    
            else:
Peter Eastman's avatar
Peter Eastman committed
3054
3055
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
                raise ValueError(outputString) 
3056
3057
3058
    
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3059
    def setPolarGroups(self, data, bonded12ParticleSets, force):
3060
3061
3062
3063
3064

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
3065
3066
3067
3068
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
3069
3070
3071
3072
            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
3073
3074
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
3075
3076
3077
3078
3079
3080
3081
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
                    
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
3082
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
3083
3084
                continue

Peter Eastman's avatar
Peter Eastman committed
3085
3086
            group = set()
            visited = set()
3087
3088
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
3089
3090
3091
3092
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
3093
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
3094
3095
                if (nextAtom not in visited):
                   visited.add(nextAtom)
3096
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
3097
3098
3099
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
3100
3101
3102

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
3103
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
3104
3105

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3106
3107
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
3108
3109
3110
3111
3112

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
3113
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
3114
3115
                continue

Peter Eastman's avatar
Peter Eastman committed
3116
            pgrp11 = set(atom.polarizationGroupSet[0])
3117
3118
3119
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3120
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
3121
3122
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3123
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
3124
3125
                
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3126
3127
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
3128
3129
3130
3131
3132

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
3133
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
3134
3135
                continue

Peter Eastman's avatar
Peter Eastman committed
3136
3137
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
3138
3139
3140
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3141
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
3142
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
3143
            pgrp13 = pgrp13 - set(pgrp11)
3144
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3145
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
3146
3147
                
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3148
3149
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
3150
3151
3152
3153
3154

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
3155
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
3156
3157
                continue

Peter Eastman's avatar
Peter Eastman committed
3158
3159
3160
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
3161
3162
3163
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3164
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
3165
3166
3167

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
3168
            pgrp14 = pgrp14 - set(pgrp11)
3169
3170

            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3171
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp14)
3172
3173
                
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3174
3175
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
3176
3177
3178

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

Peter Eastman's avatar
Peter Eastman committed
3179
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3180
3181
3182
3183
3184
3185
3186
3187
3188
3189
3190

        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)
Peter Eastman's avatar
Peter Eastman committed
3191
            if (nonbondedMethod not in methodMap): 
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3192
                raise ValueError( "AmoebaMultipoleForce: input cutoff method not available." )
3193
            else:
Peter Eastman's avatar
Peter Eastman committed
3194
                force.setNonbondedMethod(methodMap[nonbondedMethod])
3195
3196
            force.setCutoffDistance(nonbondedCutoff)

Peter Eastman's avatar
Peter Eastman committed
3197
            if ('ewaldErrorTolerance' in args):
3198
3199
                force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))

Peter Eastman's avatar
Peter Eastman committed
3200
            if ('polarization' in args):
3201
                polarizationType = args['polarization']
Peter Eastman's avatar
Peter Eastman committed
3202
                if (polarizationType.lower() == 'direct'):
3203
3204
3205
3206
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Direct)
                else:
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Mutual)

Peter Eastman's avatar
Peter Eastman committed
3207
3208
            if ('aEwald' in args):
                force.setAEwald(float(args['aEwald']))
3209

Peter Eastman's avatar
Peter Eastman committed
3210
            if ('pmeGridDimensions' in args):
3211
3212
                force.setPmeGridDimensions(args['pmeGridDimensions'])

Peter Eastman's avatar
Peter Eastman committed
3213
3214
            if ('mutualInducedMaxIterations' in args):
                force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
3215

Peter Eastman's avatar
Peter Eastman committed
3216
            if ('mutualInducedTargetEpsilon' in args):
3217
3218
3219
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230
                force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))

        else:
            force = existing[0]

        # add particles to force
        # throw error if particle type not available 

        # 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
3231
3232
3233
            bonded12ParticleSet = AmoebaVdwGenerator.getBondedParticleSet(i, data)
            bonded12ParticleSet = set(sorted(bonded12ParticleSet))
            bonded12ParticleSets.append(bonded12ParticleSet)
3234
3235
3236
3237
3238

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3239
            bonded13Set = set()
3240
3241
            bonded12ParticleSet = bonded12ParticleSets[i]
            for j in bonded12ParticleSet: 
Peter Eastman's avatar
Peter Eastman committed
3242
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
3243
3244
3245
3246

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
3247
            selfSet = set()
3248
3249
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3250
3251
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
3252
3253
3254
3255
3256

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3257
3258
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
3259
            for j in bonded13ParticleSet: 
Peter Eastman's avatar
Peter Eastman committed
3260
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
3261
3262
3263
3264
3265
           
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
3266
            selfSet = set()
3267
3268
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3269
3270
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
3271
3272
3273
3274
3275

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3276
3277
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
3278
            for j in bonded14ParticleSet: 
Peter Eastman's avatar
Peter Eastman committed
3279
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
3280
3281
3282
3283
3284
3285

            # 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
3286
            selfSet = set()
3287
3288
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3289
3290
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
3291
3292
3293
3294
3295

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

Peter Eastman's avatar
Peter Eastman committed
3296
3297
                multipoleList = self.typeMap[t]
                hit = 0
3298
3299
3300
3301
3302
3303
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3304
                    if (hit != 0):
3305
3306
                        break

Peter Eastman's avatar
Peter Eastman committed
3307
                    kIndices = multipoleDict['kIndices']
3308
    
Peter Eastman's avatar
Peter Eastman committed
3309
3310
3311
                    kz = kIndices[1]   
                    kx = kIndices[2]
                    ky = kIndices[3]
3312
3313
3314
3315
3316
3317

                    # assign multipole parameters
                    #    (1) get bonded partners
                    #    (2) match parameter types
    
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
3318
3319
3320
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3321
3322
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
3323
                       if (hit != 0):
3324
3325
3326
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
3327
3328
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
3329
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
3330
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
3331
3332
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
3333
3334
3335
3336
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
3337
3338
3339
3340
3341
3342
3343
3344
3345
3346
                                      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

3347
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3348
                                      hit = 1
3349
3350
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
3351
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
3352
3353
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
3354
3355
3356
3357
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
3358
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3359
                                              hit = 2
3360
3361
3362
3363
3364
                                         
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3365
                    if (hit != 0):
3366
3367
                        break

Peter Eastman's avatar
Peter Eastman committed
3368
                    kIndices = multipoleDict['kIndices']
3369
    
Peter Eastman's avatar
Peter Eastman committed
3370
3371
3372
                    kz = kIndices[1]   
                    kx = kIndices[2]
                    ky = kIndices[3]
3373
3374
3375
3376
3377
3378
3379
3380
    
                    # assign multipole parameters
                    #    (1) get bonded partners
                    #    (2) match parameter types
    
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
3381
3382
3383
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3384
3385
3386

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
3387
                       if (hit != 0):
3388
3389
3390
                           break

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

Peter Eastman's avatar
Peter Eastman committed
3393
                       if (bondedAtomZType == kz):
3394
3395
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
3396
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
3397
3398
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
3399
3400
3401
3402
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
3403
3404
3405
3406
3407
3408
3409
3410

                                      # 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

3411
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3412
                                      hit = 3
3413
3414
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
3415
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
3416
3417
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
3418
3419
3420
3421
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
3422
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3423
                                              hit = 4
3424
3425
3426
3427
3428
                                         
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3429
                    if (hit != 0):
3430
3431
                        break

Peter Eastman's avatar
Peter Eastman committed
3432
                    kIndices = multipoleDict['kIndices']
3433
    
Peter Eastman's avatar
Peter Eastman committed
3434
3435
                    kz = kIndices[1]   
                    kx = kIndices[2]   
3436
    
Peter Eastman's avatar
Peter Eastman committed
3437
3438
3439
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3440
3441
3442

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
3443
                        if (hit != 0):
3444
3445
3446
                            break

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

Peter Eastman's avatar
Peter Eastman committed
3449
3450
                        if (kx == 0 and kz == bondedAtomZType):
                            kz = bondedAtomZIndex
3451
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3452
                            hit = 5
3453
3454
3455
3456
3457

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3458
                    if (hit != 0):
3459
3460
                        break

Peter Eastman's avatar
Peter Eastman committed
3461
                    kIndices = multipoleDict['kIndices']
3462
    
Peter Eastman's avatar
Peter Eastman committed
3463
                    kz = kIndices[1]   
3464
    
Peter Eastman's avatar
Peter Eastman committed
3465
3466
3467
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3468

Peter Eastman's avatar
Peter Eastman committed
3469
                    if (kz == 0):
3470
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3471
                        hit = 6
3472
3473
3474
                
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
3475
                if (hit != 0):
3476

Peter Eastman's avatar
Peter Eastman committed
3477
                    atom.multipoleDict = savedMultipoleDict
3478
                    atom.polarizationGroups = dict()
3479
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
3480
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
3481
3482
3483
3484
3485
                    if (atomIndex == newIndex):
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent12, bonded12ParticleSets[atomIndex])
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent13, bonded13ParticleSets[atomIndex])
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent14, bonded14ParticleSets[atomIndex])
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent15, bonded15ParticleSets[atomIndex])
3486
                    else:
Peter Eastman's avatar
Peter Eastman committed
3487
                        raise ValueError("Atom %s of %s %d is out of synch!." %(atom.name, atom.residue.name, atom.residue.index))
3488
                else:
Peter Eastman's avatar
Peter Eastman committed
3489
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
3490
            else:
Peter Eastman's avatar
Peter Eastman committed
3491
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
3492
3493
3494

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
3495
        self.setPolarGroups(data, bonded12ParticleSets, force)
3496
3497
3498
3499
3500

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

3501
## @private
3502
3503
3504
3505
3506
3507
class AmoebaWcaDispersionGenerator:

    """A AmoebaWcaDispersionGenerator constructs a AmoebaWcaDispersionForce."""
    
    #=========================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3510
3511
3512
3513
3514
3515
3516
3517
        self.epso = epso 
        self.epsh = epsh 
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
        self.shctd = shctd 
3518

Peter Eastman's avatar
Peter Eastman committed
3519
        self.typeMap = {}
3520
3521
3522
3523
3524
3525
3526
3527
3528
3529

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

    @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" />
      
Peter Eastman's avatar
Peter Eastman committed
3530
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
3531
3532
3533
3534
3535
3536
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
                                                  element.attrib['awater'], 
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Peter Eastman's avatar
Peter Eastman committed
3537
                                                  element.attrib['shctd']) 
3538
3539
3540
3541
3542
3543
        forceField._forces.append(generator)

        # typeMap[] = [ radius, epsilon ]

        for atom in element.findall('WcaDispersion'):
            types = forceField._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
3544
            if None not in types:
3545
3546
3547
3548
3549

                values = [float(atom.attrib['radius']), float(atom.attrib['epsilon'])]
                for t in types[0]:
                    generator.typeMap[t] = values
            else:
Peter Eastman's avatar
Peter Eastman committed
3550
3551
                outputString = "AmoebaWcaDispersionGenerator: error getting type: %s" % (atom.attrib['class'])
                raise ValueError(outputString) 
3552
3553
3554
    
    #=========================================================================================
    
Peter Eastman's avatar
Peter Eastman committed
3555
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3556
3557
3558
3559
3560
3561
3562
3563
3564
3565
3566
3567
3568
3569

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

Peter Eastman's avatar
Peter Eastman committed
3570
3571
3572
3573
3574
3575
3576
3577
        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  ))
3578
3579
3580
3581
3582

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

Peter Eastman's avatar
Peter Eastman committed
3583
3584
                values = self.typeMap[t]
                force.addParticle(values[0], values[1])
3585
3586
3587
3588
3589
3590
3591
            else:
                raise ValueError('No WcaDispersion type for atom %s of %s %d' % (atom.name, atom.residue.name, atom.residue.index))

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

3592
## @private
3593
3594
3595
3596
3597
3598
class AmoebaGeneralizedKirkwoodGenerator:

    """A AmoebaGeneralizedKirkwoodGenerator constructs a AmoebaGeneralizedKirkwoodForce."""
    
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
3599
3600
3601
3602
3603
3604
3605
3606
3607
3608
3609
3610
3611
3612
3613
3614
3615
3616
3617
3618
3619
3620
3621
3622
3623
3624
3625
3626
3627
3628
3629
3630
3631
3632
3633
3634
3635
    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'] = {}
        bondiMap = self.radiusTypeMap['Bondi'] 
        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
3636
3637
3638

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

Peter Eastman's avatar
Peter Eastman committed
3639
    def getObcShct(self, data, atomIndex):
3640

Peter Eastman's avatar
Peter Eastman committed
3641
        atom = data.atoms[atomIndex]
3642
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
3643
        shct = -1.0
3644
3645
3646

        # shct
 
Peter Eastman's avatar
Peter Eastman committed
3647
        if (atomicNumber == 1):                 # H(1)
3648
            shct = 0.85       
Peter Eastman's avatar
Peter Eastman committed
3649
        elif (atomicNumber == 6):               # C(6)
3650
            shct = 0.72         
Peter Eastman's avatar
Peter Eastman committed
3651
        elif (atomicNumber == 7):               # N(7)
3652
            shct = 0.79        
Peter Eastman's avatar
Peter Eastman committed
3653
        elif (atomicNumber == 8):               # O(8)
3654
            shct = 0.85       
Peter Eastman's avatar
Peter Eastman committed
3655
        elif (atomicNumber == 9):               # F(9)
3656
            shct = 0.88    
Peter Eastman's avatar
Peter Eastman committed
3657
        elif (atomicNumber == 15):              # P(15)              
3658
            shct = 0.86 
Peter Eastman's avatar
Peter Eastman committed
3659
        elif (atomicNumber == 16):              # S(16)
3660
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
3661
        elif (atomicNumber == 26):              # Fe(26)
3662
3663
            shct = 0.88

Peter Eastman's avatar
Peter Eastman committed
3664
        if (shct < 0.0): 
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3665
            raise ValueError( "getObcShct: no GK overlap scale factor for atom %s of %s %d" % (atom.name, atom.residue.name, atom.residue.index) )
3666
3667
3668
3669
3670
 
        return shct 

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

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

Peter Eastman's avatar
Peter Eastman committed
3673
        atom = data.atoms[atomIndex]
3674
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
3675
        radius = -1.0
3676

Peter Eastman's avatar
Peter Eastman committed
3677
        if (atomicNumber == 1):                  # H(1)
3678
 
Peter Eastman's avatar
Peter Eastman committed
3679
            radius = 0.132
3680
 
Peter Eastman's avatar
Peter Eastman committed
3681
            if (len(bondedAtomIndices) < 1):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3682
                 raise ValueError( "AmoebaGeneralizedKirkwoodGenerator: error getting atom bonded to %s of %s %d " % (atom.name, atom.residue.name, atom.residue.index) )
3683
3684
 
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
3685
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
3686

Peter Eastman's avatar
Peter Eastman committed
3687
            if (bondedAtomAtomicNumber == 7):
3688
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
3689
            if (bondedAtomAtomicNumber == 8):
3690
3691
                radius = 0.105
 
Peter Eastman's avatar
Peter Eastman committed
3692
        elif (atomicNumber == 3):               # Li(3)
3693
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
3694
        elif (atomicNumber == 6):               # C(6)
3695
3696
            
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
3697
            if (len(bondedAtomIndices) == 3):
3698
3699
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
3700
            elif (len(bondedAtomIndices) == 4):
3701
3702
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
3703
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
3704
3705
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
3706
        elif (atomicNumber == 7):               # N(7)
3707
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
3708
        elif (atomicNumber == 8):               # O(8)
3709
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
3710
            if (len(bondedAtomIndices) == 2):
3711
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
3712
        elif (atomicNumber == 9):               # F(9)
3713
            radius = 0.154
Peter Eastman's avatar
Peter Eastman committed
3714
        elif (atomicNumber == 10):              
3715
            radius = 0.146
Peter Eastman's avatar
Peter Eastman committed
3716
        elif (atomicNumber == 11):              
3717
            radius = 0.209
Peter Eastman's avatar
Peter Eastman committed
3718
        elif (atomicNumber == 12):              
3719
            radius = 0.179
Peter Eastman's avatar
Peter Eastman committed
3720
        elif (atomicNumber == 14):              
3721
            radius = 0.189
Peter Eastman's avatar
Peter Eastman committed
3722
        elif (atomicNumber == 15):              # P(15)              
3723
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
3724
        elif (atomicNumber == 16):              # S(16)
3725
            radius = 0.186
Peter Eastman's avatar
Peter Eastman committed
3726
        elif (atomicNumber == 17):              
3727
            radius = 0.182
Peter Eastman's avatar
Peter Eastman committed
3728
        elif (atomicNumber == 18):              
3729
            radius = 0.179
Peter Eastman's avatar
Peter Eastman committed
3730
        elif (atomicNumber == 19):              
3731
            radius = 0.223
Peter Eastman's avatar
Peter Eastman committed
3732
        elif (atomicNumber == 20):              
3733
            radius = 0.191
Peter Eastman's avatar
Peter Eastman committed
3734
        elif (atomicNumber == 35):         
3735
            radius = 2.00
Peter Eastman's avatar
Peter Eastman committed
3736
        elif (atomicNumber == 36):   
3737
            radius = 0.190
Peter Eastman's avatar
Peter Eastman committed
3738
        elif (atomicNumber == 37):              
3739
            radius = 0.226
Peter Eastman's avatar
Peter Eastman committed
3740
        elif (atomicNumber == 53):              
3741
            radius = 0.237
Peter Eastman's avatar
Peter Eastman committed
3742
        elif (atomicNumber == 54):              
3743
            radius = 0.207
Peter Eastman's avatar
Peter Eastman committed
3744
        elif (atomicNumber == 55):              
3745
            radius = 0.263
Peter Eastman's avatar
Peter Eastman committed
3746
        elif (atomicNumber == 56):         
3747
3748
            radius = 0.230

Peter Eastman's avatar
Peter Eastman committed
3749
        if (radius < 0.0): 
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3750
3751
            outputString = "No GK radius for atom %s of %s %d" % (atom.name, atom.residue.name, atom.residue.index)
            raise ValueError( outputString )
3752
3753
3754
3755
3756
 
        return radius

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

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

Peter Eastman's avatar
Peter Eastman committed
3759
3760
        bondiMap = self.radiusTypeMap['Bondi'] 
        atom = data.atoms[atomIndex]
3761
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
3762
        if (atomicNumber in bondiMap): 
3763
3764
            radius = bondiMap[atomicNumber]
        else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3765
3766
            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 )
3767
3768
3769
3770
3771
3772
3773
3774
3775
3776
3777
3778
 
        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"  />
        
Peter Eastman's avatar
Peter Eastman committed
3779
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
3780
                                                        element.attrib['includeCavityTerm'], 
Peter Eastman's avatar
Peter Eastman committed
3781
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor']) 
3782
3783
3784
3785
        forceField._forces.append(generator)

    #=========================================================================================
    
Peter Eastman's avatar
Peter Eastman committed
3786
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3787

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

3791
3792
3793
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
3794
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3795
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
3796
        if (len(amoebaMultipoleForceList) > 0):
3797
3798
3799
3800
3801
            amoebaMultipoleForce = amoebaMultipoleForceList[0]
        else:
            # call AmoebaMultipoleForceGenerator.createForce() to ensure charges have been set

            for force in self.forceField._forces:
Peter Eastman's avatar
Peter Eastman committed
3802
3803
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'): 
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3804
          
3805
3806
3807
3808
3809
3810
3811
3812
        # 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)
 
Peter Eastman's avatar
Peter Eastman committed
3813
3814
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
3815
            else:
Peter Eastman's avatar
Peter Eastman committed
3816
                force.setSolventDielectric(   float(self.solventDielectric))
3817

Peter Eastman's avatar
Peter Eastman committed
3818
3819
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
3820
            else:
Peter Eastman's avatar
Peter Eastman committed
3821
                force.setSoluteDielectric(    float(self.soluteDielectric))
3822

Peter Eastman's avatar
Peter Eastman committed
3823
3824
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
3825
            else:
Peter Eastman's avatar
Peter Eastman committed
3826
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
3827
3828
3829
3830
3831
3832
3833

        else:
            force = existing[0]

        # add particles to force
        # throw error if particle type not available 

Peter Eastman's avatar
Peter Eastman committed
3834
3835
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
3836
3837
3838
3839
3840

        # 1-2

        bonded12ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3841
3842
3843
            bonded12ParticleSet = AmoebaVdwGenerator.getBondedParticleSet(i, data)
            bonded12ParticleSet = set(sorted(bonded12ParticleSet))
            bonded12ParticleSets.append(bonded12ParticleSet)
3844
3845

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
3846
3847
3848
3849
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
3850
            else:
Peter Eastman's avatar
Peter Eastman committed
3851
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
3852
3853
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
3854
            force.addParticle(multipoleParameters[0], radius, shct)
3855
3856
3857
3858
3859

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

3860
## @private
3861
3862
3863
3864
3865
3866
class AmoebaUreyBradleyGenerator:

    #=============================================================================================
    """An AmoebaUreyBradleyGenerator constructs a AmoebaUreyBradleyForce."""
    #=============================================================================================
    
3867
    def __init__(self):
3868

Peter Eastman's avatar
Peter Eastman committed
3869
3870
3871
        self.types1 = []
        self.types2 = []
        self.types3 = []
3872

Peter Eastman's avatar
Peter Eastman committed
3873
3874
        self.length = []
        self.k = []
3875
3876
3877
3878
3879
3880

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

    @staticmethod
    def parseElement(element, forceField):

3881
        #  <AmoebaUreyBradleyForce>
3882
3883
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" /> 

3884
        generator = AmoebaUreyBradleyGenerator()
3885
3886
3887
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
            types = forceField._findAtomTypes(bond, 3)
peastman's avatar
peastman committed
3888
            if None not in types:
3889
3890
3891
3892
3893
3894
3895
3896
3897
3898

                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
3899
3900
                                    bond.attrib['class1'], bond.attrib['class2'], bond.attrib['class3'])
                raise ValueError(outputString) 
3901
3902
3903
    
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3906
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3907
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
3908
3909

        if len(existing) == 0:
3910
            force = mm.HarmonicBondForce()
3911
3912
3913
3914
3915
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
Peter Eastman's avatar
Peter Eastman committed
3916
            if (isConstrained):
3917
3918
3919
3920
3921
3922
3923
3924
                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
3925
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
3926
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
3927
3928
3929
3930
3931
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
3932
3933
3934
3935
3936
3937
3938
3939
3940
3941
3942
3943
3944
3945
3946
3947
3948
3949
3950
3951
3952
3953
3954
3955
3956
3957
3958
3959
3960
3961
3962
3963
3964
3965
3966
3967
3968
3969
3970
3971
3972
3973
3974
3975
3976
3977
3978
3979
3980
3981
3982
3983
3984
3985
3986
3987
3988
3989
3990
3991
3992
3993
3994
3995
3996
3997
3998
3999
4000
4001
4002
4003
4004
4005
4006
4007
4008
4009
4010
4011
4012
4013
4014
4015


## @private
class DrudeGenerator:
    """A DrudeGenerator constructs a DrudeForce."""
    
    def __init__(self):
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
            generator = DrudeGenerator()
            ff._forces.append(generator)
        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'):
            types = ff._findAtomTypes(particle, 5)
            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
    
    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')
        
        # Add Drude particles.
        
        drudeMap = {}
        parentMap = {}
        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
                drudeIndex = force.addParticle(atom.index, p[0], p[1], p[2], p[3], values[4], values[5], values[6], values[7])
                drudeMap[atom.index] = p[0]
                parentMap[p[0]] = (atom.index, drudeIndex)
        sys.addForce(force)
    
    def postprocessSystem(self, sys, data, args):
        # For every nonbonded exclusion between Drude particles, add a screened pair.
        
        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)
            if charge == 0 and epsilon == 0:
                # 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]
                    type1 = data.atomType[data.atoms[drude1]]
                    type2 = data.atomType[data.atoms[drude2]]
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

parsers["DrudeForce"] = DrudeGenerator.parseElement