forcefield.py 255 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
peastman's avatar
peastman committed
43
from collections import defaultdict
44
45
import simtk.openmm as mm
import simtk.unit as unit
46
from . import element as elem
47
from simtk.openmm.app import Topology
48
from simtk.openmm.app.internal.singleton import Singleton
49

50
51
52
# Directories from which to load built in force fields.

_dataDirectories = [os.path.join(os.path.dirname(__file__), 'data')]
peastman's avatar
Bug fix  
peastman committed
53
54
55
56
57
58
try:
    from pkg_resources import iter_entry_points
    for entry in iter_entry_points(group='openmm.forcefielddir'):
        _dataDirectories.append(entry.load()())
except:
    pass # pkg_resources is not installed
59

60
61
def _convertParameterToNumber(param):
    if unit.is_quantity(param):
62
63
64
        if param.unit.is_compatible(unit.bar):
            return param / unit.bar
        return param.value_in_unit_system(unit.md_unit_system)
65
66
    return float(param)

67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
def _parseFunctions(element):
    """Parse the attributes on an XML tag to find any tabulated functions it defines."""
    functions = []
    for function in element.findall('Function'):
        values = [float(x) for x in function.text.split()]
        if 'type' in function.attrib:
            functionType = function.attrib['type']
        else:
            functionType = '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])
        functions.append((function.attrib['name'], functionType, values, params))
    return functions

def _createFunctions(force, functions):
    """Add TabulatedFunctions to a Force based on the information that was recorded by _parseFunctions()."""
    for (name, type, values, params) in 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))

101
102
# Enumerated values for nonbonded method

103
class NoCutoff(Singleton):
104
105
106
107
    def __repr__(self):
        return 'NoCutoff'
NoCutoff = NoCutoff()

108
class CutoffNonPeriodic(Singleton):
109
110
111
112
    def __repr__(self):
        return 'CutoffNonPeriodic'
CutoffNonPeriodic = CutoffNonPeriodic()

113
class CutoffPeriodic(Singleton):
114
115
116
117
    def __repr__(self):
        return 'CutoffPeriodic'
CutoffPeriodic = CutoffPeriodic()

118
class Ewald(Singleton):
119
120
121
122
    def __repr__(self):
        return 'Ewald'
Ewald = Ewald()

123
class PME(Singleton):
124
125
126
    def __repr__(self):
        return 'PME'
PME = PME()
127

Peter Eastman's avatar
Peter Eastman committed
128
129
130
131
132
class LJPME(Singleton):
    def __repr__(self):
        return 'LJPME'
LJPME = LJPME()

133
134
# Enumerated values for constraint type

135
class HBonds(Singleton):
136
137
138
139
    def __repr__(self):
        return 'HBonds'
HBonds = HBonds()

140
class AllBonds(Singleton):
141
142
143
144
    def __repr__(self):
        return 'AllBonds'
AllBonds = AllBonds()

145
class HAngles(Singleton):
146
147
148
    def __repr__(self):
        return 'HAngles'
HAngles = HAngles()
149
150
151
152
153
154
155
156
157
158

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

Robert McGibbon's avatar
Robert McGibbon committed
160
161
162
163
164
165
166
167
        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.
168
169
170
        """
        self._atomTypes = {}
        self._templates = {}
171
172
        self._patches = {}
        self._templatePatches = {}
173
        self._templateSignatures = {None:[]}
174
        self._atomClasses = {'':set()}
175
        self._forces = []
176
        self._scripts = []
177
        self._templateGenerators = []
178
        self.loadFile(files)
179

180
    def loadFile(self, files):
181
        """Load an XML file and add the definitions from it to this ForceField.
182

Robert McGibbon's avatar
Robert McGibbon committed
183
184
        Parameters
        ----------
185
186
187
        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
188
189
190
            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
191
        """
192

193
194
195
196
197
198
        if not isinstance(files, tuple):
            files = (files,)

        trees = []

        for file in files:
199
            tree = None
200
201
202
203
            try:
                # this handles either filenames or open file-like objects
                tree = etree.parse(file)
            except IOError:
204
205
206
207
208
                for dataDir in _dataDirectories:
                    f = os.path.join(dataDir, file)
                    if os.path.isfile(f):
                        tree = etree.parse(f)
                        break
209
210
211
212
213
214
215
216
217
218
219
            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)
220
221
            if tree is None:
                raise ValueError('Could not locate file "%s"' % file)
222
223
224

            trees.append(tree)

peastman's avatar
peastman committed
225
226
227
228
229
230
231
232
233
234
235
236
237
        # Process includes.
        
        for parentFile, tree in zip(files, trees):
            if isinstance(parentFile, str):
                parentDir = os.path.dirname(parentFile)
            else:
                parentDir = ''
            for include in tree.getroot().findall('Include'):
                includeFile = include.attrib['file']
                joined = os.path.join(parentDir, includeFile)
                if os.path.isfile(joined):
                    includeFile = joined
                self.loadFile(includeFile)
peastman's avatar
peastman committed
238
239
240

        # Load the atom types.

241
242
243
244
        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
245
246
247

        # Load the residue templates.

248
249
250
251
252
        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)
253
254
                    if 'override' in residue.attrib:
                        template.overrideLevel = int(residue.attrib['override'])
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
                    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']))
279
                    for patch in residue.findall('AllowPatch'):
280
                        patchName = patch.attrib['name']
281
282
                        if ':' in patchName:
                            colonIndex = patchName.find(':')
283
284
285
                            self.registerTemplatePatch(resName, patchName[:colonIndex], int(patchName[colonIndex+1:])-1)
                        else:
                            self.registerTemplatePatch(resName, patchName, 0)
286
                    self.registerResidueTemplate(template)
peastman's avatar
peastman committed
287

288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
        # 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']
311
                        patchData.addedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
312
313
314
315
316
317
318
319
320
321
322
                    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']
323
                        patchData.changedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
                    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
354
355
        # Load force definitions

356
357
358
359
        for tree in trees:
            for child in tree.getroot():
                if child.tag in parsers:
                    parsers[child.tag](child, self)
peastman's avatar
peastman committed
360
361
362

        # Load scripts

363
364
365
        for tree in trees:
            for node in tree.getroot().findall('Script'):
                self.registerScript(node.text)
366

367
368
369
    def getGenerators(self):
        """Get the list of all registered generators."""
        return self._forces
370

371
372
373
    def registerGenerator(self, generator):
        """Register a new generator."""
        self._forces.append(generator)
374

375
376
377
378
379
380
381
382
383
384
385
386
    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)
387
        self._atomTypes[name] = ForceField._AtomType(name, atomClass, mass, element)
388
389
390
391
392
393
394
        if atomClass in self._atomClasses:
            typeSet = self._atomClasses[atomClass]
        else:
            typeSet = set()
            self._atomClasses[atomClass] = typeSet
        typeSet.add(name)
        self._atomClasses[''].add(name)
395

396
397
    def registerResidueTemplate(self, template):
        """Register a new residue template."""
398
399
        if template.name in self._templates:
            # There is already a template with this name, so check the override levels.
400

401
402
403
404
405
406
407
408
409
410
411
            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))
412

413
        # Register the template.
414

415
416
417
        self._templates[template.name] = template
        signature = _createResidueSignature([atom.element for atom in template.atoms])
        if signature in self._templateSignatures:
418
            self._templateSignatures[signature].append(template)
419
420
        else:
            self._templateSignatures[signature] = [template]
421

422
423
424
    def registerPatch(self, patch):
        """Register a new patch that can be applied to templates."""
        self._patches[patch.name] = patch
425

426
427
428
429
430
431
    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))

432
433
434
    def registerScript(self, script):
        """Register a new script to be executed after building the System."""
        self._scripts.append(script)
435

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
436
    def registerTemplateGenerator(self, generator):
437
438
439
440
        """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
441
442
        .. CAUTION:: This method is experimental, and its API is subject to change.

443
444
        Parameters
        ----------
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
445
        generator : function
446
447
            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
448
        When a residue without a template is encountered, the `generator` function is called with:
449

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
450
451
        ::
           success = generator(forcefield, residue)
452
453
        ```

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
454
455
456
457
458
459
460
461
462
463
        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.
464

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
465
466
467
468
469
           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`.
470

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
471
472
473
474
475
           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.
476
477

        """
478
        self._templateGenerators.append(generator)
479

480
    def _findAtomTypes(self, attrib, num):
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
        """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.

        """
496
497
498
499
500
501
502
        types = []
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
503
            typeAttrib = 'type'+suffix
504
            if classAttrib in attrib:
505
                if typeAttrib in attrib:
506
                    raise ValueError('Specified both a type and a class for the same atom: '+str(attrib))
507
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
508
509
510
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
511
512
            elif typeAttrib in attrib:
                if attrib[typeAttrib] == '':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
513
                    types.append(self._atomClasses[''])
514
                elif attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
515
516
517
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
518
519
            else:
                types.append(None) # Unknown atom type
520
521
        return types

522
    def _parseTorsion(self, attrib):
523
        """Parse the node defining a torsion."""
524
        types = self._findAtomTypes(attrib, 4)
peastman's avatar
peastman committed
525
        if None in types:
526
527
528
529
530
            return None
        torsion = PeriodicTorsion(types)
        index = 1
        while 'phase%d'%index in attrib:
            torsion.periodicity.append(int(attrib['periodicity%d'%index]))
531
532
            torsion.phase.append(_convertParameterToNumber(attrib['phase%d'%index]))
            torsion.k.append(_convertParameterToNumber(attrib['k%d'%index]))
533
            index += 1
Justin MacCallum's avatar
Justin MacCallum committed
534
535
        return torsion

536
    class _SystemData(object):
537
538
539
        """Inner class used to encapsulate data about the system being created."""
        def __init__(self):
            self.atomType = {}
540
            self.atomParameters = {}
541
            self.atomTemplateIndexes = {}
542
            self.atoms = []
543
            self.excludeAtomWith = []
544
            self.virtualSites = {}
545
546
547
548
549
550
            self.bonds = []
            self.angles = []
            self.propers = []
            self.impropers = []
            self.atomBonds = []
            self.isAngleConstrained = []
551
552
553
554
555
556
            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:
557
                if self.constraints[key] != distance:
558
559
560
561
                    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)
562

563
564
565
566
567
568
        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
569
                self.atomTemplateIndexes[atom] = match
570
571
572
                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)
573

574
    class _TemplateData(object):
575
576
577
578
        """Inner class used to encapsulate data about a residue template definition."""
        def __init__(self, name):
            self.name = name
            self.atoms = []
579
            self.virtualSites = []
580
581
            self.bonds = []
            self.externalBonds = []
582
            self.overrideLevel = 0
583

584
585
586
587
588
589
590
        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
591
592
            msg =  "Atom name '%s' not found in residue template '%s'.\n" % (atom_name, self.name)
            msg += "Possible atom names are: %s" % str(atomIndices.keys())
593
594
            raise ValueError(msg)

595
        def addBond(self, atom1, atom2):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
596
            """Add a bond between two atoms in a template given their indices in the template."""
597
598
599
            self.bonds.append((atom1, atom2))
            self.atoms[atom1].bondedTo.append(atom2)
            self.atoms[atom2].bondedTo.append(atom1)
600

601
        def addBondByName(self, atom1_name, atom2_name):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
602
            """Add a bond between two atoms in a template given their atom names."""
603
604
605
606
607
            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
608
            """Designate that an atom in a residue template has an external bond, using atom index within template."""
609
610
611
612
            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
613
            """Designate that an atom in a residue template has an external bond, using atom name within template."""
614
615
616
            atom = self.getAtomIndexByName(atom_name)
            self.addExternalBond(atom)

617
    class _TemplateAtomData(object):
618
        """Inner class used to encapsulate data about an atom in a residue template definition."""
619
        def __init__(self, name, type, element, parameters={}):
620
621
622
            self.name = name
            self.type = type
            self.element = element
623
            self.parameters = parameters
624
625
626
            self.bondedTo = []
            self.externalBonds = 0

627
    class _BondData(object):
628
629
630
631
632
633
        """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
634

635
    class _VirtualSiteData(object):
636
        """Inner class used to encapsulate data about a virtual site."""
637
        def __init__(self, node, atomIndices):
638
639
640
            attrib = node.attrib
            self.type = attrib['type']
            if self.type == 'average2':
641
                numAtoms = 2
642
643
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
644
                numAtoms = 3
645
646
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
647
                numAtoms = 3
648
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
649
            elif self.type == 'localCoords':
650
                numAtoms = 3
651
652
653
654
                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'])]
655
656
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
657
658
659
660
661
662
            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)]
663
664
665
666
            if 'excludeWith' in attrib:
                self.excludeWith = int(attrib['excludeWith'])
            else:
                self.excludeWith = self.atoms[0]
667

668
669
670
671
672
    class _PatchData(object):
        """Inner class used to encapsulate data about a patch definition."""
        def __init__(self, name, numResidues):
            self.name = name
            self.numResidues = numResidues
