"vscode:/vscode.git/clone" did not exist on "0a93040192983292d6de80a99f149dc57ff1d61c"
forcefield.py 237 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
    """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 = []
1805
        self.ordering = 'default'
1806

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

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

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

1821
    def registerImproperTorsion(self, parameters, ordering='default'):
1822
1823
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
1824
1825
1826
1827
1828
1829
            if ordering == 'default':
                pass
            elif ordering == 'amber':
                torsion.ordering = ordering
            else:
                raise ValueError('Illegal ordering type for improper torsion %s' % torsion)
1830
1831
            self.improper.append(torsion)

1832
1833
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1834
1835
1836
1837
1838
1839
        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]
1840
        for torsion in element.findall('Proper'):
1841
            generator.registerProperTorsion(torsion.attrib)
1842
        for torsion in element.findall('Improper'):
1843
            generator.registerImproperTorsion(torsion.attrib, element.attrib['ordering'])
Justin MacCallum's avatar
Justin MacCallum committed
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
1873
1874
1875
1876
1877
1878
1879
    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]]]
1880
            match = None
1881
1882
1883
1884
1885
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
1886
                hasWildcard = (wildcard in (types1, types2, types3, types4))
1887
1888
1889
1890
                if (types1, types2, types3, types4).count(wildcard) == 2:
                    hasTwoWildcards = True
                else:
                    hasTwoWildcards = False
1891
1892
1893
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
1894
1895
1896
                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:
1897
                            # topology atom indexes
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1898
1899
                            a2 = torsion[t2[1]]
                            a3 = torsion[t3[1]]
1900
1901
1902
                            a4 = torsion[t4[1]]
                            # residue indexes
                            r2 = data.atoms[a2].residue.index
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1903
                            r3 = data.atoms[a3].residue.index
1904
1905
1906
                            r4 = data.atoms[a4].residue.index
                            # template atom indexes
                            ta2 = data.atomTemplateIndexes[data.atoms[a2]]
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1907
                            ta3 = data.atomTemplateIndexes[data.atoms[a3]]
1908
1909
                            ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                            # elements
1910
                            e2 = data.atoms[a2].element
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1911
                            e3 = data.atoms[a3].element
1912
                            e4 = data.atoms[a4].element
1913
                            if tordef.ordering == 'default':
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1914
                                    if t2[0] == t4[0] and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
1915
                                        (a2, a4) = (a4, a2)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1916
1917
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
1918
                                        ta2 = data.atomTemplateIndexes[data.atoms[a3]]
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1919
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1920
1921
                                    if t3[0] == t4[0] and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
                                        (a3, a4) = (a4, a3)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1922
1923
1924
1925
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1926
1927
                                    if t2[0] == t3[0] and (r2 > r3 or (r2 == r3 and ta2 > ta3)):
                                        (a2, a3) = (a3, a2)
1928
1929
1930
1931
1932
                                    elif hasTwoWildcards and (r2 > r3 or (r2 == r3 and ta2 > ta3)):
                                        (a2, a3) = (a3, a2)
                            elif tordef.ordering == 'amber':
                                if not hasWildcard:
                                    if t2[0] == t4[0] and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
1933
                                        (a2, a4) = (a4, a2)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1934
1935
1936
1937
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
1938
                                    if t3[0] == t4[0] and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1939
                                        (a3, a4) = (a4, a3)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1940
1941
1942
1943
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
1944
                                    if t2[0] == t3[0] and (r2 > r3 or (r2 == r3 and ta2 > ta3)):
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1945
                                        (a2, a3) = (a3, a2)
1946
1947
                                else:
                                    if e2 == e4 and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
1948
                                        (a2, a4) = (a4, a2)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1949
1950
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
1951
                                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1952
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
1953
                                    if e3 == e4 and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1954
                                        (a3, a4) = (a4, a3)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1955
1956
1957
1958
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
1959
                                    if r2 > r3 or (r2 == r3 and ta2 > ta3):
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
1960
1961
                                        (a2, a3) = (a3, a2)
                            match = (a2, a3, torsion[0], a4, tordef)
1962
                            break
1963
1964
1965
1966
1967
            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])
1968
1969
1970
1971

parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


1972
## @private
1973
class RBTorsion(object):
1974
1975
1976
1977
1978
1979
1980
1981
1982
    """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

1983
## @private
1984
class RBTorsionGenerator(object):
1985
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1986

1987
1988
    def __init__(self, forcefield):
        self.ff = forcefield
1989
1990
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
1991

1992
1993
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1994
1995
1996
1997
1998
1999
        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]
2000
        for torsion in element.findall('Proper'):
2001
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2002
            if None not in types:
2003
2004
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
2005
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2006
            if None not in types:
2007
                generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
Justin MacCallum's avatar
Justin MacCallum committed
2008

2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
    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]]]
2042
            match = None
2043
2044
2045
2046
2047
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
2048
2049
2050
2051
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
2052
2053
2054
                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:
2055
                            if hasWildcard:
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
                                # 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)
2067
                                match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
2068
2069
                            else:
                                # There are no wildcards, so the order is unambiguous.
2070
                                match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
2071
                            break
2072
2073
2074
            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])
2075
2076
2077
2078

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


