forcefield.py 175 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
Authors: Peter Eastman, Mark Friedrichs
Contributors:

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

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

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
USE OR OTHER DEALINGS IN THE SOFTWARE.
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
"""
__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

45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class NoCutoff(object):
    def __repr__(self):
        return 'NoCutoff'
NoCutoff = NoCutoff()

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

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

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

class PME(object):
    def __repr__(self):
        return 'PME'
PME = PME()
69
70
71

# Enumerated values for constraint type

72
73
74
75
76
77
78
79
80
81
82
83
84
85
class HBonds(object):
    def __repr__(self):
        return 'HBonds'
HBonds = HBonds()

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

class HAngles(object):
    def __repr__(self):
        return 'HAngles'
HAngles = HAngles()
86
87
88
89
90
91
92
93
94
95

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

parsers = {}

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

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

97
        Parameters:
98
         - files A list of XML files defining the force field.  Each entry may be an absolute file path, a path relative to the
99
100
101
102
           current working directory, or a path relative to this module's data subdirectory (for built in force fields).
        """
        self._atomTypes = {}
        self._templates = {}
103
        self._templateSignatures = {None:[]}
104
        self._atomClasses = {'':set()}
105
        self._forces = []
106
        self._scripts = []
107
108
109
110
111
112
        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()
Justin MacCallum's avatar
Justin MacCallum committed
113

114
            # Load the atom types.
Justin MacCallum's avatar
Justin MacCallum committed
115

116
117
            if tree.getroot().find('AtomTypes') is not None:
                for type in tree.getroot().find('AtomTypes').findall('Type'):
118
119
120
                    element = None
                    if 'element' in type.attrib:
                        element = elem.get_by_symbol(type.attrib['element'])
121
122
123
124
                    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)
Justin MacCallum's avatar
Justin MacCallum committed
125

126
            # Load the residue templates.
Justin MacCallum's avatar
Justin MacCallum committed
127

128
129
130
131
132
133
134
            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]))
135
136
                    for site in residue.findall('VirtualSite'):
                        template.virtualSites.append(ForceField._VirtualSiteData(site))
137
138
139
140
141
142
143
144
145
146
                    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():
147
148
149
                signature = _createResidueSignature([atom.element for atom in template.atoms])
                if signature in self._templateSignatures:
                    self._templateSignatures[signature].append(template)
150
                else:
151
                    self._templateSignatures[signature] = [template]
Justin MacCallum's avatar
Justin MacCallum committed
152

153
            # Build sets of every atom type belonging to each class
Justin MacCallum's avatar
Justin MacCallum committed
154

155
156
157
158
159
160
161
162
            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)
163
                self._atomClasses[''].add(type)
Justin MacCallum's avatar
Justin MacCallum committed
164

165
            # Load force definitions
Justin MacCallum's avatar
Justin MacCallum committed
166

167
168
169
            for child in root:
                if child.tag in parsers:
                    parsers[child.tag](child, self)
Justin MacCallum's avatar
Justin MacCallum committed
170

171
            # Load scripts
Justin MacCallum's avatar
Justin MacCallum committed
172

173
174
            for node in tree.getroot().findall('Script'):
                self._scripts.append(node.text)
175
176
177
178
179
180
181
182
183
184
185

    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
186
            typeAttrib = 'type'+suffix
187
            if classAttrib in attrib:
188
189
                if typeAttrib in attrib:
                    raise ValueError('Tag specifies both a type and a class for the same atom: '+etree.tostring(node))
190
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
191
192
193
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
194
195
            else:
                if typeAttrib not in attrib or attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
196
197
198
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
199
200
201
202
203
        return types

    def _parseTorsion(self, node):
        """Parse the node defining a torsion."""
        types = self._findAtomTypes(node, 4)
peastman's avatar
peastman committed
204
        if None in types:
205
206
207
208
209
210
211
212
213
            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
Justin MacCallum's avatar
Justin MacCallum committed
214
215
        return torsion

216
217
218
219
220
    class _SystemData:
        """Inner class used to encapsulate data about the system being created."""
        def __init__(self):
            self.atomType = {}
            self.atoms = []
221
            self.virtualSites = {}
222
223
224
225
226
227
228
229
230
231
232
233
            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 = []
234
            self.virtualSites = []
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
            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
Justin MacCallum's avatar
Justin MacCallum committed
254

255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
    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]
273
274

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

278
279
280
281
282
        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
283
         - constraints (object=None) Specifies which bonds and angles should be implemented with constraints.
284
285
           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
286
         - removeCMMotion (boolean=True) If true, a CMMotionRemover will be added to the System
287
288
         - hydrogenMass (mass=None) The mass to use for hydrogen atoms bound to heavy atoms.  Any mass added to a hydrogen is
           subtracted from the heavy atom to keep their total mass the same.
289
         - args Arbitrary additional keyword arguments may also be specified.  This allows extra parameters to be specified that are specific to