673
674
            self.addedAtoms = [[] for i in range(numResidues)]
            self.changedAtoms = [[] for i in range(numResidues)]
675
676
677
678
679
            self.deletedAtoms = []
            self.addedBonds = []
            self.deletedBonds = []
            self.addedExternalBonds = []
            self.deletedExternalBonds = []
680

681
682
683
684
        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)))
685

686
            # Construct a new version of each template.
687

688
689
690
691
            newTemplates = []
            for index, template in enumerate(templates):
                newTemplate = ForceField._TemplateData("%s-%s" % (template.name, self.name))
                newTemplates.append(newTemplate)
692

693
                # Build the list of atoms in it.
694

695
696
                for atom in template.atoms:
                    if not any(deleted.name == atom.name and deleted.residue == index for deleted in self.deletedAtoms):
697
                        newTemplate.atoms.append(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
698
                for atom in self.addedAtoms[index]:
699
700
701
                    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))
702
703
704
705
706
                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))
707
                    newTemplate.atoms[newAtomIndex[atom.name]] = ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters)
708

709
                # Copy over the virtual sites, translating the atom indices.
710

711
712
713
714
715
716
717
                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)
718

719
                # Build the lists of bonds and external bonds.
720

721
722
723
724
725
                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]
726
                    if a1 in atomMap and a2 in atomMap and (a1.name, a2.name) not in deletedBonds and (a2.name, a1.name) not in deletedBonds:
727
728
729
730
                        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:
731
                        newTemplate.addExternalBond(indexMap[atom])
732
733
734
735
736
737
738
739
740
741
                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
742

743
744
745
746
747
748
749
750
751
752
753
    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

754
    class _AtomType(object):
755
756
757
758
759
760
761
        """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

762
    class _AtomTypeParameters(object):
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
        """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))


833
    def _getResidueTemplateMatches(self, res, bondedToAtom, templateSignatures=None, ignoreExternalBonds=False):
834
835
836
837
838
839
        """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.
840
841
        bondedToAtom : list of set of int
            bondedToAtom[i] is the set of atoms bonded to atom index i
842
843
844
845
846
847
848
849
850
851
852
853

        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
854
855
        if templateSignatures is None:
            templateSignatures = self._templateSignatures
856
        signature = _createResidueSignature([atom.element for atom in res.atoms()])
857
        if signature in templateSignatures:
858
            allMatches = []
859
            for t in templateSignatures[signature]:
860
                match = _matchResidue(res, t, bondedToAtom, ignoreExternalBonds)
861
862
863
864
865
866
                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:
867
                raise Exception('Multiple matching templates found for residue %d (%s): %s.' % (res.index+1, res.name, ', '.join(match[0].name for match in allMatches)))
868
869
        return [template, matches]

870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
    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())
887
888
889
        for (atom1, atom2) in topology.bonds():
            bondedToAtom[atom1.index].add(atom2.index)
            bondedToAtom[atom2.index].add(atom1.index)
890
        return bondedToAtom
891

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

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

897
898
899
        Parameters
        ----------
        topology : Topology
900
            The Topology whose residues are to be checked against the forcefield residue templates.
901
902
903
904

        Returns
        -------
        unmatched_residues : list of Residue
905
            List of Residue objects from `topology` for which no forcefield residue templates are available.
906
907
908
909
910
            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.
911
        bondedToAtom = self._buildBondedToAtomList(topology)
912
        unmatched_residues = list() # list of unmatched residues
913
914
915
916
917
918
        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)
919
920
921

        return unmatched_residues

922
    def getMatchingTemplates(self, topology, ignoreExternalBonds=False):
923
924
925
926
927
928
929
930
        """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.
931
932
        ignoreExternalBonds : bool=False
            If true, ignore external bonds when matching residues to templates.
933
934
935
936
937
938
939
940
941
942
943
944
945
        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.
946
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
947
948
949
950
951
952
953
954
            # 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

955
956
    def generateTemplatesForUnmatchedResidues(self, topology):
        """Generate forcefield residue templates for residues in specified topology for which no forcefield templates are available.
957

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

960
961
962
        Parameters
        ----------
        topology : Topology
963
            The Topology whose residues are to be checked against the forcefield residue templates.
964
965
966

        Returns
        -------
967
968
969
970
971
972
        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]`
973
974
975
976
977

        """
        # Get a non-unique list of unmatched residues.
        unmatched_residues = self.getUnmatchedResidues(topology)
        # Generate a unique list of unmatched residues by comparing fingerprints.
978
        bondedToAtom = self._buildBondedToAtomList(topology)
979
980
        unique_unmatched_residues = list() # list of unique unmatched Residue objects from topology
        templates = list() # corresponding _TemplateData templates
981
982
983
        signatures = set()
        for residue in unmatched_residues:
            signature = _createResidueSignature([ atom.element for atom in residue.atoms() ])
984
            template = _createResidueTemplate(residue)
985
986
987
988
            is_unique = True
            if signature in signatures:
                # Signature is the same as an existing residue; check connectivity.
                for check_residue in unique_unmatched_residues:
989
                    matches = _matchResidue(check_residue, template, bondedToAtom, False)
990
991
992
993
994
995
                    if matches is not None:
                        is_unique = False
            if is_unique:
                # Residue is unique.
                unique_unmatched_residues.append(residue)
                signatures.add(signature)
996
                templates.append(template)
997

998
        return [templates, unique_unmatched_residues]
999

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

Robert McGibbon's avatar
Robert McGibbon committed
1005
1006
1007
1008
1009
1010
        Parameters
        ----------
        topology : Topology
            The Topology for which to create a System
        nonbondedMethod : object=NoCutoff
            The method to use for nonbonded interactions.  Allowed values are
Peter Eastman's avatar
Peter Eastman committed
1011
            NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME, or LJPME.
Robert McGibbon's avatar
Robert McGibbon committed
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
        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.
1026
        residueTemplates : dict=dict()
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
            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).
        ignoreExternalBonds : boolean=False
            If true, ignore external bonds when matching residues to templates.  This is
            useful when the Topology represents one piece of a larger molecule, so chains are
            not terminated properly.  This option can create ambiguities where multiple
            templates match the same residue.  If that happens, use the residueTemplates
            argument to specify which one to use.
Robert McGibbon's avatar
Robert McGibbon committed
1038
        args
1039
1040
1041
            Arbitrary additional keyword arguments may also be specified.
            This allows extra parameters to be specified that are specific to
            particular force fields.
Robert McGibbon's avatar
Robert McGibbon committed
1042
1043
1044
1045
1046

        Returns
        -------
        system
            the newly created System
1047
1048
        """
        data = ForceField._SystemData()
1049
        data.atoms = list(topology.atoms())
1050
1051
        for atom in data.atoms:
            data.excludeAtomWith.append([])
1052
1053

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

1055
        for bond in topology.bonds():
1056
            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
1057
1058

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

1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
        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
1072

1073
        unmatchedResidues = []
1074
1075
        for chain in topology.chains():
            for res in chain.residues():
1076
1077
1078
                if res in residueTemplates:
                    tname = residueTemplates[res]
                    template = self._templates[tname]
1079
                    matches = _matchResidue(res, template, bondedToAtom, ignoreExternalBonds)
1080
1081
1082
1083
                    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.
1084
                    [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1085
                if matches is None:
1086
1087
1088
                    unmatchedResidues.append(res)
                else:
                    data.recordMatchedAtomParameters(res, template, matches)
1089

1090
        # Try to apply patches to find matches for any unmatched residues.
1091

1092
        if len(unmatchedResidues) > 0:
1093
            unmatchedResidues = _applyPatchesToMatchResidues(self, data, unmatchedResidues, bondedToAtom, ignoreExternalBonds)
1094

1095
1096
        # 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).
1097

1098
1099
        for res in unmatchedResidues:
            # A template might have been generated on an earlier iteration of this loop.
1100
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1101
1102
1103
1104
1105
            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.
1106
                        [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
                        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)
1117
1118

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

1120
1121
        sys = mm.System()
        for atom in topology.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1122
            # Look up the atom type name, returning a helpful error message if it cannot be found.
1123
1124
1125
1126
            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
1127
            # Look up the type name in the list of registered atom types, returning a helpful error message if it cannot be found.
1128
1129
1130
1131
            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
1132
1133

            # Add the particle to the OpenMM system.
1134
            mass = self._atomTypes[typename].mass
1135
            sys.addParticle(mass)
1136

1137
        # Adjust hydrogen masses if requested.
1138

1139
        if hydrogenMass is not None:
1140
1141
            if not unit.is_quantity(hydrogenMass):
                hydrogenMass *= unit.dalton
1142
1143
1144
1145
1146
1147
1148
            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
1149

1150
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
1151

1152
1153
1154
        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
1155
        elif nonbondedMethod not in [NoCutoff, CutoffNonPeriodic]:
1156
1157
1158
            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
1159

1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
        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
1175

1176
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1177

1178
1179
1180
        uniquePropers = set()
        for angle in data.angles:
            for atom in bondedToAtom[angle[0]]:
pgrinaway's avatar
pgrinaway committed
1181
                if atom not in angle:
1182
1183
1184
1185
1186
                    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
1187
                if atom not in angle:
1188
1189
1190
1191
1192
                    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
1193

1194
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1195

1196
1197
1198
1199
1200
        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
1201

1202
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1203

1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
        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
1218

1219
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1220

1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
        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
1242

1243
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
1244

1245
        for atom in data.virtualSites:
1246
            (site, atoms, excludeWith) = data.virtualSites[atom]
1247
            index = atom.index
1248
            data.excludeAtomWith[excludeWith].append(index)
1249
            if site.type == 'average2':
1250
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(atoms[0], atoms[1], site.weights[0], site.weights[1]))
1251
            elif site.type == 'average3':
1252
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
1253
            elif site.type == 'outOfPlane':
1254
1255
1256
1257
1258
1259
1260
                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
1261

1262
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
1263

1264
1265
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
1266
1267
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
1268

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

peastman's avatar
peastman committed
1271
1272
1273
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
1274

1275
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
1276

1277
        for script in self._scripts:
1278
            exec(script, locals())
1279
1280
1281
        return sys


1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
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))
1306
1307
1308
    for atom in data.atoms:
        for child in data.excludeAtomWith[atom.index]:
            bondIndices.append((child, atom.index))
1309
1310
    return bondIndices

1311
1312
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
1313
1314
    counts = {}
    for element in elements:
1315
        if element in counts:
1316
1317
1318
            counts[element] += 1
        else:
            counts[element] = 1
1319
1320
1321
1322
1323
1324
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
1325
1326
    sig = []
    for c in counts:
1327
1328
        if c is not None:
            sig.append((c, counts[c]))
1329
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
1330

1331
    # Convert it to a string.
1332
1333

    s = ''
1334
    for element, count in sig:
1335
1336
1337
        s += element.symbol+str(count)
    return s

1338
def _matchResidue(res, template, bondedToAtom, ignoreExternalBonds=False):
1339
    """Determine whether a residue matches a template and return a list of corresponding atoms.
Justin MacCallum's avatar
Justin MacCallum committed
1340

Robert McGibbon's avatar
Robert McGibbon committed
1341
1342
1343
1344
1345
1346
1347
1348
    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
1349
1350
    ignoreExternalBonds : bool
        If true, ignore external bonds when matching templates
Robert McGibbon's avatar
Robert McGibbon committed
1351
1352
1353

    Returns
    -------
1354
1355
1356
    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