2079
## @private
2080
class CMAPTorsion(object):
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
    """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

2091
## @private
2092
class CMAPTorsionGenerator(object):
2093
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2094

2095
2096
    def __init__(self, forcefield):
        self.ff = forcefield
2097
2098
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
2099

2100
2101
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2102
2103
2104
2105
2106
2107
        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]
2108
2109
2110
2111
2112
2113
2114
        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'):
2115
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
2116
            if None not in types:
2117
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
2118

2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
    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
2129

2130
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
2131

2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
        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


2175
## @private
2176
class NonbondedGenerator(object):
2177
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2178

2179
2180
    SCALETOL = 1e-5

2181
2182
    def __init__(self, forcefield, coulomb14scale, lj14scale):
        self.ff = forcefield
2183
2184
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
2185
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
2186

2187
    def registerAtom(self, parameters):
2188
        self.params.registerAtom(parameters)
2189

2190
2191
2192
2193
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
        if len(existing) == 0:
2194
2195
            generator = NonbondedGenerator(ff, float(element.attrib['coulomb14scale']), float(element.attrib['lj14scale']))
            ff.registerGenerator(generator)
2196
2197
2198
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2199
2200
            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
2201
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
2202
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2203

2204
2205
2206
2207
2208
2209
2210
    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
2211
            raise ValueError('Illegal nonbonded method for NonbondedForce')
2212
2213
        force = mm.NonbondedForce()
        for atom in data.atoms:
2214
2215
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
2216
2217
2218
2219
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
2220
2221
        if 'useDispersionCorrection' in args:
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
peastman's avatar
peastman committed
2222
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
2223

peastman's avatar
peastman committed
2224
    def postprocessSystem(self, sys, data, args):
2225
2226
        # Create the exceptions.

2227
        bondIndices = _findBondsForExclusions(data, sys)
peastman's avatar
peastman committed
2228
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
2229
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
2230
2231
2232

parsers["NonbondedForce"] = NonbondedGenerator.parseElement

2233
2234

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

ChayaSt's avatar
ChayaSt committed
2238
    def __init__(self, forcefield, lj14scale):
ChayaSt's avatar
ChayaSt committed
2239
        self.ff = forcefield
2240
        self.nbfixTypes = {}
ChayaSt's avatar
ChayaSt committed
2241
        self.lj14scale = lj14scale
2242
        self.ljTypes = ForceField._AtomTypeParameters(forcefield, 'LennardJonesForce', 'Atom', ('sigma', 'epsilon'))
ChayaSt's avatar
ChayaSt committed
2243

ChayaSt's avatar
ChayaSt committed
2244
    def registerNBFIX(self, parameters):
ChayaSt's avatar
ChayaSt committed
2245
2246
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
2247
2248
            type1 = types[0][0]
            type2 = types[1][0]
ChayaSt's avatar
ChayaSt committed
2249
2250
            epsilon = _convertParameterToNumber(parameters['epsilon'])
            sigma = _convertParameterToNumber(parameters['sigma'])
2251
2252
            self.nbfixTypes[(type1, type2)] = [sigma, epsilon]
            self.nbfixTypes[(type2, type1)] = [sigma, epsilon]
2253

ChayaSt's avatar
ChayaSt committed
2254
    def registerLennardJones(self, parameters):
2255
        self.ljTypes.registerAtom(parameters)
ChayaSt's avatar
ChayaSt committed
2256
2257
2258

    @staticmethod
    def parseElement(element, ff):
ChayaSt's avatar
ChayaSt committed
2259
        existing = [f for f in ff._forces if isinstance(f, LennardJonesGenerator)]
ChayaSt's avatar
ChayaSt committed
2260
        if len(existing) == 0:
ChayaSt's avatar
ChayaSt committed
2261
            generator = LennardJonesGenerator(ff, float(element.attrib['lj14scale']))
ChayaSt's avatar
ChayaSt committed
2262
2263
            ff.registerGenerator(generator)
        else:
ChayaSt's avatar
ChayaSt committed
2264
            # Multiple <LennardJonesForce> tags were found, probably in different files
ChayaSt's avatar
ChayaSt committed
2265
            generator = existing[0]
2266
2267
            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
2268
2269
        for LJ in element.findall('Atom'):
            generator.registerLennardJones(LJ.attrib)
ChayaSt's avatar
ChayaSt committed
2270
        for Nbfix in element.findall('NBFixPair'):
ChayaSt's avatar
ChayaSt committed
2271
            generator.registerNBFIX(Nbfix.attrib)
ChayaSt's avatar
ChayaSt committed
2272

ChayaSt's avatar
ChayaSt committed
2273
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
ChayaSt's avatar
ChayaSt committed
2274
        # First derive the lookup tables
2275
2276
2277
2278
2279
2280

        nbfixTypeSet = set().union(*self.nbfixTypes)
        ljIndexList = [None]*len(data.atoms)
        numLjTypes = 0
        ljTypeList = []
        typeMap = {}
ChayaSt's avatar
ChayaSt committed
2281
        for i, atom in enumerate(data.atoms):
ChayaSt's avatar
ChayaSt committed
2282
            atype = data.atomType[atom]
2283
2284
2285
2286
            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]
2287
            else:
2288
2289
2290
2291
2292
2293
2294
                typeMap[values] = numLjTypes
                ljIndexList[i] = numLjTypes
                numLjTypes += 1
                ljTypeList.append(atype)
        reverseMap = [0]*len(typeMap)
        for typeValue in typeMap:
            reverseMap[typeMap[typeValue]] = typeValue
2295

ChayaSt's avatar
ChayaSt committed
2296
        # Now everything is assigned. Create the A- and B-coefficient arrays
2297
2298
        
        acoef = [0]*(numLjTypes*numLjTypes)
ChayaSt's avatar
ChayaSt committed
2299
        bcoef = acoef[:]
2300
2301
2302
2303
2304
2305
2306
2307
2308
        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
2309
                    continue
ChayaSt's avatar
ChayaSt committed
2310
                else:
2311
2312
2313
2314
2315
2316
2317
2318
2319
                    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
2320
        self.force.addPerParticleParameter('type')
2321
        if nonbondedMethod in [CutoffPeriodic, Ewald, PME]:
ChayaSt's avatar
ChayaSt committed
2322
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
ChayaSt's avatar
ChayaSt committed
2323
        elif nonbondedMethod is NoCutoff:
ChayaSt's avatar
ChayaSt committed
2324
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
ChayaSt's avatar
ChayaSt committed
2325
        elif nonbondedMethod is CutoffNonPeriodic:
ChayaSt's avatar
ChayaSt committed
2326
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic)
ChayaSt's avatar
ChayaSt committed
2327
        else:
2328
2329
            raise AssertionError('Unrecognized nonbonded method [%s]' % nonbondedMethod)

ChayaSt's avatar
ChayaSt committed
2330
        # Add the particles
2331
2332

        for i in ljIndexList:
ChayaSt's avatar
cleanup  
ChayaSt committed
2333
            self.force.addParticle((i,))
ChayaSt's avatar
ChayaSt committed
2334
2335
        self.force.setUseLongRangeCorrection(True)
        self.force.setCutoffDistance(nonbondedCutoff)
ChayaSt's avatar
ChayaSt committed
2336
2337
2338
        sys.addForce(self.force)

    def postprocessSystem(self, sys, data, args):
ChayaSt's avatar
ChayaSt committed
2339
        # Create the exceptions.
2340
        
2341
        bondIndices = _findBondsForExclusions(data, sys)
2342
2343
2344
        if self.lj14scale == 1:
            # Just exclude the 1-2 and 1-3 interactions.
            
peastman's avatar
Bug fix  
peastman committed
2345
            self.force.createExclusionsFromBonds(bondIndices, 2)
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
        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
2373

ChayaSt's avatar
ChayaSt committed
2374
parsers["LennardJonesForce"] = LennardJonesGenerator.parseElement
2375

2376

2377
## @private
2378
class GBSAOBCGenerator(object):
2379
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2380

2381
2382
    def __init__(self, forcefield):
        self.ff = forcefield
2383
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
2384

2385
    def registerAtom(self, parameters):
2386
        self.params.registerAtom(parameters)
2387

2388
2389
    @staticmethod
    def parseElement(element, ff):
2390
2391
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
2392
2393
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
2394
2395
2396
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2397
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2398

2399
2400
2401
2402
2403
    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
2404
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
2405
2406
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
2407
2408
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
2409
2410
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2411
2412
2413
2414
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
2415
2416
        sys.addForce(force)

2417
2418
    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
2419

2420
2421
2422
2423
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

2424
2425
2426
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


2427
## @private
2428
class CustomBondGenerator(object):
2429
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2430

2431
2432
    def __init__(self, forcefield):
        self.ff = forcefield
2433
2434
2435
2436
2437
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2438

2439
2440
    @staticmethod
    def parseElement(element, ff):
2441
2442
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
2443
2444
2445
2446
2447
2448
        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'):
2449
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2450
            if None not in types:
2451
2452
2453
                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
2454

2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
    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


2475
## @private
2476
class CustomAngleGenerator(object):
2477
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2478

2479
2480
    def __init__(self, forcefield):
        self.ff = forcefield
2481
2482
2483
2484
2485
2486
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2487

2488
2489
    @staticmethod
    def parseElement(element, ff):
2490
2491
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
2492
2493
2494
2495
2496
2497
        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'):
2498
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2499
            if None not in types:
2500
2501
2502
2503
                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
2504

2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
    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


2527
## @private
2528
class CustomTorsion(object):
2529
2530
2531
2532
2533
2534
2535
2536
2537
    """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