290
291
292
293
           particular force fields.
        Returns: the newly created System
        """
        data = ForceField._SystemData()
294
        data.atoms = list(topology.atoms())
295
296

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

298
        for bond in topology.bonds():
299
            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
300
301

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

303
304
305
306
307
308
309
310
311
312
313
314
        bondedToAtom = []
        for i in range(len(data.atoms)):
            bondedToAtom.append(set())
            data.atomBonds.append([])
        for i in range(len(data.bonds)):
            bond = data.bonds[i]
            bondedToAtom[bond.atom1].add(bond.atom2)
            bondedToAtom[bond.atom2].add(bond.atom1)
            data.atomBonds[bond.atom1].append(i)
            data.atomBonds[bond.atom2].append(i)

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

316
317
318
319
        for chain in topology.chains():
            for res in chain.residues():
                template = None
                matches = None
320
321
322
323
                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)
324
325
326
327
                        if matches is not None:
                            template = t
                            break
                if matches is None:
328
                    raise ValueError('No template found for residue %d (%s).  %s' % (res.index+1, res.name, _findMatchErrors(self, res)))
329
330
                for atom, match in zip(res.atoms(), matches):
                    data.atomType[atom] = template.atoms[match].type
331
332
333
                    for site in template.virtualSites:
                        if match == site.index:
                            data.virtualSites[atom] = site
334
335

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

337
338
339
        sys = mm.System()
        for atom in topology.atoms():
            sys.addParticle(self._atomTypes[data.atomType[atom]][1])
340
341
342
343
344
345
346
347
348
349
350
        
        # Adjust masses.
        
        if hydrogenMass is not None:
            for atom1, atom2 in topology.bonds():
                if atom1.element == elem.hydrogen:
                    (atom1, atom2) = (atom2, atom1)
                if atom2.element == elem.hydrogen and atom1.element not in (elem.hydrogen, None):
                    transferMass = hydrogenMass-sys.getParticleMass(atom2.index)
                    sys.setParticleMass(atom2.index, hydrogenMass)
                    sys.setParticleMass(atom1.index, sys.getParticleMass(atom1.index)-transferMass)
Justin MacCallum's avatar
Justin MacCallum committed
351

352
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
353

354
355
356
357
358
359
360
        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
Justin MacCallum's avatar
Justin MacCallum committed
361

362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
        uniqueAngles = set()
        for bond in data.bonds:
            for atom in bondedToAtom[bond.atom1]:
                if atom != bond.atom2:
                    if atom < bond.atom2:
                        uniqueAngles.add((atom, bond.atom1, bond.atom2))
                    else:
                        uniqueAngles.add((bond.atom2, bond.atom1, atom))
            for atom in bondedToAtom[bond.atom2]:
                if atom != bond.atom1:
                    if atom > bond.atom1:
                        uniqueAngles.add((bond.atom1, bond.atom2, atom))
                    else:
                        uniqueAngles.add((atom, bond.atom2, bond.atom1))
        data.angles = sorted(list(uniqueAngles))
Justin MacCallum's avatar
Justin MacCallum committed
377

378
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
379

380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
        uniquePropers = set()
        for angle in data.angles:
            for atom in bondedToAtom[angle[0]]:
                if atom != angle[1]:
                    if atom < angle[2]:
                        uniquePropers.add((atom, angle[0], angle[1], angle[2]))
                    else:
                        uniquePropers.add((angle[2], angle[1], angle[0], atom))
            for atom in bondedToAtom[angle[2]]:
                if atom != angle[1]:
                    if atom > angle[0]:
                        uniquePropers.add((angle[0], angle[1], angle[2], atom))
                    else:
                        uniquePropers.add((atom, angle[2], angle[1], angle[0]))
        data.propers = sorted(list(uniquePropers))
Justin MacCallum's avatar
Justin MacCallum committed
395

396
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
397

398
399
400
401
402
        for atom in range(len(bondedToAtom)):
            bondedTo = bondedToAtom[atom]
            if len(bondedTo) > 2:
                for subset in itertools.combinations(bondedTo, 3):
                    data.impropers.append((atom, subset[0], subset[1], subset[2]))
Justin MacCallum's avatar
Justin MacCallum committed
403

404
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
405

406
407
408
409
410
411
412
413
414
415
416
417
418
419
        if constraints == AllBonds or constraints == HAngles:
            for bond in data.bonds:
                bond.isConstrained = True
        elif constraints == HBonds:
            for bond in data.bonds:
                atom1 = data.atoms[bond.atom1]
                atom2 = data.atoms[bond.atom2]
                bond.isConstrained = atom1.name.startswith('H') or atom2.name.startswith('H')
        if rigidWater:
            for bond in data.bonds:
                atom1 = data.atoms[bond.atom1]
                atom2 = data.atoms[bond.atom2]
                if atom1.residue.name == 'HOH' and atom2.residue.name == 'HOH':
                    bond.isConstrained = True
Justin MacCallum's avatar
Justin MacCallum committed
420

421
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
422

423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
        if constraints == HAngles:
            for angle in data.angles:
                atom1 = data.atoms[angle[0]]
                atom2 = data.atoms[angle[1]]
                atom3 = data.atoms[angle[2]]
                numH = 0
                if atom1.name.startswith('H'):
                    numH += 1
                if atom3.name.startswith('H'):
                    numH += 1
                data.isAngleConstrained.append(numH == 2 or (numH == 1 and atom2.name.startswith('O')))
        else:
            data.isAngleConstrained = len(data.angles)*[False]
        if rigidWater:
            for i in range(len(data.angles)):
                angle = data.angles[i]
                atom1 = data.atoms[angle[0]]
                atom2 = data.atoms[angle[1]]
                atom3 = data.atoms[angle[2]]
                if atom1.residue.name == 'HOH' and atom2.residue.name == 'HOH' and atom3.residue.name == 'HOH':
                    data.isAngleConstrained[i] = True
Justin MacCallum's avatar
Justin MacCallum committed
444

445
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
446

447
448
        for atom in data.virtualSites:
            site = data.virtualSites[atom]
449
            index = atom.index
450
451
452
453
454
455
            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]))
Justin MacCallum's avatar
Justin MacCallum committed
456

457
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
458

459
460
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
461
462
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
463

peastman's avatar
peastman committed
464
        # Let generators do postprocessing
Justin MacCallum's avatar
Justin MacCallum committed
465

peastman's avatar
peastman committed
466
467
468
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
469

470
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
471

472
473
        for script in self._scripts:
            exec script
474
475
476
        return sys


477
478
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
479
480
    counts = {}
    for element in elements:
481
        if element in counts:
482
483
484
            counts[element] += 1
        else:
            counts[element] = 1
485
486
487
488
489
490
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
491
492
    sig = []
    for c in counts:
493
494
        if c is not None:
            sig.append((c, counts[c]))
495
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
496

497
    # Convert it to a string.
498
499

    s = ''
500
    for element, count in sig:
501
502
503
504
        s += element.symbol+str(count)
    return s


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

508
509
510
511
512
513
514
515
    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())
516
517
    if len(atoms) != len(template.atoms):
        return None
518
519
    matches = len(atoms)*[0]
    hasMatch = len(atoms)*[False]
Justin MacCallum's avatar
Justin MacCallum committed
520

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

523
524
    renumberAtoms = {}
    for i in range(len(atoms)):
525
        renumberAtoms[atoms[i].index] = i
526
527
528
    bondedTo = []
    externalBonds = []
    for atom in atoms:
529
        bonds = [renumberAtoms[x] for x in bondedToAtom[atom.index] if x in renumberAtoms]
530
        bondedTo.append(bonds)
531
        externalBonds.append(len([x for x in bondedToAtom[atom.index] if x not in renumberAtoms]))
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552

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

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

    # Recursively match atoms.

553
554
555
556
557
558
559
560
561
562
    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
563
    name = atoms[position].name
564
565
    for i in range(len(atoms)):
        atom = template.atoms[i]
peastman's avatar
peastman committed
566
        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]:
567
            # See if the bonds for this identification are consistent
Justin MacCallum's avatar
Justin MacCallum committed
568

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

573
574
575
576
577
578
579
580
                matches[position] = i
                hasMatch[i] = True
                if _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch, position+1):
                    return True
                hasMatch[i] = False
    return False


581
582
583
584
def _findMatchErrors(forcefield, res):
    """Try to guess why a residue failed to match any template and return an error message."""
    residueCounts = _countResidueAtoms([atom.element for atom in res.atoms()])
    numResidueAtoms = sum(residueCounts.itervalues())
585
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
    
    # Loop over templates and see how closely each one might match.
    
    bestMatchName = None
    numBestMatchAtoms = 3*numResidueAtoms
    numBestMatchHeavyAtoms = 2*numResidueHeavyAtoms
    for templateName in forcefield._templates:
        template = forcefield._templates[templateName]
        templateCounts = _countResidueAtoms([atom.element for atom in template.atoms])
        
        # Does the residue have any atoms that clearly aren't in the template?
        
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
        
        # If there are too many missing atoms, discard this template.
        
        numTemplateAtoms = sum(templateCounts.itervalues())
604
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
        
        # If this template has the same number of missing atoms as our previous best one, look at the name
        # to decide which one to use.
        
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
        
        # Accept this as our new best match.
        
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
622
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
623
624
625
626
    
    # Return an appropriate error message.
    
    if numBestMatchAtoms == numResidueAtoms:
627
628
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
629
630
631
632
            return 'The set of atoms matches %s, but the bonds are different.  Perhaps the chain is missing a terminal group?' % bestMatchName
        return 'The set of atoms matches %s, but the bonds are different.' % bestMatchName
    if bestMatchName is not None:
        if numBestMatchHeavyAtoms == numResidueHeavyAtoms:
633
634
635
636
637
            numResidueExtraParticles = len([atom for atom in res.atoms() if atom.element is None])
            if numResidueExtraParticles == 0 and numBestMatchExtraParticles == 0:
                return 'The set of atoms is similar to %s, but it is missing %d hydrogen atoms.' % (bestMatchName, numBestMatchAtoms-numResidueAtoms)
            if numBestMatchExtraParticles-numResidueExtraParticles == numBestMatchAtoms-numResidueAtoms:
                return 'The set of atoms is similar to %s, but it is missing %d extra particles.  You can add them with Modeller.addExtraParticles().' % (bestMatchName, numBestMatchAtoms-numResidueAtoms)
638
639
640
641
        return 'The set of atoms is similar to %s, but it is missing %d atoms.' % (bestMatchName, numBestMatchAtoms-numResidueAtoms)
    return 'This might mean your input topology is missing some atoms or bonds, or possibly that you are using the wrong force field.'


642
643
644
645
646
# 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.

647
## @private
648
649
class HarmonicBondGenerator:
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
650

651
652
653
654
655
    def __init__(self):
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
656

657
658
659
660
661
662
    @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
663
            if None not in types:
664
665
666
667
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.length.append(float(bond.attrib['length']))
                generator.k.append(float(bond.attrib['k']))
Justin MacCallum's avatar
Justin MacCallum committed
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
    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


694
## @private
695
696
class HarmonicAngleGenerator:
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
697

698
699
700
701
702
703
    def __init__(self):
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
704

705
706
707
708
709
710
    @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
711
            if None not in types:
712
713
714
715
716
                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']))
Justin MacCallum's avatar
Justin MacCallum committed
717

718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.HarmonicAngleForce]
        if len(existing) == 0:
            force = mm.HarmonicAngleForce()
            sys.addForce(force)
        else:
            force = existing[0]
        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
            for i in range(len(self.types1)):
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
                if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                    if isConstrained:
                        # Find the two bonds that make this angle.
Justin MacCallum's avatar
Justin MacCallum committed
737

738
739
740
741
742
743
744
745
746
                        bond1 = None
                        bond2 = None
                        for bond in data.atomBonds[angle[1]]:
                            atom1 = data.bonds[bond].atom1
                            atom2 = data.bonds[bond].atom2
                            if atom1 == angle[0] or atom2 == angle[0]:
                                bond1 = bond
                            elif atom1 == angle[2] or atom2 == angle[2]:
                                bond2 = bond
Justin MacCallum's avatar
Justin MacCallum committed
747

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

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


763
## @private
764
765
766
767
768
769
770
771
772
773
774
775
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 = []

776
## @private
777
778
class PeriodicTorsionGenerator:
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
779

780
781
782
    def __init__(self):
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
783

784
785
786
787
788
789
790
791
792
793
794
795
796
    @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)
Justin MacCallum's avatar
Justin MacCallum committed
797

798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
    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


864
## @private
865
866
867
868
869
870
871
872
873
874
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

875
## @private
876
877
class RBTorsionGenerator:
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
878

879
880
881
    def __init__(self):
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
882

883
884
885
886
887
888
889
    @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
890
            if None not in types:
891
892
893
                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
894
            if None not in types:
895
                generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
Justin MacCallum's avatar
Justin MacCallum committed
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
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
    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


959
## @private
960
961
962
963
964
965
966
967
968
969
970
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

971
## @private
972
973
class CMAPTorsionGenerator:
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
974

975
976
977
    def __init__(self):
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
978

979
980
981
982
983
984
985
986
987
988
989
990
991
    @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
992
            if None not in types:
993
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
994

995
996
997
998
999
1000
1001
1002
1003
1004
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.CMAPTorsionForce]
        if len(existing) == 0:
            force = mm.CMAPTorsionForce()
            sys.addForce(force)
        else:
            force = existing[0]
        for map in self.maps:
            force.addMap(int(sqrt(len(map))), map)
Justin MacCallum's avatar
Justin MacCallum committed
1005

1006
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
1007

1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
        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


1051
## @private
1052
1053
class NonbondedGenerator:
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1054

1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
    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']):
Justin MacCallum's avatar
Justin MacCallum committed
1070
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
1071
1072
        for atom in element.findall('Atom'):
            types = ff._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
1073
            if None not in types:
1074
1075
1076
                values = (float(atom.attrib['charge']), float(atom.attrib['sigma']), float(atom.attrib['epsilon']))
                for t in types[0]:
                    generator.typeMap[t] = values
Justin MacCallum's avatar
Justin MacCallum committed
1077

1078
1079
1080
1081
1082
1083
1084
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.NonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.NonbondedForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.NonbondedForce.CutoffPeriodic,
                     Ewald:mm.NonbondedForce.Ewald,
                     PME:mm.NonbondedForce.PME}
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
1085
            raise ValueError('Illegal nonbonded method for NonbondedForce')
1086
1087
1088
1089
1090
1091
1092
        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:
1093
                raise ValueError('No nonbonded parameters defined for atom type '+t)
peastman's avatar
peastman committed
1094
1095
1096
1097
1098
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
1099

peastman's avatar
peastman committed
1100
1101
    def postprocessSystem(self, sys, data, args):
        # Create exceptions based on bonds, virtual sites, and Drude particles.
1102
1103
1104
        bondIndices = []
        for bond in data.bonds:
            bondIndices.append((bond.atom1, bond.atom2))
1105
1106
1107
1108
1109
        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
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
        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]
Justin MacCallum's avatar
Justin MacCallum committed
1129
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
1130
1131
1132
1133

parsers["NonbondedForce"] = NonbondedGenerator.parseElement


1134
## @private
1135
1136
class GBSAOBCGenerator:
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1137

1138
1139
1140
1141
1142
    def __init__(self):
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
1143
1144
1145
1146
1147
1148
1149
        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]
1150
1151
        for atom in element.findall('Atom'):
            types = ff._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
1152
            if None not in types:
1153
1154
1155
                values = (float(atom.attrib['charge']), float(atom.attrib['radius']), float(atom.attrib['scale']))
                for t in types[0]:
                    generator.typeMap[t] = values
Justin MacCallum's avatar
Justin MacCallum committed
1156

1157
1158
1159
1160
1161
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.NonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.NonbondedForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.NonbondedForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
1162
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
1163
1164
1165
1166
1167
1168
1169
        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:
Justin MacCallum's avatar
Justin MacCallum committed
1170
                raise ValueError('No GBSAOBC parameters defined for atom type '+t)
1171
1172
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
1173
1174
1175
1176
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
1177
1178
        sys.addForce(force)

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

1182
1183
1184
1185
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

1186
1187
1188
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


1189
## @private
1190
1191
1192
class GBVIGenerator:

    """A GBVIGenerator constructs a GBVIForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1193