1357
1358
    """
    atoms = list(res.atoms())
peastman's avatar
peastman committed
1359
1360
    numAtoms = len(atoms)
    if numAtoms != len(template.atoms):
1361
        return None
Justin MacCallum's avatar
Justin MacCallum committed
1362

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

1365
    renumberAtoms = {}
peastman's avatar
peastman committed
1366
    for i in range(numAtoms):
1367
        renumberAtoms[atoms[i].index] = i
1368
1369
1370
    bondedTo = []
    externalBonds = []
    for atom in atoms:
1371
        bonds = [renumberAtoms[x] for x in bondedToAtom[atom.index] if x in renumberAtoms]
1372
        bondedTo.append(bonds)
1373
        externalBonds.append(0 if ignoreExternalBonds else len([x for x in bondedToAtom[atom.index] if x not in renumberAtoms]))
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385

    # 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):
1386
        key = (atom.element, len(atom.bondedTo), 0 if ignoreExternalBonds else atom.externalBonds)
1387
1388
1389
1390
1391
        if key not in templateTypeCount:
            templateTypeCount[key] = 1
        templateTypeCount[key] += 1
    if residueTypeCount != templateTypeCount:
        return None
1392

peastman's avatar
peastman committed
1393
    # Identify template atoms that could potentially be matches for each atom.
1394

peastman's avatar
peastman committed
1395
1396
1397
1398
1399
1400
1401
    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
1402
            if not ignoreExternalBonds and atom.externalBonds != externalBonds[i]:
peastman's avatar
peastman committed
1403
1404
1405
1406
1407
                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.
1408

peastman's avatar
peastman committed
1409
1410
    searchOrder = []
    atomsToOrder = set(range(numAtoms))
1411
1412
    efficientAtomSet = set()
    efficientAtomHeap = []
peastman's avatar
peastman committed
1413
    while len(atomsToOrder) > 0:
1414
        if len(efficientAtomSet) == 0:
peastman's avatar
peastman committed
1415
1416
1417
1418
1419
1420
            fewestNeighbors = numAtoms+1
            for i in atomsToOrder:
                if len(candidates[i]) < fewestNeighbors:
                    nextAtom = i
                    fewestNeighbors = len(candidates[i])
        else:
1421
1422
            nextAtom = heappop(efficientAtomHeap)[1]
            efficientAtomSet.remove(nextAtom)
peastman's avatar
peastman committed
1423
1424
1425
1426
        searchOrder.append(nextAtom)
        atomsToOrder.remove(nextAtom)
        for i in bondedTo[nextAtom]:
            if i in atomsToOrder:
1427
1428
1429
                if i not in efficientAtomSet:
                    efficientAtomSet.add(i)
                    heappush(efficientAtomHeap, (len(candidates[i]), i))
peastman's avatar
peastman committed
1430
1431
1432
1433
1434
    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]
1435
1436
1437

    # Recursively match atoms.

peastman's avatar
peastman committed
1438
1439
1440
1441
    matches = numAtoms*[0]
    hasMatch = numAtoms*[False]
    if _findAtomMatches(template, bondedTo, matches, hasMatch, candidates, 0):
        return [matches[inverseSearchOrder[i]] for i in range(numAtoms)]
1442
1443
1444
    return None


peastman's avatar
peastman committed
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
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):
1456
    """This is called recursively from inside _matchResidue() to identify matching atoms."""
peastman's avatar
peastman committed
1457
    if position == len(matches):
1458
        return True
peastman's avatar
peastman committed
1459
    for i in _getAtomMatchCandidates(template, bondedTo, matches, candidates, position):
1460
        atom = template.atoms[i]
peastman's avatar
peastman committed
1461
        if not hasMatch[i] and i in candidates[position]:
1462
            # See if the bonds for this identification are consistent
Justin MacCallum's avatar
Justin MacCallum committed
1463

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

1468
1469
                matches[position] = i
                hasMatch[i] = True
peastman's avatar
peastman committed
1470
                if _findAtomMatches(template, bondedTo, matches, hasMatch, candidates, position+1):
1471
1472
1473
1474
1475
                    return True
                hasMatch[i] = False
    return False


1476
def _applyPatchesToMatchResidues(forcefield, data, residues, bondedToAtom, ignoreExternalBonds):
1477
1478
1479
1480
    """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.
1481

1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
    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]
1497

1498
    # Now see if any of those templates matches any of the residues.
1499

1500
1501
    unmatchedResidues = []
    for res in residues:
1502
        [template, matches] = forcefield._getResidueTemplateMatches(res, bondedToAtom, patchedTemplateSignatures, ignoreExternalBonds)
1503
1504
1505
1506
        if matches is None:
            unmatchedResidues.append(res)
        else:
            data.recordMatchedAtomParameters(res, template, matches)
1507
1508
    if len(unmatchedResidues) == 0:
        return []
1509

1510
1511
1512
    # 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.
1513

1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
    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]
1529

1530
    # Record which unmatched residues are bonded to each other.
1531

1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
    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)
1542

1543
1544
    # 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.
1545

1546
1547
1548
1549
    clusterSize = 2
    clusters = bonds
    while clusterSize <= maxPatchSize:
        # Try to apply patches to clusters of this size.
1550

1551
1552
1553
        for patchName in patches:
            patch = forcefield._patches[patchName]
            if patch.numResidues == clusterSize:
1554
                matchedClusters = _matchToMultiResiduePatchedTemplates(data, clusters, patch, patches[patchName], bondedToAtom, ignoreExternalBonds)
1555
1556
1557
1558
1559
1560
                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.
1561

1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
        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

1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
    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
1589

1590
    # Call this function recursively to generate combinations of patches.
1591

1592
1593
1594
1595
1596
1597
    if index+1 < len(patches):
        _generatePatchedSingleResidueTemplates(template, patches, index+1, newTemplates)
        if patchedTemplate is not None:
            _generatePatchedSingleResidueTemplates(patchedTemplate, patches, index+1, newTemplates)


1598
def _matchToMultiResiduePatchedTemplates(data, clusters, patch, residueTemplates, bondedToAtom, ignoreExternalBonds):
1599
1600
1601
    """Apply a multi-residue patch to templates, then try to match them against clusters of residues."""
    matchedClusters = []
    selectedTemplates = [None]*patch.numResidues
1602
    _applyMultiResiduePatch(data, clusters, patch, residueTemplates, selectedTemplates, 0, matchedClusters, bondedToAtom, ignoreExternalBonds)
1603
1604
1605
    return matchedClusters


1606
def _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index, matchedClusters, bondedToAtom, ignoreExternalBonds):
1607
1608
1609
1610
1611
    """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
1612
            _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index+1, matchedClusters, bondedToAtom, ignoreExternalBonds)
1613
1614
1615
    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.
1616

1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
        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):
1629
                    matches = _matchResidue(residue, template, bondedToAtom, ignoreExternalBonds)
1630
1631
1632
1633
1634
1635
1636
                    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.
1637

1638
1639
1640
1641
                    for i in range(patch.numResidues):
                        data.recordMatchedAtomParameters(residues[i], patchedTemplates[i], residueMatches[i])
                    newlyMatchedClusters.append(cluster)
                    break
1642

1643
        # Record which clusters were successfully matched.
1644

1645
1646
1647
        matchedClusters += newlyMatchedClusters
        for cluster in newlyMatchedClusters:
            clusters.remove(cluster)
1648

1649

1650
1651
1652
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()])
1653
    numResidueAtoms = sum(residueCounts.values())
1654
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
1655

1656
    # Loop over templates and see how closely each one might match.
1657

1658
1659
1660
1661
1662
1663
    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])
1664

1665
        # Does the residue have any atoms that clearly aren't in the template?
1666

1667
1668
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
1669

1670
        # If there are too many missing atoms, discard this template.
1671

1672
        numTemplateAtoms = sum(templateCounts.values())
1673
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
1674
1675
1676
1677
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
1678

1679
1680
        # 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.
1681

1682
1683
1684
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
1685

1686
        # Accept this as our new best match.
1687

1688
1689
1690
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
1691
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
1692

1693
    # Return an appropriate error message.
1694

1695
    if numBestMatchAtoms == numResidueAtoms:
1696
1697
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
1698
1699
1700
1701
            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:
1702
1703
1704
1705
1706
            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)
1707
1708
1709
        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.'

1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
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
1728
        template.atoms.append(ForceField._TemplateAtomData(atom.name, None, atom.element))
1729
1730
1731
1732
1733
1734
1735
1736
1737
    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
1738

1739
1740
1741
1742
1743
# 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.

1744
## @private
1745
class HarmonicBondGenerator(object):
1746
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1747

1748
1749
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
1750
        self.bondsForAtomType = defaultdict(set)
1751
1752
1753
1754
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
1755

1756
1757
1758
    def registerBond(self, parameters):
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
peastman's avatar
peastman committed
1759
            index = len(self.types1)
1760
1761
            self.types1.append(types[0])
            self.types2.append(types[1])
peastman's avatar
peastman committed
1762
1763
1764
1765
            for t in types[0]:
                self.bondsForAtomType[t].add(index)
            for t in types[1]:
                self.bondsForAtomType[t].add(index)
1766
1767
            self.length.append(_convertParameterToNumber(parameters['length']))
            self.k.append(_convertParameterToNumber(parameters['k']))
Justin MacCallum's avatar
Justin MacCallum committed
1768

1769
1770
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1771
1772
1773
1774
1775
1776
        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]
1777
        for bond in element.findall('Bond'):
1778
            generator.registerBond(bond.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1779

1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
    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]]
peastman's avatar
peastman committed
1791
            for i in self.bondsForAtomType[type1]:
1792
1793
1794
1795
1796
                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:
1797
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
1798
1799
1800
1801
1802
1803
1804
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


1805
## @private
1806
class HarmonicAngleGenerator(object):
1807
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1808

1809
1810
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
1811
        self.anglesForAtom2Type = defaultdict(list)
1812
1813
1814
1815
1816
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
1817

1818
1819
1820
    def registerAngle(self, parameters):
        types = self.ff._findAtomTypes(parameters, 3)
        if None not in types:
peastman's avatar
peastman committed
1821
            index = len(self.types1)
1822
1823
1824
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.types3.append(types[2])
peastman's avatar
peastman committed
1825
1826
            for t in types[1]:
                self.anglesForAtom2Type[t].append(index)
1827
1828
1829
            self.angle.append(_convertParameterToNumber(parameters['angle']))
            self.k.append(_convertParameterToNumber(parameters['k']))

1830
1831
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1832
1833
1834
1835
1836
1837
        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]
1838
        for angle in element.findall('Angle'):
1839
            generator.registerAngle(angle.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1840

1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
    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]]]
peastman's avatar
peastman committed
1853
            for i in self.anglesForAtom2Type[type2]:
1854
1855
1856
1857
1858
1859
                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
1860

1861
1862
1863
1864
1865
1866
1867
1868
1869
                        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
1870

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

1873
1874
1875
1876
1877
                        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]))
1878
                                data.addConstraint(sys, angle[0], angle[2], length)
1879
1880
1881
1882
1883
1884
1885
                    elif self.k[i] != 0:
                        force.addAngle(angle[0], angle[1], angle[2], self.angle[i], self.k[i])
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


1886
## @private
1887
class PeriodicTorsion(object):
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
    """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 = []
1898
        self.ordering = 'default'
1899

1900
## @private
1901
class PeriodicTorsionGenerator(object):
1902
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1903

1904
1905
    def __init__(self, forcefield):
        self.ff = forcefield
1906
1907
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
1908

1909
1910
1911
1912
1913
    def registerProperTorsion(self, parameters):
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
            self.proper.append(torsion)

1914
    def registerImproperTorsion(self, parameters, ordering='default'):
1915
1916
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
1917
1918
1919
1920
1921
1922
            if ordering == 'default':
                pass
            elif ordering == 'amber':
                torsion.ordering = ordering
            else:
                raise ValueError('Illegal ordering type for improper torsion %s' % torsion)
1923
1924
            self.improper.append(torsion)

1925
1926
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1927
1928
1929
1930
1931
1932
        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]
1933
        for torsion in element.findall('Proper'):
1934
            generator.registerProperTorsion(torsion.attrib)
1935
        for torsion in element.findall('Improper'):
1936
1937
1938
1939
            if 'ordering' in element.attrib:
                generator.registerImproperTorsion(torsion.attrib, element.attrib['ordering'])
            else:
                generator.registerImproperTorsion(torsion.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1940

1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
    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]]]
1976
            match = None
1977
1978
1979
1980
1981
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
1982
1983
1984
1985
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
1986
1987
1988
                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:
1989
                            if tordef.ordering == 'default':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
                                # 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)
                                match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
                                break
2003
                            elif tordef.ordering == 'amber':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
                                # topology atom indexes
                                a2 = torsion[t2[1]]
                                a3 = torsion[t3[1]]
                                a4 = torsion[t4[1]]
                                # residue indexes
                                r2 = data.atoms[a2].residue.index
                                r3 = data.atoms[a3].residue.index
                                r4 = data.atoms[a4].residue.index
                                # template atom indexes
                                ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                # elements
                                e2 = data.atoms[a2].element
                                e3 = data.atoms[a3].element
                                e4 = data.atoms[a4].element
2020
2021
                                if not hasWildcard:
                                    if t2[0] == t4[0] and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
2022
                                        (a2, a4) = (a4, a2)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2023
2024
2025
2026
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
2027
                                    if t3[0] == t4[0] and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
2028
                                        (a3, a4) = (a4, a3)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2029
2030
2031
2032
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
2033
                                    if t2[0] == t3[0] and (r2 > r3 or (r2 == r3 and ta2 > ta3)):
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
2034
                                        (a2, a3) = (a3, a2)
2035
2036
                                else:
                                    if e2 == e4 and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
2037
                                        (a2, a4) = (a4, a2)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2038
2039
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
2040
                                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2041
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
2042
                                    if e3 == e4 and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
2043
                                        (a3, a4) = (a4, a3)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2044
2045
2046
2047
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
2048
                                    if r2 > r3 or (r2 == r3 and ta2 > ta3):
Rafal P. Wiewiora's avatar
fixes  
Rafal P. Wiewiora committed
2049
                                        (a2, a3) = (a3, a2)
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2050
2051
                                match = (a2, a3, torsion[0], a4, tordef)
                                break
2052
2053
2054
2055
2056
            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])
2057
2058
2059
2060

parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


2061
## @private
2062
class RBTorsion(object):
2063
2064
    """An RBTorsion records the information for a Ryckaert-Bellemans torsion definition."""