2538
## @private
2539
class CustomTorsionGenerator(object):
2540
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2541

2542
2543
    def __init__(self, forcefield):
        self.ff = forcefield
2544
2545
2546
2547
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
2548

2549
2550
    @staticmethod
    def parseElement(element, ff):
2551
2552
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
2553
2554
2555
2556
2557
2558
        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'):
2559
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2560
            if None not in types:
2561
2562
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
2563
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2564
            if None not in types:
2565
                generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
Justin MacCallum's avatar
Justin MacCallum committed
2566

2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
    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]]]
2599
            match = None
2600
2601
2602
2603
2604
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
2605
2606
2607
2608
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
2609
2610
2611
                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:
2612
                            if hasWildcard:
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
                                # 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)
2624
                                match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
2625
2626
                            else:
                                # There are no wildcards, so the order is unambiguous.
2627
                                match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
2628
                            break
2629
2630
2631
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
2632
2633
2634
2635

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


2636
## @private
2637
class CustomNonbondedGenerator(object):
2638
2639
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

2640
2641
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
2642
2643
2644
2645
2646
2647
2648
2649
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2650
2651
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
2652
2653
2654
2655
        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'])
2656
2657
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683

    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:
2684
2685
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2686
2687
2688
2689
2690
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

    def postprocessSystem(self, sys, data, args):
2691
        # Create the exclusions.
2692

2693
        bondIndices = _findBondsForExclusions(data, sys)
2694
2695
2696
2697
2698
2699
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomNonbondedForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


2700
## @private
2701
class CustomGBGenerator(object):
2702
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2703

2704
2705
    def __init__(self, forcefield):
        self.ff = forcefield
2706
2707
2708
2709
2710
2711
2712
2713
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2714
2715
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
2716
2717
2718
2719
        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'])
2720
2721
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2722
2723
2724
2725
2726
2727
2728
2729
2730
        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()]
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
            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
2742

2743
2744
2745
2746
2747
    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
2748
            raise ValueError('Illegal nonbonded method for CustomGBForce')
2749
2750
2751
2752
2753
2754
2755
2756
2757
        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])
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
        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))
2771
        for atom in data.atoms:
2772
2773
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2774
2775
2776
2777
2778
2779
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

2780
2781