1194
1195
    def __init__(self,ff):

Peter Eastman's avatar
Peter Eastman committed
1196
1197
1198
1199
1200
        self.ff = ff
        self.fixedParameters = {}
        self.fixedParameters['soluteDielectric'] = 1.0
        self.fixedParameters['solventDielectric'] = 78.3
        self.fixedParameters['scalingMethod'] = 1
1201
        self.fixedParameters['quinticUpperBornRadiusLimit'] = 5.0
Peter Eastman's avatar
Peter Eastman committed
1202
        self.fixedParameters['quinticLowerLimitFactor'] = 0.8
1203

Peter Eastman's avatar
Peter Eastman committed
1204
        self.typeMap = {}
1205
1206
1207

    @staticmethod
    def parseElement(element, ff):
Peter Eastman's avatar
Peter Eastman committed
1208
        generator = GBVIGenerator(ff)
1209
        for key in generator.fixedParameters.iterkeys():
Peter Eastman's avatar
Peter Eastman committed
1210
1211
            if (key in element.attrib):
                generator.fixedParameters[key] = float(element.attrib[key])
1212
1213
1214
1215

        ff._forces.append(generator)
        for atom in element.findall('Atom'):
            types = ff._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
1216
            if None not in types:
1217
1218
1219
                values = (float(atom.attrib['charge']), float(atom.attrib['radius']), float(atom.attrib['gamma']))
                for t in types[0]:
                    generator.typeMap[t] = values
Justin MacCallum's avatar
Justin MacCallum committed
1220

1221
1222
1223
1224
1225
1226
1227
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):

        methodMap = {NoCutoff:mm.NonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.NonbondedForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.NonbondedForce.CutoffPeriodic}

        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
1228
            raise ValueError('Illegal nonbonded method for GB/VI Force')
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238

        # 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:
Justin MacCallum's avatar
Justin MacCallum committed
1239
                raise ValueError('No GB/VI parameters defined for atom type '+t)
1240
1241
1242
1243
1244

        # get HarmonicBond generator -- exit if not found

        hbGenerator = 0
        for generator in self.ff._forces:
Justin MacCallum's avatar
Justin MacCallum committed
1245
            if (generator.__class__.__name__ == 'HarmonicBondGenerator'):
1246
1247
1248
               hbGenerator = generator
               break

Peter Eastman's avatar
Peter Eastman committed
1249
        if (hbGenerator == 0):
Justin MacCallum's avatar
Justin MacCallum committed
1250
            raise ValueError('HarmonicBondGenerator not found.')
1251
1252
1253
1254
1255
1256

        # 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
1257
            hit = 0
1258
1259
1260
1261
1262
1263
1264
1265
1266
            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
1267
1268
1269
1270
1271
        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'])
Justin MacCallum's avatar
Justin MacCallum committed
1272

1273
1274
1275
1276
        sys.addForce(force)

parsers["GBVIForce"] = GBVIGenerator.parseElement

1277
## @private
1278
1279
class CustomBondGenerator:
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1280

1281
1282
1283
1284
1285
1286
    def __init__(self):
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
1287

1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
    @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
1299
            if None not in types:
1300
1301
1302
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.paramValues.append([float(bond.attrib[param]) for param in generator.perBondParams])
Justin MacCallum's avatar
Justin MacCallum committed
1303

1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
    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


1324
## @private
1325
1326
class CustomAngleGenerator:
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1327

1328
1329
1330
1331
1332
1333
1334
    def __init__(self):
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
1335

1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
    @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
1347
            if None not in types:
1348
1349
1350
1351
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.paramValues.append([float(angle.attrib[param]) for param in generator.perAngleParams])
Justin MacCallum's avatar
Justin MacCallum committed
1352

1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
    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


1375
## @private
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
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

1386
## @private
1387
1388
class CustomTorsionGenerator:
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1389

1390
1391
1392
1393
1394
    def __init__(self):
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
1395

1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
    @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
1408
            if None not in types:
1409
1410
1411
                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
1412
            if None not in types:
1413
                generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
Justin MacCallum's avatar
Justin MacCallum committed
1414

1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
    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


1476
## @private
1477
1478
class CustomGBGenerator:
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1479

1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
    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
1499
            if None not in types:
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
                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()]
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
            if 'type' in function.attrib:
                type = function.attrib['type']
            else:
                type = 'Continuous1D'
            params = {}
            for key in function.attrib:
                if key.endswith('size'):
                    params[key] = int(function.attrib[key])
                elif key.endswith('min') or key.endswith('max'):
                    params[key] = float(function.attrib[key])
            generator.functions.append((function.attrib['name'], type, values, params))
Justin MacCallum's avatar
Justin MacCallum committed
1523

1524
1525
1526
1527
1528
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomGBForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomGBForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.CustomGBForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
1529
            raise ValueError('Illegal nonbonded method for CustomGBForce')
1530
1531
1532
1533
1534
1535
1536
1537
1538
        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])
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
        for (name, type, values, params) in self.functions:
            if type == 'Continuous1D':
                force.addTabulatedFunction(name, mm.Continuous1DFunction(values, params['min'], params['max']))
            elif type == 'Continuous2D':
                force.addTabulatedFunction(name, mm.Continuous2DFunction(params['xsize'], params['ysize'], values, params['xmin'], params['xmax'], params['ymin'], params['ymax']))
            elif type == 'Continuous3D':
                force.addTabulatedFunction(name, mm.Continuous2DFunction(params['xsize'], params['ysize'], params['zsize'], values, params['xmin'], params['xmax'], params['ymin'], params['ymax'], params['zmin'], params['zmax']))
            elif type == 'Discrete1D':
                force.addTabulatedFunction(name, mm.Discrete1DFunction(values))
            elif type == 'Discrete2D':
                force.addTabulatedFunction(name, mm.Discrete2DFunction(params['xsize'], params['ysize'], values))
            elif type == 'Discrete3D':
                force.addTabulatedFunction(name, mm.Discrete2DFunction(params['xsize'], params['ysize'], params['zsize'], values))
1552
1553
1554
1555
1556
1557
        for atom in data.atoms:
            t = data.atomType[atom]
            if t in self.typeMap:
                values = self.typeMap[t]
                force.addParticle(self.typeMap[t])
            else:
Justin MacCallum's avatar
Justin MacCallum committed
1558
                raise ValueError('No CustomGB parameters defined for atom type '+t)
1559
1560
1561
1562
1563
1564
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

Peter Eastman's avatar
Peter Eastman committed
1565
def getAtomPrint(data, atomIndex):
1566

Peter Eastman's avatar
Peter Eastman committed
1567
1568
1569
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
1570
    else:
Peter Eastman's avatar
Peter Eastman committed
1571
        returnString = "NA"
1572
1573
1574
1575
1576

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
1577
def countConstraint(data):
1578

Peter Eastman's avatar
Peter Eastman committed
1579
    bondCount = 0
1580
1581
1582
1583
1584
1585
1586
    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
1587
        if (isConstrained):
1588
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
1589

Peter Eastman's avatar
Peter Eastman committed
1590
    print "Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount))
1591

1592
## @private
1593
class AmoebaBondGenerator:
1594
1595
1596

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

1597
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
1598
1599

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

1601
1602
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
1603
1604
1605
1606
1607
1608
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
1609

1610
1611
1612
1613
1614
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

1618
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
1619
1620
1621
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
            types = forceField._findAtomTypes(bond, 2)
peastman's avatar
peastman committed
1622
            if None not in types:
1623
1624
1625
1626
1627
                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:
1628
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
1629
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
1630
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
1631
1632
                raise ValueError(outputString)

1633
1634
    #=============================================================================================

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

Mark Friedrichs's avatar
Mark Friedrichs committed
1637
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1638

Peter Eastman's avatar
Peter Eastman committed
1639
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
1640
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
1641
        if len(existing) == 0:
1642
            force = mm.AmoebaBondForce()
1643
1644
1645
1646
            sys.addForce(force)
        else:
            force = existing[0]

1647
1648
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
1649
1650
1651
1652

        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
1653
            hit = 0
1654
1655
1656
1657
1658
            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
1659
                    hit = 1
1660
1661
1662
1663
1664
1665
                    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

Justin MacCallum's avatar
Justin MacCallum committed
1666
            if (hit == 0):
1667
                outputString = "AmoebaBondGenerator missing: types=[%5s %5s] atoms=[%6d %6d] " % (type1, type2, bond.atom1, bond.atom2)
Justin MacCallum's avatar
Justin MacCallum committed
1668
                raise ValueError(outputString)
1669

1670
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
1671
1672
1673
1674

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