2065
    def __init__(self, types, c, ordering='default'):
2066
2067
2068
2069
2070
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.c = c
2071
2072
2073
2074
        if ordering == 'default' or ordering == 'amber':
            self.ordering = ordering
        else:
            raise ValueError('Illegal ordering type for RBTorsion (%s,%s,%s,%s)' % (types[0], types[1], types[2], types[3]))
2075

2076
## @private
2077
class RBTorsionGenerator(object):
2078
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2079

2080
2081
    def __init__(self, forcefield):
        self.ff = forcefield
2082
2083
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
2084

2085
2086
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2087
2088
2089
2090
2091
2092
        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]
2093
        for torsion in element.findall('Proper'):
2094
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2095
            if None not in types:
2096
2097
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
2098
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2099
            if None not in types:
2100
2101
2102
2103
                if 'ordering' in element.attrib:
                    generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)], element.attrib['ordering']))
                else:
                    generator.improper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
Justin MacCallum's avatar
Justin MacCallum committed
2104

2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
    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]]]
2138
            match = None
2139
2140
2141
2142
2143
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
2144
2145
2146
2147
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
2148
2149
2150
                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:
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
                            if tordef.ordering == 'default':
                                if hasWildcard:
                                    # 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)
                                    match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
                                else:
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2166
2167
                                    # There are no wildcards, so the order is unambiguous.
                                    match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
                                break
                            elif tordef.ordering == 'amber':
                                # topology atom indexes
                                a2 = torsion[t2[1]]
                                a3 = torsion[t3[1]]
                                a4 = torsion[t4[1]]
                                # residue indexes
                                r2 = data.atoms[a2].residue.index
                                r3 = data.atoms[a3].residue.index
                                r4 = data.atoms[a4].residue.index
                                # template atom indexes
                                ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                # elements
2183
                                e2 = data.atoms[a2].element
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
                                e3 = data.atoms[a3].element
                                e4 = data.atoms[a4].element
                                if not hasWildcard:
                                    if t2[0] == t4[0] and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                        (a2, a4) = (a4, a2)
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                    if t3[0] == t4[0] and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
                                        (a3, a4) = (a4, a3)
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                    if t2[0] == t3[0] and (r2 > r3 or (r2 == r3 and ta2 > ta3)):
                                        (a2, a3) = (a3, a2)
                                else:
                                    if e2 == e4 and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                        (a2, a4) = (a4, a2)
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                    if e3 == e4 and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
                                        (a3, a4) = (a4, a3)
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                    if r2 > r3 or (r2 == r3 and ta2 > ta3):
                                        (a2, a3) = (a3, a2)
                                match = (a2, a3, torsion[0], a4, tordef)
                                break
2218
2219
2220
            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])
2221
2222
2223
2224

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


