forcefield.py 235 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-2016 Stanford University and the Authors.
10
11
12
Authors: Peter Eastman, Mark Friedrichs
Contributors:

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

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

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

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

import os
import itertools
import xml.etree.ElementTree as etree
39
import math
40
from math import sqrt, cos
41
from copy import deepcopy
42
from heapq import heappush, heappop
43
44
import simtk.openmm as mm
import simtk.unit as unit
45
from . import element as elem
46
47
from simtk.openmm.app import Topology

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

55
56
# Enumerated values for nonbonded method

57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
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()
81
82
83

# Enumerated values for constraint type

84
85
86
87
88
89
90
91
92
93
94
95
96
97
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()
98
99
100
101
102
103
104
105
106
107

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

Robert McGibbon's avatar
Robert McGibbon committed
109
110
111
112
113
114
115
116
        Parameters
        ----------
        files : list
            A list of XML files defining the force field.  Each entry may
            be an absolute file path, a path relative to the current working
            directory, a path relative to this module's data subdirectory
            (for built in force fields), or an open file-like object with a
            read() method from which the forcefield XML data can be loaded.
117
118
119
        """
        self._atomTypes = {}
        self._templates = {}
120
121
        self._patches = {}
        self._templatePatches = {}
122
        self._templateSignatures = {None:[]}
123
        self._atomClasses = {'':set()}
124
        self._forces = []
125
        self._scripts = []
126
        self._templateGenerators = []
127
        self.loadFile(files)
128

129
    def loadFile(self, files):
130
        """Load an XML file and add the definitions from it to this ForceField.
131

Robert McGibbon's avatar
Robert McGibbon committed
132
133
        Parameters
        ----------
134
135
136
        files : string or file or tuple
            An XML file or tuple of XML files containing force field definitions.
            Each entry may be either an absolute file path, a path relative to the current working
Robert McGibbon's avatar
Robert McGibbon committed
137
138
139
            directory, a path relative to this module's data subdirectory (for
            built in force fields), or an open file-like object with a read()
            method from which the forcefield XML data can be loaded.
peastman's avatar
peastman committed
140
        """
141

142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
        if not isinstance(files, tuple):
            files = (files,)

        trees = []

        for file in files:
            try:
                # this handles either filenames or open file-like objects
                tree = etree.parse(file)
            except IOError:
                tree = etree.parse(os.path.join(os.path.dirname(__file__), 'data', file))
            except Exception as e:
                # Fail with an error message about which file could not be read.
                # TODO: Also handle case where fallback to 'data' directory encounters problems,
                # but this is much less worrisome because we control those files.
                msg  = str(e) + '\n'
                if hasattr(file, 'name'):
                    filename = file.name
                else:
                    filename = str(file)
                msg += "ForceField.loadFile() encountered an error reading file '%s'\n" % filename
                raise Exception(msg)

            trees.append(tree)

peastman's avatar
peastman committed
167
168
169

        # Load the atom types.

170
171
172
173
        for tree in trees:
            if tree.getroot().find('AtomTypes') is not None:
                for type in tree.getroot().find('AtomTypes').findall('Type'):
                    self.registerAtomType(type.attrib)
peastman's avatar
peastman committed
174
175
176

        # Load the residue templates.

177
178
179
180
181
        for tree in trees:
            if tree.getroot().find('Residues') is not None:
                for residue in tree.getroot().find('Residues').findall('Residue'):
                    resName = residue.attrib['name']
                    template = ForceField._TemplateData(resName)
182
183
                    if 'override' in residue.attrib:
                        template.overrideLevel = int(residue.attrib['override'])
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
                    atomIndices = {}
                    for atom in residue.findall('Atom'):
                        params = {}
                        for key in atom.attrib:
                            if key not in ('name', 'type'):
                                params[key] = _convertParameterToNumber(atom.attrib[key])
                        atomName = atom.attrib['name']
                        if atomName in atomIndices:
                            raise ValueError('Residue '+resName+' contains multiple atoms named '+atomName)
                        atomIndices[atomName] = len(template.atoms)
                        typeName = atom.attrib['type']
                        template.atoms.append(ForceField._TemplateAtomData(atomName, typeName, self._atomTypes[typeName].element, params))
                    for site in residue.findall('VirtualSite'):
                        template.virtualSites.append(ForceField._VirtualSiteData(site, atomIndices))
                    for bond in residue.findall('Bond'):
                        if 'atomName1' in bond.attrib:
                            template.addBondByName(bond.attrib['atomName1'], bond.attrib['atomName2'])
                        else:
                            template.addBond(int(bond.attrib['from']), int(bond.attrib['to']))
                    for bond in residue.findall('ExternalBond'):
                        if 'atomName' in bond.attrib:
                            template.addExternalBondByName(bond.attrib['atomName'])
                        else:
                            template.addExternalBond(int(bond.attrib['from']))
208
                    for patch in residue.findall('AllowPatch'):
209
                        patchName = patch.attrib['name']
210
211
                        if ':' in patchName:
                            colonIndex = patchName.find(':')
212
213
214
                            self.registerTemplatePatch(resName, patchName[:colonIndex], int(patchName[colonIndex+1:])-1)
                        else:
                            self.registerTemplatePatch(resName, patchName, 0)
215
                    self.registerResidueTemplate(template)
peastman's avatar
peastman committed
216

217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
        # Load the patch defintions.

        for tree in trees:
            if tree.getroot().find('Patches') is not None:
                for patch in tree.getroot().find('Patches').findall('Patch'):
                    patchName = patch.attrib['name']
                    if 'residues' in patch.attrib:
                        numResidues = int(patch.attrib['residues'])
                    else:
                        numResidues = 1
                    patchData = ForceField._PatchData(patchName, numResidues)
                    allAtomNames = set()
                    for atom in patch.findall('AddAtom'):
                        params = {}
                        for key in atom.attrib:
                            if key not in ('name', 'type'):
                                params[key] = _convertParameterToNumber(atom.attrib[key])
                        atomName = atom.attrib['name']
                        if atomName in allAtomNames:
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
                        allAtomNames.add(atomName)
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
240
                        patchData.addedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
241
242
243
244
245
246
247
248
249
250
251
                    for atom in patch.findall('ChangeAtom'):
                        params = {}
                        for key in atom.attrib:
                            if key not in ('name', 'type'):
                                params[key] = _convertParameterToNumber(atom.attrib[key])
                        atomName = atom.attrib['name']
                        if atomName in allAtomNames:
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
                        allAtomNames.add(atomName)
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
252
                        patchData.changedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
                    for atom in patch.findall('RemoveAtom'):
                        atomName = atom.attrib['name']
                        if atomName in allAtomNames:
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
                        allAtomNames.add(atomName)
                        atomDescription = ForceField._PatchAtomData(atomName)
                        patchData.deletedAtoms.append(atomDescription)
                    for bond in patch.findall('AddBond'):
                        atom1 = ForceField._PatchAtomData(bond.attrib['atomName1'])
                        atom2 = ForceField._PatchAtomData(bond.attrib['atomName2'])
                        patchData.addedBonds.append((atom1, atom2))
                    for bond in patch.findall('RemoveBond'):
                        atom1 = ForceField._PatchAtomData(bond.attrib['atomName1'])
                        atom2 = ForceField._PatchAtomData(bond.attrib['atomName2'])
                        patchData.deletedBonds.append((atom1, atom2))
                    for bond in patch.findall('AddExternalBond'):
                        atom = ForceField._PatchAtomData(bond.attrib['atomName'])
                        patchData.addedExternalBonds.append(atom)
                    for bond in patch.findall('RemoveExternalBond'):
                        atom = ForceField._PatchAtomData(bond.attrib['atomName'])
                        patchData.deletedExternalBonds.append(atom)
                    for residue in patch.findall('ApplyToResidue'):
                        name = residue.attrib['name']
                        if ':' in name:
                            colonIndex = name.find(':')
                            self.registerTemplatePatch(name[colonIndex+1:], patchName, int(name[:colonIndex])-1)
                        else:
                            self.registerTemplatePatch(name, patchName, 0)
                    self.registerPatch(patchData)

peastman's avatar
peastman committed
283
284
        # Load force definitions

285
286
287
288
        for tree in trees:
            for child in tree.getroot():
                if child.tag in parsers:
                    parsers[child.tag](child, self)
peastman's avatar
peastman committed
289
290
291

        # Load scripts

292
293
294
        for tree in trees:
            for node in tree.getroot().findall('Script'):
                self.registerScript(node.text)
295

296
297
298
    def getGenerators(self):
        """Get the list of all registered generators."""
        return self._forces
299

300
301
302
    def registerGenerator(self, generator):
        """Register a new generator."""
        self._forces.append(generator)
303

304
305
306
307
308
309
310
311
312
313
314
315
    def registerAtomType(self, parameters):
        """Register a new atom type."""
        name = parameters['name']
        if name in self._atomTypes:
            raise ValueError('Found multiple definitions for atom type: '+name)
        atomClass = parameters['class']
        mass = _convertParameterToNumber(parameters['mass'])
        element = None
        if 'element' in parameters:
            element = parameters['element']
            if not isinstance(element, elem.Element):
                element = elem.get_by_symbol(element)
316
        self._atomTypes[name] = ForceField._AtomType(name, atomClass, mass, element)
317
318
319
320
321
322
323
        if atomClass in self._atomClasses:
            typeSet = self._atomClasses[atomClass]
        else:
            typeSet = set()
            self._atomClasses[atomClass] = typeSet
        typeSet.add(name)
        self._atomClasses[''].add(name)
324

325
326
    def registerResidueTemplate(self, template):
        """Register a new residue template."""
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
        if template.name in self._templates:
            # There is already a template with this name, so check the override levels.
            
            existingTemplate = self._templates[template.name]
            if template.overrideLevel < existingTemplate.overrideLevel:
                # The existing one takes precedence, so just return.
                return
            if template.overrideLevel > existingTemplate.overrideLevel:
                # We need to delete the existing template.
                del self._templates[template.name]
                existingSignature = _createResidueSignature([atom.element for atom in existingTemplate.atoms])
                self._templateSignatures[existingSignature].remove(existingTemplate)
            else:
                raise ValueError('Residue template %s with the same override level %d already exists.' % (template.name, template.overrideLevel))
        
        # Register the template.
        
344
345
346
        self._templates[template.name] = template
        signature = _createResidueSignature([atom.element for atom in template.atoms])
        if signature in self._templateSignatures:
347
            self._templateSignatures[signature].append(template)
348
349
        else:
            self._templateSignatures[signature] = [template]
350

351
352
353
354
355
356
357
358
359
360
    def registerPatch(self, patch):
        """Register a new patch that can be applied to templates."""
        self._patches[patch.name] = patch
    
    def registerTemplatePatch(self, residue, patch, patchResidueIndex):
        """Register that a particular patch can be used with a particular residue."""
        if residue not in self._templatePatches:
            self._templatePatches[residue] = []
        self._templatePatches[residue].append((patch, patchResidueIndex))

361
362
363
    def registerScript(self, script):
        """Register a new script to be executed after building the System."""
        self._scripts.append(script)
364

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
365
    def registerTemplateGenerator(self, generator):
366
367
368
369
        """Register a residue template generator that can be used to parameterize residues that do not match existing forcefield templates.

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

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

372
373
        Parameters
        ----------
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
374
        generator : function
375
376
            A function that will be called when a residue is encountered that does not match an existing forcefield template.

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

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
379
380
        ::
           success = generator(forcefield, residue)