Peter Eastman's avatar
Peter Eastman committed
1676
def addAngleConstraint(angle, idealAngle, data, sys):
1677
1678

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

Peter Eastman's avatar
Peter Eastman committed
1680
1681
    bond1 = None
    bond2 = None
1682
1683
1684
1685
1686
1687
1688
    for bond in data.atomBonds[angle[1]]:
        atom1 = data.bonds[bond].atom1
        atom2 = data.bonds[bond].atom2
        if atom1 == angle[0] or atom2 == angle[0]:
            bond1 = bond
        elif atom1 == angle[2] or atom2 == angle[2]:
            bond2 = bond
Justin MacCallum's avatar
Justin MacCallum committed
1689

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

1692
1693
1694
1695
        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
1696
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
1697
1698
1699
1700
                sys.addConstraint(angle[0], angle[2], length)
                return

#=============================================================================================
1701
## @private
1702
class AmoebaAngleGenerator:
1703
1704

    #=============================================================================================
1705
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
1706
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
1707

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

Peter Eastman's avatar
Peter Eastman committed
1710
1711
1712
1713
1714
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
1715

Peter Eastman's avatar
Peter Eastman committed
1716
1717
1718
        self.types1 = []
        self.types2 = []
        self.types3 = []
1719

Peter Eastman's avatar
Peter Eastman committed
1720
1721
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
1722

1723
1724
1725
1726
1727
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

1731
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
1732
1733
1734
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
            types = forceField._findAtomTypes(angle, 3)
peastman's avatar
peastman committed
1735
            if None not in types:
1736
1737
1738
1739
1740
1741

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
1742
                angleList.append(float(angle.attrib['angle1']))
1743
1744

                try:
Peter Eastman's avatar
Peter Eastman committed
1745
                    angleList.append(float(angle.attrib['angle2']))
1746
                    try:
Peter Eastman's avatar
Peter Eastman committed
1747
                        angleList.append(float(angle.attrib['angle3']))
1748
1749
1750
1751
1752
1753
1754
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
1755
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
1756
1757
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
1758
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
1759
1760
                raise ValueError(outputString)

1761
1762
1763
1764
    #=============================================================================================
    # createForce is bypassed here since the AmoebaOutOfPlaneBendForce generator must first execute
    # and partition angles into in-plane and non-in-plane angles
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
1765

Peter Eastman's avatar
Peter Eastman committed
1766
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
1767
1768
1769
1770
1771
1772
        pass

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

Peter Eastman's avatar
Peter Eastman committed
1774
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
1775
1776
1777
1778

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
1779
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
1780
1781

        if len(existing) == 0:
1782
            force = mm.AmoebaAngleForce()
1783
1784
1785
1786
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1787
        # set scalars
1788

1789
1790
1791
1792
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
1793
1794

        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
1795
1796
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
1797

Peter Eastman's avatar
Peter Eastman committed
1798
1799
1800
1801
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
            hit = 0
1802
1803
1804
1805
1806
            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
1807
                    hit = 1
1808
1809
                    if isConstrained and self.k[i] != 0.0:
                        angleDict['idealAngle'] = self.angle[i][0]
Peter Eastman's avatar
Peter Eastman committed
1810
                        addAngleConstraint(angle, self.angle[i][0], data, sys)
1811
                    elif self.k[i] != 0:
Peter Eastman's avatar
Peter Eastman committed
1812
1813
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
1814
1815
1816
1817
1818
1819
                            # 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
1820
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
1821
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
1822
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
1823
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
1824
                            if (numberOfHydrogens < lenAngle):
1825
1826
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
1827
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
1828
                                raise ValueError(outputString)
1829
1830
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
1831

1832
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
1833
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
1834
                    break
Justin MacCallum's avatar
Justin MacCallum committed
1835
            if (hit == 0):
1836
                outputString = "AmoebaAngleGenerator missing types: [%s %s %s] for atoms: " % (type1, type2, type3)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1837
1838
1839
1840
                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])
Justin MacCallum's avatar
Justin MacCallum committed
1841
                raise ValueError(outputString)
1842
1843
1844
1845
1846

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

Peter Eastman's avatar
Peter Eastman committed
1848
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
1849
1850
1851
1852

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
1853
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
1854
1855

        if len(existing) == 0:
1856
            force = mm.AmoebaInPlaneAngleForce()
1857
1858
1859
1860
1861
1862
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

1863
1864
1865
1866
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
1867
1868

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

Peter Eastman's avatar
Peter Eastman committed
1870
1871
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
1872

Peter Eastman's avatar
Peter Eastman committed
1873
1874
1875
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
1876

Peter Eastman's avatar
Peter Eastman committed
1877
            hit = 0
1878
1879
1880
1881
1882
1883
1884
            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
1885
                    hit = 1
1886
                    angleDict['idealAngle'] = self.angle[i][0]
Peter Eastman's avatar
Peter Eastman committed
1887
1888
                    if (isConstrained and self.k[i] != 0.0):
                        addAngleConstraint(angle, self.angle[i][0], data, sys)
1889
1890
1891
1892
                    else:
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

Justin MacCallum's avatar
Justin MacCallum committed
1893
            if (hit == 0):
1894
                outputString = "AmoebaInPlaneAngleGenerator missing types: [%s %s %s] atoms: " % (type1, type2, type3)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
1895
1896
1897
1898
                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])
Justin MacCallum's avatar
Justin MacCallum committed
1899
                raise ValueError(outputString)
1900

1901
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
1902
1903
1904

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
1905
1906
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
1907
1908
#=============================================================================================

1909
## @private
1910
1911
1912
1913
1914
class AmoebaOutOfPlaneBendGenerator:

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

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

1916
1917
1918
1919
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
1920
1921
1922
1923
1924
1925
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
1926

Peter Eastman's avatar
Peter Eastman committed
1927
1928
1929
1930
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
1931

Peter Eastman's avatar
Peter Eastman committed
1932
        self.ks = []
1933
1934
1935
1936
1937

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

1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
    def findAtomTypes(self, forceField, node, num):
        """Parse the attributes on an XML tag to find the set of atom types for each atom it involves."""
        types = []
        attrib = node.attrib
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
            if classAttrib in attrib:
                if attrib[classAttrib] in forceField._atomClasses:
                    types.append(forceField._atomClasses[attrib[classAttrib]])
                else:
                    types.append(set())
        return types

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

    @staticmethod
    def parseElement(element, forceField):

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

1965
1966
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
1967
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
1968
1969
1970
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
1971
                                                   float(element.attrib['opbend-sextic']))
1972
1973
1974
1975

        forceField._forces.append(generator)

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

Peter Eastman's avatar
Peter Eastman committed
1979
1980
1981
1982
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
1983

Peter Eastman's avatar
Peter Eastman committed
1984
                generator.ks.append(float(angle.attrib['k']))
1985
1986

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

1991
1992
1993
1994
1995
1996
    #=============================================================================================
    # Get middle atom in a angle
    # return index of middle atom or -1 if no middle is found
    # This method appears not to be needed since the angle[1] entry appears to always
    # be the middle atom. However, was unsure if this is guaranteed
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
1997

Peter Eastman's avatar
Peter Eastman committed
1998
    def getMiddleAtom(self, angle, data):
1999
2000
2001

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
2002
        middleAtom = -1
Justin MacCallum's avatar
Justin MacCallum committed
2003
        for atomIndex in angle:
Peter Eastman's avatar
Peter Eastman committed
2004
            isMiddle = 0
2005
2006
2007
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
2008
                if (atom1 != atomIndex):
2009
2010
2011
                    partner = atom1
                else:
                    partner = atom2
Justin MacCallum's avatar
Justin MacCallum committed
2012
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]):
2013
2014
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
2015
            if (isMiddle == 2):
2016
2017
2018
2019
2020
                return atomIndex
        return -1

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

Peter Eastman's avatar
Peter Eastman committed
2021
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034

        # 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
2035
2036
2037
2038
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
2039
2040
2041
2042

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

Peter Eastman's avatar
Peter Eastman committed
2043
        skipAtoms = dict()
2044
2045
2046
2047

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

Peter Eastman's avatar
Peter Eastman committed
2048
2049
        inPlaneAngles = []
        nonInPlaneAngles = []
2050
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
2051
        idealAngles = []*len(data.angles)
2052
2053
2054

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

Peter Eastman's avatar
Peter Eastman committed
2055
2056
2057
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
2058
2059
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
2060
                middleType = -1
2061
2062
                middleCovalency = -1

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

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

Peter Eastman's avatar
Peter Eastman committed
2071
2072
2073
2074
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
2075
2076
2077
2078

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
2079
                    if (atom1 != middleAtom):
2080
2081
2082
2083
2084
2085
2086
2087
                        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
2088
2089
2090
2091
2092
                        if (middleType in types2 and partnerType in types1):
                            partners.append(partner)
                            partnerSet.add(partner)
                            partnerTypes.append(partnerType)
                            partnerK.append(self.ks[i])
Justin MacCallum's avatar
Justin MacCallum committed
2093

Peter Eastman's avatar
Peter Eastman committed
2094
                if (len(partners) == 3):
2095

Peter Eastman's avatar
Peter Eastman committed
2096
2097
2098
                    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])
2099

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

2102
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
2103
2104
2105
2106
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2107
2108
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
2109
2110
2111
2112
2113
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
2114
2115
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
2116
                    angleDict['isConstrained'] = 0
2117

Peter Eastman's avatar
Peter Eastman committed
2118
2119
2120
2121
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
2122
2123

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
2124
2125
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
2126

Peter Eastman's avatar
Peter Eastman committed
2127
                    inPlaneAngles.append(angleDict)
2128
2129

                else:
Peter Eastman's avatar
Peter Eastman committed
2130
2131
2132
2133
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
2134
            else:
Peter Eastman's avatar
Peter Eastman committed
2135
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
2136

Peter Eastman's avatar
Peter Eastman committed
2137
                    partnerSet = skipAtoms[middleAtom]
Justin MacCallum's avatar
Justin MacCallum committed
2138