## @private
2782
class CustomManyParticleGenerator(object):
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
    """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 = []
2795

2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
    @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(',')]))
2808
2809
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838

    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:
2839
2840
2841
            values = self.params.getAtomParameters(atom, data)
            type = int(self.params.getExtraParameters(atom, data)['filterType'])
            force.addParticle(values, type)
2842
2843
2844
2845
2846
2847
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
2859
        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)))
2860

2861
2862
        # 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.
2863

2864
2865
2866
2867
2868
2869
2870
2871
2872
        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.
2873

2874
2875
2876
2877
2878
        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
2879
def getAtomPrint(data, atomIndex):
2880

Peter Eastman's avatar
Peter Eastman committed
2881
2882
2883
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
2884
    else:
Peter Eastman's avatar
Peter Eastman committed
2885
        returnString = "NA"
2886
2887
2888
2889
2890

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
2891
def countConstraint(data):
2892

Peter Eastman's avatar
Peter Eastman committed
2893
    bondCount = 0
2894
2895
2896
2897
2898
2899
2900
    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
2901
        if (isConstrained):
2902
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
2903

2904
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
2905

2906
## @private
2907
class AmoebaBondGenerator(object):
2908
2909
2910

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

2911
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
2912
2913

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

2915
2916
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
2917
2918
2919
2920
2921
2922
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2923

2924
2925
2926
2927
2928
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

2932
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
2933
2934
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
2935
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2936
            if None not in types:
2937
2938
2939
2940
2941
                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:
2942
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
2943
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
2944
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
2945
2946
                raise ValueError(outputString)

2947
2948
    #=============================================================================================

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

Mark Friedrichs's avatar
Mark Friedrichs committed
2951
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
2952

Peter Eastman's avatar
Peter Eastman committed
2953
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
2954
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
2955
        if len(existing) == 0:
2956
            force = mm.AmoebaBondForce()
2957
2958
2959
2960
            sys.addForce(force)
        else:
            force = existing[0]

2961
2962
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972

        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:
2973
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
2974
2975
2976
2977
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

2978
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
2979
2980
2981
2982

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

Peter Eastman's avatar
Peter Eastman committed
2984
def addAngleConstraint(angle, idealAngle, data, sys):
2985
2986

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

Peter Eastman's avatar
Peter Eastman committed
2988
2989
    bond1 = None
    bond2 = None
2990
2991
2992
2993
2994
2995
2996
    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
2997

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

3000
3001
3002
3003
        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
3004
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
3005
                data.addConstraint(sys, angle[0], angle[2], length)
3006
3007
3008
                return

#=============================================================================================
3009
## @private
3010
class AmoebaAngleGenerator(object):
3011
3012

    #=============================================================================================
3013
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
3014
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3015

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

Peter Eastman's avatar
Peter Eastman committed
3018
3019
3020
3021
3022
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3023

Peter Eastman's avatar
Peter Eastman committed
3024
3025
3026
        self.types1 = []
        self.types2 = []
        self.types3 = []
3027

Peter Eastman's avatar
Peter Eastman committed
3028
3029
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3030

3031
3032
3033
3034
3035
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3039
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
3040
3041
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
3042
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
3043
            if None not in types:
3044
3045
3046
3047
3048
3049

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
3050
                angleList.append(float(angle.attrib['angle1']))
3051
3052

                try:
Peter Eastman's avatar
Peter Eastman committed
3053
                    angleList.append(float(angle.attrib['angle2']))
3054
                    try:
Peter Eastman's avatar
Peter Eastman committed
3055
                        angleList.append(float(angle.attrib['angle3']))
3056
3057
3058
3059
3060
3061
3062
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
3063
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
3064
3065
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
3066
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
3067
3068
                raise ValueError(outputString)

3069
3070
3071
3072
    #=============================================================================================
    # 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
3073

Peter Eastman's avatar
Peter Eastman committed
3074
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3075
3076
3077
3078
3079
3080
        pass

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

Peter Eastman's avatar
Peter Eastman committed
3082
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3083
3084
3085
3086

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3087
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
3088
3089

        if len(existing) == 0:
3090
            force = mm.AmoebaAngleForce()
3091
3092
3093
3094
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3095
        # set scalars
3096

3097
3098
3099
3100
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
3101
3102

        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
3103
3104
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3105

Peter Eastman's avatar
Peter Eastman committed
3106
3107
3108
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3109
3110
3111
3112
3113
3114
3115
            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]
3116
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3117
                    elif self.k[i] != 0:
Peter Eastman's avatar
Peter Eastman committed
3118
3119
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
3120
3121
3122
3123
3124
3125
                            # 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
3126
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
3127
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3128
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
3129
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3130
                            if (numberOfHydrogens < lenAngle):
3131
3132
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
3133
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
3134
                                raise ValueError(outputString)
3135
3136
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
3137

3138
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
3139
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
3140
3141
3142
3143
3144
3145
                    break

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

Peter Eastman's avatar
Peter Eastman committed
3147
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3148
3149
3150
3151

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3152
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
3153
3154

        if len(existing) == 0:
3155
            force = mm.AmoebaInPlaneAngleForce()
3156
3157
3158
3159
3160
3161
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

3162
3163
3164
3165
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
3166
3167

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

Peter Eastman's avatar
Peter Eastman committed
3169
3170
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3171

Peter Eastman's avatar
Peter Eastman committed
3172
3173
3174
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3175
3176
3177
3178
3179
3180
3181
3182
3183

            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
3184
                    if (isConstrained and self.k[i] != 0.0):
3185
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3186
3187
3188
3189
                    else:
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

3190
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
3191
3192
3193

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
3194
3195
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
3196
3197
#=============================================================================================

3198
## @private
3199
class AmoebaOutOfPlaneBendGenerator(object):
3200
3201
3202
3203

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

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

3205
3206
3207
3208
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3209
3210
3211
3212
3213
3214
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3215

Peter Eastman's avatar
Peter Eastman committed
3216
3217
3218
3219
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3220

Peter Eastman's avatar
Peter Eastman committed
3221
        self.ks = []
3222
3223
3224
3225
3226

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

3228
3229
3230
3231
3232
3233
3234
3235
3236
3237
3238
3239
3240
3241
3242
3243
3244
3245
3246
3247
3248
3249
3250
3251
3252
    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
3253

3254
3255
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
3256
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
3257
3258
3259
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
3260
                                                   float(element.attrib['opbend-sextic']))
3261
3262
3263
3264

        forceField._forces.append(generator)

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

Peter Eastman's avatar
Peter Eastman committed
3268
3269
3270
3271
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
3272

Peter Eastman's avatar
Peter Eastman committed
3273
                generator.ks.append(float(angle.attrib['k']))
3274
3275

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3276
3277
                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
3278
3279
                raise ValueError(outputString)

3280
3281
3282
3283
3284
3285
    #=============================================================================================
    # 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
3286

Peter Eastman's avatar
Peter Eastman committed
3287
    def getMiddleAtom(self, angle, data):
3288
3289
3290

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
3291
        middleAtom = -1
Justin MacCallum's avatar
Justin MacCallum committed
3292
        for atomIndex in angle:
Peter Eastman's avatar
Peter Eastman committed
3293
            isMiddle = 0
3294
3295
3296
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3297
                if (atom1 != atomIndex):
3298
3299
3300
                    partner = atom1
                else:
                    partner = atom2
Justin MacCallum's avatar
Justin MacCallum committed
3301
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]):
3302
3303
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
3304
            if (isMiddle == 2):
3305
3306
3307
3308
3309
                return atomIndex
        return -1

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

Peter Eastman's avatar
Peter Eastman committed
3310
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3311
3312
3313
3314
3315
3316
3317
3318
3319
3320
3321
3322
3323

        # 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
3324
3325
3326
3327
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
3328
3329
3330
3331

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

Peter Eastman's avatar
Peter Eastman committed
3332
        skipAtoms = dict()
3333
3334
3335
3336

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

Peter Eastman's avatar
Peter Eastman committed
3337
3338
        inPlaneAngles = []
        nonInPlaneAngles = []
3339
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
3340
        idealAngles = []*len(data.angles)
3341
3342
3343

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

Peter Eastman's avatar
Peter Eastman committed
3344
3345
3346
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
3347
3348
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
3349
                middleType = -1
3350
3351
                middleCovalency = -1

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

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

Peter Eastman's avatar
Peter Eastman committed
3360
3361
3362
3363
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
3364
3365
3366
3367

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3368
                    if (atom1 != middleAtom):
3369
3370
3371
3372
3373
3374
3375
3376
                        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
3377
3378
3379
3380
3381
                        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
3382

Peter Eastman's avatar
Peter Eastman committed
3383
                if (len(partners) == 3):
3384

Peter Eastman's avatar
Peter Eastman committed
3385
3386
3387
                    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])
3388

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

3391
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
3392
3393
3394
3395
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3396
3397
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
3398
3399
3400
3401
3402
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
3403
3404
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
3405
                    angleDict['isConstrained'] = 0
3406

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

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

Peter Eastman's avatar
Peter Eastman committed
3416
                    inPlaneAngles.append(angleDict)
3417
3418

                else:
Peter Eastman's avatar
Peter Eastman committed
3419
3420
3421
3422
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
3423
            else:
Peter Eastman's avatar
Peter Eastman committed
3424
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
3425

Peter Eastman's avatar
Peter Eastman committed
3426
                    partnerSet = skipAtoms[middleAtom]
Justin MacCallum's avatar
Justin MacCallum committed
3427

Peter Eastman's avatar
Peter Eastman committed
3428
                    angleDict = {}
3429

Peter Eastman's avatar
Peter Eastman committed
3430
3431
3432
3433
3434
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
3435

Peter Eastman's avatar
Peter Eastman committed
3436
                    angleDict['isConstrained'] = isConstrained
3437

Peter Eastman's avatar
Peter Eastman committed
3438
3439
3440
3441
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3442
3443

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3444
3445
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3446

Peter Eastman's avatar
Peter Eastman committed
3447
                    inPlaneAngles.append(angleDict)
3448
3449

                else:
Peter Eastman's avatar
Peter Eastman committed
3450
3451
                    angleDict = {}
                    angleDict['angle'] = angle
3452
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
3453
                    nonInPlaneAngles.append(angleDict)
3454

3455
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
3456
3457

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3458
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
3459
3460
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
3461
3462

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3463
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
3464
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
3465
                    nonInPlaneAngles.append(angleDict)
3466
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
3467
3468
3469
3470
3471

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

3472
## @private
3473
class AmoebaTorsionGenerator(object):
3474
3475
3476
3477
3478
3479
3480

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
3481
        self.torsionUnit = torsionUnit
3482

Peter Eastman's avatar
Peter Eastman committed
3483
3484
3485
3486
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3487

Peter Eastman's avatar
Peter Eastman committed
3488
3489
3490
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
3491

3492
3493
3494
3495
3496
3497
3498
3499
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3501
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
3502
3503
3504
3505
3506
3507
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
3508
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
3509
            if None not in types:
3510
3511
3512
3513
3514
3515
3516

                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
3517
3518
                    tInfo = []
                    suffix = str(ii)
3519
3520
3521
3522
3523
3524
                    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
3525
3526
3527
3528
3529
3530
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
3531
3532
3533

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
3534
3535
3536
3537
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
3538
3539
                raise ValueError(outputString)

3540
3541
3542
3543
3544
    #=============================================================================================

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3545
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
3546
        if len(existing) == 0:
3547
            force = mm.PeriodicTorsionForce()
3548
3549
3550
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3551

3552
3553
3554
3555
3556
3557
3558
3559
3560
3561
3562
3563
3564
3565
3566
3567
        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
3568
                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):
3569
3570
3571
3572
3573
3574
                    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])
3575
3576
3577
3578
3579
3580
                    break

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

3581
## @private
3582
class AmoebaPiTorsionGenerator(object):
3583
3584
3585
3586
3587
3588

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

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

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

3590
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
3591
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
3592
3593
3594
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3595

3596
3597
3598
3599
3600
3601
3602
3603
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
3604
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
3605
3606
3607
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
3608
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
3609
            if None not in types:
3610
3611
3612
3613
3614
3615
                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
3616
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3617
3618
                raise ValueError(outputString)

3619
3620
3621
3622
3623
3624
3625
3626
3627
3628
3629
3630
3631
3632
3633
3634
    #=============================================================================================

    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
3635

3636
3637
            atom1 = bond.atom1
            atom2 = bond.atom2
Justin MacCallum's avatar
Justin MacCallum committed
3638

3639
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
3640
3641
3642
3643
3644
3645
3646
3647
3648
3649
3650

                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
3651
3652
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
3653
3654
3655
3656
3657
3658
3659
3660
3661
3662
3663
3664

                       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
3665
                           if (bondedAtom1 != atom1):
3666
3667
3668
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
3669
3670
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3671
                                   piTorsionAtom1 = b1
3672
3673
3674
3675
3676
3677
                               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
3678
                           if (bondedAtom1 != atom2):
3679
3680
3681
3682
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3683
3684
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3685
                                   piTorsionAtom5 = b1
3686
3687
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
3688

Peter Eastman's avatar
Peter Eastman committed
3689
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
3690
3691
3692
3693
3694

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

3695
## @private
3696
class AmoebaTorsionTorsionGenerator(object):
3697
3698
3699
3700
3701

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

Peter Eastman's avatar
Peter Eastman committed
3702
    def __init__(self):
3703

Peter Eastman's avatar
Peter Eastman committed
3704
3705
3706
3707
3708
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
3709

Peter Eastman's avatar
Peter Eastman committed
3710
        self.gridIndex = []
3711

Peter Eastman's avatar
Peter Eastman committed
3712
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
3713

3714
3715
3716
3717
3718
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
3719
        generator = AmoebaTorsionTorsionGenerator()
3720
3721
3722
3723
3724
3725
3726
        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'):
3727
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
3728
            if None not in types:
3729
3730
3731
3732
3733
3734
3735

                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
3736
3737
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
3738
3739
3740
3741
3742
3743
3744
3745
3746
                    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
3747
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
3748
3749
                raise ValueError(outputString)

3750
3751
3752
3753
3754
3755
3756
3757
3758
3759
        # 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
3760
3761
3762
3763
3764
3765
        #     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
3766
3767

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

Peter Eastman's avatar
Peter Eastman committed
3771
3772
3773
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
3774

Peter Eastman's avatar
Peter Eastman committed
3775
3776
            grid = []
            gridCol = []
3777
3778
3779
3780
3781

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
3782
3783
3784
3785
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
3786
                if 'fx' in gridEntry.attrib:
3787
3788
3789
                    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
3790
                gridCol.append(gridRow)
3791
3792

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
3793
3794
3795
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
3796
3797
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
3798

Peter Eastman's avatar
Peter Eastman committed
3799
3800
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
3801
            else:
Peter Eastman's avatar
Peter Eastman committed
3802
3803
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
3804
3805
3806
3807
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
3808
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
3809
3810
3811
3812
3813
3814
3815
3816

        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
3817
        if (len(data.atomBonds[atomC]) == 4):
3818
3819
3820
3821
3822
            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
3823
3824
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
3825
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
3826
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
3827
3828
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
3829
3830
                if (hit > -1):
                    if (atomE == -1):
3831
3832
3833
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
3834

3835
3836
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
3837
3838
            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
3839
                raise ValueError(outputString)
3840
3841
3842
3843
3844

            # 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
3845
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
3846
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
3847
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
3848
                chiralAtomIndex = atomF
3849

Peter Eastman's avatar
Peter Eastman committed
3850
3851
3852
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
3853
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
3854
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
3855
                chiralAtomIndex = massF
3856
3857
3858
3859

        return chiralAtomIndex

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

3861
3862
3863
3864
3865
3866
3867
3868
3869
3870
3871
3872
3873
    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
3874
3875
            # search for bitorsions; based on TINKER subroutine bitors()

3876
3877
3878
3879
3880
3881
3882
            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
3883
                if (bondedAtom1 != ib):
3884
3885
3886
3887
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3888
                if (ia != ic and ia != id):
3889
3890
3891
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
3892
                        if (bondedAtom1 != id):
3893
3894
3895
3896
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3897
                        if (ie != ic and ie != ib and ie != ia):
3898
3899
3900
3901
3902
3903
3904
3905
3906
3907
3908
3909
3910
3911
3912
3913
3914
3915
3916
3917

                            # 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
3918
3919
3920
                                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])
3921
3922
3923

                                # match in reverse order

3924
                                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
3925
3926
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
3927
3928
3929
3930

        # set grids

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

3933
3934
3935
3936
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

3937
## @private
3938
class AmoebaStretchBendGenerator(object):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3939
3940

    #=============================================================================================
3941
3942
3943
3944
3945
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
3946
3947
3948
        self.types1 = []
        self.types2 = []
        self.types3 = []
3949

Peter Eastman's avatar
Peter Eastman committed
3950
3951
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
3952

3953
3954
3955
3956
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
3957
        generator = AmoebaStretchBendGenerator()
3958
3959
3960
3961
3962
3963
3964
        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'):
3965
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
3966
            if None not in types:
3967
3968
3969
3970
3971
3972
3973
3974
3975
3976
3977
3978

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

3982
3983
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
3984
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
3985
3986
    # 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
3987
3988
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
3989
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
3990
3991
3992

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

Peter Eastman's avatar
Peter Eastman committed
3993
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3994
3995
3996
3997
3998
3999
4000
4001
        pass

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

    # Note: request for constrained bonds is ignored.

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

4002
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
4003
4004
4005
4006
4007
4008
4009
4010
4011
4012
4013
4014

        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
4015
            if ('isConstrained' in angleDict):
4016
4017
4018
4019
4020
4021
4022
4023
                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
4024
            radian = 57.2957795130
4025
4026
4027
4028
4029
4030
4031
4032
4033
4034
            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
4035
                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
4036
4037
4038
4039
4040
4041
4042
4043
4044
4045
4046
4047
4048
4049
4050
                    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
4051

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4054
                    if ('idealAngle' not in angleDict):
4055

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4056
4057
4058
4059
4060
                       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
4061
                       raise ValueError(outputString)
4062

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4063
4064
4065
4066
4067
4068
4069
                    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
4070
                       raise ValueError(outputString)
4071

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4075
                    break
4076
4077
4078
4079
4080

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

4081
## @private
4082
class AmoebaVdwGenerator(object):
4083
4084

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

4086
4087
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4090
        self.type = type
4091

Peter Eastman's avatar
Peter Eastman committed
4092
4093
4094
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
4095

Peter Eastman's avatar
Peter Eastman committed
4096
        self.epsilonrule = epsilonrule
4097

Peter Eastman's avatar
Peter Eastman committed
4098
4099
4100
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
4101
4102
4103
4104
4105
4106
4107

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

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

4111
4112
4113
        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
4114
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
4115
4116
4117
4118
4119
4120
4121
4122
4123
4124
4125
4126
            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')
4127
        generator.params.parseDefinitions(element)
4128
4129
4130
4131
4132
        two_six = 1.122462048309372

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

    @staticmethod
4133
    def getBondedParticleSets(sys, data):
4134

4135
4136
4137
4138
4139
4140
        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
4141

4142
4143
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
4146
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
4147
4148
        epsilonMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'HARMONIC':1, 'HHG':1}

4149
4150
        force = mm.AmoebaVdwForce()
        sys.addForce(force)
4151

4152
        # sigma and epsilon combining rules
4153

4154
4155
4156
4157
        if ('sigmaCombiningRule' in args):
            sigmaRule = args['sigmaCombiningRule'].upper()
            if (sigmaRule.upper() in sigmaMap):
                force.setSigmaCombiningRule(sigmaRule.upper())
4158
            else:
4159
4160
4161
4162
                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)
4163

4164
4165
4166
4167
        if ('epsilonCombiningRule' in args):
            epsilonRule = args['epsilonCombiningRule'].upper()
            if (epsilonRule.upper() in epsilonMap):
                force.setEpsilonCombiningRule(epsilonRule.upper())
4168
            else:
4169
4170
4171
4172
                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
4173

4174
        # cutoff
4175

4176
4177
4178
4179
        if ('vdwCutoff' in args):
            force.setCutoff(args['vdwCutoff'])
        else:
            force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4180

4181
        # dispersion correction
4182

4183
4184
        if ('useDispersionCorrection' in args):
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
4185

4186
4187
        if (nonbondedMethod == PME):
            force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
4188
4189
4190

        # add particles to force

4191
4192
4193
4194
4195
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
4196
        for (i, atom) in enumerate(data.atoms):
4197
4198
4199
4200
4201
4202
4203
4204
4205
4206
            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
4207

4208
            force.addParticle(ivIndex, values[0]*sigmaScale, values[1], values[2])
4209
4210
4211
4212
4213
4214
4215

        # 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

4216
        bondedParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4217
4218

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

4220
4221
4222
4223
4224
4225
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
4226
            if (self.vdw13Scale == 0.0):
4227
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
4228
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
4229
4230
4231
4232
4233

            # self

            exclusionSet.add(i)

4234
            force.setParticleExclusions(i, tuple(exclusionSet))
4235
4236
4237
4238
4239

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

4240
## @private
4241
class AmoebaMultipoleGenerator(object):
4242
4243
4244
4245

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

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

4247
4248
    #=============================================================================================

4249
    def __init__(self, forceField):
Peter Eastman's avatar
Peter Eastman committed
4250
4251
        self.forceField = forceField
        self.typeMap = {}
4252
4253
4254
4255
4256
4257

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
4258
    def setAxisType(kIndices):
4259
4260
4261

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
4262
4263
4264
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
4265
                else:
Peter Eastman's avatar
Peter Eastman committed
4266
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
4267

Peter Eastman's avatar
Peter Eastman committed
4268
4269
                if (kIndicesLen > 2):
                    kx = kIndices[2]
4270
                else:
Peter Eastman's avatar
Peter Eastman committed
4271
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
4272

Peter Eastman's avatar
Peter Eastman committed
4273
4274
                if (kIndicesLen > 1):
                    kz = kIndices[1]
4275
                else:
Peter Eastman's avatar
Peter Eastman committed
4276
                    kz = 0
4277

Peter Eastman's avatar
Peter Eastman committed
4278
4279
                while(len(kIndices) < 4):
                    kIndices.append(0)
4280
4281

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
4282
                if (kz == 0):
4283
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
4284
                if (kz != 0 and kx == 0):
4285
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
4286
                if (kz < 0 or kx < 0):
4287
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
4288
                if (kx < 0 and ky < 0):
4289
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
4290
                if (kz < 0 and kx < 0 and ky  < 0):
4291
4292
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
4293
4294
4295
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
4296
4297
4298
4299
4300
4301
4302
4303

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
4304
        #   <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"  >
4305
4306
4307
        # <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"  />

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

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

        for atom in element.findall('Multipole'):
4319
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4320
            if None not in types:
4321
4322
4323
4324
4325
4326
4327
4328

                # 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
4329
4330
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
4331
                    except:
4332
4333
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
4334
                # set axis type based on k-Indices
4335

Peter Eastman's avatar
Peter Eastman committed
4336
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
4337
4338
4339

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
4342
4343
4344
4345
4346
4347
4348
4349
4350
4351
4352
4353
4354
                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']))
4355
4356

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4357
                    if (t not in generator.typeMap):
4358
4359
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
4360
4361
4362
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
4363
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
4364
4365
4366
4367
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
4368

4369
            else:
Peter Eastman's avatar
Peter Eastman committed
4370
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4371
4372
                raise ValueError(outputString)

4373
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
4374

4375
        for atom in element.findall('Polarize'):
4376
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4377
            if None not in types:
4378

Peter Eastman's avatar
Peter Eastman committed
4379
4380
4381
4382
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
4383
4384
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
4385
                    pdamp = pow(polarizability, 1.0/6.0)
4386

Peter Eastman's avatar
Peter Eastman committed
4387
4388
                pgrpMap = dict()
                for index in range(1, 7):
4389
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
4390
                    if (pgrp in atom.attrib):
4391
4392
4393
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4394
4395
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
4396
                        raise ValueError(outputString)
4397
4398
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
4399
4400
4401
4402
4403
4404
4405
                        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
4406
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
4407
4408
4409
4410
4411
                                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
4412
4413
                            raise ValueError(outputString)

4414
            else:
Peter Eastman's avatar
Peter Eastman committed
4415
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4416
4417
                raise ValueError(outputString)

4418
4419
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4420
    def setPolarGroups(self, data, bonded12ParticleSets, force):
4421
4422
4423
4424
4425

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
4426
4427
4428
4429
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
4430
4431
4432
4433
            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
4434
4435
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
4436
4437
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
4438

4439
4440
4441
4442
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
4443
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
4444
4445
                continue

Peter Eastman's avatar
Peter Eastman committed
4446
4447
            group = set()
            visited = set()
4448
4449
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4450
4451
4452
4453
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
4454
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
4455
4456
                if (nextAtom not in visited):
                   visited.add(nextAtom)
4457
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4458
4459
4460
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
4461
4462
4463

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
4464
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
4465
4466

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4467
4468
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
4469
4470
4471
4472
4473

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
4474
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
4475
4476
                continue

Peter Eastman's avatar
Peter Eastman committed
4477
            pgrp11 = set(atom.polarizationGroupSet[0])
4478
4479
4480
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4481
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
4482
4483
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4484
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
4485

4486
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4487
4488
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
4489
4490
4491
4492
4493

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
4494
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
4495
4496
                continue

Peter Eastman's avatar
Peter Eastman committed
4497
4498
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
4499
4500
4501
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4502
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
4503
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4504
            pgrp13 = pgrp13 - set(pgrp11)
4505
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4506
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
4507

4508
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4509
4510
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
4511
4512
4513
4514
4515

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
4516
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
4517
4518
                continue

Peter Eastman's avatar
Peter Eastman committed
4519
4520
4521
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
4522
4523
4524
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4525
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
4526
4527
4528

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4529
            pgrp14 = pgrp14 - set(pgrp11)
4530
4531

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

4534
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4535
4536
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
4537
4538
4539

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

Peter Eastman's avatar
Peter Eastman committed
4540
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4541
4542
4543
4544

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

4545
4546
4547
4548
4549
4550
4551
        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)
4552

4553
4554
        if ('ewaldErrorTolerance' in args):
            force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))
4555

4556
4557
4558
4559
4560
4561
4562
4563
        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)
4564

4565
4566
        if ('aEwald' in args):
            force.setAEwald(float(args['aEwald']))
4567

4568
4569
        if ('pmeGridDimensions' in args):
            force.setPmeGridDimensions(args['pmeGridDimensions'])
4570

4571
4572
        if ('mutualInducedMaxIterations' in args):
            force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
4573

4574
4575
        if ('mutualInducedTargetEpsilon' in args):
            force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))
4576
4577

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
4578
        # throw error if particle type not available
4579
4580
4581
4582
4583

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

        # 1-2

4584
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4585
4586
4587
4588
4589

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4590
            bonded13Set = set()
4591
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4592
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4593
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
4594
4595
4596
4597

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4598
            selfSet = set()
4599
4600
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4601
4602
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
4603
4604
4605
4606
4607

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4608
4609
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4610
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4611
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
4612

4613
4614
4615
4616
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4617
            selfSet = set()
4618
4619
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4620
4621
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
4622
4623
4624
4625
4626

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4627
4628
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4629
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4630
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
4631
4632
4633
4634
4635
4636

            # 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
4637
            selfSet = set()
4638
4639
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4640
4641
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
4642
4643
4644
4645
4646

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

Peter Eastman's avatar
Peter Eastman committed
4647
4648
                multipoleList = self.typeMap[t]
                hit = 0
4649
4650
4651
4652
4653
4654
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4655
                    if (hit != 0):
4656
4657
                        break

Peter Eastman's avatar
Peter Eastman committed
4658
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4659
4660

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4661
4662
                    kx = kIndices[2]
                    ky = kIndices[3]
4663
4664
4665
4666

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

4668
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
4669
4670
4671
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4672
4673
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
4674
                       if (hit != 0):
4675
4676
4677
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
4678
4679
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
4680
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4681
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4682
4683
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4684
4685
4686
4687
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4688
4689
4690
4691
4692
4693
4694
4695
4696
4697
                                      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

4698
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4699
                                      hit = 1
4700
4701
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4702
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4703
4704
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4705
4706
4707
4708
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4709
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4710
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
4711

4712
4713
4714
4715
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4716
                    if (hit != 0):
4717
4718
                        break

Peter Eastman's avatar
Peter Eastman committed
4719
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4720
4721

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4722
4723
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
4724

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

4729
4730
4731
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
4732
4733
4734
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4735
4736
4737

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4738
                       if (hit != 0):
4739
4740
4741
                           break

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

Peter Eastman's avatar
Peter Eastman committed
4744
                       if (bondedAtomZType == kz):
4745
4746
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
4747
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4748
4749
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4750
4751
4752
4753
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4754
4755
4756
4757
4758
4759
4760
4761

                                      # 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

4762
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4763
                                      hit = 3
4764
4765
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
4766
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4767
4768
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4769
4770
4771
4772
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4773
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4774
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
4775

4776
4777
4778
4779
                # assign multipole parameters via only a z-defining atom

                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
4787

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

Peter Eastman's avatar
Peter Eastman committed
4788
4789
4790
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4791
4792
4793

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4794
                        if (hit != 0):
4795
4796
4797
                            break

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

Peter Eastman's avatar
Peter Eastman committed
4800
4801
                        if (kx == 0 and kz == bondedAtomZType):
                            kz = bondedAtomZIndex
4802
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4803
                            hit = 5
4804
4805
4806
4807
4808

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4809
                    if (hit != 0):
4810
4811
                        break

Peter Eastman's avatar
Peter Eastman committed
4812
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4813
4814
4815

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
4816
4817
4818
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4819

Peter Eastman's avatar
Peter Eastman committed
4820
                    if (kz == 0):
4821
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4822
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
4823

4824
4825
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
4826
                if (hit != 0):
4827

Peter Eastman's avatar
Peter Eastman committed
4828
                    atom.multipoleDict = savedMultipoleDict
4829
                    atom.polarizationGroups = dict()
4830
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
4831
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
4832
                    if (atomIndex == newIndex):
4833
4834
4835
4836
                        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]))
4837
                    else:
4838
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
4839
                else:
Peter Eastman's avatar
Peter Eastman committed
4840
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
4841
            else:
Peter Eastman's avatar
Peter Eastman committed
4842
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
4843
4844
4845

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
4846
        self.setPolarGroups(data, bonded12ParticleSets, force)
4847
4848
4849
4850
4851

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

4852
## @private
4853
class AmoebaWcaDispersionGenerator(object):
4854
4855

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

4857
4858
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4861
4862
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
4863
4864
4865
4866
4867
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
4868
        self.shctd = shctd
4869
4870
4871
4872
4873
4874
4875
4876
4877

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

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

Peter Eastman's avatar
Peter Eastman committed
4879
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
4880
4881
4882
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
Justin MacCallum's avatar
Justin MacCallum committed
4883
                                                  element.attrib['awater'],
4884
4885
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Justin MacCallum's avatar
Justin MacCallum committed
4886
                                                  element.attrib['shctd'])
4887
        forceField._forces.append(generator)
4888
4889
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaWcaDispersionForce', 'WcaDispersion', ('radius', 'epsilon'))
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
4890

4891
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
4892

Peter Eastman's avatar
Peter Eastman committed
4893
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4894
4895
4896
4897
4898
4899
4900
4901
4902
4903
4904
4905

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

Peter Eastman's avatar
Peter Eastman committed
4908
4909
4910
4911
4912
4913
4914
4915
        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  ))
4916

4917
4918
4919
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
4920
4921
4922
4923
4924

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

4925
## @private
4926
class AmoebaGeneralizedKirkwoodGenerator(object):
4927
4928

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

4930
4931
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
4932
4933
4934
4935
4936
4937
4938
4939
4940
4941
4942
    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
4943
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
4944
4945
4946
4947
4948
4949
4950
4951
4952
4953
4954
4955
4956
4957
4958
4959
4960
4961
4962
4963
4964
4965
4966
4967
4968
        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
4969
4970
4971

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

Peter Eastman's avatar
Peter Eastman committed
4972
    def getObcShct(self, data, atomIndex):
4973

Peter Eastman's avatar
Peter Eastman committed
4974
        atom = data.atoms[atomIndex]
4975
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
4976
        shct = -1.0
4977
4978

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
4979

Peter Eastman's avatar
Peter Eastman committed
4980
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
4981
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
4982
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
4983
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
4984
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
4985
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
4986
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
4987
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
4988
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
4989
4990
4991
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
4992
        elif (atomicNumber == 16):              # S(16)
4993
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
4994
        elif (atomicNumber == 26):              # Fe(26)
4995
4996
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
4997
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4998
            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
4999
5000

        return shct
5001
5002
5003

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

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

Peter Eastman's avatar
Peter Eastman committed
5006
        atom = data.atoms[atomIndex]
5007
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5008
        radius = -1.0
5009

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

Peter Eastman's avatar
Peter Eastman committed
5012
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
5013

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

5017
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5018
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
5019

Peter Eastman's avatar
Peter Eastman committed
5020
            if (bondedAtomAtomicNumber == 7):
5021
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
5022
            if (bondedAtomAtomicNumber == 8):
5023
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
5024

Peter Eastman's avatar
Peter Eastman committed
5025
        elif (atomicNumber == 3):               # Li(3)
5026
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
5027
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5028

5029
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
5030
            if (len(bondedAtomIndices) == 3):
5031
5032
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
5033
            elif (len(bondedAtomIndices) == 4):
5034
5035
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5036
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
5037
5038
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
5039
        elif (atomicNumber == 7):               # N(7)
5040
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
5041
        elif (atomicNumber == 8):               # O(8)
5042
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
5043
            if (len(bondedAtomIndices) == 2):
5044
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
5045
        elif (atomicNumber == 9):               # F(9)
5046
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
5047
        elif (atomicNumber == 10):
5048
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
5049
        elif (atomicNumber == 11):
5050
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
5051
        elif (atomicNumber == 12):
5052
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5053
        elif (atomicNumber == 14):
5054
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
5055
        elif (atomicNumber == 15):              # P(15)
5056
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
5057
        elif (atomicNumber == 16):              # S(16)
5058
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
5059
        elif (atomicNumber == 17):
5060
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
5061
        elif (atomicNumber == 18):
5062
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5063
        elif (atomicNumber == 19):
5064
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
5065
        elif (atomicNumber == 20):
5066
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
5067
        elif (atomicNumber == 35):
5068
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
5069
        elif (atomicNumber == 36):
5070
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
5071
        elif (atomicNumber == 37):
5072
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
5073
        elif (atomicNumber == 53):
5074
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
5075
        elif (atomicNumber == 54):
5076
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
5077
        elif (atomicNumber == 55):
5078
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
5079
        elif (atomicNumber == 56):
5080
5081
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
5082
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5083
5084
            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
5085

5086
5087
5088
5089
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
5092
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5093
        atom = data.atoms[atomIndex]
5094
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
5095
        if (atomicNumber in bondiMap):
5096
5097
            radius = bondiMap[atomicNumber]
        else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5098
5099
            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
5100

5101
5102
5103
5104
5105
5106
5107
5108
5109
5110
        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
5111

Peter Eastman's avatar
Peter Eastman committed
5112
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
Justin MacCallum's avatar
Justin MacCallum committed
5113
5114
                                                        element.attrib['includeCavityTerm'],
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor'])
5115
5116
5117
        forceField._forces.append(generator)

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

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

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

5124
5125
5126
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
5127
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5128
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
5129
        if (len(amoebaMultipoleForceList) > 0):
5130
5131
5132
5133
5134
            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
5135
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
5136
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
5137

5138
5139
5140
5141
5142
5143
5144
        # 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
5145

Peter Eastman's avatar
Peter Eastman committed
5146
5147
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
5148
            else:
Peter Eastman's avatar
Peter Eastman committed
5149
                force.setSolventDielectric(   float(self.solventDielectric))
5150

Peter Eastman's avatar
Peter Eastman committed
5151
5152
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
5153
            else:
Peter Eastman's avatar
Peter Eastman committed
5154
                force.setSoluteDielectric(    float(self.soluteDielectric))
5155

Peter Eastman's avatar
Peter Eastman committed
5156
5157
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
5158
            else:
Peter Eastman's avatar
Peter Eastman committed
5159
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
5160
5161
5162
5163
5164

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
5167
5168
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
5169
5170
5171

        # 1-2

5172
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5173
5174

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
5175
5176
5177
5178
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5179
            else:
Peter Eastman's avatar
Peter Eastman committed
5180
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5181
5182
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
5183
            force.addParticle(multipoleParameters[0], radius, shct)
5184
5185
5186
5187
5188

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

5189
## @private
5190
class AmoebaUreyBradleyGenerator(object):
5191
5192
5193
5194

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

5196
    def __init__(self):
5197

Peter Eastman's avatar
Peter Eastman committed
5198
5199
5200
        self.types1 = []
        self.types2 = []
        self.types3 = []
5201

Peter Eastman's avatar
Peter Eastman committed
5202
5203
        self.length = []
        self.k = []
5204
5205
5206
5207
5208
5209

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

    @staticmethod
    def parseElement(element, forceField):

5210
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
5211
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
5212

5213
        generator = AmoebaUreyBradleyGenerator()
5214
5215
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
5216
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
5217
            if None not in types:
5218
5219
5220
5221
5222
5223
5224
5225
5226
5227

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

5231
5232
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
5235
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5236
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
5237
5238

        if len(existing) == 0:
5239
            force = mm.HarmonicBondForce()
5240
5241
5242
5243
5244
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
Peter Eastman's avatar
Peter Eastman committed
5245
            if (isConstrained):
5246
5247
5248
5249
5250
5251
5252
5253
                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
5254
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
5255
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
5256
5257
5258
5259
5260
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
5261
5262
5263


## @private
5264
class DrudeGenerator(object):
peastman's avatar
peastman committed
5265
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
5266

5267
5268
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
5269
5270
5271
5272
5273
5274
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
5275
5276
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
5277
5278
5279
5280
        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'):
5281
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
5282
5283
5284
5285
5286
5287
5288
5289
5290
5291
            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
5292

peastman's avatar
peastman committed
5293
5294
5295
5296
    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
5297

peastman's avatar
peastman committed
5298
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
5299

peastman's avatar
peastman committed
5300
5301
5302
5303
5304
5305
5306
5307
5308
5309
5310
5311
5312
5313
5314
5315
        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
5316
5317
                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
5318
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
5319

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

peastman's avatar
peastman committed
5323
5324
5325
5326
5327
5328
5329
        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)
5330
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
5331
5332
5333
5334
5335
                # 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]
5336
5337
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
5338
5339
5340
5341
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

Justin MacCallum's avatar
Justin MacCallum committed
5342
parsers["DrudeForce"] = DrudeGenerator.parseElement
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
5343
5344

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