2225
## @private
2226
class CMAPTorsion(object):
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
    """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

2237
## @private
2238
class CMAPTorsionGenerator(object):
2239
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2240

2241
2242
    def __init__(self, forcefield):
        self.ff = forcefield
2243
2244
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
2245

2246
2247
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2248
2249
2250
2251
2252
2253
        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]
2254
2255
2256
2257
2258
2259
2260
        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'):
2261
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
2262
            if None not in types:
2263
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
2264

2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
    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
2275

2276
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
2277

2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
        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


2321
## @private
2322
class NonbondedGenerator(object):
2323
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2324

2325
2326
    SCALETOL = 1e-5

2327
2328
    def __init__(self, forcefield, coulomb14scale, lj14scale):
        self.ff = forcefield
2329
2330
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
2331
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
2332

2333
    def registerAtom(self, parameters):
2334
        self.params.registerAtom(parameters)
2335

2336
2337
2338
2339
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
        if len(existing) == 0:
2340
2341
            generator = NonbondedGenerator(ff, float(element.attrib['coulomb14scale']), float(element.attrib['lj14scale']))
            ff.registerGenerator(generator)
2342
2343
2344
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2345
2346
            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
2347
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
2348
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2349

2350
2351
2352
2353
2354
    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,
Peter Eastman's avatar
Peter Eastman committed
2355
2356
                     PME:mm.NonbondedForce.PME,
                     LJPME:mm.NonbondedForce.LJPME}
2357
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2358
            raise ValueError('Illegal nonbonded method for NonbondedForce')
2359
2360
        force = mm.NonbondedForce()
        for atom in data.atoms:
2361
2362
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
2363
2364
2365
2366
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
2367
2368
        if 'useDispersionCorrection' in args:
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
peastman's avatar
peastman committed
2369
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
2370

peastman's avatar
peastman committed
2371
    def postprocessSystem(self, sys, data, args):
2372
2373
        # Create the exceptions.

2374
        bondIndices = _findBondsForExclusions(data, sys)
peastman's avatar
peastman committed
2375
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
2376
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
2377
2378
2379

parsers["NonbondedForce"] = NonbondedGenerator.parseElement

2380
2381

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

ChayaSt's avatar
ChayaSt committed
2385
    def __init__(self, forcefield, lj14scale):
ChayaSt's avatar
ChayaSt committed
2386
        self.ff = forcefield
2387
        self.nbfixTypes = {}
ChayaSt's avatar
ChayaSt committed
2388
        self.lj14scale = lj14scale
2389
        self.ljTypes = ForceField._AtomTypeParameters(forcefield, 'LennardJonesForce', 'Atom', ('sigma', 'epsilon'))
ChayaSt's avatar
ChayaSt committed
2390

ChayaSt's avatar
ChayaSt committed
2391
    def registerNBFIX(self, parameters):
ChayaSt's avatar
ChayaSt committed
2392
2393
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
2394
2395
            type1 = types[0][0]
            type2 = types[1][0]
ChayaSt's avatar
ChayaSt committed
2396
2397
            epsilon = _convertParameterToNumber(parameters['epsilon'])
            sigma = _convertParameterToNumber(parameters['sigma'])
2398
2399
            self.nbfixTypes[(type1, type2)] = [sigma, epsilon]
            self.nbfixTypes[(type2, type1)] = [sigma, epsilon]
2400

ChayaSt's avatar
ChayaSt committed
2401
    def registerLennardJones(self, parameters):
2402
        self.ljTypes.registerAtom(parameters)
ChayaSt's avatar
ChayaSt committed
2403
2404
2405

    @staticmethod
    def parseElement(element, ff):
ChayaSt's avatar
ChayaSt committed
2406
        existing = [f for f in ff._forces if isinstance(f, LennardJonesGenerator)]
ChayaSt's avatar
ChayaSt committed
2407
        if len(existing) == 0:
ChayaSt's avatar
ChayaSt committed
2408
            generator = LennardJonesGenerator(ff, float(element.attrib['lj14scale']))
ChayaSt's avatar
ChayaSt committed
2409
2410
            ff.registerGenerator(generator)
        else:
ChayaSt's avatar
ChayaSt committed
2411
            # Multiple <LennardJonesForce> tags were found, probably in different files
ChayaSt's avatar
ChayaSt committed
2412
            generator = existing[0]
2413
2414
            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
2415
2416
        for LJ in element.findall('Atom'):
            generator.registerLennardJones(LJ.attrib)
ChayaSt's avatar
ChayaSt committed
2417
        for Nbfix in element.findall('NBFixPair'):
ChayaSt's avatar
ChayaSt committed
2418
            generator.registerNBFIX(Nbfix.attrib)
ChayaSt's avatar
ChayaSt committed
2419

ChayaSt's avatar
ChayaSt committed
2420
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
ChayaSt's avatar
ChayaSt committed
2421
        # First derive the lookup tables
2422
2423
2424
2425
2426
2427

        nbfixTypeSet = set().union(*self.nbfixTypes)
        ljIndexList = [None]*len(data.atoms)
        numLjTypes = 0
        ljTypeList = []
        typeMap = {}
ChayaSt's avatar
ChayaSt committed
2428
        for i, atom in enumerate(data.atoms):
ChayaSt's avatar
ChayaSt committed
2429
            atype = data.atomType[atom]
2430
2431
2432
2433
            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]
2434
            else:
2435
2436
2437
2438
2439
2440
2441
                typeMap[values] = numLjTypes
                ljIndexList[i] = numLjTypes
                numLjTypes += 1
                ljTypeList.append(atype)
        reverseMap = [0]*len(typeMap)
        for typeValue in typeMap:
            reverseMap[typeMap[typeValue]] = typeValue
2442

ChayaSt's avatar
ChayaSt committed
2443
        # Now everything is assigned. Create the A- and B-coefficient arrays
2444

2445
        acoef = [0]*(numLjTypes*numLjTypes)
ChayaSt's avatar
ChayaSt committed
2446
        bcoef = acoef[:]
2447
2448
2449
2450
2451
2452
2453
2454
2455
        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
2456
                    continue
ChayaSt's avatar
ChayaSt committed
2457
                else:
2458
2459
2460
2461
2462
2463
2464
2465
2466
                    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
2467
        self.force.addPerParticleParameter('type')
Peter Eastman's avatar
Peter Eastman committed
2468
        if nonbondedMethod in [CutoffPeriodic, Ewald, PME, LJPME]:
ChayaSt's avatar
ChayaSt committed
2469
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
ChayaSt's avatar
ChayaSt committed
2470
        elif nonbondedMethod is NoCutoff:
ChayaSt's avatar
ChayaSt committed
2471
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
ChayaSt's avatar
ChayaSt committed
2472
        elif nonbondedMethod is CutoffNonPeriodic:
ChayaSt's avatar
ChayaSt committed
2473
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic)
ChayaSt's avatar
ChayaSt committed
2474
        else:
2475
2476
            raise AssertionError('Unrecognized nonbonded method [%s]' % nonbondedMethod)

ChayaSt's avatar
ChayaSt committed
2477
        # Add the particles
2478
2479

        for i in ljIndexList:
ChayaSt's avatar
cleanup  
ChayaSt committed
2480
            self.force.addParticle((i,))
ChayaSt's avatar
ChayaSt committed
2481
2482
        self.force.setUseLongRangeCorrection(True)
        self.force.setCutoffDistance(nonbondedCutoff)
ChayaSt's avatar
ChayaSt committed
2483
2484
2485
        sys.addForce(self.force)

    def postprocessSystem(self, sys, data, args):
ChayaSt's avatar
ChayaSt committed
2486
        # Create the exceptions.
2487

2488
        bondIndices = _findBondsForExclusions(data, sys)
2489
2490
        if self.lj14scale == 1:
            # Just exclude the 1-2 and 1-3 interactions.
2491

peastman's avatar
Bug fix  
peastman committed
2492
            self.force.createExclusionsFromBonds(bondIndices, 2)
2493
2494
2495
2496
2497
2498
        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.
2499

2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
                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
2520

ChayaSt's avatar
ChayaSt committed
2521
parsers["LennardJonesForce"] = LennardJonesGenerator.parseElement
2522

2523

2524
## @private
2525
class GBSAOBCGenerator(object):
2526
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2527

2528
2529
    def __init__(self, forcefield):
        self.ff = forcefield
2530
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
2531

2532
    def registerAtom(self, parameters):
2533
        self.params.registerAtom(parameters)
2534

2535
2536
    @staticmethod
    def parseElement(element, ff):
2537
2538
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
2539
2540
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
2541
2542
2543
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2544
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2545

2546
2547
2548
2549
2550
    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
2551
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
2552
2553
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
2554
2555
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
2556
2557
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2558
2559
2560
2561
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
2562
2563
        sys.addForce(force)

2564
2565
    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
2566

2567
2568
2569
2570
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

2571
2572
2573
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


2574
## @private
2575
class CustomBondGenerator(object):
2576
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2577

2578
2579
    def __init__(self, forcefield):
        self.ff = forcefield
2580
2581
2582
2583
2584
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2585

2586
2587
    @staticmethod
    def parseElement(element, ff):
2588
2589
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
2590
2591
2592
2593
2594
2595
        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'):
2596
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2597
            if None not in types:
2598
2599
2600
                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
2601

2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
    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


2622
## @private
2623
class CustomAngleGenerator(object):
2624
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2625

2626
2627
    def __init__(self, forcefield):
        self.ff = forcefield
2628
2629
2630
2631
2632
2633
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2634

2635
2636
    @staticmethod
    def parseElement(element, ff):
2637
2638
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
2639
2640
2641
2642
2643
2644
        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'):
2645
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2646
            if None not in types:
2647
2648
2649
2650
                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
2651

2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
    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


2674
## @private
2675
class CustomTorsion(object):
2676
2677
    """A CustomTorsion records the information for a custom torsion definition."""

2678
    def __init__(self, types, paramValues, ordering='default'):
2679
2680
2681
2682
2683
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.paramValues = paramValues
2684
2685
2686
2687
        if ordering == 'default' or ordering == 'amber':
            self.ordering = ordering
        else:
            raise ValueError('Illegal ordering type for CustomTorsion (%s,%s,%s,%s)' % (types[0], types[1], types[2], types[3]))
2688

2689
## @private
2690
class CustomTorsionGenerator(object):
2691
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2692

2693
2694
    def __init__(self, forcefield):
        self.ff = forcefield
2695
2696
2697
2698
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
2699

2700
2701
    @staticmethod
    def parseElement(element, ff):
2702
2703
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
2704
2705
2706
2707
2708
2709
        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'):
2710
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2711
            if None not in types:
2712
2713
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
2714
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2715
            if None not in types:
2716
2717
2718
2719
                if 'ordering' in element.attrib:
                    generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams], element.attrib['ordering']))
                else:
                    generator.improper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
Justin MacCallum's avatar
Justin MacCallum committed
2720

2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
    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]]]
2753
            match = None
2754
2755
2756
2757
2758
            for tordef in self.improper:
                types1 = tordef.types1
                types2 = tordef.types2
                types3 = tordef.types3
                types4 = tordef.types4
2759
2760
2761
2762
                hasWildcard = (wildcard in (types1, types2, types3, types4))
                if match is not None and hasWildcard:
                    # Prefer specific definitions over ones with wildcards
                    continue
2763
2764
2765
                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:
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
                            if tordef.ordering == 'default':
                                if hasWildcard:
                                    # 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)
                                    match = (a1, a2, torsion[0], torsion[t4[1]], tordef)
                                else:
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2781
2782
                                    # There are no wildcards, so the order is unambiguous.
                                    match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
                                break
                            elif tordef.ordering == 'amber':
                                # topology atom indexes
                                a2 = torsion[t2[1]]
                                a3 = torsion[t3[1]]
                                a4 = torsion[t4[1]]
                                # residue indexes
                                r2 = data.atoms[a2].residue.index
                                r3 = data.atoms[a3].residue.index
                                r4 = data.atoms[a4].residue.index
                                # template atom indexes
                                ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                # elements
2798
                                e2 = data.atoms[a2].element
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
                                e3 = data.atoms[a3].element
                                e4 = data.atoms[a4].element
                                if not hasWildcard:
                                    if t2[0] == t4[0] and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                        (a2, a4) = (a4, a2)
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                    if t3[0] == t4[0] and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
                                        (a3, a4) = (a4, a3)
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                    if t2[0] == t3[0] and (r2 > r3 or (r2 == r3 and ta2 > ta3)):
                                        (a2, a3) = (a3, a2)
                                else:
                                    if e2 == e4 and (r2 > r4 or (r2 == r4 and ta2 > ta4)):
                                        (a2, a4) = (a4, a2)
                                        r2 = data.atoms[a2].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta2 = data.atomTemplateIndexes[data.atoms[a2]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                    if e3 == e4 and (r3 > r4 or (r3 == r4 and ta3 > ta4)):
                                        (a3, a4) = (a4, a3)
                                        r3 = data.atoms[a3].residue.index
                                        r4 = data.atoms[a4].residue.index
                                        ta3 = data.atomTemplateIndexes[data.atoms[a3]]
                                        ta4 = data.atomTemplateIndexes[data.atoms[a4]]
                                    if r2 > r3 or (r2 == r3 and ta2 > ta3):
                                        (a2, a3) = (a3, a2)
                                match = (a2, a3, torsion[0], a4, tordef)
                                break
2833
2834
2835
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
2836
2837
2838
2839

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


2840
## @private
2841
class CustomNonbondedGenerator(object):
2842
2843
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

2844
2845
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
2846
2847
2848
2849
2850
2851
2852
2853
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2854
2855
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
2856
2857
2858
2859
        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'])
2860
2861
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2862
        generator.functions += _parseFunctions(element)
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874

    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)
2875
        _createFunctions(force, self.functions)
2876
        for atom in data.atoms:
2877
2878
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2879
2880
2881
2882
2883
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

    def postprocessSystem(self, sys, data, args):
2884
        # Create the exclusions.
2885

2886
        bondIndices = _findBondsForExclusions(data, sys)
2887
2888
2889
2890
2891
2892
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomNonbondedForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


2893
## @private
2894
class CustomGBGenerator(object):
2895
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2896

2897
2898
    def __init__(self, forcefield):
        self.ff = forcefield
2899
2900
2901
2902
2903
2904
2905
2906
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2907
2908
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
2909
2910
2911
2912
        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'])
2913
2914
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2915
2916
2917
2918
2919
2920
2921
        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']]))
2922
        generator.functions += _parseFunctions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2923

2924
2925
2926
2927
2928
    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
2929
            raise ValueError('Illegal nonbonded method for CustomGBForce')
2930
2931
2932
2933
2934
2935
2936
2937
2938
        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])
2939
        _createFunctions(force, self.functions)
2940
        for atom in data.atoms:
2941
2942
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2943
2944
2945
2946
2947
2948
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

2949

2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
## @private
class CustomHbondGenerator(object):
    """A CustomHbondGenerator constructs a CustomHbondForce."""

    def __init__(self, forcefield):
        self.ff = forcefield
        self.donorTypes1 = []
        self.donorTypes2 = []
        self.donorTypes3 = []
        self.acceptorTypes1 = []
        self.acceptorTypes2 = []
        self.acceptorTypes3 = []
        self.globalParams = {}
        self.perDonorParams = []
        self.perAcceptorParams = []
        self.donorParamValues = []
        self.acceptorParamValues = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
        generator = CustomHbondGenerator(ff)
        ff.registerGenerator(generator)
        generator.energy = element.attrib['energy']
        generator.bondCutoff = int(element.attrib['bondCutoff'])
        generator.particlesPerDonor = int(element.attrib['particlesPerDonor'])
        generator.particlesPerAcceptor = int(element.attrib['particlesPerAcceptor'])
        if generator.particlesPerDonor < 1 or generator.particlesPerDonor > 3:
            raise ValueError('Illegal value for particlesPerDonor for CustomHbondForce')
        if generator.particlesPerAcceptor < 1 or generator.particlesPerAcceptor > 3:
            raise ValueError('Illegal value for particlesPerAcceptor for CustomHbondForce')
        for param in element.findall('GlobalParameter'):
            generator.globalParams[param.attrib['name']] = float(param.attrib['defaultValue'])
        for param in element.findall('PerDonorParameter'):
            generator.perDonorParams.append(param.attrib['name'])
        for param in element.findall('PerAcceptorParameter'):
            generator.perAcceptorParams.append(param.attrib['name'])
        for donor in element.findall('Donor'):
2988
            types = ff._findAtomTypes(donor.attrib, 3)[:generator.particlesPerDonor]
2989
2990
2991
2992
2993
2994
2995
2996
            if None not in types:
                generator.donorTypes1.append(types[0])
                if len(types) > 1:
                    generator.donorTypes2.append(types[1])
                if len(types) > 2:
                    generator.donorTypes3.append(types[2])
                generator.donorParamValues.append([float(donor.attrib[param]) for param in generator.perDonorParams])
        for acceptor in element.findall('Acceptor'):
2997
            types = ff._findAtomTypes(acceptor.attrib, 3)[:generator.particlesPerAcceptor]
2998
2999
3000
3001
3002
3003
3004
3005
3006
3007
3008
3009
3010
3011
3012
3013
3014
3015
3016
3017
3018
3019
3020
3021
3022
3023
3024
3025
3026
3027
3028
3029
3030
3031
3032
3033
3034
3035
3036
3037
3038
3039
3040
            if None not in types:
                generator.acceptorTypes1.append(types[0])
                if len(types) > 1:
                    generator.acceptorTypes2.append(types[1])
                if len(types) > 2:
                    generator.acceptorTypes3.append(types[2])
                generator.acceptorParamValues.append([float(acceptor.attrib[param]) for param in generator.perAcceptorParams])
        generator.functions += _parseFunctions(element)

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomHbondForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomHbondForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.CustomHbondForce.CutoffPeriodic}
        if nonbondedMethod not in methodMap:
            raise ValueError('Illegal nonbonded method for CustomNonbondedForce')
        force = mm.CustomHbondForce(self.energy)
        sys.addForce(force)
        for param in self.globalParams:
            force.addGlobalParameter(param, self.globalParams[param])
        for param in self.perDonorParams:
            force.addPerDonorParameter(param)
        for param in self.perAcceptorParams:
            force.addPerAcceptorParameter(param)
        _createFunctions(force, self.functions)
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)

        # Add donors.

        if self.particlesPerDonor == 1:
            for atom in data.atoms:
                type1 = data.atomType[atom]
                for i in range(len(self.donorTypes1)):
                    types1 = self.donorTypes1[i]
                    if type1 in self.donorTypes1[i]:
                        force.addDonor(atom.index, -1, -1, self.donorParamValues[i])
        elif self.particlesPerDonor == 2:
            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.donorTypes1)):
                    types1 = self.donorTypes1[i]
                    types2 = self.donorTypes2[i]
3041
3042
3043
3044
                    if type1 in types1 and type2 in types2:
                        force.addDonor(bond.atom1, bond.atom2, -1, self.donorParamValues[i])
                    elif type1 in types2 and type2 in types1:
                        force.addDonor(bond.atom2, bond.atom1, -1, self.donorParamValues[i])
3045
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
        else:
            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.donorTypes1)):
                    types1 = self.donorTypes1[i]
                    types2 = self.donorTypes2[i]
                    types3 = self.donorTypes3[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.addDonor(angle[0], angle[1], angle[2], self.donorParamValues[i])

        # Add acceptors.

        if self.particlesPerAcceptor == 1:
            for atom in data.atoms:
                type1 = data.atomType[atom]
                for i in range(len(self.acceptorTypes1)):
                    types1 = self.acceptorTypes1[i]
                    if type1 in self.acceptorTypes1[i]:
                        force.addAcceptor(atom.index, -1, -1, self.acceptorParamValues[i])
        elif self.particlesPerAcceptor == 2:
            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.acceptorTypes1)):
                    types1 = self.acceptorTypes1[i]
                    types2 = self.acceptorTypes2[i]
3073
                    if type1 in types1 and type2 in types2:
3074
                        force.addAcceptor(bond.atom1, bond.atom2, -1, self.acceptorParamValues[i])
3075
3076
                    elif type1 in types2 and type2 in types1:
                        force.addAcceptor(bond.atom2, bond.atom1, -1, self.acceptorParamValues[i])
3077
3078
3079
3080
3081
3082
3083
3084
3085
3086
3087
3088
3089
3090
3091
3092
3093
3094
3095
3096
3097
3098
3099
3100
3101
3102
3103
3104
3105
3106
3107
3108
3109
3110
3111
3112
3113
3114
        else:
            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.acceptorTypes1)):
                    types1 = self.acceptorTypes1[i]
                    types2 = self.acceptorTypes2[i]
                    types3 = self.acceptorTypes3[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.addAcceptor(angle[0], angle[1], angle[2], self.acceptorParamValues[i])

        # Add exclusions.

        for donor in range(force.getNumDonors()):
            (d1, d2, d3, params) = force.getDonorParameters(donor)
            outerAtoms = set((d1, d2, d3))
            if -1 in outerAtoms:
                outerAtoms.remove(-1)
            excludedAtoms = set(outerAtoms)
            for i in range(self.bondCutoff):
                newOuterAtoms = set()
                for atom in outerAtoms:
                    for bond in data.atomBonds[atom]:
                        b = data.bonds[bond]
                        bondedAtom = (b.atom2 if b.atom1 == atom else b.atom1)
                        if bondedAtom not in excludedAtoms:
                            newOuterAtoms.add(bondedAtom)
                            excludedAtoms.add(bondedAtom)
                outerAtoms = newOuterAtoms
            for acceptor in range(force.getNumAcceptors()):
                (a1, a2, a3, params) = force.getAcceptorParameters(acceptor)
                if a1 in excludedAtoms or a2 in excludedAtoms or a3 in excludedAtoms:
                    force.addExclusion(donor, acceptor)

parsers["CustomHbondForce"] = CustomHbondGenerator.parseElement


3115
## @private
3116
class CustomManyParticleGenerator(object):
3117
3118
3119
3120
3121
3122
3123
3124
3125
3126
3127
3128
    """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 = []
3129

3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140
3141
    @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(',')]))
3142
3143
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172

    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:
3173
3174
3175
            values = self.params.getAtomParameters(atom, data)
            type = int(self.params.getExtraParameters(atom, data)['filterType'])
            force.addParticle(values, type)
3176
3177
3178
3179
3180
3181
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

3183
3184
3185
3186
3187
3188
3189
3190
3191
3192
3193
        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)))
3194

3195
3196
        # 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.
3197

3198
3199
3200
3201
3202
3203
3204
3205
3206
        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.
3207

3208
3209
3210
3211
3212
        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
3213
def getAtomPrint(data, atomIndex):
3214

Peter Eastman's avatar
Peter Eastman committed
3215
3216
3217
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
3218
    else:
Peter Eastman's avatar
Peter Eastman committed
3219
        returnString = "NA"
3220
3221
3222
3223
3224

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
3225
def countConstraint(data):
3226