Peter Eastman's avatar
Peter Eastman committed
2139
                    angleDict = {}
2140

Peter Eastman's avatar
Peter Eastman committed
2141
2142
2143
2144
2145
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
2146

Peter Eastman's avatar
Peter Eastman committed
2147
                    angleDict['isConstrained'] = isConstrained
2148

Peter Eastman's avatar
Peter Eastman committed
2149
2150
2151
2152
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
2153
2154

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
2155
2156
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
2157

Peter Eastman's avatar
Peter Eastman committed
2158
                    inPlaneAngles.append(angleDict)
2159
2160

                else:
Peter Eastman's avatar
Peter Eastman committed
2161
2162
                    angleDict = {}
                    angleDict['angle'] = angle
2163
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
2164
                    nonInPlaneAngles.append(angleDict)
2165

2166
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
2167
2168

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
2169
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
2170
2171
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
2172
2173

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
2174
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
2175
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
2176
                    nonInPlaneAngles.append(angleDict)
2177
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
2178
2179
2180
2181
2182

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

2183
## @private
2184
2185
2186
2187
2188
2189
2190
2191
class AmoebaTorsionGenerator:

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
2192
        self.torsionUnit = torsionUnit
2193

Peter Eastman's avatar
Peter Eastman committed
2194
2195
2196
2197
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
2198

Peter Eastman's avatar
Peter Eastman committed
2199
2200
2201
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
2202

2203
2204
2205
2206
2207
2208
2209
2210
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
2212
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
2213
2214
2215
2216
2217
2218
2219
        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
2220
            if None not in types:
2221
2222
2223
2224
2225
2226
2227

                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
2228
2229
                    tInfo = []
                    suffix = str(ii)
2230
2231
2232
2233
2234
2235
                    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
2236
2237
2238
2239
2240
2241
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
2242
2243
2244
2245
2246
2247

            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
2248
                                    stretchBend.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
2249
2250
                raise ValueError(outputString)

2251
2252
2253
2254
2255
    #=============================================================================================

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2256
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
2257
        if len(existing) == 0:
2258
            force = mm.PeriodicTorsionForce()
2259
2260
2261
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2262

2263
2264
2265
2266
2267
2268
2269
        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
2270
            hit = 0
2271
2272
2273
2274
2275
2276
2277
2278
2279
            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
2280
2281
                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
2282
2283
2284
2285
2286
2287
                    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])
2288
2289
                    break

Justin MacCallum's avatar
Justin MacCallum committed
2290
            if (hit == 0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2291
2292
2293
2294
2295
2296
                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])
Justin MacCallum's avatar
Justin MacCallum committed
2297
                raise ValueError(outputString)
2298
2299
2300
2301
2302

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

2303
## @private
2304
2305
2306
2307
2308
2309
2310
class AmoebaPiTorsionGenerator:

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

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

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

2312
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
2313
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
2314
2315
2316
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2317

2318
2319
2320
2321
2322
2323
2324
2325
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
2326
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
2327
2328
2329
2330
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
            types = forceField._findAtomTypes(piTorsion, 2)
peastman's avatar
peastman committed
2331
            if None not in types:
2332
2333
2334
2335
2336
2337
                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
2338
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
2339
2340
                raise ValueError(outputString)

2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
    #=============================================================================================

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

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

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

        for bond in data.bonds:

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

2358
2359
            atom1 = bond.atom1
            atom2 = bond.atom2
Justin MacCallum's avatar
Justin MacCallum committed
2360

Peter Eastman's avatar
Peter Eastman committed
2361
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom1]) == 3):
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372

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

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

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

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

Justin MacCallum's avatar
Justin MacCallum committed
2373
2374
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386

                       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
2387
                           if (bondedAtom1 != atom1):
2388
2389
2390
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
2391
2392
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
2393
                                   piTorsionAtom1 = b1
2394
2395
2396
2397
2398
2399
                               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
2400
                           if (bondedAtom1 != atom2):
2401
2402
2403
2404
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
2405
2406
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
2407
                                   piTorsionAtom5 = b1
2408
2409
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
2410

Peter Eastman's avatar
Peter Eastman committed
2411
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
2412
2413
2414
2415
2416

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

2417
## @private
2418
2419
2420
2421
2422
2423
class AmoebaTorsionTorsionGenerator:

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

Peter Eastman's avatar
Peter Eastman committed
2424
    def __init__(self):
2425

Peter Eastman's avatar
Peter Eastman committed
2426
2427
2428
2429
2430
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
2431

Peter Eastman's avatar
Peter Eastman committed
2432
        self.gridIndex = []
2433

Peter Eastman's avatar
Peter Eastman committed
2434
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
2435

2436
2437
2438
2439
2440
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
2441
        generator = AmoebaTorsionTorsionGenerator()
2442
2443
2444
2445
2446
2447
2448
        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
2449
            types = forceField._findAtomTypes(torsionTorsion, 5)
peastman's avatar
peastman committed
2450
            if None not in types:
2451
2452
2453
2454
2455
2456
2457

                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
2458
2459
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
2460
2461
2462
2463
2464
2465
2466
2467
2468
                    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
2469
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
2470
2471
                raise ValueError(outputString)

2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
        # load grid

        # xml source

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

        # output grid:

Justin MacCallum's avatar
Justin MacCallum committed
2482
2483
2484
2485
2486
2487
        #     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
2488
2489

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

Peter Eastman's avatar
Peter Eastman committed
2493
2494
2495
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
2496

Peter Eastman's avatar
Peter Eastman committed
2497
2498
            grid = []
            gridCol = []
2499
2500
2501
2502
2503

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
2504
2505
2506
2507
2508
2509
2510
2511
                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)
2512
2513

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
2514
2515
2516
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
2517
2518
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
2519

Peter Eastman's avatar
Peter Eastman committed
2520
2521
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
2522
            else:
Peter Eastman's avatar
Peter Eastman committed
2523
2524
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
2525
2526
2527
2528
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
2529
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
2530
2531
2532
2533
2534
2535
2536
2537

        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
2538
        if (len(data.atomBonds[atomC]) == 4):
2539
2540
2541
2542
2543
            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
2544
2545
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
2546
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
2547
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
2548
2549
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
2550
2551
                if (hit > -1):
                    if (atomE == -1):
2552
2553
2554
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
2555

2556
2557
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
2558
2559
            if (atomE == -1 or atomF == -1):
                outputString = "getChiralAtomIndex: error getting bonded partners of atomC=%s %d %s" % (atomC.name, atomC.resiude.index, atomC.resiude.name,)
Justin MacCallum's avatar
Justin MacCallum committed
2560
                raise ValueError(outputString)
2561
2562
2563
2564
2565

            # 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
2566
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
2567
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
2568
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
2569
                chiralAtomIndex = atomF
2570

Peter Eastman's avatar
Peter Eastman committed
2571
2572
2573
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
2574
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
2575
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
2576
                chiralAtomIndex = massF
2577
2578
2579
2580

        return chiralAtomIndex

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

2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
    def createForce(self, sys, data, nonpiTorsionedMethod, nonpiTorsionedCutoff, args):

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

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

        for angle in data.angles:

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

2597
2598
2599
2600
2601
2602
2603
            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
2604
                if (bondedAtom1 != ib):
2605
2606
2607
2608
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
2609
                if (ia != ic and ia != id):
2610
2611
2612
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
2613
                        if (bondedAtom1 != id):
2614
2615
2616
2617
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
2618
                        if (ie != ic and ie != ib and ie != ia):
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638

                            # 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
2639
2640
2641
                                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])
2642
2643
2644

                                # match in reverse order

Peter Eastman's avatar
Peter Eastman committed
2645
2646
2647
                                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])
2648
2649
2650
2651

        # set grids

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

2654
2655
2656
2657
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

2658
## @private
2659
class AmoebaStretchBendGenerator:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2660
2661

    #=============================================================================================
2662
2663
2664
2665
2666
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
2667
2668
2669
        self.types1 = []
        self.types2 = []
        self.types3 = []
2670

Peter Eastman's avatar
Peter Eastman committed
2671
2672
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
2673

2674
2675
2676
2677
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
2678
        generator = AmoebaStretchBendGenerator()
2679
2680
2681
2682
2683
2684
2685
2686
        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
2687
            if None not in types:
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699

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

2703
2704
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
2705
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
2706
2707
    # having been called since the ideal bond lengths and angle are needed here.
    # As a conseqeunce, createForce() is not implemented since it is not guaranteed that the generator for
Justin MacCallum's avatar
Justin MacCallum committed
2708
2709
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
2710
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
2711
2712
2713

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

Peter Eastman's avatar
Peter Eastman committed
2714
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2715
2716
2717
2718
2719
2720
2721
2722
        pass

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

    # Note: request for constrained bonds is ignored.

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

2723
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735

        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
2736
            if ('isConstrained' in angleDict):
2737
2738
2739
2740
2741
2742
2743
2744
                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
2745
            radian = 57.2957795130
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
            for i in range(len(self.types1)):

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

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

Justin MacCallum's avatar
Justin MacCallum committed
2756
                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
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
                    bondAB = -1.0
                    bondCB = -1.0
                    swap = 0
                    for bond in data.atomBonds[angle[1]]:
                        atom1 = data.bonds[bond].atom1
                        atom2 = data.bonds[bond].atom2
                        length = data.bonds[bond].length
                        if (atom1 == angle[0]):
                            bondAB = length
                        if (atom1 == angle[2]):
                            bondCB = length
                        if (atom2 == angle[2]):
                            bondCB = length
                        if (atom2 == angle[0]):
                            bondAB = length
Justin MacCallum's avatar
Justin MacCallum committed
2772

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2775
                    if ('idealAngle' not in angleDict):
2776

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2777
2778
2779
2780
2781
                       outputString = "AmoebaStretchBendGenerator: ideal angle is not set for following entry:\n"
                       outputString += "   types: %5s %5s %5s atoms: " % (type1, type2, type3)
                       outputString += getAtomPrint( data, angle[0] ) + ' '
                       outputString += getAtomPrint( data, angle[1] ) + ' '
                       outputString += getAtomPrint( data, angle[2] )