381
382
        ```

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
383
384
385
386
387
388
389
390
391
392
        where `forcefield` is the calling `ForceField` object and `residue` is a simtk.openmm.app.topology.Residue object.

        `generator` must conform to the following API:
        ::
          Parameters
           ----------
           forcefield : simtk.openmm.app.ForceField
               The ForceField object to which residue templates and/or parameters are to be added.
           residue : simtk.openmm.app.Topology.Residue
               The residue topology for which a template is to be generated.
393

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
394
395
396
397
398
           Returns
           -------
           success : bool
               If the generator is able to successfully parameterize the residue, `True` is returned.
               If the generator cannot parameterize the residue, it should return `False` and not modify `forcefield`.
399

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
400
401
402
403
404
           The generator should either register a residue template directly with `forcefield.registerResidueTemplate(template)`
           or it should call `forcefield.loadFile(file)` to load residue definitions from an ffxml file.

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

        """
407
        self._templateGenerators.append(generator)
408

409
    def _findAtomTypes(self, attrib, num):
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
        """Parse the attributes on an XML tag to find the set of atom types for each atom it involves.

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

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

        """
425
426
427
428
429
430
431
        types = []
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
432
            typeAttrib = 'type'+suffix
433
            if classAttrib in attrib:
434
                if typeAttrib in attrib:
435
                    raise ValueError('Specified both a type and a class for the same atom: '+str(attrib))
436
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
437
438
439
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
440
441
            elif typeAttrib in attrib:
                if attrib[typeAttrib] == '':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
442
                    types.append(self._atomClasses[''])
443
                elif attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
444
445
446
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
447
448
            else:
                types.append(None) # Unknown atom type
449
450
        return types

451
    def _parseTorsion(self, attrib):
452
        """Parse the node defining a torsion."""
453
        types = self._findAtomTypes(attrib, 4)
peastman's avatar
peastman committed
454
        if None in types:
455
456
457
458
459
            return None
        torsion = PeriodicTorsion(types)
        index = 1
        while 'phase%d'%index in attrib:
            torsion.periodicity.append(int(attrib['periodicity%d'%index]))
460
461
            torsion.phase.append(_convertParameterToNumber(attrib['phase%d'%index]))
            torsion.k.append(_convertParameterToNumber(attrib['k%d'%index]))
462
            index += 1
Justin MacCallum's avatar
Justin MacCallum committed
463
464
        return torsion

465
    class _SystemData(object):
466
467
468
        """Inner class used to encapsulate data about the system being created."""
        def __init__(self):
            self.atomType = {}
469
            self.atomParameters = {}
470
            self.atomTemplateIndexes = {}
471
            self.atoms = []
472
            self.excludeAtomWith = []
473
            self.virtualSites = {}
474
475
476
477
478
479
            self.bonds = []
            self.angles = []
            self.propers = []
            self.impropers = []
            self.atomBonds = []
            self.isAngleConstrained = []
480
481
482
483
484
485
            self.constraints = {}

        def addConstraint(self, system, atom1, atom2, distance):
            """Add a constraint to the system, avoiding duplicate constraints."""
            key = (min(atom1, atom2), max(atom1, atom2))
            if key in self.constraints:
486
                if self.constraints[key] != distance:
487
488
489
490
                    raise ValueError('Two constraints were specified between atoms %d and %d with different distances' % (atom1, atom2))
            else:
                self.constraints[key] = distance
                system.addConstraint(atom1, atom2, distance)
491
492
493
494
495
496
497
        
        def recordMatchedAtomParameters(self, residue, template, matches):
            """Record parameters for atoms based on having matched a residue to a template."""
            matchAtoms = dict(zip(matches, residue.atoms()))
            for atom, match in zip(residue.atoms(), matches):
                self.atomType[atom] = template.atoms[match].type
                self.atomParameters[atom] = template.atoms[match].parameters
498
                self.atomTemplateIndexes[atom] = match
499
500
501
                for site in template.virtualSites:
                    if match == site.index:
                        self.virtualSites[atom] = (site, [matchAtoms[i].index for i in site.atoms], matchAtoms[site.excludeWith].index)
502

503
    class _TemplateData(object):
504
505
506
507
        """Inner class used to encapsulate data about a residue template definition."""
        def __init__(self, name):
            self.name = name
            self.atoms = []
508
            self.virtualSites = []
509
510
            self.bonds = []
            self.externalBonds = []
511
            self.overrideLevel = 0
512

513
514
515
516
517
518
519
        def getAtomIndexByName(self, atom_name):
            """Look up an atom index by atom name, providing a helpful error message if not found."""
            for (index, atom) in enumerate(self.atoms):
                if atom.name == atom_name:
                    return index

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

524
        def addBond(self, atom1, atom2):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
525
            """Add a bond between two atoms in a template given their indices in the template."""
526
527
528
            self.bonds.append((atom1, atom2))
            self.atoms[atom1].bondedTo.append(atom2)
            self.atoms[atom2].bondedTo.append(atom1)
529

530
        def addBondByName(self, atom1_name, atom2_name):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
531
            """Add a bond between two atoms in a template given their atom names."""
532
533
534
535
536
            atom1 = self.getAtomIndexByName(atom1_name)
            atom2 = self.getAtomIndexByName(atom2_name)
            self.addBond(atom1, atom2)

        def addExternalBond(self, atom_index):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
537
            """Designate that an atom in a residue template has an external bond, using atom index within template."""
538
539
540
541
            self.externalBonds.append(atom_index)
            self.atoms[atom_index].externalBonds += 1

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

546
    class _TemplateAtomData(object):
547
        """Inner class used to encapsulate data about an atom in a residue template definition."""
548
        def __init__(self, name, type, element, parameters={}):
549
550
551
            self.name = name
            self.type = type
            self.element = element
552
            self.parameters = parameters
553
554
555
            self.bondedTo = []
            self.externalBonds = 0

556
    class _BondData(object):
557
558
559
560
561
562
        """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
563

564
    class _VirtualSiteData(object):
565
        """Inner class used to encapsulate data about a virtual site."""
566
        def __init__(self, node, atomIndices):
567
568
569
            attrib = node.attrib
            self.type = attrib['type']
            if self.type == 'average2':
570
                numAtoms = 2
571
572
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
573
                numAtoms = 3
574
575
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
576
                numAtoms = 3
577
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
578
            elif self.type == 'localCoords':
579
                numAtoms = 3
580
581
582
583
                self.originWeights = [float(attrib['wo1']), float(attrib['wo2']), float(attrib['wo3'])]
                self.xWeights = [float(attrib['wx1']), float(attrib['wx2']), float(attrib['wx3'])]
                self.yWeights = [float(attrib['wy1']), float(attrib['wy2']), float(attrib['wy3'])]
                self.localPos = [float(attrib['p1']), float(attrib['p2']), float(attrib['p3'])]
584
585
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
586
587
588
589
590
591
            if 'siteName' in attrib:
                self.index = atomIndices[attrib['siteName']]
                self.atoms = [atomIndices[attrib['atomName%d'%(i+1)]] for i in range(numAtoms)]
            else:
                self.index = int(attrib['index'])
                self.atoms = [int(attrib['atom%d'%(i+1)]) for i in range(numAtoms)]
592
593
594
595
            if 'excludeWith' in attrib:
                self.excludeWith = int(attrib['excludeWith'])
            else:
                self.excludeWith = self.atoms[0]
596

597
598
599
600
601
    class _PatchData(object):
        """Inner class used to encapsulate data about a patch definition."""
        def __init__(self, name, numResidues):
            self.name = name
            self.numResidues = numResidues
602
603
            self.addedAtoms = [[] for i in range(numResidues)]
            self.changedAtoms = [[] for i in range(numResidues)]
604
605
606
607
608
            self.deletedAtoms = []
            self.addedBonds = []
            self.deletedBonds = []
            self.addedExternalBonds = []
            self.deletedExternalBonds = []
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
        
        def createPatchedTemplates(self, templates):
            """Apply this patch to a set of templates, creating new modified ones."""
            if len(templates) != self.numResidues:
                raise ValueError("Patch '%s' expected %d templates, received %d", (self.name, self.numResidues, len(templates)))
            
            # Construct a new version of each template.
            
            newTemplates = []
            for index, template in enumerate(templates):
                newTemplate = ForceField._TemplateData("%s-%s" % (template.name, self.name))
                newTemplates.append(newTemplate)
                
                # Build the list of atoms in it.
                
                for atom in template.atoms:
                    if not any(deleted.name == atom.name and deleted.residue == index for deleted in self.deletedAtoms):
626
                        newTemplate.atoms.append(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
627
                for atom in self.addedAtoms[index]:
628
629
630
                    if any(a.name == atom.name for a in newTemplate.atoms):
                        raise ValueError("Patch '%s' adds an atom with the same name as an existing atom: %s" % (self.name, atom.name))
                    newTemplate.atoms.append(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
631
632
633
634
635
                oldAtomIndex = dict([(atom.name, i) for i, atom in enumerate(template.atoms)])
                newAtomIndex = dict([(atom.name, i) for i, atom in enumerate(newTemplate.atoms)])
                for atom in self.changedAtoms[index]:
                    if atom.name not in newAtomIndex:
                        raise ValueError("Patch '%s' modifies nonexistent atom '%s' in template '%s'" % (self.name, atom.name, template.name))
636
                    newTemplate.atoms[newAtomIndex[atom.name]] = ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters)
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
                
                # Copy over the virtual sites, translating the atom indices.
                
                indexMap = dict([(oldAtomIndex[name], newAtomIndex[name]) for name in newAtomIndex if name in oldAtomIndex])
                for site in template.virtualSites:
                    if site.index in indexMap and all(i in indexMap for i in site.atoms):
                        newSite = deepcopy(site)
                        newSite.index = indexMap[site.index]
                        newSite.atoms = [indexMap[i] for i in site.atoms]
                        newTemplate.virtualSites.append(newSite)
                
                # Build the lists of bonds and external bonds.
                
                atomMap = dict([(template.atoms[i], indexMap[i]) for i in indexMap])
                deletedBonds = [(atom1.name, atom2.name) for atom1, atom2 in self.deletedBonds if atom1.residue == index and atom2.residue == index]
                for atom1, atom2 in template.bonds:
                    a1 = template.atoms[atom1]
                    a2 = template.atoms[atom2]
655
                    if a1 in atomMap and a2 in atomMap and (a1.name, a2.name) not in deletedBonds and (a2.name, a1.name) not in deletedBonds:
656
657
658
659
                        newTemplate.addBond(atomMap[a1], atomMap[a2])
                deletedExternalBonds = [atom.name for atom in self.deletedExternalBonds if atom.residue == index]
                for atom in template.externalBonds:
                    if template.atoms[atom].name not in deletedExternalBonds:
660
                        newTemplate.addExternalBond(indexMap[atom])
661
662
663
664
665
666
667
668
669
670
                for atom1, atom2 in self.addedBonds:
                    if atom1.residue == index and atom2.residue == index:
                        newTemplate.addBondByName(atom1.name, atom2.name)
                    elif atom1.residue == index:
                        newTemplate.addExternalBondByName(atom1.name)
                    elif atom2.residue == index:
                        newTemplate.addExternalBondByName(atom2.name)
                for atom in self.addedExternalBonds:
                    newTemplate.addExternalBondByName(atom.name)
            return newTemplates
671
672
673
674
675
676
677
678
679
680
681
682
            
    class _PatchAtomData(object):
        """Inner class used to encapsulate data about an atom in a patch definition."""
        def __init__(self, description):
            if ':' in description:
                colonIndex = description.find(':')
                self.residue = int(description[:colonIndex])-1
                self.name = description[colonIndex+1:]
            else:
                self.residue = 0
                self.name = description

683
    class _AtomType(object):
684
685
686
687
688
689
690
        """Inner class used to record atom types and associated properties."""
        def __init__(self, name, atomClass, mass, element):
            self.name = name
            self.atomClass = atomClass
            self.mass = mass
            self.element = element

691
    class _AtomTypeParameters(object):
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
        """Inner class used to record parameter values for atom types."""
        def __init__(self, forcefield, forceName, atomTag, paramNames):
            self.ff = forcefield
            self.forceName = forceName
            self.atomTag = atomTag
            self.paramNames = paramNames
            self.paramsForType = {}
            self.extraParamsForType = {}

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

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

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

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


762
    def _getResidueTemplateMatches(self, res, bondedToAtom, templateSignatures=None):
763
764
765
766
767
768
        """Return the residue template matches, or None if none are found.

        Parameters
        ----------
        res : Topology.Residue
            The residue for which template matches are to be retrieved.
769
770
        bondedToAtom : list of set of int
            bondedToAtom[i] is the set of atoms bonded to atom index i
771
772
773
774
775
776
777
778
779
780
781
782

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

        """
        template = None
        matches = None
783
784
        if templateSignatures is None:
            templateSignatures = self._templateSignatures
785
        signature = _createResidueSignature([atom.element for atom in res.atoms()])
786
        if signature in templateSignatures:
787
            allMatches = []
788
            for t in templateSignatures[signature]:
789
790
791
792
793
794
795
796
                match = _matchResidue(res, t, bondedToAtom)
                if match is not None:
                    allMatches.append((t, match))
            if len(allMatches) == 1:
                template = allMatches[0][0]
                matches = allMatches[0][1]
            elif len(allMatches) > 1:
                raise Exception('Multiple matching templates found for residue %d (%s).' % (res.index+1, res.name))
797
798
        return [template, matches]

799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
    def _buildBondedToAtomList(self, topology):
        """Build a list of which atom indices are bonded to each atom.

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

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

        """
        bondedToAtom = []
        for atom in topology.atoms():
            bondedToAtom.append(set())
816
817
818
        for (atom1, atom2) in topology.bonds():
            bondedToAtom[atom1.index].add(atom2.index)
            bondedToAtom[atom2.index].add(atom1.index)
819
        return bondedToAtom
820

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

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

826
827
828
        Parameters
        ----------
        topology : Topology
829
            The Topology whose residues are to be checked against the forcefield residue templates.
830
831
832
833

        Returns
        -------
        unmatched_residues : list of Residue
834
            List of Residue objects from `topology` for which no forcefield residue templates are available.
835
836
837
838
839
            Note that multiple instances of the same residue appearing at different points in the topology may be returned.

        This method may be of use in generating missing residue templates or diagnosing parameterization failures.
        """
        # Find the template matching each residue, compiling a list of residues for which no templates are available.
840
        bondedToAtom = self._buildBondedToAtomList(topology)
841
        unmatched_residues = list() # list of unmatched residues
842
843
844
845
846
847
        for res in topology.residues():
            # Attempt to match one of the existing templates.
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
            if matches is None:
                # No existing templates match.
                unmatched_residues.append(res)
848
849
850

        return unmatched_residues

851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
    def getMatchingTemplates(self, topology):
        """Return a list of forcefield residue templates matching residues in the specified topology.

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

        Parameters
        ----------
        topology : Topology
            The Topology whose residues are to be checked against the forcefield residue templates.

        Returns
        -------
        templates : list of _TemplateData
            List of forcefield residue templates corresponding to residues in the topology.
            templates[index] is template corresponding to residue `index` in topology.residues()

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

        return templates

883
884
    def generateTemplatesForUnmatchedResidues(self, topology):
        """Generate forcefield residue templates for residues in specified topology for which no forcefield templates are available.
885

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

888
889
890
        Parameters
        ----------
        topology : Topology
891
            The Topology whose residues are to be checked against the forcefield residue templates.
892
893
894

        Returns
        -------
895
896
897
898
899
900
        templates : list of _TemplateData
            List of forcefield residue templates corresponding to residues in `topology` for which no forcefield templates are currently available.
            Atom types will be set to `None`, but template name, atom names, elements, and connectivity will be taken from corresponding Residue objects.
        residues : list of Residue
            List of Residue objects that were used to generate the templates.
            `residues[index]` is the Residue that was used to generate the template `templates[index]`
901
902
903
904
905

        """
        # Get a non-unique list of unmatched residues.
        unmatched_residues = self.getUnmatchedResidues(topology)
        # Generate a unique list of unmatched residues by comparing fingerprints.
906
        bondedToAtom = self._buildBondedToAtomList(topology)
907
908
        unique_unmatched_residues = list() # list of unique unmatched Residue objects from topology
        templates = list() # corresponding _TemplateData templates
909
910
911
        signatures = set()
        for residue in unmatched_residues:
            signature = _createResidueSignature([ atom.element for atom in residue.atoms() ])
912
            template = _createResidueTemplate(residue)
913
914
915
916
917
918
919
920
921
922
923
            is_unique = True
            if signature in signatures:
                # Signature is the same as an existing residue; check connectivity.
                for check_residue in unique_unmatched_residues:
                    matches = _matchResidue(check_residue, template, bondedToAtom)
                    if matches is not None:
                        is_unique = False
            if is_unique:
                # Residue is unique.
                unique_unmatched_residues.append(residue)
                signatures.add(signature)
924
                templates.append(template)
925

926
        return [templates, unique_unmatched_residues]
927

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

Robert McGibbon's avatar
Robert McGibbon committed
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
        Parameters
        ----------
        topology : Topology
            The Topology for which to create a System
        nonbondedMethod : object=NoCutoff
            The method to use for nonbonded interactions.  Allowed values are
            NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, or PME.
        nonbondedCutoff : distance=1*nanometer
            The cutoff distance to use for nonbonded interactions
        constraints : object=None
            Specifies which bonds and angles should be implemented with constraints.
            Allowed values are None, HBonds, AllBonds, or HAngles.
        rigidWater : boolean=True
            If true, water molecules will be fully rigid regardless of the value
            passed for the constraints argument
        removeCMMotion : boolean=True
            If true, a CMMotionRemover will be added to the System
        hydrogenMass : mass=None
            The mass to use for hydrogen atoms bound to heavy atoms.  Any mass
            added to a hydrogen is subtracted from the heavy atom to keep
            their total mass the same.
953
954
955
956
957
958
959
        residueTemplates : dict=dict()
           Key: Topology Residue object
           Value: string, name of _TemplateData residue template object to use for
                  (Key) residue
           This allows user to specify which template to apply to particular Residues
           in the event that multiple matching templates are available (e.g Fe2+ and Fe3+
           templates in the ForceField for a monoatomic iron ion in the topology).
Robert McGibbon's avatar
Robert McGibbon committed
960
961
962
963
964
965
966
967
968
        args
             Arbitrary additional keyword arguments may also be specified.
             This allows extra parameters to be specified that are specific to
             particular force fields.

        Returns
        -------
        system
            the newly created System
969
970
        """
        data = ForceField._SystemData()
971
        data.atoms = list(topology.atoms())
972
973
        for atom in data.atoms:
            data.excludeAtomWith.append([])
974
975

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

977
        for bond in topology.bonds():
978
            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
979
980

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

982
983
984
985
986
987
988
989
990
991
992
993
        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
994

995
        unmatchedResidues = []
996
997
        for chain in topology.chains():
            for res in chain.residues():
998
999
1000
1001
1002
1003
1004
1005
1006
                if res in residueTemplates:
                    tname = residueTemplates[res]
                    template = self._templates[tname]
                    matches = _matchResidue(res, template, bondedToAtom)
                    if matches is None:
                        raise Exception('User-supplied template %s does not match the residue %d (%s)' % (tname, res.index+1, res.name))
                else:
                    # Attempt to match one of the existing templates.
                    [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
1007
                if matches is None:
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
                    unmatchedResidues.append(res)
                else:
                    data.recordMatchedAtomParameters(res, template, matches)
        
        # Try to apply patches to find matches for any unmatched residues.
        
        if len(unmatchedResidues) > 0:
            unmatchedResidues = _applyPatchesToMatchResidues(self, data, unmatchedResidues, bondedToAtom)
        
        # If we still haven't found a match for a residue, attempt to use residue template generators to create
        # new templates (and potentially atom types/parameters).
1019

1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
        for res in unmatchedResidues:
            # A template might have been generated on an earlier iteration of this loop.
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
            if matches is None:
                # Try all generators.
                for generator in self._templateGenerators:
                    if generator(self, res):
                        # This generator has registered a new residue template that should match.
                        [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
                        if matches is None:
                            # Something went wrong because the generated template does not match the residue signature.
                            raise Exception('The residue handler %s indicated it had correctly parameterized residue %s, but the generated template did not match the residue signature.' % (generator.__class__.__name__, str(res)))
                        else:
                            # We successfully generated a residue template.  Break out of the for loop.
                            break
            if matches is None:
                raise ValueError('No template found for residue %d (%s).  %s' % (res.index+1, res.name, _findMatchErrors(self, res)))
            else:
                data.recordMatchedAtomParameters(res, template, matches)
1039
1040

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

1042
1043
        sys = mm.System()
        for atom in topology.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1044
            # Look up the atom type name, returning a helpful error message if it cannot be found.
1045
1046
1047
1048
            if atom not in data.atomType:
                raise Exception("Could not identify atom type for atom '%s'." % str(atom))
            typename = data.atomType[atom]

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1049
            # Look up the type name in the list of registered atom types, returning a helpful error message if it cannot be found.
1050
1051
1052
1053
            if typename not in self._atomTypes:
                msg  = "Could not find typename '%s' for atom '%s' in list of known atom types.\n" % (typename, str(atom))
                msg += "Known atom types are: %s" % str(self._atomTypes.keys())
                raise Exception(msg)
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1054
1055

            # Add the particle to the OpenMM system.
1056
            mass = self._atomTypes[typename].mass
1057
            sys.addParticle(mass)
1058

1059
        # Adjust hydrogen masses if requested.
1060

1061
        if hydrogenMass is not None:
1062
1063
            if not unit.is_quantity(hydrogenMass):
                hydrogenMass *= unit.dalton
1064
1065
1066
1067
1068
1069
1070
            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
1071

1072
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
1073

1074
1075
1076
        boxVectors = topology.getPeriodicBoxVectors()
        if boxVectors is not None:
            sys.setDefaultPeriodicBoxVectors(boxVectors[0], boxVectors[1], boxVectors[2])
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1077
        elif nonbondedMethod not in [NoCutoff, CutoffNonPeriodic]:
1078
1079
1080
            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
1081

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

1098
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1099

1100
1101
1102
        uniquePropers = set()
        for angle in data.angles:
            for atom in bondedToAtom[angle[0]]:
pgrinaway's avatar
pgrinaway committed
1103
                if atom not in angle:
1104
1105
1106
1107
1108
                    if atom < angle[2]:
                        uniquePropers.add((atom, angle[0], angle[1], angle[2]))
                    else:
                        uniquePropers.add((angle[2], angle[1], angle[0], atom))
            for atom in bondedToAtom[angle[2]]:
pgrinaway's avatar
pgrinaway committed
1109
                if atom not in angle:
1110
1111
1112
1113
1114
                    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
1115

1116
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1117

1118
1119
1120
1121
1122
        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
1123

1124
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1125

1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
        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
1140

1141
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1142

1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
        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
1164

1165
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
1166

1167
        for atom in data.virtualSites:
1168
            (site, atoms, excludeWith) = data.virtualSites[atom]
1169
            index = atom.index
1170
            data.excludeAtomWith[excludeWith].append(index)
1171
            if site.type == 'average2':
1172
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(atoms[0], atoms[1], site.weights[0], site.weights[1]))
1173
            elif site.type == 'average3':
1174
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
1175
            elif site.type == 'outOfPlane':
1176
1177
1178
1179
1180
1181
1182
                sys.setVirtualSite(index, mm.OutOfPlaneSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
            elif site.type == 'localCoords':
                sys.setVirtualSite(index, mm.LocalCoordinatesSite(atoms[0], atoms[1], atoms[2],
                                                                  mm.Vec3(site.originWeights[0], site.originWeights[1], site.originWeights[2]),
                                                                  mm.Vec3(site.xWeights[0], site.xWeights[1], site.xWeights[2]),
                                                                  mm.Vec3(site.yWeights[0], site.yWeights[1], site.yWeights[2]),
                                                                  mm.Vec3(site.localPos[0], site.localPos[1], site.localPos[2])))
Justin MacCallum's avatar
Justin MacCallum committed
1183

1184
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
1185

1186
1187
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
1188
1189
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
1190

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

peastman's avatar
peastman committed
1193
1194
1195
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
1196

1197
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
1198

1199
        for script in self._scripts:
1200
            exec(script, locals())
1201
1202
1203
        return sys


1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
def _findBondsForExclusions(data, sys):
    """Create a list of bonds to use when identifying exclusions."""
    bondIndices = []
    for bond in data.bonds:
        bondIndices.append((bond.atom1, bond.atom2))

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

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

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

    for atom1, atom2 in bondIndices:
        for child1 in data.excludeAtomWith[atom1]:
            bondIndices.append((child1, atom2))
            for child2 in data.excludeAtomWith[atom2]:
                bondIndices.append((child1, child2))
        for child2 in data.excludeAtomWith[atom2]:
            bondIndices.append((atom1, child2))
    return bondIndices

1230
1231
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
1232
1233
    counts = {}
    for element in elements:
1234
        if element in counts:
1235
1236
1237
            counts[element] += 1
        else:
            counts[element] = 1
1238
1239
1240
1241
1242
1243
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
1244
1245
    sig = []
    for c in counts:
1246
1247
        if c is not None:
            sig.append((c, counts[c]))
1248
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
1249

1250
    # Convert it to a string.
1251
1252

    s = ''
1253
    for element, count in sig:
1254
1255
1256
        s += element.symbol+str(count)
    return s

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

Robert McGibbon's avatar
Robert McGibbon committed
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
    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
    -------
1271
1272
1273
    list
        a list specifying which atom of the template each atom of the residue
        corresponds to, or None if it does not match the template
1274
1275
    """
    atoms = list(res.atoms())
peastman's avatar
peastman committed
1276
1277
    numAtoms = len(atoms)
    if numAtoms != len(template.atoms):
1278
        return None
Justin MacCallum's avatar
Justin MacCallum committed
1279

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

1282
    renumberAtoms = {}
peastman's avatar
peastman committed
1283
    for i in range(numAtoms):
1284
        renumberAtoms[atoms[i].index] = i
1285
1286
1287
    bondedTo = []
    externalBonds = []
    for atom in atoms:
1288
        bonds = [renumberAtoms[x] for x in bondedToAtom[atom.index] if x in renumberAtoms]
1289
        bondedTo.append(bonds)
1290
        externalBonds.append(len([x for x in bondedToAtom[atom.index] if x not in renumberAtoms]))
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308

    # 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
peastman's avatar
peastman committed
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
        
    # Identify template atoms that could potentially be matches for each atom.
    
    candidates = [[] for i in range(numAtoms)]
    for i in range(numAtoms):
        for j, atom in enumerate(template.atoms):
            if (atom.element is not None and atom.element != atoms[i].element) or (atom.element is None and atom.name != atoms[i].name):
                continue
            if len(atom.bondedTo) != len(bondedTo[i]):
                continue
            if atom.externalBonds != externalBonds[i]:
                continue
            candidates[i].append(j)

    # Find an optimal ordering for matching atoms.  This means 1) start with the one that has the fewest options,
    # and 2) follow with ones that are bonded to an already matched atom.
    
    searchOrder = []
    atomsToOrder = set(range(numAtoms))
1328
1329
    efficientAtomSet = set()
    efficientAtomHeap = []
peastman's avatar
peastman committed
1330
    while len(atomsToOrder) > 0:
1331
        if len(efficientAtomSet) == 0:
peastman's avatar
peastman committed
1332
1333
1334
1335
1336
1337
            fewestNeighbors = numAtoms+1
            for i in atomsToOrder:
                if len(candidates[i]) < fewestNeighbors:
                    nextAtom = i
                    fewestNeighbors = len(candidates[i])
        else:
1338
1339
            nextAtom = heappop(efficientAtomHeap)[1]
            efficientAtomSet.remove(nextAtom)
peastman's avatar
peastman committed
1340
1341
1342
1343
        searchOrder.append(nextAtom)
        atomsToOrder.remove(nextAtom)
        for i in bondedTo[nextAtom]:
            if i in atomsToOrder:
1344
1345
1346
                if i not in efficientAtomSet:
                    efficientAtomSet.add(i)
                    heappush(efficientAtomHeap, (len(candidates[i]), i))
peastman's avatar
peastman committed
1347
1348
1349
1350
1351
    inverseSearchOrder = [0]*numAtoms
    for i in range(numAtoms):
        inverseSearchOrder[searchOrder[i]] = i
    bondedTo = [[inverseSearchOrder[bondedTo[i][j]] for j in range(len(bondedTo[i]))] for i in searchOrder]
    candidates = [candidates[i] for i in searchOrder]
1352
1353
1354

    # Recursively match atoms.

peastman's avatar
peastman committed
1355
1356
1357
1358
    matches = numAtoms*[0]
    hasMatch = numAtoms*[False]
    if _findAtomMatches(template, bondedTo, matches, hasMatch, candidates, 0):
        return [matches[inverseSearchOrder[i]] for i in range(numAtoms)]
1359
1360
1361
    return None


peastman's avatar
peastman committed
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
def _getAtomMatchCandidates(template, bondedTo, matches, candidates, position):
    """Get a list of template atoms that are potential matches for the next atom."""
    for bonded in bondedTo[position]:
        if bonded < position:
            # This atom is bonded to another one for which we already have a match, so only consider
            # template atoms that *that* one is bonded to.
            return template.atoms[matches[bonded]].bondedTo
    return candidates[position]


def _findAtomMatches(template, bondedTo, matches, hasMatch, candidates, position):
1373
    """This is called recursively from inside _matchResidue() to identify matching atoms."""
peastman's avatar
peastman committed
1374
    if position == len(matches):
1375
        return True
peastman's avatar
peastman committed
1376
    for i in _getAtomMatchCandidates(template, bondedTo, matches, candidates, position):
1377
        atom = template.atoms[i]
peastman's avatar
peastman committed
1378
        if not hasMatch[i] and i in candidates[position]:
1379
            # See if the bonds for this identification are consistent
Justin MacCallum's avatar
Justin MacCallum committed
1380

1381
1382
            allBondsMatch = all((bonded > position or matches[bonded] in atom.bondedTo for bonded in bondedTo[position]))
            if allBondsMatch:
peastman's avatar
peastman committed
1383
                # This is a possible match, so try matching the rest of the residue.
Justin MacCallum's avatar
Justin MacCallum committed
1384

1385
1386
                matches[position] = i
                hasMatch[i] = True
peastman's avatar
peastman committed
1387
                if _findAtomMatches(template, bondedTo, matches, hasMatch, candidates, position+1):
1388
1389
1390
1391
1392
                    return True
                hasMatch[i] = False
    return False


1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
def _applyPatchesToMatchResidues(forcefield, data, residues, bondedToAtom):
    """Try to apply patches to find matches for residues."""
    # Start by creating all templates than can be created by applying a combination of one-residue patches
    # to a single template.  The number of these is usually not too large, and they often cover a large fraction
    # of residues.
    
    patchedTemplateSignatures = {}
    patchedTemplates = {}
    for name, template in forcefield._templates.items():
        if name in forcefield._templatePatches:
            patches = [forcefield._patches[patchName] for patchName, patchResidueIndex in forcefield._templatePatches[name] if forcefield._patches[patchName].numResidues == 1]
            if len(patches) > 0:
                newTemplates = []
                patchedTemplates[name] = newTemplates
                _generatePatchedSingleResidueTemplates(template, patches, 0, newTemplates)
                for patchedTemplate in newTemplates:
                    signature = _createResidueSignature([atom.element for atom in patchedTemplate.atoms])
                    if signature in patchedTemplateSignatures:
                        patchedTemplateSignatures[signature].append(patchedTemplate)
                    else:
                        patchedTemplateSignatures[signature] = [patchedTemplate]
    
    # Now see if any of those templates matches any of the residues.
    
    unmatchedResidues = []
    for res in residues:
        [template, matches] = forcefield._getResidueTemplateMatches(res, bondedToAtom, patchedTemplateSignatures)
        if matches is None:
            unmatchedResidues.append(res)
        else:
            data.recordMatchedAtomParameters(res, template, matches)
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
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
    if len(unmatchedResidues) == 0:
        return []
    
    # We need to consider multi-residue patches.  This can easily lead to a combinatorial explosion, so we make a simplifying
    # assumption: that no residue is affected by more than one multi-residue patch (in addition to any number of single-residue
    # patches).  Record all multi-residue patches, and the templates they can be applied to.
    
    patches = {}
    maxPatchSize = 0
    for patch in forcefield._patches.values():
        if patch.numResidues > 1:
            patches[patch.name] = [[] for i in range(patch.numResidues)]
            maxPatchSize = max(maxPatchSize, patch.numResidues)
    if maxPatchSize == 0:
        return unmatchedResidues # There aren't any multi-residue patches
    for templateName in forcefield._templatePatches:
        for patchName, patchResidueIndex in forcefield._templatePatches[templateName]:
            if patchName in patches:
                # The patch should accept this template, *and* all patched versions of it generated above.
                patches[patchName][patchResidueIndex].append(forcefield._templates[templateName])
                if templateName in patchedTemplates:
                    patches[patchName][patchResidueIndex] += patchedTemplates[templateName]
    
    # Record which unmatched residues are bonded to each other.
    
    bonds = set()
    topology = residues[0].chain.topology
    for atom1, atom2 in topology.bonds():
        if atom1.residue != atom2.residue:
            res1 = atom1.residue
            res2 = atom2.residue
            if res1 in unmatchedResidues and res2 in unmatchedResidues:
                bond = tuple(sorted((res1, res2), key=lambda x: x.index))
                if bond not in bonds:
                    bonds.add(bond)
    
    # Identify clusters of unmatched residues that are all bonded to each other.  These are the ones we'll
    # try to apply multi-residue patches to.
    
    clusterSize = 2
    clusters = bonds
    while clusterSize <= maxPatchSize:
        # Try to apply patches to clusters of this size.
        
        for patchName in patches:
            patch = forcefield._patches[patchName]
            if patch.numResidues == clusterSize:
                matchedClusters = _matchToMultiResiduePatchedTemplates(data, clusters, patch, patches[patchName], bondedToAtom)
                for cluster in matchedClusters:
                    for residue in cluster:
                        unmatchedResidues.remove(residue)
                bonds = set(bond for bond in bonds if bond[0] in unmatchedResidues and bond[1] in unmatchedResidues)

        # Now extend the clusters to find ones of the next size up.
        
        largerClusters = set()
        for cluster in clusters:
            for bond in bonds:
                if bond[0] in cluster and bond[1] not in cluster:
                    newCluster = tuple(sorted(cluster+(bond[1],), key=lambda x: x.index))
                    largerClusters.add(newCluster)
                elif bond[1] in cluster and bond[0] not in cluster:
                    newCluster = tuple(sorted(cluster+(bond[0],), key=lambda x: x.index))
                    largerClusters.add(newCluster)
        if len(largerClusters) == 0:
            # There are no clusters of this size or larger
            break
        clusters = largerClusters
        clusterSize += 1

1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
    return unmatchedResidues


def _generatePatchedSingleResidueTemplates(template, patches, index, newTemplates):
    """Apply all possible combinations of a set of single-residue patches to a template."""
    try:
        patchedTemplate = patches[index].createPatchedTemplates([template])[0]
        newTemplates.append(patchedTemplate)
    except:
        # This probably means the patch is inconsistent with another one that has already been applied,
        # so just ignore it.
        patchedTemplate = None
    
    # Call this function recursively to generate combinations of patches.
    
    if index+1 < len(patches):
        _generatePatchedSingleResidueTemplates(template, patches, index+1, newTemplates)
        if patchedTemplate is not None:
            _generatePatchedSingleResidueTemplates(patchedTemplate, patches, index+1, newTemplates)


1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
def _matchToMultiResiduePatchedTemplates(data, clusters, patch, residueTemplates, bondedToAtom):
    """Apply a multi-residue patch to templates, then try to match them against clusters of residues."""
    matchedClusters = []
    selectedTemplates = [None]*patch.numResidues
    _applyMultiResiduePatch(data, clusters, patch, residueTemplates, selectedTemplates, 0, matchedClusters, bondedToAtom)
    return matchedClusters


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

    if index < patch.numResidues:
        for template in candidateTemplates[index]:
            selectedTemplates[index] = template
            _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index+1, matchedClusters, bondedToAtom)
    else:
        # We're at the deepest level of the recursion.  We've selected a template for each residue, so apply the patch,
        # then try to match it against clusters.
        
        try:
            patchedTemplates = patch.createPatchedTemplates(selectedTemplates)
        except:
            # This probably means the patch is inconsistent with another one that has already been applied,
            # so just ignore it.
            raise
            return
        newlyMatchedClusters = []
        for cluster in clusters:
            for residues in itertools.permutations(cluster):
                residueMatches = []
                for residue, template in zip(residues, patchedTemplates):
                    matches = _matchResidue(residue, template, bondedToAtom)
                    if matches is None:
                        residueMatches = None
                        break
                    else:
                        residueMatches.append(matches)
                if residueMatches is not None:
                    # We successfully matched the template to the residues.  Record the parameters.
                    
                    for i in range(patch.numResidues):
                        data.recordMatchedAtomParameters(residues[i], patchedTemplates[i], residueMatches[i])
                    newlyMatchedClusters.append(cluster)
                    break
        
        # Record which clusters were successfully matched.
        
        matchedClusters += newlyMatchedClusters
        for cluster in newlyMatchedClusters:
            clusters.remove(cluster)
        

1567
1568
1569
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()])
1570
    numResidueAtoms = sum(residueCounts.values())
1571
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
1572

1573
    # Loop over templates and see how closely each one might match.
1574

1575
1576
1577
1578
1579
1580
    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])
1581

1582
        # Does the residue have any atoms that clearly aren't in the template?
1583

1584
1585
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
1586

1587
        # If there are too many missing atoms, discard this template.
1588

1589
        numTemplateAtoms = sum(templateCounts.values())
1590
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
1591
1592
1593
1594
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
1595

1596
1597
        # 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.
1598

1599
1600
1601
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
1602

1603
        # Accept this as our new best match.
1604

1605
1606
1607
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
1608
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
1609

1610
    # Return an appropriate error message.
1611

1612
    if numBestMatchAtoms == numResidueAtoms:
1613
1614
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
1615
1616
1617
1618
            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:
1619
1620
1621
1622
1623
            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)
1624
1625
1626
        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.'

1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
def _createResidueTemplate(residue):
    """Create a _TemplateData template from a Residue object.

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

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

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

    """
    template = ForceField._TemplateData(residue.name)
    for atom in residue.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1645
        template.atoms.append(ForceField._TemplateAtomData(atom.name, None, atom.element))
1646
1647
1648
1649
1650
1651
1652
1653
1654
    for (atom1,atom2) in residue.internal_bonds():
        template.addBondByName(atom1.name, atom2.name)
    residue_atoms = [ atom for atom in residue.atoms() ]
    for (atom1,atom2) in residue.external_bonds():
        if atom1 in residue_atoms:
            template.addExternalBondByName(atom1.name)
        elif atom2 in residue_atoms:
            template.addExternalBondByName(atom2.name)
    return template
1655

1656
1657
1658
1659
1660
# 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.

1661
## @private
1662
class HarmonicBondGenerator(object):
1663
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1664

1665
1666
    def __init__(self, forcefield):
        self.ff = forcefield
1667
1668
1669
1670
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
1671

1672
1673
1674
1675
1676
1677
1678
    def registerBond(self, parameters):
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.length.append(_convertParameterToNumber(parameters['length']))
            self.k.append(_convertParameterToNumber(parameters['k']))
Justin MacCallum's avatar
Justin MacCallum committed
1679

1680
1681
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1682
1683
1684
1685
1686
1687
        existing = [f for f in ff._forces if isinstance(f, HarmonicBondGenerator)]
        if len(existing) == 0:
            generator = HarmonicBondGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
1688
        for bond in element.findall('Bond'):
1689
            generator.registerBond(bond.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1690

1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
    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:
1708
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
1709
1710
1711
1712
1713
1714
1715
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


1716
## @private
1717
class HarmonicAngleGenerator(object):
1718
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1719

1720
1721
    def __init__(self, forcefield):
        self.ff = forcefield
1722
1723
1724
1725
1726
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
1727

1728
1729
1730
1731
1732
1733
1734
1735
1736
    def registerAngle(self, parameters):
        types = self.ff._findAtomTypes(parameters, 3)
        if None not in types:
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.types3.append(types[2])
            self.angle.append(_convertParameterToNumber(parameters['angle']))
            self.k.append(_convertParameterToNumber(parameters['k']))

1737
1738
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1739
1740
1741
1742
1743
1744
        existing = [f for f in ff._forces if isinstance(f, HarmonicAngleGenerator)]
        if len(existing) == 0:
            generator = HarmonicAngleGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
1745
        for angle in element.findall('Angle'):
1746
            generator.registerAngle(angle.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1747

1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
    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
1767

1768
1769
1770
1771
1772
1773
1774
1775
1776
                        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
1777

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

1780
1781
1782
1783
1784
                        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]))
1785
                                data.addConstraint(sys, angle[0], angle[2], length)
1786
1787
1788
1789
1790
1791
1792
                    elif self.k[i] != 0:
                        force.addAngle(angle[0], angle[1], angle[2], self.angle[i], self.k[i])
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


1793
## @private
1794
class PeriodicTorsion(object):
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
    """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 = []

1806
## @private
1807
class PeriodicTorsionGenerator(object):
1808
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1809

1810
1811
    def __init__(self, forcefield):
        self.ff = forcefield
1812
1813
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
1814

1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
    def registerProperTorsion(self, parameters):
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
            self.proper.append(torsion)

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

1825
1826
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1827
1828
1829
1830
1831
1832
        existing = [f for f in ff._forces if isinstance(f, PeriodicTorsionGenerator)]
        if len(existing) == 0:
            generator = PeriodicTorsionGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
1833
        for torsion in element.findall('Proper'):
1834
            generator.registerProperTorsion(torsion.attrib)
1835
        for torsion in element.findall('Improper'):
1836
            generator.registerImproperTorsion(torsion.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1837

1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
    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]]]
1873
            match = None
1874
1875
1876
1877
1878
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
1879
                hasWildcard = (wildcard in (types1, types2, types3, types4))
1880
1881
1882
1883
                if (types1, types2, types3, types4).count(wildcard) == 2:
                    hasTwoWildcards = True
                else:
                    hasTwoWildcards = False
1884
1885
1886
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
1887
1888
1889
                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:
1890
                            # topology atom indexes
1891
1892
                            a1 = torsion[t2[1]]
                            a2 = torsion[t3[1]]
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
                            a4 = torsion[t4[1]]
                            # residue indexes
                            r1 = data.atoms[a1].residue.index
                            r2 = data.atoms[a2].residue.index
                            r4 = data.atoms[a4].residue.index
                            # template atom indexes
                            ta1 = data.atomTemplateIndexes[data.atoms[a1]]
                            ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                            ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                            # elements
1903
1904
                            e1 = data.atoms[a1].element
                            e2 = data.atoms[a2].element
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
                            e4 = data.atoms[a4].element
                            # the following for AMBER only - TODO: decide how to pass this in ffxml's
                            if isAmber:
                                if not hasWildcard:
                                    if t2[0] == t4[0] and (r1 > r4 or (r1 == r4 and ta1 > ta4)):
                                        (a1, a4) = (a4, a1)
                                    if t3[0] == t4[0] and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                        (a2, a4) = (a4, a2)
                                    if t2[0] == t3[0] and (r1 > r2 or (r1 == r2 and ta1 > ta2)):
                                        (a1, a2) = (a2, a1)
                                else:
                                    if e1 == e4 and (r1 > r4 or (r1 == r4 and ta1 > ta4)):
                                        (a1, a4) = (a4, a1)
                                    if e2 == e4 and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                        (a2, a4) = (a4, a2)
                                    if (r1 > r2 or (r1 == r2 and ta1 > ta2)):
                                        (a1, a2) = (a2, a1)
                            # the following is OpenMM default
                            else:
                                    if t2[0] == t4[0] and (r1 > r4 or (r1 == r4 and ta1 > ta4)):
                                        (a1, a4) = (a4, a1)
                                    if t3[0] == t4[0] and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                        (a2, a4) = (a4, a2)
                                    if t2[0] == t3[0] and (r1 > r2 or (r1 == r2 and ta1 > ta2)):
                                        (a1, a2) = (a2, a1)
                                    if hasTwoWildcards and (r1 > r2 or (r1 == r2 and ta1 > ta2)):
                                        (a1, a2) = (a2, a1)
                            match = (a1, a2, torsion[0], a4, tordef)
1933
                            break
1934
1935
1936
1937
1938
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                for i in range(len(tordef.phase)):
                    if tordef.k[i] != 0:
                        force.addTorsion(a1, a2, a3, a4, tordef.periodicity[i], tordef.phase[i], tordef.k[i])
1939
1940
1941
1942

parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


1943
## @private
1944
class RBTorsion(object):
1945
1946
1947
1948
1949
1950
1951
1952
1953
    """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

1954
## @private
1955
class RBTorsionGenerator(object):
1956
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1957

1958
1959
    def __init__(self, forcefield):
        self.ff = forcefield
1960
1961
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
1962

1963
1964
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1965
1966
1967
1968
1969
1970
        existing = [f for f in ff._forces if isinstance(f, RBTorsionGenerator)]
        if len(existing) == 0:
            generator = RBTorsionGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
1971
        for torsion in element.findall('Proper'):
1972
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1973
            if None not in types:
1974
1975
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
1976
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
1977
            if None not in types:
1978
                generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
Justin MacCallum's avatar
Justin MacCallum committed
1979

1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
    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]]]
2013
            match = None
2014
2015
2016
2017
2018
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
2019
2020
2021
2022
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
2023
2024
2025
                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:
2026
                            if hasWildcard:
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
                                # Workaround to be more consistent with AMBER.  It uses wildcards to define most of its
                                # impropers, which leaves the ordering ambiguous.  It then follows some bizarre rules
                                # to pick the order.
                                a1 = torsion[t2[1]]
                                a2 = torsion[t3[1]]
                                e1 = data.atoms[a1].element
                                e2 = data.atoms[a2].element
                                if e1 == e2 and a1 > a2:
                                    (a1, a2) = (a2, a1)
                                elif e1 != elem.carbon and (e2 == elem.carbon or e1.mass < e2.mass):
                                    (a1, a2) = (a2, a1)
2038
                                match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
2039
2040
                            else:
                                # There are no wildcards, so the order is unambiguous.
2041
                                match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
2042
                            break
2043
2044
2045
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.c[0], tordef.c[1], tordef.c[2], tordef.c[3], tordef.c[4], tordef.c[5])
2046
2047
2048
2049

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


2050
## @private
2051
class CMAPTorsion(object):
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
    """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

2062
## @private
2063
class CMAPTorsionGenerator(object):
2064
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2065

2066
2067
    def __init__(self, forcefield):
        self.ff = forcefield
2068
2069
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
2070

2071
2072
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2073
2074
2075
2076
2077
2078
        existing = [f for f in ff._forces if isinstance(f, CMAPTorsionGenerator)]
        if len(existing) == 0:
            generator = CMAPTorsionGenerator(ff)
            ff.registerGenerator(generator)
        else:
            generator = existing[0]
2079
2080
2081
2082
2083
2084
2085
        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'):
2086
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
2087
            if None not in types:
2088
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
2089

2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
    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
2100

2101
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
2102

2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
        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


2146
## @private
2147
class NonbondedGenerator(object):
2148
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2149

2150
2151
    SCALETOL = 1e-5

2152
2153
    def __init__(self, forcefield, coulomb14scale, lj14scale):
        self.ff = forcefield
2154
2155
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
2156
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
2157

2158
    def registerAtom(self, parameters):
2159
        self.params.registerAtom(parameters)
2160

2161
2162
2163
2164
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
        if len(existing) == 0:
2165
2166
            generator = NonbondedGenerator(ff, float(element.attrib['coulomb14scale']), float(element.attrib['lj14scale']))
            ff.registerGenerator(generator)
2167
2168
2169
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2170
2171
            if abs(generator.coulomb14scale - float(element.attrib['coulomb14scale'])) > NonbondedGenerator.SCALETOL or \
                    abs(generator.lj14scale - float(element.attrib['lj14scale'])) > NonbondedGenerator.SCALETOL:
Justin MacCallum's avatar
Justin MacCallum committed
2172
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
2173
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2174

2175
2176
2177
2178
2179
2180
2181
    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
2182
            raise ValueError('Illegal nonbonded method for NonbondedForce')
2183
2184
        force = mm.NonbondedForce()
        for atom in data.atoms:
2185
2186
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
2187
2188
2189
2190
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
2191
2192
        if 'useDispersionCorrection' in args:
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
peastman's avatar
peastman committed
2193
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
2194

peastman's avatar
peastman committed
2195
    def postprocessSystem(self, sys, data, args):
2196
2197
        # Create the exceptions.

2198
        bondIndices = _findBondsForExclusions(data, sys)
peastman's avatar
peastman committed
2199
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
2200
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
2201
2202
2203

parsers["NonbondedForce"] = NonbondedGenerator.parseElement

2204
2205

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

ChayaSt's avatar
ChayaSt committed
2209
    def __init__(self, forcefield, lj14scale):
ChayaSt's avatar
ChayaSt committed
2210
        self.ff = forcefield
2211
        self.nbfixTypes = {}
ChayaSt's avatar
ChayaSt committed
2212
        self.lj14scale = lj14scale
2213
        self.ljTypes = ForceField._AtomTypeParameters(forcefield, 'LennardJonesForce', 'Atom', ('sigma', 'epsilon'))
ChayaSt's avatar
ChayaSt committed
2214

ChayaSt's avatar
ChayaSt committed
2215
    def registerNBFIX(self, parameters):
ChayaSt's avatar
ChayaSt committed
2216
2217
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
2218
2219
            type1 = types[0][0]
            type2 = types[1][0]
ChayaSt's avatar
ChayaSt committed
2220
2221
            epsilon = _convertParameterToNumber(parameters['epsilon'])
            sigma = _convertParameterToNumber(parameters['sigma'])
2222
2223
            self.nbfixTypes[(type1, type2)] = [sigma, epsilon]
            self.nbfixTypes[(type2, type1)] = [sigma, epsilon]
2224

ChayaSt's avatar
ChayaSt committed
2225
    def registerLennardJones(self, parameters):
2226
        self.ljTypes.registerAtom(parameters)
ChayaSt's avatar
ChayaSt committed
2227
2228
2229

    @staticmethod
    def parseElement(element, ff):
ChayaSt's avatar
ChayaSt committed
2230
        existing = [f for f in ff._forces if isinstance(f, LennardJonesGenerator)]
ChayaSt's avatar
ChayaSt committed
2231
        if len(existing) == 0:
ChayaSt's avatar
ChayaSt committed
2232
            generator = LennardJonesGenerator(ff, float(element.attrib['lj14scale']))
ChayaSt's avatar
ChayaSt committed
2233
2234
            ff.registerGenerator(generator)
        else:
ChayaSt's avatar
ChayaSt committed
2235
            # Multiple <LennardJonesForce> tags were found, probably in different files
ChayaSt's avatar
ChayaSt committed
2236
            generator = existing[0]
2237
2238
            if abs(generator.lj14scale - float(element.attrib['lj14scale'])) > NonbondedGenerator.SCALETOL:
                raise ValueError('Found multiple LennardJonesForce tags with different 1-4 scales')
ChayaSt's avatar
ChayaSt committed
2239
2240
        for LJ in element.findall('Atom'):
            generator.registerLennardJones(LJ.attrib)
ChayaSt's avatar
ChayaSt committed
2241
        for Nbfix in element.findall('NBFixPair'):
ChayaSt's avatar
ChayaSt committed
2242
            generator.registerNBFIX(Nbfix.attrib)
ChayaSt's avatar
ChayaSt committed
2243

ChayaSt's avatar
ChayaSt committed
2244
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
ChayaSt's avatar
ChayaSt committed
2245
        # First derive the lookup tables
2246
2247
2248
2249
2250
2251

        nbfixTypeSet = set().union(*self.nbfixTypes)
        ljIndexList = [None]*len(data.atoms)
        numLjTypes = 0
        ljTypeList = []
        typeMap = {}
ChayaSt's avatar
ChayaSt committed
2252
        for i, atom in enumerate(data.atoms):
ChayaSt's avatar
ChayaSt committed
2253
            atype = data.atomType[atom]
2254
2255
2256
2257
            values = tuple(self.ljTypes.getAtomParameters(atom, data))
            if values in typeMap and atype not in nbfixTypeSet:
                # Only non-NBFIX types can be compressed
                ljIndexList[i] = typeMap[values]
2258
            else:
2259
2260
2261
2262
2263
2264
2265
                typeMap[values] = numLjTypes
                ljIndexList[i] = numLjTypes
                numLjTypes += 1
                ljTypeList.append(atype)
        reverseMap = [0]*len(typeMap)
        for typeValue in typeMap:
            reverseMap[typeMap[typeValue]] = typeValue
2266

ChayaSt's avatar
ChayaSt committed
2267
        # Now everything is assigned. Create the A- and B-coefficient arrays
2268
2269
        
        acoef = [0]*(numLjTypes*numLjTypes)
ChayaSt's avatar
ChayaSt committed
2270
        bcoef = acoef[:]
2271
2272
2273
2274
2275
2276
2277
2278
2279
        for m in range(numLjTypes):
            for n in range(numLjTypes):
                pair = (ljTypeList[m], ljTypeList[n])
                if pair in self.nbfixTypes:
                    epsilon = self.nbfixTypes[pair][1]
                    sigma = self.nbfixTypes[pair][0]
                    sigma6 = sigma**6
                    acoef[m+numLjTypes*n] = 4*epsilon*sigma6*sigma6
                    bcoef[m+numLjTypes*n] = 4*epsilon*sigma6
ChayaSt's avatar
cleanup  
ChayaSt committed
2280
                    continue
ChayaSt's avatar
ChayaSt committed
2281
                else:
2282
2283
2284
2285
2286
2287
2288
2289
2290
                    sigma = 0.5*(reverseMap[m][0]+reverseMap[n][0])
                    sigma6 = sigma**6
                    epsilon = math.sqrt(reverseMap[m][-1]*reverseMap[n][-1])
                    acoef[m+numLjTypes*n] = 4*epsilon*sigma6*sigma6
                    bcoef[m+numLjTypes*n] = 4*epsilon*sigma6

        self.force = mm.CustomNonbondedForce('acoef(type1, type2)/r^12 - bcoef(type1, type2)/r^6;')
        self.force.addTabulatedFunction('acoef', mm.Discrete2DFunction(numLjTypes, numLjTypes, acoef))
        self.force.addTabulatedFunction('bcoef', mm.Discrete2DFunction(numLjTypes, numLjTypes, bcoef))
ChayaSt's avatar
ChayaSt committed
2291
        self.force.addPerParticleParameter('type')
2292
        if nonbondedMethod in [CutoffPeriodic, Ewald, PME]:
ChayaSt's avatar
ChayaSt committed
2293
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
ChayaSt's avatar
ChayaSt committed
2294
        elif nonbondedMethod is NoCutoff:
ChayaSt's avatar
ChayaSt committed
2295
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
ChayaSt's avatar
ChayaSt committed
2296
        elif nonbondedMethod is CutoffNonPeriodic:
ChayaSt's avatar
ChayaSt committed
2297
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic)
ChayaSt's avatar
ChayaSt committed
2298
        else:
2299
2300
            raise AssertionError('Unrecognized nonbonded method [%s]' % nonbondedMethod)

ChayaSt's avatar
ChayaSt committed
2301
        # Add the particles
2302
2303

        for i in ljIndexList:
ChayaSt's avatar
cleanup  
ChayaSt committed
2304
            self.force.addParticle((i,))
ChayaSt's avatar
ChayaSt committed
2305
2306
        self.force.setUseLongRangeCorrection(True)
        self.force.setCutoffDistance(nonbondedCutoff)
ChayaSt's avatar
ChayaSt committed
2307
2308
2309
        sys.addForce(self.force)

    def postprocessSystem(self, sys, data, args):
ChayaSt's avatar
ChayaSt committed
2310
        # Create the exceptions.
2311
        
2312
        bondIndices = _findBondsForExclusions(data, sys)
2313
2314
2315
        if self.lj14scale == 1:
            # Just exclude the 1-2 and 1-3 interactions.
            
peastman's avatar
Bug fix  
peastman committed
2316
            self.force.createExclusionsFromBonds(bondIndices, 2)
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
        else:
            forceCopy = deepcopy(self.force)
            forceCopy.createExclusionsFromBonds(bondIndices, 2)
            self.force.createExclusionsFromBonds(bondIndices, 3)
            if self.force.getNumExclusions() > forceCopy.getNumExclusions() and self.lj14scale != 0:
                # We need to create a CustomBondForce and use it to implement the scaled 1-4 interactions.
                
                bonded = mm.CustomBondForce('%g*epsilon*((sigma/r)^12-(sigma/r)^6)' % (4*self.lj14scale))
                bonded.addPerBondParameter('sigma')
                bonded.addPerBondParameter('epsilon')
                sys.addForce(bonded)
                skip = set(tuple(forceCopy.getExclusionParticles(i)) for i in range(forceCopy.getNumExclusions()))
                for i in range(self.force.getNumExclusions()):
                    p1,p2 = self.force.getExclusionParticles(i)
                    a1 = data.atoms[p1]
                    a2 = data.atoms[p2]
                    if (p1,p2) not in skip and (p2,p1) not in skip:
                        type1 = data.atomType[a1]
                        type2 = data.atomType[a2]
                        if (type1, type2) in self.nbfixTypes:
                            sigma, epsilon = self.nbfixTypes[(type1, type2)]
                        else:
                            values1 = self.ljTypes.getAtomParameters(a1, data)
                            values2 = self.ljTypes.getAtomParameters(a2, data)
                            sigma = 0.5*(values1[0]+values2[0])
                            epsilon = sqrt(values1[1]*values2[1])
                        bonded.addBond(p1, p2, (sigma, epsilon))
ChayaSt's avatar
ChayaSt committed
2344

ChayaSt's avatar
ChayaSt committed
2345
parsers["LennardJonesForce"] = LennardJonesGenerator.parseElement
2346

2347

2348
## @private
2349
class GBSAOBCGenerator(object):
2350
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2351

2352
2353
    def __init__(self, forcefield):
        self.ff = forcefield
2354
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
2355

2356
    def registerAtom(self, parameters):
2357
        self.params.registerAtom(parameters)
2358

2359
2360
    @staticmethod
    def parseElement(element, ff):
2361
2362
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
2363
2364
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
2365
2366
2367
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2368
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2369

2370
2371
2372
2373
2374
    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
2375
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
2376
2377
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
2378
2379
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
2380
2381
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2382
2383
2384
2385
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
2386
2387
        sys.addForce(force)

2388
2389
    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
2390

2391
2392
2393
2394
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

2395
2396
2397
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


2398
## @private
2399
class CustomBondGenerator(object):
2400
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2401

2402
2403
    def __init__(self, forcefield):
        self.ff = forcefield
2404
2405
2406
2407
2408
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2409

2410
2411
    @staticmethod
    def parseElement(element, ff):
2412
2413
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
2414
2415
2416
2417
2418
2419
        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'):
2420
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2421
            if None not in types:
2422
2423
2424
                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
2425

2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
    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


2446
## @private
2447
class CustomAngleGenerator(object):
2448
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2449

2450
2451
    def __init__(self, forcefield):
        self.ff = forcefield
2452
2453
2454
2455
2456
2457
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2458

2459
2460
    @staticmethod
    def parseElement(element, ff):
2461
2462
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
2463
2464
2465
2466
2467
2468
        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'):
2469
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2470
            if None not in types:
2471
2472
2473
2474
                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
2475

2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
    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


2498
## @private
2499
class CustomTorsion(object):
2500
2501
2502
2503
2504
2505
2506
2507
2508
    """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

2509
## @private
2510
class CustomTorsionGenerator(object):
2511
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2512

2513
2514
    def __init__(self, forcefield):
        self.ff = forcefield
2515
2516
2517
2518
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
2519

2520
2521
    @staticmethod
    def parseElement(element, ff):
2522
2523
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
2524
2525
2526
2527
2528
2529
        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'):
2530
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2531
            if None not in types:
2532
2533
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
2534
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2535
            if None not in types:
2536
                generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
Justin MacCallum's avatar
Justin MacCallum committed
2537

2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
    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]]]
2570
            match = None
2571
2572
2573
2574
2575
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
2576
2577
2578
2579
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
2580
2581
2582
                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:
2583
                            if hasWildcard:
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
                                # Workaround to be more consistent with AMBER.  It uses wildcards to define most of its
                                # impropers, which leaves the ordering ambiguous.  It then follows some bizarre rules
                                # to pick the order.
                                a1 = torsion[t2[1]]
                                a2 = torsion[t3[1]]
                                e1 = data.atoms[a1].element
                                e2 = data.atoms[a2].element
                                if e1 == e2 and a1 > a2:
                                    (a1, a2) = (a2, a1)
                                elif e1 != elem.carbon and (e2 == elem.carbon or e1.mass < e2.mass):
                                    (a1, a2) = (a2, a1)
2595
                                match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
2596
2597
                            else:
                                # There are no wildcards, so the order is unambiguous.
2598
                                match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
2599
                            break
2600
2601
2602
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
2603
2604
2605
2606

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


2607
## @private
2608
class CustomNonbondedGenerator(object):
2609
2610
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

2611
2612
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
2613
2614
2615
2616
2617
2618
2619
2620
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2621
2622
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
2623
2624
2625
2626
        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'])
2627
2628
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomNonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomNonbondedForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.CustomNonbondedForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
            raise ValueError('Illegal nonbonded method for CustomNonbondedForce')
        force = mm.CustomNonbondedForce(self.energy)
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perParticleParams:
            force.addPerParticleParameter(param)
        for (name, type, values, params) in self.functions:
            if type == 'Continuous1D':
                force.addTabulatedFunction(name, mm.Continuous1DFunction(values, params['min'], params['max']))
            elif type == 'Continuous2D':
                force.addTabulatedFunction(name, mm.Continuous2DFunction(params['xsize'], params['ysize'], values, params['xmin'], params['xmax'], params['ymin'], params['ymax']))
            elif type == 'Continuous3D':
                force.addTabulatedFunction(name, mm.Continuous2DFunction(params['xsize'], params['ysize'], params['zsize'], values, params['xmin'], params['xmax'], params['ymin'], params['ymax'], params['zmin'], params['zmax']))
            elif type == 'Discrete1D':
                force.addTabulatedFunction(name, mm.Discrete1DFunction(values))
            elif type == 'Discrete2D':
                force.addTabulatedFunction(name, mm.Discrete2DFunction(params['xsize'], params['ysize'], values))
            elif type == 'Discrete3D':
                force.addTabulatedFunction(name, mm.Discrete2DFunction(params['xsize'], params['ysize'], params['zsize'], values))
        for atom in data.atoms:
2655
2656
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2657
2658
2659
2660
2661
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

    def postprocessSystem(self, sys, data, args):
2662
        # Create the exclusions.
2663

2664
        bondIndices = _findBondsForExclusions(data, sys)
2665
2666
2667
2668
2669
2670
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomNonbondedForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


2671
## @private
2672
class CustomGBGenerator(object):
2673
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2674

2675
2676
    def __init__(self, forcefield):
        self.ff = forcefield
2677
2678
2679
2680
2681
2682
2683
2684
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2685
2686
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
2687
2688
2689
2690
        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'])
2691
2692
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2693
2694
2695
2696
2697
2698
2699
2700
2701
        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()]
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
            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
2713

2714
2715
2716
2717
2718
    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
2719
            raise ValueError('Illegal nonbonded method for CustomGBForce')
2720
2721
2722
2723
2724
2725
2726
2727
2728
        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])
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
        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))
2742
        for atom in data.atoms:
2743
2744
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2745
2746
2747
2748
2749
2750
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

2751
2752

## @private
2753
class CustomManyParticleGenerator(object):
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
    """A CustomManyParticleGenerator constructs a CustomManyParticleForce."""

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

2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
    @staticmethod
    def parseElement(element, ff):
        permutationMap = {"SinglePermutation" : mm.CustomManyParticleForce.SinglePermutation,
                          "UniqueCentralParticle" : mm.CustomManyParticleForce.UniqueCentralParticle}
        generator = CustomManyParticleGenerator(ff, int(element.attrib['particlesPerSet']), element.attrib['energy'], permutationMap[element.attrib['permutationMode']], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerParticleParameter'):
            generator.perParticleParams.append(param.attrib['name'])
        for param in element.findall('TypeFilter'):
            generator.typeFilters.append((int(param.attrib['index']), [int(x) for x in param.attrib['types'].split(',')]))
2779
2780
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809

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

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

2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
        bondIndices = []
        for bond in data.bonds:
            bondIndices.append((bond.atom1, bond.atom2))

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

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

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

2835
2836
2837
2838
2839
2840
2841
2842
2843
        for atom1, atom2 in bondIndices:
            for child1 in data.excludeAtomWith[atom1]:
                bondIndices.append((child1, atom2))
                for child2 in data.excludeAtomWith[atom2]:
                    bondIndices.append((child1, child2))
            for child2 in data.excludeAtomWith[atom2]:
                bondIndices.append((atom1, child2))

        # Create the exclusions.
2844

2845
2846
2847
2848
2849
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomManyParticleForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomManyParticleForce"] = CustomManyParticleGenerator.parseElement

Peter Eastman's avatar
Peter Eastman committed
2850
def getAtomPrint(data, atomIndex):
2851

Peter Eastman's avatar
Peter Eastman committed
2852
2853
2854
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
2855
    else:
Peter Eastman's avatar
Peter Eastman committed
2856
        returnString = "NA"
2857
2858
2859
2860
2861

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
2862
def countConstraint(data):
2863

Peter Eastman's avatar
Peter Eastman committed
2864
    bondCount = 0
2865
2866
2867
2868
2869
2870
2871
    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
2872
        if (isConstrained):
2873
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
2874

2875
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
2876

2877
## @private
2878
class AmoebaBondGenerator(object):
2879
2880
2881

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

2882
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
2883
2884

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

2886
2887
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
2888
2889
2890
2891
2892
2893
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2894

2895
2896
2897
2898
2899
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

2903
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
2904
2905
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
2906
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2907
            if None not in types:
2908
2909
2910
2911
2912
                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:
2913
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
2914
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
2915
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
2916
2917
                raise ValueError(outputString)

2918
2919
    #=============================================================================================

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

Mark Friedrichs's avatar
Mark Friedrichs committed
2922
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2923

Peter Eastman's avatar
Peter Eastman committed
2924
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2925
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
2926
        if len(existing) == 0:
2927
            force = mm.AmoebaBondForce()
2928
2929
2930
2931
            sys.addForce(force)
        else:
            force = existing[0]

2932
2933
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943

        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:
2944
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
2945
2946
2947
2948
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

2949
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
2950
2951
2952
2953

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

Peter Eastman's avatar
Peter Eastman committed
2955
def addAngleConstraint(angle, idealAngle, data, sys):
2956
2957

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

Peter Eastman's avatar
Peter Eastman committed
2959
2960
    bond1 = None
    bond2 = None
2961
2962
2963
2964
2965
2966
2967
    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
2968

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

2971
2972
2973
2974
        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
2975
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
2976
                data.addConstraint(sys, angle[0], angle[2], length)
2977
2978
2979
                return

#=============================================================================================
2980
## @private
2981
class AmoebaAngleGenerator(object):
2982
2983

    #=============================================================================================
2984
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
2985
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
2986

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

Peter Eastman's avatar
Peter Eastman committed
2989
2990
2991
2992
2993
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
2994

Peter Eastman's avatar
Peter Eastman committed
2995
2996
2997
        self.types1 = []
        self.types2 = []
        self.types3 = []
2998

Peter Eastman's avatar
Peter Eastman committed
2999
3000
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3001

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

    @staticmethod
    def parseElement(element, forceField):

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

3010
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
3011
3012
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
3013
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
3014
            if None not in types:
3015
3016
3017
3018
3019
3020

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
3021
                angleList.append(float(angle.attrib['angle1']))
3022
3023

                try:
Peter Eastman's avatar
Peter Eastman committed
3024
                    angleList.append(float(angle.attrib['angle2']))
3025
                    try:
Peter Eastman's avatar
Peter Eastman committed
3026
                        angleList.append(float(angle.attrib['angle3']))
3027
3028
3029
3030
3031
3032
3033
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
3034
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
3035
3036
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
3037
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
3038
3039
                raise ValueError(outputString)

3040
3041
3042
3043
    #=============================================================================================
    # 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
3044

Peter Eastman's avatar
Peter Eastman committed
3045
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3046
3047
3048
3049
3050
3051
        pass

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

Peter Eastman's avatar
Peter Eastman committed
3053
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3054
3055
3056
3057

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3058
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
3059
3060

        if len(existing) == 0:
3061
            force = mm.AmoebaAngleForce()
3062
3063
3064
3065
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3066
        # set scalars
3067

3068
3069
3070
3071
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
3072
3073

        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
3074
3075
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3076

Peter Eastman's avatar
Peter Eastman committed
3077
3078
3079
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3080
3081
3082
3083
3084
3085
3086
            for i in range(len(self.types1)):
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
                if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                    if isConstrained and self.k[i] != 0.0:
                        angleDict['idealAngle'] = self.angle[i][0]
3087
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3088
                    elif self.k[i] != 0:
Peter Eastman's avatar
Peter Eastman committed
3089
3090
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
3091
3092
3093
3094
3095
3096
                            # 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
3097
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
3098
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3099
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
3100
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3101
                            if (numberOfHydrogens < lenAngle):
3102
3103
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
3104
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
3105
                                raise ValueError(outputString)
3106
3107
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
3108

3109
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
3110
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
3111
3112
3113
3114
3115
3116
                    break

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

Peter Eastman's avatar
Peter Eastman committed
3118
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3119
3120
3121
3122

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3123
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
3124
3125

        if len(existing) == 0:
3126
            force = mm.AmoebaInPlaneAngleForce()
3127
3128
3129
3130
3131
3132
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

3133
3134
3135
3136
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
3137
3138

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

Peter Eastman's avatar
Peter Eastman committed
3140
3141
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3142

Peter Eastman's avatar
Peter Eastman committed
3143
3144
3145
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3146
3147
3148
3149
3150
3151
3152
3153
3154

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

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

                if (type1 in types1 and type2 in types2 and type3 in types3) or (type1 in types3 and type2 in types2 and type3 in types1):
                    angleDict['idealAngle'] = self.angle[i][0]
Peter Eastman's avatar
Peter Eastman committed
3155
                    if (isConstrained and self.k[i] != 0.0):
3156
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3157
3158
3159
3160
                    else:
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

3161
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
3162
3163
3164

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
3165
3166
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
3167
3168
#=============================================================================================

3169
## @private
3170
class AmoebaOutOfPlaneBendGenerator(object):
3171
3172
3173
3174

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

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

3176
3177
3178
3179
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3180
3181
3182
3183
3184
3185
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3186

Peter Eastman's avatar
Peter Eastman committed
3187
3188
3189
3190
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3191

Peter Eastman's avatar
Peter Eastman committed
3192
        self.ks = []
3193
3194
3195
3196
3197

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

3199
3200
3201
3202
3203
3204
3205
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
3218
3219
3220
3221
3222
3223
    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
3224

3225
3226
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
3227
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
3228
3229
3230
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
3231
                                                   float(element.attrib['opbend-sextic']))
3232
3233
3234
3235

        forceField._forces.append(generator)

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

Peter Eastman's avatar
Peter Eastman committed
3239
3240
3241
3242
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
3243

Peter Eastman's avatar
Peter Eastman committed
3244
                generator.ks.append(float(angle.attrib['k']))
3245
3246

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3247
3248
                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
3249
3250
                raise ValueError(outputString)

3251
3252
3253
3254
3255
3256
    #=============================================================================================
    # 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
3257

Peter Eastman's avatar
Peter Eastman committed
3258
    def getMiddleAtom(self, angle, data):
3259
3260
3261

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
3262
        middleAtom = -1
Justin MacCallum's avatar
Justin MacCallum committed
3263
        for atomIndex in angle:
Peter Eastman's avatar
Peter Eastman committed
3264
            isMiddle = 0
3265
3266
3267
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3268
                if (atom1 != atomIndex):
3269
3270
3271
                    partner = atom1
                else:
                    partner = atom2
Justin MacCallum's avatar
Justin MacCallum committed
3272
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]):
3273
3274
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
3275
            if (isMiddle == 2):
3276
3277
3278
3279
3280
                return atomIndex
        return -1

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

Peter Eastman's avatar
Peter Eastman committed
3281
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3282
3283
3284
3285
3286
3287
3288
3289
3290
3291
3292
3293
3294

        # 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
3295
3296
3297
3298
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
3299
3300
3301
3302

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

Peter Eastman's avatar
Peter Eastman committed
3303
        skipAtoms = dict()
3304
3305
3306
3307

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

Peter Eastman's avatar
Peter Eastman committed
3308
3309
        inPlaneAngles = []
        nonInPlaneAngles = []
3310
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
3311
        idealAngles = []*len(data.angles)
3312
3313
3314

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

Peter Eastman's avatar
Peter Eastman committed
3315
3316
3317
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
3318
3319
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
3320
                middleType = -1
3321
3322
                middleCovalency = -1

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

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

Peter Eastman's avatar
Peter Eastman committed
3331
3332
3333
3334
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
3335
3336
3337
3338

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3339
                    if (atom1 != middleAtom):
3340
3341
3342
3343
3344
3345
3346
3347
                        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
3348
3349
3350
3351
3352
                        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
3353

Peter Eastman's avatar
Peter Eastman committed
3354
                if (len(partners) == 3):
3355

Peter Eastman's avatar
Peter Eastman committed
3356
3357
3358
                    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])
3359

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

3362
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
3363
3364
3365
3366
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3367
3368
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
3369
3370
3371
3372
3373
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
3374
3375
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
3376
                    angleDict['isConstrained'] = 0
3377

Peter Eastman's avatar
Peter Eastman committed
3378
3379
3380
3381
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3382
3383

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3384
3385
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3386

Peter Eastman's avatar
Peter Eastman committed
3387
                    inPlaneAngles.append(angleDict)
3388
3389

                else:
Peter Eastman's avatar
Peter Eastman committed
3390
3391
3392
3393
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
3394
            else:
Peter Eastman's avatar
Peter Eastman committed
3395
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
3396

Peter Eastman's avatar
Peter Eastman committed
3397
                    partnerSet = skipAtoms[middleAtom]
Justin MacCallum's avatar
Justin MacCallum committed
3398

Peter Eastman's avatar
Peter Eastman committed
3399
                    angleDict = {}
3400

Peter Eastman's avatar
Peter Eastman committed
3401
3402
3403
3404
3405
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
3406

Peter Eastman's avatar
Peter Eastman committed
3407
                    angleDict['isConstrained'] = isConstrained
3408

Peter Eastman's avatar
Peter Eastman committed
3409
3410
3411
3412
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3413
3414

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3415
3416
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3417

Peter Eastman's avatar
Peter Eastman committed
3418
                    inPlaneAngles.append(angleDict)
3419
3420

                else:
Peter Eastman's avatar
Peter Eastman committed
3421
3422
                    angleDict = {}
                    angleDict['angle'] = angle
3423
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
3424
                    nonInPlaneAngles.append(angleDict)
3425

3426
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
3427
3428

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3429
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
3430
3431
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
3432
3433

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3434
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
3435
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
3436
                    nonInPlaneAngles.append(angleDict)
3437
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
3438
3439
3440
3441
3442

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

3443
## @private
3444
class AmoebaTorsionGenerator(object):
3445
3446
3447
3448
3449
3450
3451

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
3452
        self.torsionUnit = torsionUnit
3453

Peter Eastman's avatar
Peter Eastman committed
3454
3455
3456
3457
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3458

Peter Eastman's avatar
Peter Eastman committed
3459
3460
3461
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
3462

3463
3464
3465
3466
3467
3468
3469
3470
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3472
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
3473
3474
3475
3476
3477
3478
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
3479
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
3480
            if None not in types:
3481
3482
3483
3484
3485
3486
3487

                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
3488
3489
                    tInfo = []
                    suffix = str(ii)
3490
3491
3492
3493
3494
3495
                    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
3496
3497
3498
3499
3500
3501
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
3502
3503
3504

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
3505
3506
3507
3508
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
3509
3510
                raise ValueError(outputString)

3511
3512
3513
3514
3515
    #=============================================================================================

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3516
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
3517
        if len(existing) == 0:
3518
            force = mm.PeriodicTorsionForce()
3519
3520
3521
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3522

3523
3524
3525
3526
3527
3528
3529
3530
3531
3532
3533
3534
3535
3536
3537
3538
        for torsion in data.propers:

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

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

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

                # match types in forward or reverse direction

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

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

3552
## @private
3553
class AmoebaPiTorsionGenerator(object):
3554
3555
3556
3557
3558
3559

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

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

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

3561
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
3562
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
3563
3564
3565
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3566

3567
3568
3569
3570
3571
3572
3573
3574
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
3575
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
3576
3577
3578
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
3579
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
3580
            if None not in types:
3581
3582
3583
3584
3585
3586
                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
3587
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3588
3589
                raise ValueError(outputString)

3590
3591
3592
3593
3594
3595
3596
3597
3598
3599
3600
3601
3602
3603
3604
3605
    #=============================================================================================

    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
3606

3607
3608
            atom1 = bond.atom1
            atom2 = bond.atom2
Justin MacCallum's avatar
Justin MacCallum committed
3609

3610
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
3611
3612
3613
3614
3615
3616
3617
3618
3619
3620
3621

                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
3622
3623
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
3624
3625
3626
3627
3628
3629
3630
3631
3632
3633
3634
3635

                       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
3636
                           if (bondedAtom1 != atom1):
3637
3638
3639
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
3640
3641
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3642
                                   piTorsionAtom1 = b1
3643
3644
3645
3646
3647
3648
                               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
3649
                           if (bondedAtom1 != atom2):
3650
3651
3652
3653
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3654
3655
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3656
                                   piTorsionAtom5 = b1
3657
3658
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
3659

Peter Eastman's avatar
Peter Eastman committed
3660
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
3661
3662
3663
3664
3665

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

3666
## @private
3667
class AmoebaTorsionTorsionGenerator(object):
3668
3669
3670
3671
3672

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

Peter Eastman's avatar
Peter Eastman committed
3673
    def __init__(self):
3674

Peter Eastman's avatar
Peter Eastman committed
3675
3676
3677
3678
3679
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
3680

Peter Eastman's avatar
Peter Eastman committed
3681
        self.gridIndex = []
3682

Peter Eastman's avatar
Peter Eastman committed
3683
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
3684

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

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
3690
        generator = AmoebaTorsionTorsionGenerator()
3691
3692
3693
3694
3695
3696
3697
        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'):
3698
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
3699
            if None not in types:
3700
3701
3702
3703
3704
3705
3706

                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
3707
3708
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
3709
3710
3711
3712
3713
3714
3715
3716
3717
                    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
3718
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
3719
3720
                raise ValueError(outputString)

3721
3722
3723
3724
3725
3726
3727
3728
3729
3730
        # 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
3731
3732
3733
3734
3735
3736
        #     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
3737
3738

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

Peter Eastman's avatar
Peter Eastman committed
3742
3743
3744
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
3745

Peter Eastman's avatar
Peter Eastman committed
3746
3747
            grid = []
            gridCol = []
3748
3749
3750
3751
3752

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
3753
3754
3755
3756
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
3757
                if 'fx' in gridEntry.attrib:
3758
3759
3760
                    gridRow.append(float(gridEntry.attrib['fx']))
                    gridRow.append(float(gridEntry.attrib['fy']))
                    gridRow.append(float(gridEntry.attrib['fxy']))
Peter Eastman's avatar
Peter Eastman committed
3761
                gridCol.append(gridRow)
3762
3763

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
3764
3765
3766
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
3767
3768
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
3769

Peter Eastman's avatar
Peter Eastman committed
3770
3771
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
3772
            else:
Peter Eastman's avatar
Peter Eastman committed
3773
3774
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
3775
3776
3777
3778
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
3779
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
3780
3781
3782
3783
3784
3785
3786
3787

        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
3788
        if (len(data.atomBonds[atomC]) == 4):
3789
3790
3791
3792
3793
            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
3794
3795
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
3796
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
3797
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
3798
3799
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
3800
3801
                if (hit > -1):
                    if (atomE == -1):
3802
3803
3804
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
3805

3806
3807
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
3808
3809
            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
3810
                raise ValueError(outputString)
3811
3812
3813
3814
3815

            # 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
3816
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
3817
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
3818
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
3819
                chiralAtomIndex = atomF
3820

Peter Eastman's avatar
Peter Eastman committed
3821
3822
3823
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
3824
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
3825
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
3826
                chiralAtomIndex = massF
3827
3828
3829
3830

        return chiralAtomIndex

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

3832
3833
3834
3835
3836
3837
3838
3839
3840
3841
3842
3843
3844
    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
3845
3846
            # search for bitorsions; based on TINKER subroutine bitors()

3847
3848
3849
3850
3851
3852
3853
            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
3854
                if (bondedAtom1 != ib):
3855
3856
3857
3858
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3859
                if (ia != ic and ia != id):
3860
3861
3862
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
3863
                        if (bondedAtom1 != id):
3864
3865
3866
3867
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3868
                        if (ie != ic and ie != ib and ie != ia):
3869
3870
3871
3872
3873
3874
3875
3876
3877
3878
3879
3880
3881
3882
3883
3884
3885
3886
3887
3888

                            # 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
3889
3890
3891
                                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])
3892
3893
3894

                                # match in reverse order

3895
                                elif (type5 in types1 and type4 in types2 and type3 in types3 and type2 in types4 and type1 in types5):
Peter Eastman's avatar
Peter Eastman committed
3896
3897
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
3898
3899
3900
3901

        # set grids

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

3904
3905
3906
3907
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

3908
## @private
3909
class AmoebaStretchBendGenerator(object):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3910
3911

    #=============================================================================================
3912
3913
3914
3915
3916
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
3917
3918
3919
        self.types1 = []
        self.types2 = []
        self.types3 = []
3920

Peter Eastman's avatar
Peter Eastman committed
3921
3922
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
3923

3924
3925
3926
3927
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
3928
        generator = AmoebaStretchBendGenerator()
3929
3930
3931
3932
3933
3934
3935
        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'):
3936
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
3937
            if None not in types:
3938
3939
3940
3941
3942
3943
3944
3945
3946
3947
3948
3949

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

3953
3954
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
3955
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
3956
3957
    # 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
3958
3959
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
3960
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
3961
3962
3963

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

Peter Eastman's avatar
Peter Eastman committed
3964
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3965
3966
3967
3968
3969
3970
3971
3972
        pass

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

    # Note: request for constrained bonds is ignored.

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

3973
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3974
3975
3976
3977
3978
3979
3980
3981
3982
3983
3984
3985

        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
3986
            if ('isConstrained' in angleDict):
3987
3988
3989
3990
3991
3992
3993
3994
                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
3995
            radian = 57.2957795130
3996
3997
3998
3999
4000
4001
4002
4003
4004
4005
            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
4006
                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
4007
4008
4009
4010
4011
4012
4013
4014
4015
4016
4017
4018
4019
4020
4021
                    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
4022

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4025
                    if ('idealAngle' not in angleDict):
4026

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4027
4028
4029
4030
4031
                       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
4032
                       raise ValueError(outputString)
4033

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4034
4035
4036
4037
4038
4039
4040
                    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
4041
                       raise ValueError(outputString)
4042

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4046
                    break
4047
4048
4049
4050
4051

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

4052
## @private
4053
class AmoebaVdwGenerator(object):
4054
4055

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

4057
4058
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4061
        self.type = type
4062

Peter Eastman's avatar
Peter Eastman committed
4063
4064
4065
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
4066

Peter Eastman's avatar
Peter Eastman committed
4067
        self.epsilonrule = epsilonrule
4068

Peter Eastman's avatar
Peter Eastman committed
4069
4070
4071
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
4072
4073
4074
4075
4076
4077
4078

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

    @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
4079
4080
4081
        #   <Vdw class="1" sigma="0.371" epsilon="0.46024" reduction="1.0" />
        #   <Vdw class="2" sigma="0.382" epsilon="0.422584" reduction="1.0" />

4082
4083
4084
        existing = [f for f in forceField._forces if isinstance(f, AmoebaVdwGenerator)]
        if len(existing) == 0:
            generator = AmoebaVdwGenerator(element.attrib['type'], element.attrib['radiusrule'], element.attrib['radiustype'], element.attrib['radiussize'], element.attrib['epsilonrule'],
Justin MacCallum's avatar
Justin MacCallum committed
4085
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
4086
4087
4088
4089
4090
4091
4092
4093
4094
4095
4096
4097
            forceField.registerGenerator(generator)
            generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaVdwForce', 'Vdw', ('sigma', 'epsilon', 'reduction'))
        else:
            # Multiple <AmoebaVdwForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
            if abs(generator.vdw13Scale - float(element.attrib['vdw-13-scale'])) > NonbondedGenerator.SCALETOL or \
                    abs(generator.vdw14Scale - float(element.attrib['vdw-14-scale'])) > NonbondedGenerator.SCALETOL or \
                    abs(generator.vdw15Scale - float(element.attrib['vdw-15-scale'])) > NonbondedGenerator.SCALETOL:
                raise ValueError('Found multiple AmoebaVdwForce tags with different scale factors')
            if generator.radiusrule != element.attrib['radiusrule'] or generator.epsilonrule != element.attrib['epsilonrule'] or \
                    generator.radiustype != element.attrib['radiustype'] or generator.radiussize != element.attrib['radiussize']:
                raise ValueError('Found multiple AmoebaVdwForce tags with different combining rules')
4098
        generator.params.parseDefinitions(element)
4099
4100
4101
4102
4103
        two_six = 1.122462048309372

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

    @staticmethod
4104
    def getBondedParticleSets(sys, data):
4105

4106
4107
4108
4109
4110
4111
        bondedParticleSets = [set() for i in range(len(data.atoms))]
        bondIndices = _findBondsForExclusions(data, sys)
        for atom1, atom2 in bondIndices:
            bondedParticleSets[atom1].add(atom2)
            bondedParticleSets[atom2].add(atom1)
        return bondedParticleSets
Justin MacCallum's avatar
Justin MacCallum committed
4112

4113
4114
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
4117
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
4118
4119
        epsilonMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'HARMONIC':1, 'HHG':1}

4120
4121
        force = mm.AmoebaVdwForce()
        sys.addForce(force)
4122

4123
        # sigma and epsilon combining rules
4124

4125
4126
4127
4128
        if ('sigmaCombiningRule' in args):
            sigmaRule = args['sigmaCombiningRule'].upper()
            if (sigmaRule.upper() in sigmaMap):
                force.setSigmaCombiningRule(sigmaRule.upper())
4129
            else:
4130
4131
4132
4133
                stringList = ' ' . join(str(x) for x in sigmaMap.keys())
                raise ValueError( "AmoebaVdwGenerator: sigma combining rule %s not recognized; valid values are %s; using default." % (sigmaRule, stringList) )
        else:
            force.setSigmaCombiningRule(self.radiusrule)
4134

4135
4136
4137
4138
        if ('epsilonCombiningRule' in args):
            epsilonRule = args['epsilonCombiningRule'].upper()
            if (epsilonRule.upper() in epsilonMap):
                force.setEpsilonCombiningRule(epsilonRule.upper())
4139
            else:
4140
4141
4142
4143
                stringList = ' ' . join(str(x) for x in epsilonMap.keys())
                raise ValueError( "AmoebaVdwGenerator: epsilon combining rule %s not recognized; valid values are %s; using default." % (epsilonRule, stringList) )
        else:
            force.setEpsilonCombiningRule(self.epsilonrule)
Justin MacCallum's avatar
Justin MacCallum committed
4144

4145
        # cutoff
4146

4147
4148
4149
4150
        if ('vdwCutoff' in args):
            force.setCutoff(args['vdwCutoff'])
        else:
            force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4151

4152
        # dispersion correction
4153

4154
4155
        if ('useDispersionCorrection' in args):
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
4156

4157
4158
        if (nonbondedMethod == PME):
            force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
4159
4160
4161

        # add particles to force

4162
4163
4164
4165
4166
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
4167
        for (i, atom) in enumerate(data.atoms):
4168
4169
4170
4171
4172
4173
4174
4175
4176
4177
            values = self.params.getAtomParameters(atom, data)
            # ivIndex = index of bonded partner for hydrogens; otherwise ivIndex = particle index

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

4179
            force.addParticle(ivIndex, values[0]*sigmaScale, values[1], values[2])
4180
4181
4182
4183
4184
4185
4186

        # 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

4187
        bondedParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4188
4189

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

4191
4192
4193
4194
4195
4196
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
4197
            if (self.vdw13Scale == 0.0):
4198
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
4199
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
4200
4201
4202
4203
4204

            # self

            exclusionSet.add(i)

4205
            force.setParticleExclusions(i, tuple(exclusionSet))
4206
4207
4208
4209
4210

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

4211
## @private
4212
class AmoebaMultipoleGenerator(object):
4213
4214
4215
4216

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

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

4218
4219
    #=============================================================================================

4220
    def __init__(self, forceField):
Peter Eastman's avatar
Peter Eastman committed
4221
4222
        self.forceField = forceField
        self.typeMap = {}
4223
4224
4225
4226
4227
4228

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
4229
    def setAxisType(kIndices):
4230
4231
4232

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
4233
4234
4235
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
4236
                else:
Peter Eastman's avatar
Peter Eastman committed
4237
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
4238

Peter Eastman's avatar
Peter Eastman committed
4239
4240
                if (kIndicesLen > 2):
                    kx = kIndices[2]
4241
                else:
Peter Eastman's avatar
Peter Eastman committed
4242
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
4243

Peter Eastman's avatar
Peter Eastman committed
4244
4245
                if (kIndicesLen > 1):
                    kz = kIndices[1]
4246
                else:
Peter Eastman's avatar
Peter Eastman committed
4247
                    kz = 0
4248

Peter Eastman's avatar
Peter Eastman committed
4249
4250
                while(len(kIndices) < 4):
                    kIndices.append(0)
4251
4252

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
4253
                if (kz == 0):
4254
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
4255
                if (kz != 0 and kx == 0):
4256
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
4257
                if (kz < 0 or kx < 0):
4258
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
4259
                if (kx < 0 and ky < 0):
4260
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
4261
                if (kz < 0 and kx < 0 and ky  < 0):
4262
4263
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
4264
4265
4266
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
4267
4268
4269
4270
4271
4272
4273
4274

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
4275
        #   <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"  >
4276
4277
4278
        # <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"  />

4279
4280
4281
4282
4283
4284
4285
        existing = [f for f in forceField._forces if isinstance(f, AmoebaMultipoleGenerator)]
        if len(existing) == 0:
            generator = AmoebaMultipoleGenerator(forceField)
            forceField.registerGenerator(generator)
        else:
            # Multiple <AmoebaMultipoleForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
4286
4287
4288
4289

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

        for atom in element.findall('Multipole'):
4290
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4291
            if None not in types:
4292
4293
4294
4295
4296
4297
4298
4299

                # 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
4300
4301
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
4302
                    except:
4303
4304
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
4305
                # set axis type based on k-Indices
4306

Peter Eastman's avatar
Peter Eastman committed
4307
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
4308
4309
4310

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
4313
4314
4315
4316
4317
4318
4319
4320
4321
4322
4323
4324
4325
                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']))
4326
4327

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4328
                    if (t not in generator.typeMap):
4329
4330
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
4331
4332
4333
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
4334
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
4335
4336
4337
4338
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
4339

4340
            else:
Peter Eastman's avatar
Peter Eastman committed
4341
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4342
4343
                raise ValueError(outputString)

4344
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
4345

4346
        for atom in element.findall('Polarize'):
4347
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4348
            if None not in types:
4349

Peter Eastman's avatar
Peter Eastman committed
4350
4351
4352
4353
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
4354
4355
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
4356
                    pdamp = pow(polarizability, 1.0/6.0)
4357

Peter Eastman's avatar
Peter Eastman committed
4358
4359
                pgrpMap = dict()
                for index in range(1, 7):
4360
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
4361
                    if (pgrp in atom.attrib):
4362
4363
4364
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4365
4366
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
4367
                        raise ValueError(outputString)
4368
4369
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
4370
4371
4372
4373
4374
4375
4376
                        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
4377
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
4378
4379
4380
4381
4382
                                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
4383
4384
                            raise ValueError(outputString)

4385
            else:
Peter Eastman's avatar
Peter Eastman committed
4386
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4387
4388
                raise ValueError(outputString)

4389
4390
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4391
    def setPolarGroups(self, data, bonded12ParticleSets, force):
4392
4393
4394
4395
4396

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
4397
4398
4399
4400
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
4401
4402
4403
4404
            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
4405
4406
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
4407
4408
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
4409

4410
4411
4412
4413
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
4414
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
4415
4416
                continue

Peter Eastman's avatar
Peter Eastman committed
4417
4418
            group = set()
            visited = set()
4419
4420
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4421
4422
4423
4424
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
4425
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
4426
4427
                if (nextAtom not in visited):
                   visited.add(nextAtom)
4428
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4429
4430
4431
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
4432
4433
4434

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
4435
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
4436
4437

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4438
4439
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
4440
4441
4442
4443
4444

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
4445
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
4446
4447
                continue

Peter Eastman's avatar
Peter Eastman committed
4448
            pgrp11 = set(atom.polarizationGroupSet[0])
4449
4450
4451
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4452
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
4453
4454
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4455
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
4456

4457
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4458
4459
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
4460
4461
4462
4463
4464

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
4465
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
4466
4467
                continue

Peter Eastman's avatar
Peter Eastman committed
4468
4469
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
4470
4471
4472
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4473
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
4474
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4475
            pgrp13 = pgrp13 - set(pgrp11)
4476
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4477
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
4478

4479
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4480
4481
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
4482
4483
4484
4485
4486

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
4487
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
4488
4489
                continue

Peter Eastman's avatar
Peter Eastman committed
4490
4491
4492
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
4493
4494
4495
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4496
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
4497
4498
4499

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4500
            pgrp14 = pgrp14 - set(pgrp11)
4501
4502

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

4505
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4506
4507
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
4508
4509
4510

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

Peter Eastman's avatar
Peter Eastman committed
4511
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4512
4513
4514
4515

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

4516
4517
4518
4519
4520
4521
4522
        force = mm.AmoebaMultipoleForce()
        sys.addForce(force)
        if (nonbondedMethod not in methodMap):
            raise ValueError( "AmoebaMultipoleForce: input cutoff method not available." )
        else:
            force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
4523

4524
4525
        if ('ewaldErrorTolerance' in args):
            force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))
4526

4527
4528
4529
4530
4531
4532
4533
4534
        if ('polarization' in args):
            polarizationType = args['polarization']
            if (polarizationType.lower() == 'direct'):
                force.setPolarizationType(mm.AmoebaMultipoleForce.Direct)
            elif (polarizationType.lower() == 'extrapolated'):
                force.setPolarizationType(mm.AmoebaMultipoleForce.Extrapolated)
            else:
                force.setPolarizationType(mm.AmoebaMultipoleForce.Mutual)
4535

4536
4537
        if ('aEwald' in args):
            force.setAEwald(float(args['aEwald']))
4538

4539
4540
        if ('pmeGridDimensions' in args):
            force.setPmeGridDimensions(args['pmeGridDimensions'])
4541

4542
4543
        if ('mutualInducedMaxIterations' in args):
            force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
4544

4545
4546
        if ('mutualInducedTargetEpsilon' in args):
            force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))
4547
4548

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

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

        # 1-2

4555
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4556
4557
4558
4559
4560

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4561
            bonded13Set = set()
4562
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4563
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4564
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
4565
4566
4567
4568

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4569
            selfSet = set()
4570
4571
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4572
4573
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
4574
4575
4576
4577
4578

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4579
4580
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4581
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4582
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
4583

4584
4585
4586
4587
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4588
            selfSet = set()
4589
4590
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4591
4592
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
4593
4594
4595
4596
4597

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4598
4599
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4600
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4601
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
4602
4603
4604
4605
4606
4607

            # 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
4608
            selfSet = set()
4609
4610
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4611
4612
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
4613
4614
4615
4616
4617

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

Peter Eastman's avatar
Peter Eastman committed
4618
4619
                multipoleList = self.typeMap[t]
                hit = 0
4620
4621
4622
4623
4624
4625
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4626
                    if (hit != 0):
4627
4628
                        break

Peter Eastman's avatar
Peter Eastman committed
4629
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4630
4631

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4632
4633
                    kx = kIndices[2]
                    ky = kIndices[3]
4634
4635
4636
4637

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

4639
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
4640
4641
4642
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4643
4644
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
4645
                       if (hit != 0):
4646
4647
4648
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
4649
4650
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
4651
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4652
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4653
4654
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4655
4656
4657
4658
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4659
4660
4661
4662
4663
4664
4665
4666
4667
4668
                                      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

4669
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4670
                                      hit = 1
4671
4672
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4673
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4674
4675
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4676
4677
4678
4679
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4680
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4681
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
4682

4683
4684
4685
4686
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4687
                    if (hit != 0):
4688
4689
                        break

Peter Eastman's avatar
Peter Eastman committed
4690
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4691
4692

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4693
4694
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
4695

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

4700
4701
4702
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
4703
4704
4705
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4706
4707
4708

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4709
                       if (hit != 0):
4710
4711
4712
                           break

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

Peter Eastman's avatar
Peter Eastman committed
4715
                       if (bondedAtomZType == kz):
4716
4717
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
4718
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4719
4720
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4721
4722
4723
4724
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4725
4726
4727
4728
4729
4730
4731
4732

                                      # 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

4733
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4734
                                      hit = 3
4735
4736
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
4737
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4738
4739
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4740
4741
4742
4743
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4744
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4745
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
4746

4747
4748
4749
4750
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4751
                    if (hit != 0):
4752
4753
                        break

Peter Eastman's avatar
Peter Eastman committed
4754
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4755
4756
4757
4758

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

Peter Eastman's avatar
Peter Eastman committed
4759
4760
4761
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4762
4763
4764

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4765
                        if (hit != 0):
4766
4767
4768
                            break

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

Peter Eastman's avatar
Peter Eastman committed
4771
4772
                        if (kx == 0 and kz == bondedAtomZType):
                            kz = bondedAtomZIndex
4773
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4774
                            hit = 5
4775
4776
4777
4778
4779

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4780
                    if (hit != 0):
4781
4782
                        break

Peter Eastman's avatar
Peter Eastman committed
4783
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4784
4785
4786

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
4787
4788
4789
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4790

Peter Eastman's avatar
Peter Eastman committed
4791
                    if (kz == 0):
4792
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4793
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
4794

4795
4796
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
4797
                if (hit != 0):
4798

Peter Eastman's avatar
Peter Eastman committed
4799
                    atom.multipoleDict = savedMultipoleDict
4800
                    atom.polarizationGroups = dict()
4801
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
4802
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
4803
                    if (atomIndex == newIndex):
4804
4805
4806
4807
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent12, tuple(bonded12ParticleSets[atomIndex]))
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent13, tuple(bonded13ParticleSets[atomIndex]))
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent14, tuple(bonded14ParticleSets[atomIndex]))
                        force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.Covalent15, tuple(bonded15ParticleSets[atomIndex]))
4808
                    else:
4809
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
4810
                else:
Peter Eastman's avatar
Peter Eastman committed
4811
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
4812
            else:
Peter Eastman's avatar
Peter Eastman committed
4813
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
4814
4815
4816

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
4817
        self.setPolarGroups(data, bonded12ParticleSets, force)
4818
4819
4820
4821
4822

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

4823
## @private
4824
class AmoebaWcaDispersionGenerator(object):
4825
4826

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

4828
4829
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4832
4833
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
4834
4835
4836
4837
4838
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
4839
        self.shctd = shctd
4840
4841
4842
4843
4844
4845
4846
4847
4848

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

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

Peter Eastman's avatar
Peter Eastman committed
4850
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
4851
4852
4853
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
Justin MacCallum's avatar
Justin MacCallum committed
4854
                                                  element.attrib['awater'],
4855
4856
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Justin MacCallum's avatar
Justin MacCallum committed
4857
                                                  element.attrib['shctd'])
4858
        forceField._forces.append(generator)
4859
4860
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaWcaDispersionForce', 'WcaDispersion', ('radius', 'epsilon'))
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
4861

4862
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
4863

Peter Eastman's avatar
Peter Eastman committed
4864
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4865
4866
4867
4868
4869
4870
4871
4872
4873
4874
4875
4876

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

Peter Eastman's avatar
Peter Eastman committed
4879
4880
4881
4882
4883
4884
4885
4886
        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  ))
4887

4888
4889
4890
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
4891
4892
4893
4894
4895

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

4896
## @private
4897
class AmoebaGeneralizedKirkwoodGenerator(object):
4898
4899

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

4901
4902
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
4903
4904
4905
4906
4907
4908
4909
4910
4911
4912
4913
    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
4914
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
4915
4916
4917
4918
4919
4920
4921
4922
4923
4924
4925
4926
4927
4928
4929
4930
4931
4932
4933
4934
4935
4936
4937
4938
4939
        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
4940
4941
4942

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

Peter Eastman's avatar
Peter Eastman committed
4943
    def getObcShct(self, data, atomIndex):
4944

Peter Eastman's avatar
Peter Eastman committed
4945
        atom = data.atoms[atomIndex]
4946
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4947
        shct = -1.0
4948
4949

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
4950

Peter Eastman's avatar
Peter Eastman committed
4951
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
4952
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
4953
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
4954
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
4955
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
4956
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
4957
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
4958
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
4959
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
4960
4961
4962
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
4963
        elif (atomicNumber == 16):              # S(16)
4964
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
4965
        elif (atomicNumber == 26):              # Fe(26)
4966
4967
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
4968
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4969
            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
4970
4971

        return shct
4972
4973
4974

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

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

Peter Eastman's avatar
Peter Eastman committed
4977
        atom = data.atoms[atomIndex]
4978
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4979
        radius = -1.0
4980

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

Peter Eastman's avatar
Peter Eastman committed
4983
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
4984

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

4988
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4989
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
4990

Peter Eastman's avatar
Peter Eastman committed
4991
            if (bondedAtomAtomicNumber == 7):
4992
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
4993
            if (bondedAtomAtomicNumber == 8):
4994
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
4995

Peter Eastman's avatar
Peter Eastman committed
4996
        elif (atomicNumber == 3):               # Li(3)
4997
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
4998
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
4999

5000
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
5001
            if (len(bondedAtomIndices) == 3):
5002
5003
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
5004
            elif (len(bondedAtomIndices) == 4):
5005
5006
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5007
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
5008
5009
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
5010
        elif (atomicNumber == 7):               # N(7)
5011
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
5012
        elif (atomicNumber == 8):               # O(8)
5013
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
5014
            if (len(bondedAtomIndices) == 2):
5015
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
5016
        elif (atomicNumber == 9):               # F(9)
5017
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
5018
        elif (atomicNumber == 10):
5019
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
5020
        elif (atomicNumber == 11):
5021
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
5022
        elif (atomicNumber == 12):
5023
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5024
        elif (atomicNumber == 14):
5025
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
5026
        elif (atomicNumber == 15):              # P(15)
5027
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
5028
        elif (atomicNumber == 16):              # S(16)
5029
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
5030
        elif (atomicNumber == 17):
5031
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
5032
        elif (atomicNumber == 18):
5033
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5034
        elif (atomicNumber == 19):
5035
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
5036
        elif (atomicNumber == 20):
5037
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
5038
        elif (atomicNumber == 35):
5039
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
5040
        elif (atomicNumber == 36):
5041
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
5042
        elif (atomicNumber == 37):
5043
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
5044
        elif (atomicNumber == 53):
5045
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
5046
        elif (atomicNumber == 54):
5047
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
5048
        elif (atomicNumber == 55):
5049
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
5050
        elif (atomicNumber == 56):
5051
5052
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
5053
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5054
5055
            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
5056

5057
5058
5059
5060
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
5063
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5064
        atom = data.atoms[atomIndex]
5065
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
5066
        if (atomicNumber in bondiMap):
5067
5068
            radius = bondiMap[atomicNumber]
        else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5069
5070
            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
5071

5072
5073
5074
5075
5076
5077
5078
5079
5080
5081
        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
5082

Peter Eastman's avatar
Peter Eastman committed
5083
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
Justin MacCallum's avatar
Justin MacCallum committed
5084
5085
                                                        element.attrib['includeCavityTerm'],
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor'])
5086
5087
5088
        forceField._forces.append(generator)

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

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

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

5095
5096
5097
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
5098
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5099
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
5100
        if (len(amoebaMultipoleForceList) > 0):
5101
5102
5103
5104
5105
            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
5106
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
5107
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
5108

5109
5110
5111
5112
5113
5114
5115
        # 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
5116

Peter Eastman's avatar
Peter Eastman committed
5117
5118
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
5119
            else:
Peter Eastman's avatar
Peter Eastman committed
5120
                force.setSolventDielectric(   float(self.solventDielectric))
5121

Peter Eastman's avatar
Peter Eastman committed
5122
5123
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
5124
            else:
Peter Eastman's avatar
Peter Eastman committed
5125
                force.setSoluteDielectric(    float(self.soluteDielectric))
5126

Peter Eastman's avatar
Peter Eastman committed
5127
5128
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
5129
            else:
Peter Eastman's avatar
Peter Eastman committed
5130
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
5131
5132
5133
5134
5135

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
5138
5139
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
5140
5141
5142

        # 1-2

5143
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5144
5145

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
5146
5147
5148
5149
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5150
            else:
Peter Eastman's avatar
Peter Eastman committed
5151
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5152
5153
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
5154
            force.addParticle(multipoleParameters[0], radius, shct)
5155
5156
5157
5158
5159

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

5160
## @private
5161
class AmoebaUreyBradleyGenerator(object):
5162
5163
5164
5165

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

5167
    def __init__(self):
5168

Peter Eastman's avatar
Peter Eastman committed
5169
5170
5171
        self.types1 = []
        self.types2 = []
        self.types3 = []
5172

Peter Eastman's avatar
Peter Eastman committed
5173
5174
        self.length = []
        self.k = []
5175
5176
5177
5178
5179
5180

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

    @staticmethod
    def parseElement(element, forceField):

5181
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
5182
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
5183

5184
        generator = AmoebaUreyBradleyGenerator()
5185
5186
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
5187
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
5188
            if None not in types:
5189
5190
5191
5192
5193
5194
5195
5196
5197
5198

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

5202
5203
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
5206
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5207
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
5208
5209

        if len(existing) == 0:
5210
            force = mm.HarmonicBondForce()
5211
5212
5213
5214
5215
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
Peter Eastman's avatar
Peter Eastman committed
5216
            if (isConstrained):
5217
5218
5219
5220
5221
5222
5223
5224
                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
5225
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
5226
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
5227
5228
5229
5230
5231
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
5232
5233
5234


## @private
5235
class DrudeGenerator(object):
peastman's avatar
peastman committed
5236
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
5237

5238
5239
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
5240
5241
5242
5243
5244
5245
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
5246
5247
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
5248
5249
5250
5251
        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'):
5252
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
5253
5254
5255
5256
5257
5258
5259
5260
5261
5262
            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
5263

peastman's avatar
peastman committed
5264
5265
5266
5267
    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
5268

peastman's avatar
peastman committed
5269
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
5270

peastman's avatar
peastman committed
5271
5272
5273
5274
5275
5276
5277
5278
5279
5280
5281
5282
5283
5284
5285
5286
        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
5287
5288
                force.addParticle(atom.index, p[0], p[1], p[2], p[3], values[4], values[5], values[6], values[7])
                data.excludeAtomWith[p[0]].append(atom.index)
peastman's avatar
peastman committed
5289
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
5290

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

peastman's avatar
peastman committed
5294
5295
5296
5297
5298
5299
5300
        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)
5301
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
5302
5303
5304
5305
5306
                # 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]
5307
5308
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
5309
5310
5311
5312
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

Justin MacCallum's avatar
Justin MacCallum committed
5313
parsers["DrudeForce"] = DrudeGenerator.parseElement
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
5314
5315

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