Peter Eastman's avatar
Peter Eastman committed
3227
    bondCount = 0
3228
3229
3230
3231
3232
3233
3234
    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
3235
        if (isConstrained):
3236
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
3237

3238
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
3239

3240
## @private
3241
class AmoebaBondGenerator(object):
3242
3243
3244

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

3245
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
3246
3247

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

3249
3250
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
3251
3252
3253
3254
3255
3256
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3257

3258
3259
3260
3261
3262
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3266
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
3267
3268
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
3269
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
3270
            if None not in types:
3271
3272
3273
3274
3275
                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:
3276
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
3277
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
3278
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3279
3280
                raise ValueError(outputString)

3281
3282
    #=============================================================================================

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

Mark Friedrichs's avatar
Mark Friedrichs committed
3285
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3286

Peter Eastman's avatar
Peter Eastman committed
3287
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3288
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
3289
        if len(existing) == 0:
3290
            force = mm.AmoebaBondForce()
3291
3292
3293
3294
            sys.addForce(force)
        else:
            force = existing[0]

3295
3296
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
3297
3298
3299
3300
3301
3302
3303
3304
3305
3306

        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:
3307
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
3308
3309
3310
3311
                    elif self.k[i] != 0:
                        force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
                    break

3312
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
3313
3314
3315
3316

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

Peter Eastman's avatar
Peter Eastman committed
3318
def addAngleConstraint(angle, idealAngle, data, sys):
3319
3320

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

Peter Eastman's avatar
Peter Eastman committed
3322
3323
    bond1 = None
    bond2 = None
3324
3325
3326
3327
3328
3329
3330
    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
3331

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

3334
3335
3336
3337
        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
3338
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
3339
                data.addConstraint(sys, angle[0], angle[2], length)
3340
3341
3342
                return

#=============================================================================================
3343
## @private
3344
class AmoebaAngleGenerator(object):
3345
3346

    #=============================================================================================
3347
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
3348
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3349

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

Peter Eastman's avatar
Peter Eastman committed
3352
3353
3354
3355
3356
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3357

Peter Eastman's avatar
Peter Eastman committed
3358
3359
3360
        self.types1 = []
        self.types2 = []
        self.types3 = []
3361

Peter Eastman's avatar
Peter Eastman committed
3362
3363
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3364

3365
3366
3367
3368
3369
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3373
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
3374
3375
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
3376
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
3377
            if None not in types:
3378
3379
3380
3381
3382
3383

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
3384
                angleList.append(float(angle.attrib['angle1']))
3385
3386

                try:
Peter Eastman's avatar
Peter Eastman committed
3387
                    angleList.append(float(angle.attrib['angle2']))
3388
                    try:
Peter Eastman's avatar
Peter Eastman committed
3389
                        angleList.append(float(angle.attrib['angle3']))
3390
3391
3392
3393
3394
3395
3396
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
3397
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
3398
3399
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
3400
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
3401
3402
                raise ValueError(outputString)

3403
3404
3405
3406
    #=============================================================================================
    # 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
3407

Peter Eastman's avatar
Peter Eastman committed
3408
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3409
3410
3411
3412
3413
3414
        pass

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

Peter Eastman's avatar
Peter Eastman committed
3416
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3417
3418
3419
3420

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3421
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
3422
3423

        if len(existing) == 0:
3424
            force = mm.AmoebaAngleForce()
3425
3426
3427
3428
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3429
        # set scalars
3430

3431
3432
3433
3434
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
3435
3436

        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
3437
3438
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3439

Peter Eastman's avatar
Peter Eastman committed
3440
3441
3442
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3443
3444
3445
3446
3447
3448
3449
            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]
3450
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3451
                    elif self.k[i] != 0:
Peter Eastman's avatar
Peter Eastman committed
3452
3453
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
3454
3455
3456
3457
3458
3459
                            # 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
3460
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
3461
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3462
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
3463
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3464
                            if (numberOfHydrogens < lenAngle):
3465
3466
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
3467
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
3468
                                raise ValueError(outputString)
3469
3470
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
3471

3472
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
3473
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
3474
3475
3476
3477
3478
3479
                    break

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

Peter Eastman's avatar
Peter Eastman committed
3481
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3482
3483
3484
3485

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3486
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
3487
3488

        if len(existing) == 0:
3489
            force = mm.AmoebaInPlaneAngleForce()
3490
3491
3492
3493
3494
3495
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

3496
3497
3498
3499
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
3500
3501

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

Peter Eastman's avatar
Peter Eastman committed
3503
3504
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3505

Peter Eastman's avatar
Peter Eastman committed
3506
3507
3508
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3509
3510
3511
3512
3513
3514
3515
3516
3517

            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
3518
                    if (isConstrained and self.k[i] != 0.0):
3519
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3520
3521
3522
3523
                    else:
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

3524
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
3525
3526
3527

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
3528
3529
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
3530
3531
#=============================================================================================

3532
## @private
3533
class AmoebaOutOfPlaneBendGenerator(object):
3534
3535
3536
3537

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

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

3539
3540
3541
3542
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3543
3544
3545
3546
3547
3548
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3549

Peter Eastman's avatar
Peter Eastman committed
3550
3551
3552
3553
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3554

Peter Eastman's avatar
Peter Eastman committed
3555
        self.ks = []
3556
3557
3558
3559
3560

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

3562
3563
3564
3565
3566
3567
3568
3569
3570
3571
3572
3573
3574
3575
3576
3577
3578
3579
3580
3581
3582
3583
3584
3585
3586
    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
3587

3588
3589
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
3590
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
3591
3592
3593
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
3594
                                                   float(element.attrib['opbend-sextic']))
3595
3596
3597
3598

        forceField._forces.append(generator)

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

Peter Eastman's avatar
Peter Eastman committed
3602
3603
3604
3605
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
3606

Peter Eastman's avatar
Peter Eastman committed
3607
                generator.ks.append(float(angle.attrib['k']))
3608
3609

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3610
3611
                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
3612
3613
                raise ValueError(outputString)

3614
3615
3616
3617
3618
3619
    #=============================================================================================
    # 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
3620

Peter Eastman's avatar
Peter Eastman committed
3621
    def getMiddleAtom(self, angle, data):
3622
3623
3624

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
3625
        middleAtom = -1
Justin MacCallum's avatar
Justin MacCallum committed
3626
        for atomIndex in angle:
Peter Eastman's avatar
Peter Eastman committed
3627
            isMiddle = 0
3628
3629
3630
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3631
                if (atom1 != atomIndex):
3632
3633
3634
                    partner = atom1
                else:
                    partner = atom2
Justin MacCallum's avatar
Justin MacCallum committed
3635
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]):
3636
3637
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
3638
            if (isMiddle == 2):
3639
3640
3641
3642
3643
                return atomIndex
        return -1

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

Peter Eastman's avatar
Peter Eastman committed
3644
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3645
3646
3647
3648
3649
3650
3651
3652
3653
3654
3655
3656
3657

        # 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
3658
3659
3660
3661
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
3662
3663
3664
3665

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

Peter Eastman's avatar
Peter Eastman committed
3666
        skipAtoms = dict()
3667
3668
3669
3670

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

Peter Eastman's avatar
Peter Eastman committed
3671
3672
        inPlaneAngles = []
        nonInPlaneAngles = []
3673
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
3674
        idealAngles = []*len(data.angles)
3675
3676
3677

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

Peter Eastman's avatar
Peter Eastman committed
3678
3679
3680
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
3681
3682
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
3683
                middleType = -1
3684
3685
                middleCovalency = -1

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

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

Peter Eastman's avatar
Peter Eastman committed
3694
3695
3696
3697
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
3698
3699
3700
3701

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3702
                    if (atom1 != middleAtom):
3703
3704
3705
3706
3707
3708
3709
3710
                        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
3711
3712
3713
3714
3715
                        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
3716

Peter Eastman's avatar
Peter Eastman committed
3717
                if (len(partners) == 3):
3718

Peter Eastman's avatar
Peter Eastman committed
3719
3720
3721
                    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])
3722

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

3725
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
3726
3727
3728
3729
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3730
3731
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
3732
3733
3734
3735
3736
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
3737
3738
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
3739
                    angleDict['isConstrained'] = 0
3740

Peter Eastman's avatar
Peter Eastman committed
3741
3742
3743
3744
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3745
3746

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3747
3748
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3749

Peter Eastman's avatar
Peter Eastman committed
3750
                    inPlaneAngles.append(angleDict)
3751
3752

                else:
Peter Eastman's avatar
Peter Eastman committed
3753
3754
3755
3756
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
3757
            else:
Peter Eastman's avatar
Peter Eastman committed
3758
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
3759

Peter Eastman's avatar
Peter Eastman committed
3760
                    partnerSet = skipAtoms[middleAtom]
Justin MacCallum's avatar
Justin MacCallum committed
3761

Peter Eastman's avatar
Peter Eastman committed
3762
                    angleDict = {}
3763

Peter Eastman's avatar
Peter Eastman committed
3764
3765
3766
3767
3768
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
3769

Peter Eastman's avatar
Peter Eastman committed
3770
                    angleDict['isConstrained'] = isConstrained
3771

Peter Eastman's avatar
Peter Eastman committed
3772
3773
3774
3775
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3776
3777

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3778
3779
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3780

Peter Eastman's avatar
Peter Eastman committed
3781
                    inPlaneAngles.append(angleDict)
3782
3783

                else:
Peter Eastman's avatar
Peter Eastman committed
3784
3785
                    angleDict = {}
                    angleDict['angle'] = angle
3786
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
3787
                    nonInPlaneAngles.append(angleDict)
3788

3789
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
3790
3791

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3792
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
3793
3794
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
3795
3796

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3797
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
3798
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
3799
                    nonInPlaneAngles.append(angleDict)
3800
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
3801
3802
3803
3804
3805

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

3806
## @private
3807
class AmoebaTorsionGenerator(object):
3808
3809
3810
3811
3812
3813
3814

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
3815
        self.torsionUnit = torsionUnit
3816

Peter Eastman's avatar
Peter Eastman committed
3817
3818
3819
3820
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3821

Peter Eastman's avatar
Peter Eastman committed
3822
3823
3824
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
3825

3826
3827
3828
3829
3830
3831
3832
3833
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3835
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
3836
3837
3838
3839
3840
3841
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
3842
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
3843
            if None not in types:
3844
3845
3846
3847
3848
3849
3850

                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
3851
3852
                    tInfo = []
                    suffix = str(ii)
3853
3854
3855
3856
3857
3858
                    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
3859
3860
3861
3862
3863
3864
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
3865
3866
3867

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
3868
3869
3870
3871
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
3872
3873
                raise ValueError(outputString)

3874
3875
3876
3877
3878
    #=============================================================================================

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3879
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
3880
        if len(existing) == 0:
3881
            force = mm.PeriodicTorsionForce()
3882
3883
3884
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3885

3886
3887
3888
3889
3890
3891
3892
3893
3894
3895
3896
3897
3898
3899
3900
3901
        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
3902
                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):
3903
3904
3905
3906
3907
3908
                    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])
3909
3910
3911
3912
3913
3914
                    break

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

3915
## @private
3916
class AmoebaPiTorsionGenerator(object):
3917
3918
3919
3920
3921
3922

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

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

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

3924
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
3925
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
3926
3927
3928
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3929

3930
3931
3932
3933
3934
3935
3936
3937
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
3938
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
3939
3940
3941
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
3942
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
3943
            if None not in types:
3944
3945
3946
3947
3948
3949
                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
3950
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3951
3952
                raise ValueError(outputString)

3953
3954
3955
3956
3957
3958
3959
3960
3961
3962
3963
3964
3965
3966
3967
3968
    #=============================================================================================

    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
3969

3970
3971
            atom1 = bond.atom1
            atom2 = bond.atom2
Justin MacCallum's avatar
Justin MacCallum committed
3972

3973
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
3974
3975
3976
3977
3978
3979
3980
3981
3982
3983
3984

                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
3985
3986
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
3987
3988
3989
3990
3991
3992
3993
3994
3995
3996
3997
3998

                       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
3999
                           if (bondedAtom1 != atom1):
4000
4001
4002
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
4003
4004
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
4005
                                   piTorsionAtom1 = b1
4006
4007
4008
4009
4010
4011
                               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
4012
                           if (bondedAtom1 != atom2):
4013
4014
4015
4016
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4017
4018
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
4019
                                   piTorsionAtom5 = b1
4020
4021
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
4022

Peter Eastman's avatar
Peter Eastman committed
4023
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
4024
4025
4026
4027
4028

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

4029
## @private
4030
class AmoebaTorsionTorsionGenerator(object):
4031
4032
4033
4034
4035

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

Peter Eastman's avatar
Peter Eastman committed
4036
    def __init__(self):
4037

Peter Eastman's avatar
Peter Eastman committed
4038
4039
4040
4041
4042
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
4043

Peter Eastman's avatar
Peter Eastman committed
4044
        self.gridIndex = []