Justin MacCallum's avatar
Justin MacCallum committed
2782
                       raise ValueError(outputString)
2783

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2784
2785
2786
2787
2788
2789
2790
                    elif (bondAB < 0 or bondCB < 0):

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

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2796
                    break
2797
2798
2799
2800
2801

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

2802
## @private
2803
2804
2805
class AmoebaVdwGenerator:

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

2807
2808
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
2811
        self.type = type
2812

Peter Eastman's avatar
Peter Eastman committed
2813
2814
2815
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
2816

Peter Eastman's avatar
Peter Eastman committed
2817
        self.epsilonrule = epsilonrule
2818

Peter Eastman's avatar
Peter Eastman committed
2819
2820
2821
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
2822

Peter Eastman's avatar
Peter Eastman committed
2823
        self.typeMap = {}
2824
2825
2826
2827
2828
2829
2830

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

    @staticmethod
    def parseElement(element, forceField):

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

        generator = AmoebaVdwGenerator(element.attrib['type'], element.attrib['radiusrule'], element.attrib['radiustype'], element.attrib['radiussize'], element.attrib['epsilonrule'],
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
2836
2837
2838
2839
2840
2841
2842
2843
        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
2844
            if None not in types:
2845

Peter Eastman's avatar
Peter Eastman committed
2846
                values = [float(atom.attrib['sigma']), float(atom.attrib['epsilon']), float(atom.attrib['reduction'])]
2847
2848


Peter Eastman's avatar
Peter Eastman committed
2849
                if (generator.radiustype == 'SIGMA'):
2850
                    values[0] *= two_six
Justin MacCallum's avatar
Justin MacCallum committed
2851

Peter Eastman's avatar
Peter Eastman committed
2852
                if (generator.radiussize == 'DIAMETER'):
2853
2854
2855
2856
                    values[0] *= 0.5

                for t in types[0]:
                    generator.typeMap[t] = values
Justin MacCallum's avatar
Justin MacCallum committed
2857

2858
            else:
Peter Eastman's avatar
Peter Eastman committed
2859
                outputString = "AmoebaVdwGenerator: error getting type: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
2860
2861
                raise ValueError(outputString)

2862
2863
2864
2865
2866
2867
2868
    #=============================================================================================

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

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
2869
    def getBondedParticleSet(particleIndex, data):
2870
2871
2872
2873
2874
2875

        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
2876
2877
            if (atom1 != particleIndex):
                bondedParticleSet.add(atom1)
2878
            else:
Peter Eastman's avatar
Peter Eastman committed
2879
                bondedParticleSet.add(atom2)
2880
2881

        return bondedParticleSet
Justin MacCallum's avatar
Justin MacCallum committed
2882

2883
2884
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
2887
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
        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
2900
            if ('sigmaCombiningRule' in args):
2901
                sigmaRule = args['sigmaCombiningRule'].upper()
Peter Eastman's avatar
Peter Eastman committed
2902
2903
                if (sigmaRule.upper() in sigmaMap):
                    force.setSigmaCombiningRule(sigmaRule.upper())
2904
                else:
Peter Eastman's avatar
Peter Eastman committed
2905
                    stringList = ' ' . join(str(x) for x in sigmaMap.keys())
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2906
                    raise ValueError( "AmoebaVdwGenerator: sigma combining rule %s not recognized; valid values are %s; using default." % (sigmaRule, stringList) )
2907
            else:
Peter Eastman's avatar
Peter Eastman committed
2908
                force.setSigmaCombiningRule(self.radiusrule)
2909

Peter Eastman's avatar
Peter Eastman committed
2910
            if ('epsilonCombiningRule' in args):
2911
                epsilonRule = args['epsilonCombiningRule'].upper()
Peter Eastman's avatar
Peter Eastman committed
2912
2913
                if (epsilonRule.upper() in epsilonMap):
                    force.setEpsilonCombiningRule(epsilonRule.upper())
2914
                else:
Peter Eastman's avatar
Peter Eastman committed
2915
                    stringList = ' ' . join(str(x) for x in epsilonMap.keys())
Justin MacCallum's avatar
Justin MacCallum committed
2916
                    raise ValueError( "AmoebaVdwGenerator: epsilon combining rule %s not recognized; valid values are %s; using default." % (epsilonRule, stringList) )
2917
            else:
Peter Eastman's avatar
Peter Eastman committed
2918
                force.setEpsilonCombiningRule(self.epsilonrule)
Justin MacCallum's avatar
Justin MacCallum committed
2919

2920
2921
            # cutoff

Peter Eastman's avatar
Peter Eastman committed
2922
            if ('vdwCutoff' in args):
2923
                force.setCutoff(args['vdwCutoff'])
2924
            else:
Peter Eastman's avatar
Peter Eastman committed
2925
                force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2926

2927
2928
2929
2930
2931
            # dispersion correction

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

Peter Eastman's avatar
Peter Eastman committed
2932
            if (nonbondedMethod == PME):
2933
                force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
Justin MacCallum's avatar
Justin MacCallum committed
2934

2935
2936
2937
2938
        else:
            force = existing[0]

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
2939
        # throw error if particle type not available
2940
2941
2942
2943
2944

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

Peter Eastman's avatar
Peter Eastman committed
2945
                values = self.typeMap[t]
Justin MacCallum's avatar
Justin MacCallum committed
2946

2947
2948
                # ivIndex = index of bonded partner for hydrogens; otherwise ivIndex = particle index

Peter Eastman's avatar
Peter Eastman committed
2949
2950
2951
                ivIndex = i
                mass = sys.getParticleMass(i)/unit.dalton
                if (mass < 1.9 and len(data.atomBonds[i]) == 1):
2952
                    bondIndex = data.atomBonds[i][0]
Peter Eastman's avatar
Peter Eastman committed
2953
                    if (data.bonds[bondIndex].atom1 == i):
2954
2955
2956
2957
                        ivIndex = data.bonds[bondIndex].atom2
                    else:
                        ivIndex = data.bonds[bondIndex].atom1

2958
                force.addParticle(ivIndex, values[0], values[1], values[2])
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
            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
2970
            bondedParticleSets.append(AmoebaVdwGenerator.getBondedParticleSet(i, data))
2971
2972

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

2974
2975
2976
2977
2978
2979
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
2980
            if (self.vdw13Scale == 0.0):
2981
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
2982
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
2983
2984
2985
2986
2987

            # self

            exclusionSet.add(i)

Peter Eastman's avatar
Peter Eastman committed
2988
            force.setParticleExclusions(i, exclusionSet)
2989
2990
2991
2992
2993

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

2994
## @private
2995
2996
2997
2998
2999
class AmoebaMultipoleGenerator:

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

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

3001
3002
3003
3004
3005
3006
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3009
        self.forceField = forceField
3010

Justin MacCallum's avatar
Justin MacCallum committed
3011
3012
3013
3014
        self.direct11Scale = direct11Scale
        self.direct12Scale = direct12Scale
        self.direct13Scale = direct13Scale
        self.direct14Scale = direct14Scale
3015

Justin MacCallum's avatar
Justin MacCallum committed
3016
3017
3018
3019
        self.mpole12Scale = mpole12Scale
        self.mpole13Scale = mpole13Scale
        self.mpole14Scale = mpole14Scale
        self.mpole15Scale = mpole15Scale
3020

Justin MacCallum's avatar
Justin MacCallum committed
3021
3022
3023
3024
        self.mutual11Scale = mutual11Scale
        self.mutual12Scale = mutual12Scale
        self.mutual13Scale = mutual13Scale
        self.mutual14Scale = mutual14Scale
3025

Justin MacCallum's avatar
Justin MacCallum committed
3026
3027
3028
3029
        self.polar12Scale = polar12Scale
        self.polar13Scale = polar13Scale
        self.polar14Scale = polar14Scale
        self.polar15Scale = polar15Scale
3030

Peter Eastman's avatar
Peter Eastman committed
3031
        self.typeMap = {}
3032
3033
3034
3035
3036
3037

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
3038
    def setAxisType(kIndices):
3039
3040
3041

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
3042
3043
3044
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
3045
                else:
Peter Eastman's avatar
Peter Eastman committed
3046
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
3047

Peter Eastman's avatar
Peter Eastman committed
3048
3049
                if (kIndicesLen > 2):
                    kx = kIndices[2]
3050
                else:
Peter Eastman's avatar
Peter Eastman committed
3051
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
3052

Peter Eastman's avatar
Peter Eastman committed
3053
3054
                if (kIndicesLen > 1):
                    kz = kIndices[1]
3055
                else:
Peter Eastman's avatar
Peter Eastman committed
3056
                    kz = 0
3057

Peter Eastman's avatar
Peter Eastman committed
3058
3059
                while(len(kIndices) < 4):
                    kIndices.append(0)
3060
3061

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
3062
                if (kz == 0):
3063
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
3064
                if (kz != 0 and kx == 0):
3065
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
3066
                if (kz < 0 or kx < 0):
3067
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
3068
                if (kx < 0 and ky < 0):
3069
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
3070
                if (kz < 0 and kx < 0 and ky  < 0):
3071
3072
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
3073
3074
3075
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
3076
3077
3078
3079
3080
3081
3082
3083

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
3084
        #   <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"  >
3085
3086
3087
        # <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
3088
        generator = AmoebaMultipoleGenerator(forceField,
3089
3090
3091
3092
3093
3094
3095
3096
3097
3098
3099
3100
3101
3102
3103
3104
3105
3106
                                              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
3107
                                              element.attrib['polar15Scale'])
3108
3109
3110
3111
3112
3113
3114
3115
3116



        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
3117
            if None not in types:
3118
3119
3120
3121
3122
3123
3124
3125

                # 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
3126
3127
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
3128
                    except:
3129
3130
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
3131
                # set axis type based on k-Indices
3132

Peter Eastman's avatar
Peter Eastman committed
3133
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
3134
3135
3136

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
3139
3140
3141
3142
3143
3144
3145
3146
3147
3148
3149
3150
3151
                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']))
3152
3153

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
3154
                    if (t not in generator.typeMap):
3155
3156
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
3157
3158
3159
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
3160
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
3161
3162
3163
3164
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
3165

3166
            else:
Peter Eastman's avatar
Peter Eastman committed
3167
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
3168
3169
                raise ValueError(outputString)

3170
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
3171

3172
3173
        for atom in element.findall('Polarize'):
            types = forceField._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
3174
            if None not in types:
3175

Peter Eastman's avatar
Peter Eastman committed
3176
3177
3178
3179
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
3180
3181
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
3182
                    pdamp = pow(polarizability, 1.0/6.0)
3183

Peter Eastman's avatar
Peter Eastman committed
3184
3185
                pgrpMap = dict()
                for index in range(1, 7):
3186
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
3187
                    if (pgrp in atom.attrib):
3188
3189
3190
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
3191
3192
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
3193
                        raise ValueError(outputString)
3194
3195
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
3196
3197
3198
3199
3200
3201
3202
                        hit = 0
                        for (ii, typeMap) in enumerate(typeMapList):

                            if (typeMap['classIndex'] == classIndex):
                                typeMap['polarizability'] = polarizability
                                typeMap['thole'] = thole
                                typeMap['pdamp'] = pdamp
Justin MacCallum's avatar
Justin MacCallum committed
3203
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
3204
3205
3206
3207
3208
                                typeMapList[ii] = typeMap
                                hit = 1

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

3211
            else:
Peter Eastman's avatar
Peter Eastman committed
3212
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
3213
3214
                raise ValueError(outputString)

3215
3216
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3217
    def setPolarGroups(self, data, bonded12ParticleSets, force):
3218
3219
3220
3221
3222

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
3223
3224
3225
3226
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
3227
3228
3229
3230
            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
3231
3232
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
3233
3234
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
3235

3236
3237
3238
3239
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
3240
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
3241
3242
                continue

Peter Eastman's avatar
Peter Eastman committed
3243
3244
            group = set()
            visited = set()
3245
3246
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
3247
3248
3249
3250
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
3251
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
3252
3253
                if (nextAtom not in visited):
                   visited.add(nextAtom)
3254
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
3255
3256
3257
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
3258
3259
3260

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
3261
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
3262
3263

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3264
3265
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
3266
3267
3268
3269
3270

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
3271
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
3272
3273
                continue

Peter Eastman's avatar
Peter Eastman committed
3274
            pgrp11 = set(atom.polarizationGroupSet[0])
3275
3276
3277
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3278
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
3279
3280
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3281
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
3282

3283
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3284
3285
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
3286
3287
3288
3289
3290

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
3291
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
3292
3293
                continue

Peter Eastman's avatar
Peter Eastman committed
3294
3295
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
3296
3297
3298
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3299
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
3300
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
3301
            pgrp13 = pgrp13 - set(pgrp11)
3302
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
3303
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
3304

3305
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3306
3307
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
3308
3309
3310
3311
3312

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
3313
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
3314
3315
                continue

Peter Eastman's avatar
Peter Eastman committed
3316
3317
3318
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
3319
3320
3321
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
3322
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
3323
3324
3325

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
3326
            pgrp14 = pgrp14 - set(pgrp11)
3327
3328

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

3331
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
3332
3333
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
3334
3335
3336

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

Peter Eastman's avatar
Peter Eastman committed
3337
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3338
3339
3340
3341
3342
3343
3344
3345
3346
3347
3348

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

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
        existing = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
        if len(existing) == 0:
            force = mm.AmoebaMultipoleForce()
            sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
3349
            if (nonbondedMethod not in methodMap):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3350
                raise ValueError( "AmoebaMultipoleForce: input cutoff method not available." )
3351
            else:
Peter Eastman's avatar
Peter Eastman committed
3352
                force.setNonbondedMethod(methodMap[nonbondedMethod])
3353
3354
            force.setCutoffDistance(nonbondedCutoff)

Peter Eastman's avatar
Peter Eastman committed
3355
            if ('ewaldErrorTolerance' in args):
3356
3357
                force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))

Peter Eastman's avatar
Peter Eastman committed
3358
            if ('polarization' in args):
3359
                polarizationType = args['polarization']
Peter Eastman's avatar
Peter Eastman committed
3360
                if (polarizationType.lower() == 'direct'):
3361
3362
3363
3364
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Direct)
                else:
                    force.setPolarizationType(mm.AmoebaMultipoleForce.Mutual)

Peter Eastman's avatar
Peter Eastman committed
3365
3366
            if ('aEwald' in args):
                force.setAEwald(float(args['aEwald']))
3367

Peter Eastman's avatar
Peter Eastman committed
3368
            if ('pmeGridDimensions' in args):
3369
3370
                force.setPmeGridDimensions(args['pmeGridDimensions'])

Peter Eastman's avatar
Peter Eastman committed
3371
3372
            if ('mutualInducedMaxIterations' in args):
                force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
3373

Peter Eastman's avatar
Peter Eastman committed
3374
            if ('mutualInducedTargetEpsilon' in args):
3375
3376
3377
3378
3379
3380
                force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))

        else:
            force = existing[0]

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
3381
        # throw error if particle type not available
3382
3383
3384
3385
3386
3387
3388

        # 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
3389
3390
3391
            bonded12ParticleSet = AmoebaVdwGenerator.getBondedParticleSet(i, data)
            bonded12ParticleSet = set(sorted(bonded12ParticleSet))
            bonded12ParticleSets.append(bonded12ParticleSet)
3392
3393
3394
3395
3396

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3397
            bonded13Set = set()
3398
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
3399
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
3400
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
3401
3402
3403
3404

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
3405
            selfSet = set()
3406
3407
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3408
3409
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
3410
3411
3412
3413
3414

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3415
3416
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
3417
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
3418
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
3419

3420
3421
3422
3423
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
3424
            selfSet = set()
3425
3426
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3427
3428
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
3429
3430
3431
3432
3433

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3434
3435
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
3436
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
3437
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
3438
3439
3440
3441
3442
3443

            # 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
3444
            selfSet = set()
3445
3446
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
3447
3448
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
3449
3450
3451
3452
3453

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

Peter Eastman's avatar
Peter Eastman committed
3454
3455
                multipoleList = self.typeMap[t]
                hit = 0
3456
3457
3458
3459
3460
3461
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3462
                    if (hit != 0):
3463
3464
                        break

Peter Eastman's avatar
Peter Eastman committed
3465
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
3466
3467

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
3468
3469
                    kx = kIndices[2]
                    ky = kIndices[3]
3470
3471
3472
3473

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

3475
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
3476
3477
3478
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3479
3480
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
3481
                       if (hit != 0):
3482
3483
3484
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
3485
3486
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
3487
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
3488
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
3489
3490
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
3491
3492
3493
3494
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
3495
3496
3497
3498
3499
3500
3501
3502
3503
3504
                                      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

3505
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3506
                                      hit = 1
3507
3508
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
3509
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
3510
3511
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
3512
3513
3514
3515
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
3516
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3517
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
3518

3519
3520
3521
3522
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3523
                    if (hit != 0):
3524
3525
                        break

Peter Eastman's avatar
Peter Eastman committed
3526
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
3527
3528

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
3529
3530
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
3531

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

3536
3537
3538
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
3539
3540
3541
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3542
3543
3544

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
3545
                       if (hit != 0):
3546
3547
3548
                           break

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

Peter Eastman's avatar
Peter Eastman committed
3551
                       if (bondedAtomZType == kz):
3552
3553
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
3554
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
3555
3556
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
3557
3558
3559
3560
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
3561
3562
3563
3564
3565
3566
3567
3568

                                      # 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

3569
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3570
                                      hit = 3
3571
3572
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
3573
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
3574
3575
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
3576
3577
3578
3579
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
3580
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3581
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
3582

3583
3584
3585
3586
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3587
                    if (hit != 0):
3588
3589
                        break

Peter Eastman's avatar
Peter Eastman committed
3590
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
3591
3592
3593
3594

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

Peter Eastman's avatar
Peter Eastman committed
3595
3596
3597
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3598
3599
3600

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
3601
                        if (hit != 0):
3602
3603
3604
                            break

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

Peter Eastman's avatar
Peter Eastman committed
3607
3608
                        if (kx == 0 and kz == bondedAtomZType):
                            kz = bondedAtomZIndex
3609
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3610
                            hit = 5
3611
3612
3613
3614
3615

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
3616
                    if (hit != 0):
3617
3618
                        break

Peter Eastman's avatar
Peter Eastman committed
3619
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
3620
3621
3622

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
3623
3624
3625
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
3626

Peter Eastman's avatar
Peter Eastman committed
3627
                    if (kz == 0):
3628
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
3629
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
3630

3631
3632
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
3633
                if (hit != 0):
3634

Peter Eastman's avatar
Peter Eastman committed
3635
                    atom.multipoleDict = savedMultipoleDict
3636
                    atom.polarizationGroups = dict()
3637
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
3638
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
3639
3640
3641
3642
3643
                    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])
3644
                    else:
Peter Eastman's avatar
Peter Eastman committed
3645
                        raise ValueError("Atom %s of %s %d is out of synch!." %(atom.name, atom.residue.name, atom.residue.index))
3646
                else:
Peter Eastman's avatar
Peter Eastman committed
3647
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
3648
            else:
Peter Eastman's avatar
Peter Eastman committed
3649
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
3650
3651
3652

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
3653
        self.setPolarGroups(data, bonded12ParticleSets, force)
3654
3655
3656
3657
3658

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

3659
## @private
3660
3661
3662
class AmoebaWcaDispersionGenerator:

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

3664
3665
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
3668
3669
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
3670
3671
3672
3673
3674
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
3675
        self.shctd = shctd
3676

Peter Eastman's avatar
Peter Eastman committed
3677
        self.typeMap = {}