4045

Peter Eastman's avatar
Peter Eastman committed
4046
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
4047

4048
4049
4050
4051
4052
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
4053
        generator = AmoebaTorsionTorsionGenerator()
4054
4055
4056
4057
4058
4059
4060
        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'):
4061
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
4062
            if None not in types:
4063
4064
4065
4066
4067
4068
4069

                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
4070
4071
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
4072
4073
4074
4075
4076
4077
4078
4079
4080
                    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
4081
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
4082
4083
                raise ValueError(outputString)

4084
4085
4086
4087
4088
4089
4090
4091
4092
4093
        # 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
4094
4095
4096
4097
4098
4099
        #     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
4100
4101

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

Peter Eastman's avatar
Peter Eastman committed
4105
4106
4107
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
4108

Peter Eastman's avatar
Peter Eastman committed
4109
4110
            grid = []
            gridCol = []
4111
4112
4113
4114
4115

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
4116
4117
4118
4119
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
4120
                if 'fx' in gridEntry.attrib:
4121
4122
4123
                    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
4124
                gridCol.append(gridRow)
4125
4126

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
4127
4128
4129
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
4130
4131
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
4132

Peter Eastman's avatar
Peter Eastman committed
4133
4134
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
4135
            else:
Peter Eastman's avatar
Peter Eastman committed
4136
4137
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
4138
4139
4140
4141
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
4142
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
4143
4144
4145
4146
4147
4148
4149
4150

        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
4151
        if (len(data.atomBonds[atomC]) == 4):
4152
4153
4154
4155
4156
            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
4157
4158
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
4159
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
4160
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
4161
4162
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
4163
4164
                if (hit > -1):
                    if (atomE == -1):
4165
4166
4167
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
4168

4169
4170
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
4171
4172
            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
4173
                raise ValueError(outputString)
4174
4175
4176
4177
4178

            # 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
4179
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
4180
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
4181
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
4182
                chiralAtomIndex = atomF
4183

Peter Eastman's avatar
Peter Eastman committed
4184
4185
4186
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
4187
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
4188
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
4189
                chiralAtomIndex = massF
4190
4191
4192
4193

        return chiralAtomIndex

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

4195
4196
4197
4198
4199
4200
4201
4202
4203
4204
4205
4206
4207
    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
4208
4209
            # search for bitorsions; based on TINKER subroutine bitors()

4210
4211
4212
4213
4214
4215
4216
            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
4217
                if (bondedAtom1 != ib):
4218
4219
4220
4221
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4222
                if (ia != ic and ia != id):
4223
4224
4225
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
4226
                        if (bondedAtom1 != id):
4227
4228
4229
4230
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4231
                        if (ie != ic and ie != ib and ie != ia):
4232
4233
4234
4235
4236
4237
4238
4239
4240
4241
4242
4243
4244
4245
4246
4247
4248
4249
4250
4251

                            # 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
4252
4253
4254
                                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])
4255
4256
4257

                                # match in reverse order

4258
                                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
4259
4260
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
4261
4262
4263
4264

        # set grids

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

4267
4268
4269
4270
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

4271
## @private
4272
class AmoebaStretchBendGenerator(object):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4273
4274

    #=============================================================================================
4275
4276
4277
4278
4279
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
4280
4281
4282
        self.types1 = []
        self.types2 = []
        self.types3 = []
4283

Peter Eastman's avatar
Peter Eastman committed
4284
4285
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
4286

4287
4288
4289
4290
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
4291
        generator = AmoebaStretchBendGenerator()
4292
4293
4294
4295
4296
4297
4298
        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'):
4299
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
4300
            if None not in types:
4301
4302
4303
4304
4305
4306
4307
4308
4309
4310
4311
4312

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

4316
4317
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
4318
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
4319
4320
    # 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
4321
4322
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
4323
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
4324
4325
4326

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

Peter Eastman's avatar
Peter Eastman committed
4327
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4328
4329
4330
4331
4332
4333
4334
4335
        pass

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

    # Note: request for constrained bonds is ignored.

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

4336
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
4337
4338
4339
4340
4341
4342
4343
4344
4345
4346
4347
4348

        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
4349
            if ('isConstrained' in angleDict):
4350
4351
4352
4353
4354
4355
4356
4357
                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
4358
            radian = 57.2957795130
4359
4360
4361
4362
4363
4364
4365
4366
4367
4368
            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
4369
                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
4370
4371
4372
4373
4374
4375
4376
4377
4378
4379
4380
4381
4382
4383
4384
                    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
4385

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4388
                    if ('idealAngle' not in angleDict):
4389

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4390
4391
4392
4393
4394
                       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
4395
                       raise ValueError(outputString)
4396

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4397
4398
4399
4400
4401
4402
4403
                    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
4404
                       raise ValueError(outputString)
4405

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4409
                    break
4410
4411
4412
4413
4414

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

4415
## @private
4416
class AmoebaVdwGenerator(object):
4417
4418

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

4420
4421
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4424
        self.type = type
4425

Peter Eastman's avatar
Peter Eastman committed
4426
4427
4428
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
4429

Peter Eastman's avatar
Peter Eastman committed
4430
        self.epsilonrule = epsilonrule
4431

Peter Eastman's avatar
Peter Eastman committed
4432
4433
4434
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
4435
4436
4437
4438
4439
4440
4441

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

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

4445
4446
4447
        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
4448
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
4449
4450
4451
4452
4453
4454
4455
4456
4457
4458
4459
4460
            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')
4461
        generator.params.parseDefinitions(element)
4462
4463
4464
4465
4466
        two_six = 1.122462048309372

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

    @staticmethod
4467
    def getBondedParticleSets(sys, data):
4468

4469
4470
4471
4472
4473
4474
        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
4475

4476
4477
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
4480
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
4481
4482
        epsilonMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'HARMONIC':1, 'HHG':1}

4483
4484
        force = mm.AmoebaVdwForce()
        sys.addForce(force)
4485

4486
        # sigma and epsilon combining rules
4487

4488
4489
4490
4491
        if ('sigmaCombiningRule' in args):
            sigmaRule = args['sigmaCombiningRule'].upper()
            if (sigmaRule.upper() in sigmaMap):
                force.setSigmaCombiningRule(sigmaRule.upper())
4492
            else:
4493
4494
4495
4496
                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)
4497

4498
4499
4500
4501
        if ('epsilonCombiningRule' in args):
            epsilonRule = args['epsilonCombiningRule'].upper()
            if (epsilonRule.upper() in epsilonMap):
                force.setEpsilonCombiningRule(epsilonRule.upper())
4502
            else:
4503
4504
4505
4506
                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
4507

4508
        # cutoff
4509

4510
4511
4512
4513
        if ('vdwCutoff' in args):
            force.setCutoff(args['vdwCutoff'])
        else:
            force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4514

4515
        # dispersion correction
4516

4517
4518
        if ('useDispersionCorrection' in args):
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
4519

4520
4521
        if (nonbondedMethod == PME):
            force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
4522
4523
4524

        # add particles to force

4525
4526
4527
4528
4529
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
4530
        for (i, atom) in enumerate(data.atoms):
4531
4532
4533
4534
4535
4536
4537
4538
4539
4540
            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
4541

4542
            force.addParticle(ivIndex, values[0]*sigmaScale, values[1], values[2])
4543
4544
4545
4546
4547
4548
4549

        # 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

4550
        bondedParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4551
4552

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

4554
4555
4556
4557
4558
4559
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
4560
            if (self.vdw13Scale == 0.0):
4561
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
4562
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
4563
4564
4565
4566
4567

            # self

            exclusionSet.add(i)

4568
            force.setParticleExclusions(i, tuple(exclusionSet))
4569
4570
4571
4572
4573

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

4574
## @private
4575
class AmoebaMultipoleGenerator(object):
4576
4577
4578
4579

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

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

4581
4582
    #=============================================================================================

4583
    def __init__(self, forceField):
Peter Eastman's avatar
Peter Eastman committed
4584
4585
        self.forceField = forceField
        self.typeMap = {}
4586
4587
4588
4589
4590
4591

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
4592
    def setAxisType(kIndices):
4593
4594
4595

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
4596
4597
4598
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
4599
                else:
Peter Eastman's avatar
Peter Eastman committed
4600
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
4601

Peter Eastman's avatar
Peter Eastman committed
4602
4603
                if (kIndicesLen > 2):
                    kx = kIndices[2]
4604
                else:
Peter Eastman's avatar
Peter Eastman committed
4605
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
4606

Peter Eastman's avatar
Peter Eastman committed
4607
4608
                if (kIndicesLen > 1):
                    kz = kIndices[1]
4609
                else:
Peter Eastman's avatar
Peter Eastman committed
4610
                    kz = 0
4611

Peter Eastman's avatar
Peter Eastman committed
4612
4613
                while(len(kIndices) < 4):
                    kIndices.append(0)
4614
4615

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
4616
                if (kz == 0):
4617
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
4618
                if (kz != 0 and kx == 0):
4619
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
4620
                if (kz < 0 or kx < 0):
4621
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
4622
                if (kx < 0 and ky < 0):
4623
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
4624
                if (kz < 0 and kx < 0 and ky  < 0):
4625
4626
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
4627
4628
4629
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
4630
4631
4632
4633
4634
4635
4636
4637

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
4638
        #   <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"  >
4639
4640
4641
        # <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"  />

4642
4643
4644
4645
4646
4647
4648
        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]
4649
4650
4651
4652

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

        for atom in element.findall('Multipole'):
4653
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4654
            if None not in types:
4655
4656
4657
4658
4659
4660
4661
4662

                # 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
4663
4664
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
4665
                    except:
4666
4667
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
4668
                # set axis type based on k-Indices
4669

Peter Eastman's avatar
Peter Eastman committed
4670
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
4671
4672
4673

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
4676
4677
4678
4679
4680
4681
4682
4683
4684
4685
4686
4687
4688
                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']))
4689
4690

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4691
                    if (t not in generator.typeMap):
4692
4693
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
4694
4695
4696
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
4697
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
4698
4699
4700
4701
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
4702

4703
            else:
Peter Eastman's avatar
Peter Eastman committed
4704
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4705
4706
                raise ValueError(outputString)

4707
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
4708

4709
        for atom in element.findall('Polarize'):
4710
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4711
            if None not in types:
4712

Peter Eastman's avatar
Peter Eastman committed
4713
4714
4715
4716
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
4717
4718
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
4719
                    pdamp = pow(polarizability, 1.0/6.0)
4720

Peter Eastman's avatar
Peter Eastman committed
4721
4722
                pgrpMap = dict()
                for index in range(1, 7):
4723
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
4724
                    if (pgrp in atom.attrib):
4725
4726
4727
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4728
4729
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
4730
                        raise ValueError(outputString)
4731
4732
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
4733
4734
4735
4736
4737
4738
4739
                        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
4740
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
4741
4742
4743
4744
4745
                                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
4746
4747
                            raise ValueError(outputString)

4748
            else:
Peter Eastman's avatar
Peter Eastman committed
4749
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4750
4751
                raise ValueError(outputString)

4752
4753
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4754
    def setPolarGroups(self, data, bonded12ParticleSets, force):
4755
4756
4757
4758
4759

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
4760
4761
4762
4763
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
4764
4765
4766
4767
            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
4768
4769
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
4770
4771
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
4772

4773
4774
4775
4776
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
4777
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
4778
4779
                continue

Peter Eastman's avatar
Peter Eastman committed
4780
4781
            group = set()
            visited = set()
4782
4783
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4784
4785
4786
4787
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
4788
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
4789
4790
                if (nextAtom not in visited):
                   visited.add(nextAtom)
4791
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4792
4793
4794
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
4795
4796
4797

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
4798
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
4799
4800

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4801
4802
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
4803
4804
4805
4806
4807

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
4808
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
4809
4810
                continue

Peter Eastman's avatar
Peter Eastman committed
4811
            pgrp11 = set(atom.polarizationGroupSet[0])
4812
4813
4814
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4815
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
4816
4817
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4818
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
4819

4820
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4821
4822
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
4823
4824
4825
4826
4827

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
4828
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
4829
4830
                continue

Peter Eastman's avatar
Peter Eastman committed
4831
4832
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
4833
4834
4835
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4836
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
4837
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4838
            pgrp13 = pgrp13 - set(pgrp11)
4839
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4840
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
4841

4842
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4843
4844
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
4845
4846
4847
4848
4849

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
4850
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
4851
4852
                continue

Peter Eastman's avatar
Peter Eastman committed
4853
4854
4855
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
4856
4857
4858
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4859
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
4860
4861
4862

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4863
            pgrp14 = pgrp14 - set(pgrp11)
4864
4865

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

4868
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4869
4870
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
4871
4872
4873

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

Peter Eastman's avatar
Peter Eastman committed
4874
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4875
4876
4877
4878

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

4879
4880
4881
4882
4883
4884
4885
        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)
4886

4887
4888
        if ('ewaldErrorTolerance' in args):
            force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))
4889

4890
4891
4892
4893
4894
4895
4896
4897
        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)
4898

4899
4900
        if ('aEwald' in args):
            force.setAEwald(float(args['aEwald']))
4901

4902
4903
        if ('pmeGridDimensions' in args):
            force.setPmeGridDimensions(args['pmeGridDimensions'])
4904

4905
4906
        if ('mutualInducedMaxIterations' in args):
            force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
4907

4908
4909
        if ('mutualInducedTargetEpsilon' in args):
            force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))
4910
4911

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
4912
        # throw error if particle type not available
4913
4914
4915
4916
4917

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

        # 1-2

4918
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4919
4920
4921
4922
4923

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4924
            bonded13Set = set()
4925
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4926
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4927
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
4928
4929
4930
4931

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4932
            selfSet = set()
4933
4934
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4935
4936
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
4937
4938
4939
4940
4941

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4942
4943
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4944
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4945
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
4946

4947
4948
4949
4950
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4951
            selfSet = set()
4952
4953
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4954
4955
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
4956
4957
4958
4959
4960

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4961
4962
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4963
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4964
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
4965
4966
4967
4968
4969
4970

            # 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
4971
            selfSet = set()
4972
4973
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4974
4975
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
4976
4977
4978
4979
4980

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

Peter Eastman's avatar
Peter Eastman committed
4981
4982
                multipoleList = self.typeMap[t]
                hit = 0
4983
4984
4985
4986
4987
4988
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4989
                    if (hit != 0):
4990
4991
                        break

Peter Eastman's avatar
Peter Eastman committed
4992
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4993
4994

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4995
4996
                    kx = kIndices[2]
                    ky = kIndices[3]
4997
4998
4999
5000

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

5002
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
5003
5004
5005
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5006
5007
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
5008
                       if (hit != 0):
5009
5010
5011
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
5012
5013
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
5014
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5015
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
5016
5017
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
5018
5019
5020
5021
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
5022
5023
5024
5025
5026
5027
5028
5029
5030
5031
                                      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

5032
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5033
                                      hit = 1
5034
5035
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5036
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
5037
5038
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
5039
5040
5041
5042
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
5043
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5044
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
5045

5046
5047
5048
5049
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5050
                    if (hit != 0):
5051
5052
                        break

Peter Eastman's avatar
Peter Eastman committed
5053
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5054
5055

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
5056
5057
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
5058

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

5063
5064
5065
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
5066
5067
5068
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5069
5070
5071

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
5072
                       if (hit != 0):
5073
5074
5075
                           break

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

Peter Eastman's avatar
Peter Eastman committed
5078
                       if (bondedAtomZType == kz):
5079
5080
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
5081
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
5082
5083
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
5084
5085
5086
5087
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
5088
5089
5090
5091
5092
5093
5094
5095

                                      # 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

5096
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5097
                                      hit = 3
5098
5099
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
5100
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
5101
5102
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
5103
5104
5105
5106
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
5107
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5108
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
5109

5110
5111
5112
5113
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5114
                    if (hit != 0):
5115
5116
                        break

Peter Eastman's avatar
Peter Eastman committed
5117
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5118
5119
5120
5121

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

Peter Eastman's avatar
Peter Eastman committed
5122
5123
5124
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5125
5126
5127

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
5128
                        if (hit != 0):
5129
5130
5131
                            break

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

Peter Eastman's avatar
Peter Eastman committed
5134
                        if (kx == 0 and kz == bondedAtomZType):
5135
                            zaxis = bondedAtomZIndex
5136
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5137
                            hit = 5
5138
5139
5140
5141
5142

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5143
                    if (hit != 0):
5144
5145
                        break

Peter Eastman's avatar
Peter Eastman committed
5146
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5147
5148
5149

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
5150
5151
5152
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5153

Peter Eastman's avatar
Peter Eastman committed
5154
                    if (kz == 0):
5155
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5156
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
5157

5158
5159
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
5160
                if (hit != 0):
5161

Peter Eastman's avatar
Peter Eastman committed
5162
                    atom.multipoleDict = savedMultipoleDict
5163
                    atom.polarizationGroups = dict()
5164
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
5165
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
5166
                    if (atomIndex == newIndex):
5167
5168
5169
5170
                        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]))
5171
                    else:
5172
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
5173
                else:
Peter Eastman's avatar
Peter Eastman committed
5174
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
5175
            else:
Peter Eastman's avatar
Peter Eastman committed
5176
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
5177
5178
5179

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
5180
        self.setPolarGroups(data, bonded12ParticleSets, force)
5181
5182
5183
5184
5185

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

5186
## @private
5187
class AmoebaWcaDispersionGenerator(object):
5188
5189

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

5191
5192
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
5195
5196
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
5197
5198
5199
5200
5201
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
5202
        self.shctd = shctd
5203
5204
5205
5206
5207
5208
5209
5210
5211

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

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

Peter Eastman's avatar
Peter Eastman committed
5213
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
5214
5215
5216
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
Justin MacCallum's avatar
Justin MacCallum committed
5217
                                                  element.attrib['awater'],
5218
5219
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Justin MacCallum's avatar
Justin MacCallum committed
5220
                                                  element.attrib['shctd'])
5221
        forceField._forces.append(generator)
5222
5223
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaWcaDispersionForce', 'WcaDispersion', ('radius', 'epsilon'))
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
5224

5225
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
5226

Peter Eastman's avatar
Peter Eastman committed
5227
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
5228
5229
5230
5231
5232
5233
5234
5235
5236
5237
5238
5239

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

Peter Eastman's avatar
Peter Eastman committed
5242
5243
5244
5245
5246
5247
5248
5249
        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  ))
5250

5251
5252
5253
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
5254
5255
5256
5257
5258

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

5259
## @private
5260
class AmoebaGeneralizedKirkwoodGenerator(object):
5261
5262

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

5264
5265
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5266
5267
5268
5269
5270
5271
5272
5273
5274
5275
5276
    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
5277
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5278
5279
5280
5281
5282
5283
5284
5285
5286
5287
5288
5289
5290
5291
5292
5293
5294
5295
5296
5297
5298
5299
5300
5301
5302
        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
5303
5304
5305

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

Peter Eastman's avatar
Peter Eastman committed
5306
    def getObcShct(self, data, atomIndex):
5307

Peter Eastman's avatar
Peter Eastman committed
5308
        atom = data.atoms[atomIndex]
5309
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5310
        shct = -1.0
5311
5312

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
5313

Peter Eastman's avatar
Peter Eastman committed
5314
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
5315
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5316
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5317
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
5318
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
5319
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
5320
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
5321
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5322
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
5323
5324
5325
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
5326
        elif (atomicNumber == 16):              # S(16)
5327
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
5328
        elif (atomicNumber == 26):              # Fe(26)
5329
5330
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
5331
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5332
            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
5333
5334

        return shct
5335
5336
5337

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

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

Peter Eastman's avatar
Peter Eastman committed
5340
        atom = data.atoms[atomIndex]
5341
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5342
        radius = -1.0
5343

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

Peter Eastman's avatar
Peter Eastman committed
5346
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
5347

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

5351
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5352
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
5353

Peter Eastman's avatar
Peter Eastman committed
5354
            if (bondedAtomAtomicNumber == 7):
5355
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
5356
            if (bondedAtomAtomicNumber == 8):
5357
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
5358

Peter Eastman's avatar
Peter Eastman committed
5359
        elif (atomicNumber == 3):               # Li(3)
5360
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
5361
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5362

5363
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
5364
            if (len(bondedAtomIndices) == 3):
5365
5366
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
5367
            elif (len(bondedAtomIndices) == 4):
5368
5369
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5370
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
5371
5372
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
5373
        elif (atomicNumber == 7):               # N(7)
5374
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
5375
        elif (atomicNumber == 8):               # O(8)
5376
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
5377
            if (len(bondedAtomIndices) == 2):
5378
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
5379
        elif (atomicNumber == 9):               # F(9)
5380
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
5381
        elif (atomicNumber == 10):
5382
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
5383
        elif (atomicNumber == 11):
5384
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
5385
        elif (atomicNumber == 12):
5386
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5387
        elif (atomicNumber == 14):
5388
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
5389
        elif (atomicNumber == 15):              # P(15)
5390
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
5391
        elif (atomicNumber == 16):              # S(16)
5392
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
5393
        elif (atomicNumber == 17):
5394
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
5395
        elif (atomicNumber == 18):
5396
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5397
        elif (atomicNumber == 19):
5398
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
5399
        elif (atomicNumber == 20):
5400
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
5401
        elif (atomicNumber == 35):
5402
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
5403
        elif (atomicNumber == 36):
5404
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
5405
        elif (atomicNumber == 37):
5406
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
5407
        elif (atomicNumber == 53):
5408
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
5409
        elif (atomicNumber == 54):
5410
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
5411
        elif (atomicNumber == 55):
5412
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
5413
        elif (atomicNumber == 56):
5414
5415
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
5416
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5417
5418
            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
5419

5420
5421
5422
5423
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
5426
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5427
        atom = data.atoms[atomIndex]
5428
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
5429
        if (atomicNumber in bondiMap):
5430
5431
            radius = bondiMap[atomicNumber]
        else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5432
5433
            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
5434

5435
5436
5437
5438
5439
5440
5441
5442
5443
5444
        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
5445

Peter Eastman's avatar
Peter Eastman committed
5446
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
Justin MacCallum's avatar
Justin MacCallum committed
5447
5448
                                                        element.attrib['includeCavityTerm'],
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor'])
5449
5450
5451
        forceField._forces.append(generator)

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

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

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

5458
5459
5460
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
5461
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5462
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
5463
        if (len(amoebaMultipoleForceList) > 0):
5464
5465
5466
5467
5468
            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
5469
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
5470
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
5471

5472
5473
5474
5475
5476
5477
5478
        # 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
5479

Peter Eastman's avatar
Peter Eastman committed
5480
5481
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
5482
            else:
Peter Eastman's avatar
Peter Eastman committed
5483
                force.setSolventDielectric(   float(self.solventDielectric))
5484

Peter Eastman's avatar
Peter Eastman committed
5485
5486
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
5487
            else:
Peter Eastman's avatar
Peter Eastman committed
5488
                force.setSoluteDielectric(    float(self.soluteDielectric))
5489

Peter Eastman's avatar
Peter Eastman committed
5490
5491
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
5492
            else:
Peter Eastman's avatar
Peter Eastman committed
5493
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
5494
5495
5496
5497
5498

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
5501
5502
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
5503
5504
5505

        # 1-2

5506
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5507
5508

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
5509
5510
5511
5512
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5513
            else:
Peter Eastman's avatar
Peter Eastman committed
5514
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5515
5516
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
5517
            force.addParticle(multipoleParameters[0], radius, shct)
5518
5519
5520
5521
5522

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

5523
## @private
5524
class AmoebaUreyBradleyGenerator(object):
5525
5526
5527
5528

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

5530
    def __init__(self):
5531

Peter Eastman's avatar
Peter Eastman committed
5532
5533
5534
        self.types1 = []
        self.types2 = []
        self.types3 = []
5535

Peter Eastman's avatar
Peter Eastman committed
5536
5537
        self.length = []
        self.k = []
5538
5539
5540
5541
5542
5543

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

    @staticmethod
    def parseElement(element, forceField):

5544
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
5545
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
5546

5547
        generator = AmoebaUreyBradleyGenerator()
5548
5549
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
5550
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
5551
            if None not in types:
5552
5553
5554
5555
5556
5557
5558
5559
5560
5561

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

5565
5566
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
5569
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5570
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
5571
5572

        if len(existing) == 0:
5573
            force = mm.HarmonicBondForce()
5574
5575
5576
5577
5578
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
Peter Eastman's avatar
Peter Eastman committed
5579
            if (isConstrained):
5580
5581
5582
5583
5584
5585
5586
5587
                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
5588
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
5589
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
5590
5591
5592
5593
5594
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
5595
5596
5597


## @private
5598
class DrudeGenerator(object):
peastman's avatar
peastman committed
5599
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
5600

5601
5602
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
5603
5604
5605
5606
5607
5608
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
5609
5610
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
5611
5612
5613
5614
        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'):
5615
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
5616
5617
5618
5619
5620
5621
5622
5623
5624
5625
            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
5626

peastman's avatar
peastman committed
5627
5628
5629
5630
    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
5631

peastman's avatar
peastman committed
5632
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
5633

peastman's avatar
peastman committed
5634
5635
5636
5637
5638
5639
5640
5641
5642
5643
5644
5645
5646
5647
5648
5649
        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
5650
5651
                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
5652
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
5653

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

peastman's avatar
peastman committed
5657
5658
5659
5660
5661
5662
5663
        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)
5664
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
5665
5666
5667
5668
5669
                # 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]
5670
5671
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
5672
5673
5674
5675
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

Justin MacCallum's avatar
Justin MacCallum committed
5676
parsers["DrudeForce"] = DrudeGenerator.parseElement
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
5677
5678

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