3678
3679
3680
3681
3682
3683
3684
3685
3686

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

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
3688
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
3689
3690
3691
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
Justin MacCallum's avatar
Justin MacCallum committed
3692
                                                  element.attrib['awater'],
3693
3694
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Justin MacCallum's avatar
Justin MacCallum committed
3695
                                                  element.attrib['shctd'])
3696
3697
3698
3699
3700
3701
        forceField._forces.append(generator)

        # typeMap[] = [ radius, epsilon ]

        for atom in element.findall('WcaDispersion'):
            types = forceField._findAtomTypes(atom, 1)
peastman's avatar
peastman committed
3702
            if None not in types:
3703
3704
3705
3706
3707

                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
3708
                outputString = "AmoebaWcaDispersionGenerator: error getting type: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
3709
3710
                raise ValueError(outputString)

3711
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3712

Peter Eastman's avatar
Peter Eastman committed
3713
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3714
3715
3716
3717
3718
3719
3720
3721
3722
3723
3724
3725

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

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

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

Peter Eastman's avatar
Peter Eastman committed
3728
3729
3730
3731
3732
3733
3734
3735
        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  ))
3736
3737
3738
3739
3740

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

Peter Eastman's avatar
Peter Eastman committed
3741
3742
                values = self.typeMap[t]
                force.addParticle(values[0], values[1])
3743
3744
3745
3746
3747
3748
3749
            else:
                raise ValueError('No WcaDispersion type for atom %s of %s %d' % (atom.name, atom.residue.name, atom.residue.index))

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

3750
## @private
3751
3752
3753
class AmoebaGeneralizedKirkwoodGenerator:

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

3755
3756
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
3757
3758
3759
3760
3761
3762
3763
3764
3765
3766
3767
    def __init__(self, forceField, solventDielectric, soluteDielectric, includeCavityTerm, probeRadius, surfaceAreaFactor):

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

        self.radiusTypeMap = {}
        self.radiusTypeMap['Bondi'] = {}
Justin MacCallum's avatar
Justin MacCallum committed
3768
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
3769
3770
3771
3772
3773
3774
3775
3776
3777
3778
3779
3780
3781
3782
3783
3784
3785
3786
3787
3788
3789
3790
3791
3792
3793
        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
3794
3795
3796

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

Peter Eastman's avatar
Peter Eastman committed
3797
    def getObcShct(self, data, atomIndex):
3798

Peter Eastman's avatar
Peter Eastman committed
3799
        atom = data.atoms[atomIndex]
3800
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
3801
        shct = -1.0
3802
3803

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
3804

Peter Eastman's avatar
Peter Eastman committed
3805
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
3806
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
3807
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
3808
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
3809
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
3810
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
3811
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
3812
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
3813
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
3814
3815
3816
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
3817
        elif (atomicNumber == 16):              # S(16)
3818
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
3819
        elif (atomicNumber == 26):              # Fe(26)
3820
3821
            shct = 0.88

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

        return shct
3826
3827
3828

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

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

Peter Eastman's avatar
Peter Eastman committed
3831
        atom = data.atoms[atomIndex]
3832
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
3833
        radius = -1.0
3834

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

Peter Eastman's avatar
Peter Eastman committed
3837
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
3838

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

3842
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
3843
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
3844

Peter Eastman's avatar
Peter Eastman committed
3845
            if (bondedAtomAtomicNumber == 7):
3846
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
3847
            if (bondedAtomAtomicNumber == 8):
3848
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
3849

Peter Eastman's avatar
Peter Eastman committed
3850
        elif (atomicNumber == 3):               # Li(3)
3851
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
3852
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
3853

3854
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
3855
            if (len(bondedAtomIndices) == 3):
3856
3857
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
3858
            elif (len(bondedAtomIndices) == 4):
3859
3860
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
3861
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
3862
3863
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
3864
        elif (atomicNumber == 7):               # N(7)
3865
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
3866
        elif (atomicNumber == 8):               # O(8)
3867
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
3868
            if (len(bondedAtomIndices) == 2):
3869
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
3870
        elif (atomicNumber == 9):               # F(9)
3871
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
3872
        elif (atomicNumber == 10):
3873
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
3874
        elif (atomicNumber == 11):
3875
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
3876
        elif (atomicNumber == 12):
3877
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
3878
        elif (atomicNumber == 14):
3879
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
3880
        elif (atomicNumber == 15):              # P(15)
3881
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
3882
        elif (atomicNumber == 16):              # S(16)
3883
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
3884
        elif (atomicNumber == 17):
3885
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
3886
        elif (atomicNumber == 18):
3887
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
3888
        elif (atomicNumber == 19):
3889
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
3890
        elif (atomicNumber == 20):
3891
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
3892
        elif (atomicNumber == 35):
3893
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
3894
        elif (atomicNumber == 36):
3895
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
3896
        elif (atomicNumber == 37):
3897
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
3898
        elif (atomicNumber == 53):
3899
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
3900
        elif (atomicNumber == 54):
3901
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
3902
        elif (atomicNumber == 55):
3903
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
3904
        elif (atomicNumber == 56):
3905
3906
            radius = 0.230

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

3911
3912
3913
3914
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
3917
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
3918
        atom = data.atoms[atomIndex]
3919
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
3920
        if (atomicNumber in bondiMap):
3921
3922
            radius = bondiMap[atomicNumber]
        else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3923
3924
            outputString = "Warning no Bondi radius for atom %s of %s %d using default value=%f" % (atom.name, atom.residue.name, atom.residue.index, radius)
            raise ValueError( outputString )
Justin MacCallum's avatar
Justin MacCallum committed
3925

3926
3927
3928
3929
3930
3931
3932
3933
3934
3935
        return radius

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

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
3937
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
Justin MacCallum's avatar
Justin MacCallum committed
3938
3939
                                                        element.attrib['includeCavityTerm'],
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor'])
3940
3941
3942
        forceField._forces.append(generator)

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

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

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

3949
3950
3951
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
3952
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3953
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
3954
        if (len(amoebaMultipoleForceList) > 0):
3955
3956
3957
3958
3959
            amoebaMultipoleForce = amoebaMultipoleForceList[0]
        else:
            # call AmoebaMultipoleForceGenerator.createForce() to ensure charges have been set

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

3963
3964
3965
3966
3967
3968
3969
        # get or create force depending on whether it has already been added to the system

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

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

Peter Eastman's avatar
Peter Eastman committed
3971
3972
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
3973
            else:
Peter Eastman's avatar
Peter Eastman committed
3974
                force.setSolventDielectric(   float(self.solventDielectric))
3975

Peter Eastman's avatar
Peter Eastman committed
3976
3977
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
3978
            else:
Peter Eastman's avatar
Peter Eastman committed
3979
                force.setSoluteDielectric(    float(self.soluteDielectric))
3980

Peter Eastman's avatar
Peter Eastman committed
3981
3982
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
3983
            else:
Peter Eastman's avatar
Peter Eastman committed
3984
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
3985
3986
3987
3988
3989

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
3992
3993
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
3994
3995
3996
3997
3998

        # 1-2

        bonded12ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
3999
4000
4001
            bonded12ParticleSet = AmoebaVdwGenerator.getBondedParticleSet(i, data)
            bonded12ParticleSet = set(sorted(bonded12ParticleSet))
            bonded12ParticleSets.append(bonded12ParticleSet)
4002
4003

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
4004
4005
4006
4007
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
4008
            else:
Peter Eastman's avatar
Peter Eastman committed
4009
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
4010
4011
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
4012
            force.addParticle(multipoleParameters[0], radius, shct)
4013
4014
4015
4016
4017

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

4018
## @private
4019
4020
4021
4022
4023
class AmoebaUreyBradleyGenerator:

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

4025
    def __init__(self):
4026

Peter Eastman's avatar
Peter Eastman committed
4027
4028
4029
        self.types1 = []
        self.types2 = []
        self.types3 = []
4030

Peter Eastman's avatar
Peter Eastman committed
4031
4032
        self.length = []
        self.k = []
4033
4034
4035
4036
4037
4038

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

    @staticmethod
    def parseElement(element, forceField):

4039
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
4040
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
4041

4042
        generator = AmoebaUreyBradleyGenerator()
4043
4044
4045
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
            types = forceField._findAtomTypes(bond, 3)
peastman's avatar
peastman committed
4046
            if None not in types:
4047
4048
4049
4050
4051
4052
4053
4054
4055
4056

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

4060
4061
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
4064
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
4065
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
4066
4067

        if len(existing) == 0:
4068
            force = mm.HarmonicBondForce()
4069
4070
4071
4072
4073
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
Peter Eastman's avatar
Peter Eastman committed
4074
            if (isConstrained):
4075
4076
4077
4078
4079
4080
4081
4082
                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
4083
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
4084
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
4085
4086
4087
4088
4089
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
4090
4091
4092
4093
4094


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

peastman's avatar
peastman committed
4096
4097
4098
4099
4100
4101
4102
4103
4104
4105
4106
4107
4108
4109
4110
4111
4112
4113
4114
4115
4116
4117
4118
4119
    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
Justin MacCallum's avatar
Justin MacCallum committed
4120

peastman's avatar
peastman committed
4121
4122
4123
4124
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        force = mm.DrudeForce()
        if not any(isinstance(f, mm.NonbondedForce) for f in sys.getForces()):
            raise ValueError('<DrudeForce> must come after <NonbondedForce> in XML file')
Justin MacCallum's avatar
Justin MacCallum committed
4125

peastman's avatar
peastman committed
4126
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
4127

peastman's avatar
peastman committed
4128
4129
4130
4131
4132
4133
4134
4135
4136
4137
4138
4139
4140
4141
4142
4143
4144
4145
4146
4147
4148
4149
        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)
Justin MacCallum's avatar
Justin MacCallum committed
4150

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

peastman's avatar
peastman committed
4154
4155
4156
4157
4158
4159
4160
4161
4162
4163
4164
4165
4166
4167
4168
4169
4170
4171
4172
        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)

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