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

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

52
53
54
55
56
57
58
59
60
61
62
63
64
_dataDirectories = None

def _getDataDirectories():
    global _dataDirectories
    if _dataDirectories is None:
        _dataDirectories = [os.path.join(os.path.dirname(__file__), 'data')]
        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
    return _dataDirectories
65

66
67
def _convertParameterToNumber(param):
    if unit.is_quantity(param):
68
69
70
        if param.unit.is_compatible(unit.bar):
            return param / unit.bar
        return param.value_in_unit_system(unit.md_unit_system)
71
72
    return float(param)

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
101
102
103
104
105
106
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))

107
108
# Enumerated values for nonbonded method

109
class NoCutoff(Singleton):
110
111
112
113
    def __repr__(self):
        return 'NoCutoff'
NoCutoff = NoCutoff()

114
class CutoffNonPeriodic(Singleton):
115
116
117
118
    def __repr__(self):
        return 'CutoffNonPeriodic'
CutoffNonPeriodic = CutoffNonPeriodic()

119
class CutoffPeriodic(Singleton):
120
121
122
123
    def __repr__(self):
        return 'CutoffPeriodic'
CutoffPeriodic = CutoffPeriodic()

124
class Ewald(Singleton):
125
126
127
128
    def __repr__(self):
        return 'Ewald'
Ewald = Ewald()

129
class PME(Singleton):
130
131
132
    def __repr__(self):
        return 'PME'
PME = PME()
133

Peter Eastman's avatar
Peter Eastman committed
134
135
136
137
138
class LJPME(Singleton):
    def __repr__(self):
        return 'LJPME'
LJPME = LJPME()

139
140
# Enumerated values for constraint type

141
class HBonds(Singleton):
142
143
144
145
    def __repr__(self):
        return 'HBonds'
HBonds = HBonds()

146
class AllBonds(Singleton):
147
148
149
150
    def __repr__(self):
        return 'AllBonds'
AllBonds = AllBonds()

151
class HAngles(Singleton):
152
153
154
    def __repr__(self):
        return 'HAngles'
HAngles = HAngles()
155
156
157
158
159
160
161
162
163
164

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

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

186
    def loadFile(self, files):
187
        """Load an XML file and add the definitions from it to this ForceField.
188

Robert McGibbon's avatar
Robert McGibbon committed
189
190
        Parameters
        ----------
191
192
193
        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
194
195
196
            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
197
        """
198

199
200
201
202
        if isinstance(files, tuple):
            files = list(files)
        else:
            files = [files]
203
204
205

        trees = []

206
207
208
        i = 0
        while i < len(files):
            file = files[i]
209
            tree = None
210
211
212
213
            try:
                # this handles either filenames or open file-like objects
                tree = etree.parse(file)
            except IOError:
214
                for dataDir in _getDataDirectories():
215
216
217
218
                    f = os.path.join(dataDir, file)
                    if os.path.isfile(f):
                        tree = etree.parse(f)
                        break
219
220
221
222
223
224
225
226
227
228
229
            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)
230
231
            if tree is None:
                raise ValueError('Could not locate file "%s"' % file)
232
233

            trees.append(tree)
234
            i += 1
235

236
            # Process includes in this file.
237

238
239
            if isinstance(file, str):
                parentDir = os.path.dirname(file)
peastman's avatar
peastman committed
240
241
            else:
                parentDir = ''
peastman's avatar
peastman committed
242
243
            for included in tree.getroot().findall('Include'):
                includeFile = included.attrib['file']
peastman's avatar
peastman committed
244
245
246
                joined = os.path.join(parentDir, includeFile)
                if os.path.isfile(joined):
                    includeFile = joined
247
248
                if includeFile not in files:
                    files.append(includeFile)
peastman's avatar
peastman committed
249
250
251

        # Load the atom types.

252
253
254
255
        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
256
257
258

        # Load the residue templates.

259
260
261
262
263
        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)
264
265
                    if 'override' in residue.attrib:
                        template.overrideLevel = int(residue.attrib['override'])
266
267
                    atomIndices = template.atomIndices
                    for ia, atom in enumerate(residue.findall('Atom')):
268
                        params = {}
269
270
271
272
                        for key in atom.attrib:
                            if key not in ('name', 'type'):
                                params[key] = _convertParameterToNumber(atom.attrib[key])
                        atomName = atom.attrib['name']
273
274
                        if atomName in atomIndices:
                            raise ValueError('Residue '+resName+' contains multiple atoms named '+atomName)
275
                        typeName = atom.attrib['type']
276
                        atomIndices[atomName] = ia
277
278
279
280
281
282
283
284
285
286
287
288
289
                        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']))
290
                    for patch in residue.findall('AllowPatch'):
291
                        patchName = patch.attrib['name']
292
293
                        if ':' in patchName:
                            colonIndex = patchName.find(':')
294
295
296
                            self.registerTemplatePatch(resName, patchName[:colonIndex], int(patchName[colonIndex+1:])-1)
                        else:
                            self.registerTemplatePatch(resName, patchName, 0)
297
                    self.registerResidueTemplate(template)
peastman's avatar
peastman committed
298

299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
        # 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)
                    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']
peastman's avatar
peastman committed
316
                        if atomName in patchData.allAtomNames:
317
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
318
                        patchData.allAtomNames.add(atomName)
319
320
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
321
                        patchData.addedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
322
323
324
325
326
327
                    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']
peastman's avatar
peastman committed
328
                        if atomName in patchData.allAtomNames:
329
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
330
                        patchData.allAtomNames.add(atomName)
331
332
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
333
                        patchData.changedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
334
335
                    for atom in patch.findall('RemoveAtom'):
                        atomName = atom.attrib['name']
peastman's avatar
peastman committed
336
                        if atomName in patchData.allAtomNames:
337
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
338
                        patchData.allAtomNames.add(atomName)
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
                        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
364
365
        # Load force definitions

366
367
368
369
        for tree in trees:
            for child in tree.getroot():
                if child.tag in parsers:
                    parsers[child.tag](child, self)
peastman's avatar
peastman committed
370
371
372

        # Load scripts

373
374
375
        for tree in trees:
            for node in tree.getroot().findall('Script'):
                self.registerScript(node.text)
376

377
378
379
    def getGenerators(self):
        """Get the list of all registered generators."""
        return self._forces
380

381
382
383
    def registerGenerator(self, generator):
        """Register a new generator."""
        self._forces.append(generator)
384

385
386
387
388
389
390
391
392
393
394
395
396
    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)
397
        self._atomTypes[name] = ForceField._AtomType(name, atomClass, mass, element)
398
399
400
401
402
403
404
        if atomClass in self._atomClasses:
            typeSet = self._atomClasses[atomClass]
        else:
            typeSet = set()
            self._atomClasses[atomClass] = typeSet
        typeSet.add(name)
        self._atomClasses[''].add(name)
405

406
407
    def registerResidueTemplate(self, template):
        """Register a new residue template."""
408
409
        if template.name in self._templates:
            # There is already a template with this name, so check the override levels.
410

411
412
413
414
415
416
417
418
419
420
421
            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))
422

423
        # Register the template.
424

425
426
427
        self._templates[template.name] = template
        signature = _createResidueSignature([atom.element for atom in template.atoms])
        if signature in self._templateSignatures:
428
            self._templateSignatures[signature].append(template)
429
430
        else:
            self._templateSignatures[signature] = [template]
431

432
433
434
    def registerPatch(self, patch):
        """Register a new patch that can be applied to templates."""
        self._patches[patch.name] = patch
435

436
437
438
    def registerTemplatePatch(self, residue, patch, patchResidueIndex):
        """Register that a particular patch can be used with a particular residue."""
        if residue not in self._templatePatches:
439
440
            self._templatePatches[residue] = set()
        self._templatePatches[residue].add((patch, patchResidueIndex))
441

442
443
444
    def registerScript(self, script):
        """Register a new script to be executed after building the System."""
        self._scripts.append(script)
445

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
446
    def registerTemplateGenerator(self, generator):
447
448
449
450
        """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
451
452
        .. CAUTION:: This method is experimental, and its API is subject to change.

453
454
        Parameters
        ----------
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
455
        generator : function
456
457
            A function that will be called when a residue is encountered that does not match an existing forcefield template.

John Chodera's avatar
John Chodera committed
458
        When a residue without a template is encountered, the ``generator`` function is called with:
459

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
460
461
        ::
           success = generator(forcefield, residue)
462

John Chodera's avatar
John Chodera committed
463
464
465
        where ``forcefield`` is the calling ``ForceField`` object and ``residue`` is a simtk.openmm.app.topology.Residue object.

        ``generator`` must conform to the following API:
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
466
467

        ::
John Chodera's avatar
John Chodera committed
468
469
470
           generator API

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

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
477
478
479
480
481
           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`.
482

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
483
484
485
486
487
           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.
488
489

        """
490
        self._templateGenerators.append(generator)
491

492
    def _findAtomTypes(self, attrib, num):
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
        """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.

        """
508
509
510
511
512
513
514
        types = []
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
515
            typeAttrib = 'type'+suffix
516
            if classAttrib in attrib:
517
                if typeAttrib in attrib:
518
                    raise ValueError('Specified both a type and a class for the same atom: '+str(attrib))
519
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
520
521
522
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
523
524
            elif typeAttrib in attrib:
                if attrib[typeAttrib] == '':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
525
                    types.append(self._atomClasses[''])
526
                elif attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
527
528
529
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
530
531
            else:
                types.append(None) # Unknown atom type
532
533
        return types

534
    def _parseTorsion(self, attrib):
535
        """Parse the node defining a torsion."""
536
        types = self._findAtomTypes(attrib, 4)
peastman's avatar
peastman committed
537
        if None in types:
538
539
540
541
542
            return None
        torsion = PeriodicTorsion(types)
        index = 1
        while 'phase%d'%index in attrib:
            torsion.periodicity.append(int(attrib['periodicity%d'%index]))
543
544
            torsion.phase.append(_convertParameterToNumber(attrib['phase%d'%index]))
            torsion.k.append(_convertParameterToNumber(attrib['k%d'%index]))
545
            index += 1
Justin MacCallum's avatar
Justin MacCallum committed
546
547
        return torsion

548
    class _SystemData(object):
549
550
551
        """Inner class used to encapsulate data about the system being created."""
        def __init__(self):
            self.atomType = {}
552
            self.atomParameters = {}
553
            self.atomTemplateIndexes = {}
554
            self.atoms = []
555
            self.excludeAtomWith = []
556
            self.virtualSites = {}
557
558
559
560
561
562
            self.bonds = []
            self.angles = []
            self.propers = []
            self.impropers = []
            self.atomBonds = []
            self.isAngleConstrained = []
563
564
565
566
567
568
            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:
569
                if self.constraints[key] != distance:
570
571
572
573
                    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)
574

575
576
577
578
579
580
        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
581
                self.atomTemplateIndexes[atom] = match
582
583
584
                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)
585

586
    class _TemplateData(object):
587
588
589
590
        """Inner class used to encapsulate data about a residue template definition."""
        def __init__(self, name):
            self.name = name
            self.atoms = []
591
            self.atomIndices = {}
592
            self.virtualSites = []
593
594
            self.bonds = []
            self.externalBonds = []
595
            self.overrideLevel = 0
596

597
598
        def getAtomIndexByName(self, atom_name):
            """Look up an atom index by atom name, providing a helpful error message if not found."""
599
600
601
            index = self.atomIndices.get(atom_name, None)
            if index is not None:
                return index
602
603

            # Provide a helpful error message if atom name not found.
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
604
            msg =  "Atom name '%s' not found in residue template '%s'.\n" % (atom_name, self.name)
605
            msg += "Possible atom names are: %s" % str(list(map(lambda x: x.name, self.atoms)))
606
607
            raise ValueError(msg)

608
609
        def addAtom(self, atom):
            self.atoms.append(atom)
610
            self.atomIndices[atom.name] = len(self.atoms)-1
611

612
        def addBond(self, atom1, atom2):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
613
            """Add a bond between two atoms in a template given their indices in the template."""
614
615
616
            self.bonds.append((atom1, atom2))
            self.atoms[atom1].bondedTo.append(atom2)
            self.atoms[atom2].bondedTo.append(atom1)
617

618
        def addBondByName(self, atom1_name, atom2_name):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
619
            """Add a bond between two atoms in a template given their atom names."""
620
621
622
623
624
            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
625
            """Designate that an atom in a residue template has an external bond, using atom index within template."""
626
627
628
629
            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
630
            """Designate that an atom in a residue template has an external bond, using atom name within template."""
631
632
633
            atom = self.getAtomIndexByName(atom_name)
            self.addExternalBond(atom)

634
    class _TemplateAtomData(object):
635
        """Inner class used to encapsulate data about an atom in a residue template definition."""
636
        def __init__(self, name, type, element, parameters={}):
637
638
639
            self.name = name
            self.type = type
            self.element = element
640
            self.parameters = parameters
641
642
643
            self.bondedTo = []
            self.externalBonds = 0

644
    class _BondData(object):
645
646
647
648
649
650
        """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
651

652
    class _VirtualSiteData(object):
653
        """Inner class used to encapsulate data about a virtual site."""
654
        def __init__(self, node, atomIndices):
655
656
657
            attrib = node.attrib
            self.type = attrib['type']
            if self.type == 'average2':
658
                numAtoms = 2
659
660
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
661
                numAtoms = 3
662
663
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
664
                numAtoms = 3
665
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
666
            elif self.type == 'localCoords':
667
668
669
670
671
672
673
674
675
                numAtoms = 0
                self.originWeights = []
                self.xWeights = []
                self.yWeights = []
                while ('wo%d' % (numAtoms+1)) in attrib:
                    numAtoms += 1
                    self.originWeights.append(float(attrib['wo%d' % numAtoms]))
                    self.xWeights.append(float(attrib['wx%d' % numAtoms]))
                    self.yWeights.append(float(attrib['wy%d' % numAtoms]))
676
                self.localPos = [float(attrib['p1']), float(attrib['p2']), float(attrib['p3'])]
677
678
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
679
680
681
682
683
684
            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)]
685
686
687
688
            if 'excludeWith' in attrib:
                self.excludeWith = int(attrib['excludeWith'])
            else:
                self.excludeWith = self.atoms[0]
689

690
691
692
693
694
    class _PatchData(object):
        """Inner class used to encapsulate data about a patch definition."""
        def __init__(self, name, numResidues):
            self.name = name
            self.numResidues = numResidues
695
696
            self.addedAtoms = [[] for i in range(numResidues)]
            self.changedAtoms = [[] for i in range(numResidues)]
697
698
699
700
701
            self.deletedAtoms = []
            self.addedBonds = []
            self.deletedBonds = []
            self.addedExternalBonds = []
            self.deletedExternalBonds = []
peastman's avatar
peastman committed
702
            self.allAtomNames = set()
703

704
705
706
707
        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)))
708

709
            # Construct a new version of each template.
710

711
712
713
714
            newTemplates = []
            for index, template in enumerate(templates):
                newTemplate = ForceField._TemplateData("%s-%s" % (template.name, self.name))
                newTemplates.append(newTemplate)
715

716
                # Build the list of atoms in it.
717

718
719
                for atom in template.atoms:
                    if not any(deleted.name == atom.name and deleted.residue == index for deleted in self.deletedAtoms):
720
                        newTemplate.addAtom(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
721
                for atom in self.addedAtoms[index]:
722
723
                    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))
724
                    newTemplate.addAtom(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
725
726
727
728
729
                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))
730
                    newTemplate.atoms[newAtomIndex[atom.name]] = ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters)
731

732
                # Copy over the virtual sites, translating the atom indices.
733

734
735
736
737
738
739
740
                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)
741

742
                # Build the lists of bonds and external bonds.
743

744
745
746
747
748
                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]
749
                    if a1 in atomMap and a2 in atomMap and (a1.name, a2.name) not in deletedBonds and (a2.name, a1.name) not in deletedBonds:
750
751
752
753
                        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:
754
                        newTemplate.addExternalBond(indexMap[atom])
755
756
757
758
759
760
761
762
763
764
                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
765

766
767
768
769
770
771
772
773
774
775
776
    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

777
    class _AtomType(object):
778
779
780
781
782
783
784
        """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

785
    class _AtomTypeParameters(object):
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
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
        """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))


856
    def _getResidueTemplateMatches(self, res, bondedToAtom, templateSignatures=None, ignoreExternalBonds=False):
857
858
859
860
861
862
        """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.
863
864
        bondedToAtom : list of set of int
            bondedToAtom[i] is the set of atoms bonded to atom index i
865
866
867
868
869
870
871
872
873
874
875
876

        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
877
878
        if templateSignatures is None:
            templateSignatures = self._templateSignatures
879
        signature = _createResidueSignature([atom.element for atom in res.atoms()])
880
        if signature in templateSignatures:
881
            allMatches = []
882
            for t in templateSignatures[signature]:
peastman's avatar
peastman committed
883
                match = compiled.matchResidueToTemplate(res, t, bondedToAtom, ignoreExternalBonds)
884
885
886
887
888
889
                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:
890
                raise Exception('Multiple matching templates found for residue %d (%s): %s.' % (res.index+1, res.name, ', '.join(match[0].name for match in allMatches)))
891
892
        return [template, matches]

893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
    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())
910
911
912
        for (atom1, atom2) in topology.bonds():
            bondedToAtom[atom1.index].add(atom2.index)
            bondedToAtom[atom2.index].add(atom1.index)
913
        return bondedToAtom
914

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

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

920
921
922
        Parameters
        ----------
        topology : Topology
923
            The Topology whose residues are to be checked against the forcefield residue templates.
924
925
926
927

        Returns
        -------
        unmatched_residues : list of Residue
928
            List of Residue objects from `topology` for which no forcefield residue templates are available.
929
930
931
932
933
            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.
934
        bondedToAtom = self._buildBondedToAtomList(topology)
935
        unmatched_residues = list() # list of unmatched residues
936
937
938
939
940
941
        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)
942
943
944

        return unmatched_residues

945
    def getMatchingTemplates(self, topology, ignoreExternalBonds=False):
946
947
948
949
950
951
952
953
        """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.
954
955
        ignoreExternalBonds : bool=False
            If true, ignore external bonds when matching residues to templates.
956
957
958
959
960
961
962
963
964
965
966
967
968
        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.
969
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
970
971
972
973
974
975
976
977
            # 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

978
979
    def generateTemplatesForUnmatchedResidues(self, topology):
        """Generate forcefield residue templates for residues in specified topology for which no forcefield templates are available.
980

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

983
984
985
        Parameters
        ----------
        topology : Topology
986
            The Topology whose residues are to be checked against the forcefield residue templates.
987
988
989

        Returns
        -------
990
991
992
993
994
995
        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]`
996
997
998
999
1000

        """
        # Get a non-unique list of unmatched residues.
        unmatched_residues = self.getUnmatchedResidues(topology)
        # Generate a unique list of unmatched residues by comparing fingerprints.
1001
        bondedToAtom = self._buildBondedToAtomList(topology)
1002
1003
        unique_unmatched_residues = list() # list of unique unmatched Residue objects from topology
        templates = list() # corresponding _TemplateData templates
1004
1005
1006
        signatures = set()
        for residue in unmatched_residues:
            signature = _createResidueSignature([ atom.element for atom in residue.atoms() ])
1007
            template = _createResidueTemplate(residue)
1008
1009
1010
1011
            is_unique = True
            if signature in signatures:
                # Signature is the same as an existing residue; check connectivity.
                for check_residue in unique_unmatched_residues:
peastman's avatar
peastman committed
1012
                    matches = compiled.matchResidueToTemplate(check_residue, template, bondedToAtom, False)
1013
1014
1015
1016
1017
1018
                    if matches is not None:
                        is_unique = False
            if is_unique:
                # Residue is unique.
                unique_unmatched_residues.append(residue)
                signatures.add(signature)
1019
                templates.append(template)
1020

1021
        return [templates, unique_unmatched_residues]
1022

1023
    def createSystem(self, topology, nonbondedMethod=NoCutoff, nonbondedCutoff=1.0*unit.nanometer,
1024
                     constraints=None, rigidWater=True, removeCMMotion=True, hydrogenMass=None, residueTemplates=dict(),
1025
                     ignoreExternalBonds=False, switchDistance=None, flexibleConstraints=False, **args):
1026
        """Construct an OpenMM System representing a Topology with this force field.
Justin MacCallum's avatar
Justin MacCallum committed
1027

Robert McGibbon's avatar
Robert McGibbon committed
1028
1029
1030
1031
1032
1033
        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
1034
            NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME, or LJPME.
Robert McGibbon's avatar
Robert McGibbon committed
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
        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.
1049
        residueTemplates : dict=dict()
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
            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.
1061
1062
1063
        switchDistance : float=None
            The distance at which the potential energy switching function is turned on for
            Lennard-Jones interactions. If this is None, no switching function will be used.
1064
1065
        flexibleConstraints : boolean=False
            If True, parameters for constrained degrees of freedom will be added to the System
Robert McGibbon's avatar
Robert McGibbon committed
1066
        args
1067
1068
1069
            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
1070
1071
1072
1073
1074

        Returns
        -------
        system
            the newly created System
1075
        """
1076
1077
        args['switchDistance'] = switchDistance
        args['flexibleConstraints'] = flexibleConstraints
1078
        data = ForceField._SystemData()
1079
        data.atoms = list(topology.atoms())
1080
1081
        for atom in data.atoms:
            data.excludeAtomWith.append([])
1082
1083

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

1085
        for bond in topology.bonds():
1086
            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
1087
1088

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

1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
        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
1102

1103
        unmatchedResidues = []
1104
1105
        for chain in topology.chains():
            for res in chain.residues():
1106
1107
1108
                if res in residueTemplates:
                    tname = residueTemplates[res]
                    template = self._templates[tname]
peastman's avatar
peastman committed
1109
                    matches = compiled.matchResidueToTemplate(res, template, bondedToAtom, ignoreExternalBonds)
1110
1111
1112
1113
                    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.
1114
                    [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1115
                if matches is None:
1116
1117
1118
                    unmatchedResidues.append(res)
                else:
                    data.recordMatchedAtomParameters(res, template, matches)
1119

1120
        # Try to apply patches to find matches for any unmatched residues.
1121

1122
        if len(unmatchedResidues) > 0:
1123
            unmatchedResidues = _applyPatchesToMatchResidues(self, data, unmatchedResidues, bondedToAtom, ignoreExternalBonds)
1124

1125
1126
        # 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).
1127

1128
1129
        for res in unmatchedResidues:
            # A template might have been generated on an earlier iteration of this loop.
1130
            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1131
1132
1133
1134
1135
            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.
1136
                        [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
                        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)
1147
1148

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

1150
1151
        sys = mm.System()
        for atom in topology.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1152
            # Look up the atom type name, returning a helpful error message if it cannot be found.
1153
1154
1155
1156
            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
1157
            # Look up the type name in the list of registered atom types, returning a helpful error message if it cannot be found.
1158
1159
1160
1161
            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
1162
1163

            # Add the particle to the OpenMM system.
1164
            mass = self._atomTypes[typename].mass
1165
            sys.addParticle(mass)
1166

1167
        # Adjust hydrogen masses if requested.
1168

1169
        if hydrogenMass is not None:
1170
1171
            if not unit.is_quantity(hydrogenMass):
                hydrogenMass *= unit.dalton
1172
            for atom1, atom2 in topology.bonds():
1173
                if atom1.element is elem.hydrogen:
1174
                    (atom1, atom2) = (atom2, atom1)
1175
                if atom2.element is elem.hydrogen and atom1.element not in (elem.hydrogen, None):
1176
1177
1178
                    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
1179

1180
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
1181

1182
1183
1184
        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
1185
        elif nonbondedMethod not in [NoCutoff, CutoffNonPeriodic]:
1186
1187
1188
            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
1189

1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
        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
1205

1206
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1207

1208
1209
1210
        uniquePropers = set()
        for angle in data.angles:
            for atom in bondedToAtom[angle[0]]:
pgrinaway's avatar
pgrinaway committed
1211
                if atom not in angle:
1212
1213
1214
1215
1216
                    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
1217
                if atom not in angle:
1218
1219
1220
1221
1222
                    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
1223

1224
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1225

1226
1227
1228
1229
1230
        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
1231

1232
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1233

1234
1235
1236
1237
1238
1239
1240
        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]
1241
                bond.isConstrained = atom1.element is elem.hydrogen or atom2.element is elem.hydrogen
1242
1243
1244
1245
1246
1247
        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
1248

1249
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1250

1251
1252
1253
1254
1255
1256
        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
1257
                if atom1.element is elem.hydrogen:
1258
                    numH += 1
1259
                if atom3.element is elem.hydrogen:
1260
                    numH += 1
1261
                data.isAngleConstrained.append(numH == 2 or (numH == 1 and atom2.element is elem.oxygen))
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
        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
1272

1273
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
1274

1275
        for atom in data.virtualSites:
1276
            (site, atoms, excludeWith) = data.virtualSites[atom]
1277
            index = atom.index
1278
            data.excludeAtomWith[excludeWith].append(index)
1279
            if site.type == 'average2':
1280
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(atoms[0], atoms[1], site.weights[0], site.weights[1]))
1281
            elif site.type == 'average3':
1282
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
1283
            elif site.type == 'outOfPlane':
1284
1285
                sys.setVirtualSite(index, mm.OutOfPlaneSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
            elif site.type == 'localCoords':
1286
                sys.setVirtualSite(index, mm.LocalCoordinatesSite(atoms, site.originWeights, site.xWeights, site.yWeights, site.localPos))
Justin MacCallum's avatar
Justin MacCallum committed
1287

1288
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
1289

1290
1291
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
1292
1293
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
1294

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

peastman's avatar
peastman committed
1297
1298
1299
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
1300

1301
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
1302

1303
        for script in self._scripts:
1304
            exec(script, locals())
1305
1306
1307
        return sys


1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
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))
1332
1333
1334
    for atom in data.atoms:
        for child in data.excludeAtomWith[atom.index]:
            bondIndices.append((child, atom.index))
1335
1336
    return bondIndices

1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
def _findExclusions(bondIndices, maxSeparation, numAtoms):
    """Identify pairs of atoms in the same molecule separated by no more than maxSeparation bonds."""
    bondedTo = [set() for i in range(numAtoms)]
    for i, j in bondIndices:
        bondedTo[i].add(j)
        bondedTo[j].add(i)

    # Identify all neighbors of each atom with each separation.

    bondedWithSeparation = [bondedTo]
    for i in range(maxSeparation-1):
        lastBonds = bondedWithSeparation[-1]
        newBonds = deepcopy(lastBonds)
        for atom in range(numAtoms):
            for a1 in lastBonds[atom]:
                for a2 in bondedTo[a1]:
                    newBonds[atom].add(a2)
        bondedWithSeparation.append(newBonds)

    # Build the list of pairs.

    pairs = []
    for atom in range(numAtoms):
        for otherAtom in bondedWithSeparation[-1][atom]:
            if otherAtom > atom:
                # Determine the minimum number of bonds between them.
                sep = maxSeparation
                for i in reversed(range(maxSeparation-1)):
                    if otherAtom in bondedWithSeparation[i][atom]:
                        sep -= 1
                    else:
                        break
                pairs.append((atom, otherAtom, sep))
    return pairs


def _findGroups(bondedTo):
    """Given bonds that connect atoms, identify the connected groups."""
    atomGroup = [None]*len(bondedTo)
    numGroups = 0
    for i in range(len(bondedTo)):
        if atomGroup[i] is None:
            # Start a new group.

            atomStack = [i]
            neighborStack = [0]
            group = numGroups
            numGroups += 1

            # Recursively tag all the bonded atoms.

            while len(atomStack) > 0:
                atom = atomStack[-1]
                atomGroup[atom] = group
                while neighborStack[-1] < len(bondedTo[atom]) and atomGroup[bondedTo[atom][neighborStack[-1]]] is not None:
                    neighborStack[-1] += 1
                if neighborStack[-1] < len(bondedTo[atom]):
                    atomStack.append(bondedTo[atom][neighborStack[-1]])
                    neighborStack.append(0)
                else:
                    atomStack.pop()
                    neighborStack.pop()
    return atomGroup

1401
1402
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
1403
1404
    counts = {}
    for element in elements:
1405
        if element in counts:
1406
1407
1408
            counts[element] += 1
        else:
            counts[element] = 1
1409
1410
1411
1412
1413
1414
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
1415
1416
    sig = []
    for c in counts:
1417
1418
        if c is not None:
            sig.append((c, counts[c]))
1419
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
1420

1421
    # Convert it to a string.
1422
1423

    s = ''
1424
    for element, count in sig:
1425
1426
1427
1428
        s += element.symbol+str(count)
    return s


1429
def _applyPatchesToMatchResidues(forcefield, data, residues, bondedToAtom, ignoreExternalBonds):
1430
1431
1432
1433
    """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.
1434

1435
1436
1437
1438
1439
1440
1441
1442
    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
peastman's avatar
peastman committed
1443
                _generatePatchedSingleResidueTemplates(template, patches, 0, newTemplates, set())
1444
1445
1446
1447
1448
1449
                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]
1450

1451
    # Now see if any of those templates matches any of the residues.
1452

1453
1454
    unmatchedResidues = []
    for res in residues:
1455
        [template, matches] = forcefield._getResidueTemplateMatches(res, bondedToAtom, patchedTemplateSignatures, ignoreExternalBonds)
1456
1457
1458
1459
        if matches is None:
            unmatchedResidues.append(res)
        else:
            data.recordMatchedAtomParameters(res, template, matches)
1460
1461
    if len(unmatchedResidues) == 0:
        return []
1462

1463
1464
1465
    # 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.
1466

1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
    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]
1482

1483
    # Record which unmatched residues are bonded to each other.
1484

1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
    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)
1495

1496
1497
    # 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.
1498

1499
1500
1501
1502
    clusterSize = 2
    clusters = bonds
    while clusterSize <= maxPatchSize:
        # Try to apply patches to clusters of this size.
1503

1504
1505
1506
        for patchName in patches:
            patch = forcefield._patches[patchName]
            if patch.numResidues == clusterSize:
1507
                matchedClusters = _matchToMultiResiduePatchedTemplates(data, clusters, patch, patches[patchName], bondedToAtom, ignoreExternalBonds)
1508
1509
1510
1511
1512
1513
                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.
1514

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

1530
1531
1532
    return unmatchedResidues


peastman's avatar
peastman committed
1533
def _generatePatchedSingleResidueTemplates(template, patches, index, newTemplates, alteredAtoms):
1534
1535
    """Apply all possible combinations of a set of single-residue patches to a template."""
    try:
peastman's avatar
peastman committed
1536
1537
1538
1539
1540
1541
1542
        if len(alteredAtoms.intersection(patches[index].allAtomNames)) > 0:
            # This patch would alter an atom that another patch has already altered,
            # so don't apply it.
            patchedTemplate = None
        else:
            patchedTemplate = patches[index].createPatchedTemplates([template])[0]
            newTemplates.append(patchedTemplate)
1543
1544
1545
1546
    except:
        # This probably means the patch is inconsistent with another one that has already been applied,
        # so just ignore it.
        patchedTemplate = None
1547

1548
    # Call this function recursively to generate combinations of patches.
1549

1550
    if index+1 < len(patches):
peastman's avatar
peastman committed
1551
        _generatePatchedSingleResidueTemplates(template, patches, index+1, newTemplates, alteredAtoms)
1552
        if patchedTemplate is not None:
peastman's avatar
peastman committed
1553
1554
            newAlteredAtoms = alteredAtoms.union(patches[index].allAtomNames)
            _generatePatchedSingleResidueTemplates(patchedTemplate, patches, index+1, newTemplates, newAlteredAtoms)
1555
1556


1557
def _matchToMultiResiduePatchedTemplates(data, clusters, patch, residueTemplates, bondedToAtom, ignoreExternalBonds):
1558
1559
1560
    """Apply a multi-residue patch to templates, then try to match them against clusters of residues."""
    matchedClusters = []
    selectedTemplates = [None]*patch.numResidues
1561
    _applyMultiResiduePatch(data, clusters, patch, residueTemplates, selectedTemplates, 0, matchedClusters, bondedToAtom, ignoreExternalBonds)
1562
1563
1564
    return matchedClusters


1565
def _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index, matchedClusters, bondedToAtom, ignoreExternalBonds):
1566
1567
1568
1569
1570
    """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
1571
            _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index+1, matchedClusters, bondedToAtom, ignoreExternalBonds)
1572
1573
1574
    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.
1575

1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
        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.
            return
        newlyMatchedClusters = []
        for cluster in clusters:
            for residues in itertools.permutations(cluster):
                residueMatches = []
                for residue, template in zip(residues, patchedTemplates):
peastman's avatar
peastman committed
1587
                    matches = compiled.matchResidueToTemplate(residue, template, bondedToAtom, ignoreExternalBonds)
1588
1589
1590
1591
1592
1593
                    if matches is None:
                        residueMatches = None
                        break
                    else:
                        residueMatches.append(matches)
                if residueMatches is not None:
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
                    # Each residue individually matches.  Now make sure they're bonded in the correct way.

                    bondsMatch = True
                    for a1, a2 in patch.addedBonds:
                        res1 = a1.residue
                        res2 = a2.residue
                        if res1 != res2:
                            # The patch adds a bond between residues.  Make sure that bond exists.

                            atoms1 = patchedTemplates[res1].atoms
                            atoms2 = patchedTemplates[res2].atoms
                            index1 = next(i for i in range(len(atoms1)) if atoms1[residueMatches[res1][i]].name == a1.name)
                            index2 = next(i for i in range(len(atoms2)) if atoms2[residueMatches[res2][i]].name == a2.name)
                            atom1 = list(residues[res1].atoms())[index1]
                            atom2 = list(residues[res2].atoms())[index2]
                            bondsMatch &= atom2.index in bondedToAtom[atom1.index]
                    if bondsMatch:
                        # We successfully matched the template to the residues.  Record the parameters.
1612

1613
1614
1615
1616
                        for i in range(patch.numResidues):
                            data.recordMatchedAtomParameters(residues[i], patchedTemplates[i], residueMatches[i])
                        newlyMatchedClusters.append(cluster)
                        break
1617

1618
        # Record which clusters were successfully matched.
1619

1620
1621
1622
        matchedClusters += newlyMatchedClusters
        for cluster in newlyMatchedClusters:
            clusters.remove(cluster)
1623

1624

1625
1626
1627
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()])
1628
    numResidueAtoms = sum(residueCounts.values())
1629
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
1630

1631
    # Loop over templates and see how closely each one might match.
1632

1633
1634
1635
1636
1637
1638
    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])
1639

1640
        # Does the residue have any atoms that clearly aren't in the template?
1641

1642
1643
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
1644

1645
        # If there are too many missing atoms, discard this template.
1646

1647
        numTemplateAtoms = sum(templateCounts.values())
1648
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
1649
1650
1651
1652
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
1653

1654
1655
        # 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.
1656

1657
1658
1659
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
1660

1661
        # Accept this as our new best match.
1662

1663
1664
1665
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
1666
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
1667

1668
    # Return an appropriate error message.
1669

1670
    if numBestMatchAtoms == numResidueAtoms:
1671
1672
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
1673
1674
1675
1676
            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:
1677
1678
1679
1680
1681
            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)
1682
1683
1684
        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.'

1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
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():
1703
        template.addAtom(ForceField._TemplateAtomData(atom.name, None, atom.element))
1704
1705
1706
1707
1708
1709
1710
1711
1712
    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
1713

1714
1715
1716
1717
1718
def _matchImproper(data, torsion, generator):
    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]]]
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1719
    wildcard = generator.ff._atomClasses['']
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
    match = None
    for tordef in generator.improper:
        types1 = tordef.types1
        types2 = tordef.types2
        types3 = tordef.types3
        types4 = tordef.types4
        hasWildcard = (wildcard in (types1, types2, types3, types4))
        if match is not None and hasWildcard:
            # Prefer specific definitions over ones with wildcards
            continue
        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:
                    if tordef.ordering == 'default':
                        # 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
                    elif tordef.ordering == 'charmm':
                        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:
                            # There are no wildcards, so the order is unambiguous.
                            match = (torsion[0], torsion[t2[1]], torsion[t3[1]], torsion[t4[1]], tordef)
                        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
                        e2 = data.atoms[a2].element
                        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
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1814
    return match
1815

1816

1817
1818
1819
1820
1821
# 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.

1822
## @private
1823
class HarmonicBondGenerator(object):
1824
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1825

1826
1827
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
1828
        self.bondsForAtomType = defaultdict(set)
1829
1830
1831
1832
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
1833

1834
1835
1836
    def registerBond(self, parameters):
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
peastman's avatar
peastman committed
1837
            index = len(self.types1)
1838
1839
            self.types1.append(types[0])
            self.types2.append(types[1])
peastman's avatar
peastman committed
1840
1841
1842
1843
            for t in types[0]:
                self.bondsForAtomType[t].add(index)
            for t in types[1]:
                self.bondsForAtomType[t].add(index)
1844
1845
            self.length.append(_convertParameterToNumber(parameters['length']))
            self.k.append(_convertParameterToNumber(parameters['k']))
Justin MacCallum's avatar
Justin MacCallum committed
1846

1847
1848
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1849
1850
1851
1852
1853
1854
        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]
1855
        for bond in element.findall('Bond'):
1856
            generator.registerBond(bond.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1857

1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
    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
1869
            for i in self.bondsForAtomType[type1]:
1870
1871
1872
1873
1874
                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:
1875
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
1876
1877
1878
1879
1880
                    if self.k[i] != 0:
                        # flexibleConstraints allows us to add parameters even if the DOF is
                        # constrained
                        if not bond.isConstrained or args.get('flexibleConstraints', False):
                            force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
1881
1882
1883
1884
1885
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


1886
## @private
1887
class HarmonicAngleGenerator(object):
1888
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1889

1890
1891
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
1892
        self.anglesForAtom2Type = defaultdict(list)
1893
1894
1895
1896
1897
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
1898

1899
1900
1901
    def registerAngle(self, parameters):
        types = self.ff._findAtomTypes(parameters, 3)
        if None not in types:
peastman's avatar
peastman committed
1902
            index = len(self.types1)
1903
1904
1905
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.types3.append(types[2])
peastman's avatar
peastman committed
1906
1907
            for t in types[1]:
                self.anglesForAtom2Type[t].append(index)
1908
1909
1910
            self.angle.append(_convertParameterToNumber(parameters['angle']))
            self.k.append(_convertParameterToNumber(parameters['k']))

1911
1912
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1913
1914
1915
1916
1917
1918
        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]
1919
        for angle in element.findall('Angle'):
1920
            generator.registerAngle(angle.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
1921

1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
    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
1934
            for i in self.anglesForAtom2Type[type2]:
1935
1936
1937
1938
1939
1940
                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
1941

1942
1943
1944
1945
1946
1947
1948
1949
1950
                        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
1951

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

1954
1955
1956
1957
1958
                        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]))
1959
                                data.addConstraint(sys, angle[0], angle[2], length)
1960
1961
1962
                    if self.k[i] != 0:
                        if not isConstrained or args.get('flexibleConstraints', False):
                            force.addAngle(angle[0], angle[1], angle[2], self.angle[i], self.k[i])
1963
1964
1965
1966
1967
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


1968
## @private
1969
class PeriodicTorsion(object):
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
    """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 = []
1980
        self.ordering = 'default'
1981

1982
## @private
1983
class PeriodicTorsionGenerator(object):
1984
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1985

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

1992
1993
1994
    def registerProperTorsion(self, parameters):
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
1995
            index = len(self.proper)
1996
            self.proper.append(torsion)
1997
1998
1999
2000
            for t in torsion.types2:
                self.propersForAtomType[t].add(index)
            for t in torsion.types3:
                self.propersForAtomType[t].add(index)
2001

2002
    def registerImproperTorsion(self, parameters, ordering='default'):
2003
2004
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
2005
            if ordering in ['default', 'charmm', 'amber']:
2006
2007
                torsion.ordering = ordering
            else:
2008
                raise ValueError('Illegal ordering type %s for improper torsion %s' % (ordering, torsion))
2009
2010
            self.improper.append(torsion)

2011
2012
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2013
2014
2015
2016
2017
2018
        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]
2019
        for torsion in element.findall('Proper'):
2020
            generator.registerProperTorsion(torsion.attrib)
2021
        for torsion in element.findall('Improper'):
2022
2023
2024
2025
            if 'ordering' in element.attrib:
                generator.registerImproperTorsion(torsion.attrib, element.attrib['ordering'])
            else:
                generator.registerImproperTorsion(torsion.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
2026

2027
2028
2029
2030
2031
2032
2033
2034
2035
    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['']
tic20's avatar
tic20 committed
2036
        proper_cache = {}
2037
        for torsion in data.propers:
tic20's avatar
tic20 committed
2038
2039
            type1, type2, type3, type4 = [data.atomType[data.atoms[torsion[i]]] for i in range(4)]
            sig = (type1, type2, type3, type4)
2040
            sig = frozenset((sig, sig[::-1]))
tic20's avatar
tic20 committed
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
            match = proper_cache.get(sig, None)
            if match == -1:
                continue
            if match is None:
                for index in self.propersForAtomType[type2]:
                    tordef = self.proper[index]
                    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 None:
                    proper_cache[sig] = -1
                else:
                    proper_cache[sig] = match
2061
2062
2063
2064
            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])
2065
        impr_cache = {}
2066
        for torsion in data.impropers:
2067
2068
            t1, t2, t3, t4 = tatoms = [data.atomType[data.atoms[torsion[i]]] for i in range(4)]
            sig = (t1, t2, t3, t4)
2069
2070
2071
2072
            match = impr_cache.get(sig, None)
            if match == -1:
                # Previously checked, and doesn't appear in the database
                continue
2073
2074
2075
2076
            elif match:
                i1, i2, i3, i4, tordef = match
                a1, a2, a3, a4 = (torsion[i] for i in (i1, i2, i3, i4))
                match = (a1, a2, a3, a4, tordef)
2077
2078
2079
            if match is None:
                match = _matchImproper(data, torsion, self)
                if match is not None:
2080
2081
2082
                    order = match[:4]
                    i1, i2, i3, i4 = tuple(torsion.index(a) for a in order)
                    impr_cache[sig] = (i1, i2, i3, i4, match[-1])
2083
2084
                else:
                    impr_cache[sig] = -1
2085
2086
2087
2088
2089
            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])
2090
2091
2092
parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


2093
## @private
2094
class RBTorsion(object):
2095
2096
    """An RBTorsion records the information for a Ryckaert-Bellemans torsion definition."""

2097
    def __init__(self, types, c, ordering='charmm'):
2098
2099
2100
2101
2102
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.c = c
2103
        if ordering in ['default', 'charmm', 'amber']:
2104
2105
            self.ordering = ordering
        else:
2106
            raise ValueError('Illegal ordering type %s for RBTorsion (%s,%s,%s,%s)' % (ordering, types[0], types[1], types[2], types[3]))
2107

2108
## @private
2109
class RBTorsionGenerator(object):
2110
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2111

2112
2113
    def __init__(self, forcefield):
        self.ff = forcefield
2114
2115
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
2116

2117
2118
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2119
2120
2121
2122
2123
2124
        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]
2125
        for torsion in element.findall('Proper'):
2126
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2127
            if None not in types:
2128
2129
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
2130
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2131
            if None not in types:
2132
2133
2134
2135
                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
2136

2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
    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:
2166
            match = _matchImproper(data, torsion, self)
2167
2168
2169
            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])
2170
2171
2172
2173

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


2174
## @private
2175
class CMAPTorsion(object):
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
    """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

2186
## @private
2187
class CMAPTorsionGenerator(object):
2188
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2189

2190
2191
    def __init__(self, forcefield):
        self.ff = forcefield
2192
2193
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
2194

2195
2196
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2197
2198
2199
2200
2201
2202
        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]
2203
2204
2205
2206
2207
2208
2209
        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'):
2210
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
2211
            if None not in types:
2212
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
2213

2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
    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
2224

2225
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
2226

2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
        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


2270
## @private
2271
class NonbondedGenerator(object):
2272
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2273

2274
2275
    SCALETOL = 1e-5

2276
2277
    def __init__(self, forcefield, coulomb14scale, lj14scale):
        self.ff = forcefield
2278
2279
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
2280
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
2281

2282
    def registerAtom(self, parameters):
2283
        self.params.registerAtom(parameters)
2284

2285
2286
2287
2288
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
        if len(existing) == 0:
2289
2290
            generator = NonbondedGenerator(ff, float(element.attrib['coulomb14scale']), float(element.attrib['lj14scale']))
            ff.registerGenerator(generator)
2291
2292
2293
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2294
2295
            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
2296
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
2297
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2298

2299
2300
2301
2302
2303
    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
2304
2305
                     PME:mm.NonbondedForce.PME,
                     LJPME:mm.NonbondedForce.LJPME}
2306
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2307
            raise ValueError('Illegal nonbonded method for NonbondedForce')
2308
2309
        force = mm.NonbondedForce()
        for atom in data.atoms:
2310
2311
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
2312
2313
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2314
2315
2316
        if args['switchDistance'] is not None:
            force.setUseSwitchingFunction(True)
            force.setSwitchingDistance(args['switchDistance'])
peastman's avatar
peastman committed
2317
2318
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
2319
2320
        if 'useDispersionCorrection' in args:
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
peastman's avatar
peastman committed
2321
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
2322

peastman's avatar
peastman committed
2323
    def postprocessSystem(self, sys, data, args):
2324
2325
        # Create the exceptions.

2326
        bondIndices = _findBondsForExclusions(data, sys)
peastman's avatar
peastman committed
2327
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
2328
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
2329
2330
2331

parsers["NonbondedForce"] = NonbondedGenerator.parseElement

2332
2333

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

ChayaSt's avatar
ChayaSt committed
2337
    def __init__(self, forcefield, lj14scale):
ChayaSt's avatar
ChayaSt committed
2338
        self.ff = forcefield
2339
        self.nbfixTypes = {}
ChayaSt's avatar
ChayaSt committed
2340
        self.lj14scale = lj14scale
2341
        self.ljTypes = ForceField._AtomTypeParameters(forcefield, 'LennardJonesForce', 'Atom', ('sigma', 'epsilon'))
ChayaSt's avatar
ChayaSt committed
2342

ChayaSt's avatar
ChayaSt committed
2343
    def registerNBFIX(self, parameters):
ChayaSt's avatar
ChayaSt committed
2344
2345
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
2346
2347
            type1 = types[0][0]
            type2 = types[1][0]
ChayaSt's avatar
ChayaSt committed
2348
2349
            epsilon = _convertParameterToNumber(parameters['epsilon'])
            sigma = _convertParameterToNumber(parameters['sigma'])
2350
2351
            self.nbfixTypes[(type1, type2)] = [sigma, epsilon]
            self.nbfixTypes[(type2, type1)] = [sigma, epsilon]
2352

ChayaSt's avatar
ChayaSt committed
2353
    def registerLennardJones(self, parameters):
2354
        self.ljTypes.registerAtom(parameters)
ChayaSt's avatar
ChayaSt committed
2355
2356
2357

    @staticmethod
    def parseElement(element, ff):
ChayaSt's avatar
ChayaSt committed
2358
        existing = [f for f in ff._forces if isinstance(f, LennardJonesGenerator)]
ChayaSt's avatar
ChayaSt committed
2359
        if len(existing) == 0:
ChayaSt's avatar
ChayaSt committed
2360
            generator = LennardJonesGenerator(ff, float(element.attrib['lj14scale']))
ChayaSt's avatar
ChayaSt committed
2361
2362
            ff.registerGenerator(generator)
        else:
ChayaSt's avatar
ChayaSt committed
2363
            # Multiple <LennardJonesForce> tags were found, probably in different files
ChayaSt's avatar
ChayaSt committed
2364
            generator = existing[0]
2365
2366
            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
2367
2368
        for LJ in element.findall('Atom'):
            generator.registerLennardJones(LJ.attrib)
ChayaSt's avatar
ChayaSt committed
2369
        for Nbfix in element.findall('NBFixPair'):
ChayaSt's avatar
ChayaSt committed
2370
            generator.registerNBFIX(Nbfix.attrib)
ChayaSt's avatar
ChayaSt committed
2371

ChayaSt's avatar
ChayaSt committed
2372
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
peastman's avatar
peastman committed
2373
2374
        # First derive the lookup tables.  We need to include entries for every type
        # that a) appears in the system and b) has unique parameters.
2375
2376

        nbfixTypeSet = set().union(*self.nbfixTypes)
peastman's avatar
peastman committed
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
        allTypes = set(data.atomType[atom] for atom in data.atoms)
        mergedTypes = []
        mergedTypeParams = []
        paramsToMergedType = {}
        typeToMergedType = {}
        for t in allTypes:
            typeParams = self.ljTypes.paramsForType[t]
            params = (typeParams['sigma'], typeParams['epsilon'])
            if t in nbfixTypeSet:
                # NBFIX types cannot be merged.
                typeToMergedType[t] = len(mergedTypes)
                mergedTypes.append(t)
                mergedTypeParams.append(params)
            elif params in paramsToMergedType:
                # We can merge this with another type.
                typeToMergedType[t] = paramsToMergedType[params]
2393
            else:
peastman's avatar
peastman committed
2394
2395
2396
2397
2398
                # This is a new type.
                typeToMergedType[t] = len(mergedTypes)
                paramsToMergedType[params] = len(mergedTypes)
                mergedTypes.append(t)
                mergedTypeParams.append(params)
2399

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

peastman's avatar
peastman committed
2402
        numLjTypes = len(mergedTypes)
2403
        acoef = [0]*(numLjTypes*numLjTypes)
ChayaSt's avatar
ChayaSt committed
2404
        bcoef = acoef[:]
2405
2406
        for m in range(numLjTypes):
            for n in range(numLjTypes):
peastman's avatar
peastman committed
2407
                pair = (mergedTypes[m], mergedTypes[n])
2408
2409
2410
2411
2412
2413
                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
2414
                    continue
ChayaSt's avatar
ChayaSt committed
2415
                else:
peastman's avatar
peastman committed
2416
                    sigma = 0.5*(mergedTypeParams[m][0]+mergedTypeParams[n][0])
2417
                    sigma6 = sigma**6
peastman's avatar
peastman committed
2418
                    epsilon = math.sqrt(mergedTypeParams[m][1]*mergedTypeParams[n][1])
2419
2420
2421
2422
2423
2424
                    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
2425
        self.force.addPerParticleParameter('type')
Peter Eastman's avatar
Peter Eastman committed
2426
        if nonbondedMethod in [CutoffPeriodic, Ewald, PME, LJPME]:
ChayaSt's avatar
ChayaSt committed
2427
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
ChayaSt's avatar
ChayaSt committed
2428
        elif nonbondedMethod is NoCutoff:
ChayaSt's avatar
ChayaSt committed
2429
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
ChayaSt's avatar
ChayaSt committed
2430
        elif nonbondedMethod is CutoffNonPeriodic:
ChayaSt's avatar
ChayaSt committed
2431
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic)
ChayaSt's avatar
ChayaSt committed
2432
        else:
2433
            raise AssertionError('Unrecognized nonbonded method [%s]' % nonbondedMethod)
2434
        if args['switchDistance'] is not None:
2435
2436
            self.force.setUseSwitchingFunction(True)
            self.force.setSwitchingDistance(args['switchDistance'])
2437

ChayaSt's avatar
ChayaSt committed
2438
        # Add the particles
2439

peastman's avatar
peastman committed
2440
2441
        for atom in data.atoms:
            self.force.addParticle((typeToMergedType[data.atomType[atom]],))
2442
        self.force.setUseLongRangeCorrection(bool(args.get('useLongRangeCorrection', False)))
ChayaSt's avatar
ChayaSt committed
2443
        self.force.setCutoffDistance(nonbondedCutoff)
ChayaSt's avatar
ChayaSt committed
2444
2445
2446
        sys.addForce(self.force)

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

2449
        bondIndices = _findBondsForExclusions(data, sys)
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
        forceCopy = deepcopy(self.force)
        forceCopy.createExclusionsFromBonds(bondIndices, 2)
        self.force.createExclusionsFromBonds(bondIndices, 3)
        if self.force.getNumExclusions() > forceCopy.getNumExclusions() and self.lj14scale != 0:
            # We need to create a CustomBondForce and use it to implement the scaled 1-4 interactions.

            bonded = mm.CustomBondForce('%g*epsilon*((sigma/r)^12-(sigma/r)^6)' % (4*self.lj14scale))
            bonded.addPerBondParameter('sigma')
            bonded.addPerBondParameter('epsilon')
            sys.addForce(bonded)
            skip = set(tuple(forceCopy.getExclusionParticles(i)) for i in range(forceCopy.getNumExclusions()))
            for i in range(self.force.getNumExclusions()):
                p1,p2 = self.force.getExclusionParticles(i)
                a1 = data.atoms[p1]
                a2 = data.atoms[p2]
                if (p1,p2) not in skip and (p2,p1) not in skip:
                    type1 = data.atomType[a1]
                    type2 = data.atomType[a2]
                    if (type1, type2) in self.nbfixTypes:
                        sigma, epsilon = self.nbfixTypes[(type1, type2)]
                    else:
                        values1 = self.ljTypes.getAtomParameters(a1, data)
                        values2 = self.ljTypes.getAtomParameters(a2, data)
2473
2474
2475
2476
2477
2478
2479
2480
                        extra1 = self.ljTypes.getExtraParameters(a1, data)
                        extra2 = self.ljTypes.getExtraParameters(a2, data)
                        sigma1 = float(extra1['sigma14']) if 'sigma14' in extra1 else values1[0]
                        sigma2 = float(extra2['sigma14']) if 'sigma14' in extra2 else values2[0]
                        epsilon1 = float(extra1['epsilon14']) if 'epsilon14' in extra1 else values1[1]
                        epsilon2 = float(extra2['epsilon14']) if 'epsilon14' in extra2 else values2[1]
                        sigma = 0.5*(sigma1+sigma2)
                        epsilon = sqrt(epsilon1*epsilon2)
2481
                    bonded.addBond(p1, p2, (sigma, epsilon))
ChayaSt's avatar
ChayaSt committed
2482

ChayaSt's avatar
ChayaSt committed
2483
parsers["LennardJonesForce"] = LennardJonesGenerator.parseElement
2484

2485

2486
## @private
2487
class GBSAOBCGenerator(object):
2488
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2489

2490
2491
    def __init__(self, forcefield):
        self.ff = forcefield
2492
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
2493

2494
    def registerAtom(self, parameters):
2495
        self.params.registerAtom(parameters)
2496

2497
2498
    @staticmethod
    def parseElement(element, ff):
2499
2500
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
2501
2502
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
2503
2504
2505
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2506
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2507

2508
2509
2510
2511
2512
    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
2513
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
2514
2515
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
2516
2517
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
2518
2519
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2520
2521
2522
2523
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
2524
2525
        sys.addForce(force)

2526
2527
    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
2528

2529
2530
2531
2532
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

2533
2534
2535
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


2536
## @private
2537
class CustomBondGenerator(object):
2538
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2539

2540
2541
    def __init__(self, forcefield):
        self.ff = forcefield
2542
2543
2544
2545
2546
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2547

2548
2549
    @staticmethod
    def parseElement(element, ff):
2550
2551
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
2552
2553
2554
2555
2556
2557
        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'):
2558
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2559
            if None not in types:
2560
2561
2562
                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
2563

2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
    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


2584
## @private
2585
class CustomAngleGenerator(object):
2586
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2587

2588
2589
    def __init__(self, forcefield):
        self.ff = forcefield
2590
2591
2592
2593
2594
2595
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2596

2597
2598
    @staticmethod
    def parseElement(element, ff):
2599
2600
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
2601
2602
2603
2604
2605
2606
        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'):
2607
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2608
            if None not in types:
2609
2610
2611
2612
                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
2613

2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
    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


2636
## @private
2637
class CustomTorsion(object):
2638
2639
    """A CustomTorsion records the information for a custom torsion definition."""

2640
    def __init__(self, types, paramValues, ordering='charmm'):
2641
2642
2643
2644
2645
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.paramValues = paramValues
2646
        if ordering in ['default', 'charmm', 'amber']:
2647
2648
            self.ordering = ordering
        else:
2649
            raise ValueError('Illegal ordering type %s for CustomTorsion (%s,%s,%s,%s)' % (ordering, types[0], types[1], types[2], types[3]))
2650

2651
## @private
2652
class CustomTorsionGenerator(object):
2653
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2654

2655
2656
    def __init__(self, forcefield):
        self.ff = forcefield
2657
2658
2659
2660
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
2661

2662
2663
    @staticmethod
    def parseElement(element, ff):
2664
2665
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
2666
2667
2668
2669
2670
2671
        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'):
2672
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2673
            if None not in types:
2674
2675
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
2676
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2677
            if None not in types:
2678
2679
2680
2681
                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
2682

2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
    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:
2711
            match = _matchImproper(data, torsion, self)
2712
2713
2714
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
2715
2716
2717
2718

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


2719
## @private
2720
class CustomNonbondedGenerator(object):
2721
2722
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

2723
2724
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
2725
2726
2727
2728
2729
2730
2731
2732
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2733
2734
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
2735
2736
2737
2738
        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'])
2739
2740
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2741
        generator.functions += _parseFunctions(element)
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753

    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)
2754
        _createFunctions(force, self.functions)
2755
        for atom in data.atoms:
2756
2757
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2758
2759
2760
2761
2762
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

    def postprocessSystem(self, sys, data, args):
2763
        # Create the exclusions.
2764

2765
        bondIndices = _findBondsForExclusions(data, sys)
2766
2767
2768
2769
2770
2771
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomNonbondedForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


2772
## @private
2773
class CustomGBGenerator(object):
2774
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2775

2776
2777
    def __init__(self, forcefield):
        self.ff = forcefield
2778
2779
2780
2781
2782
2783
2784
2785
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2786
2787
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
2788
2789
2790
2791
        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'])
2792
2793
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2794
2795
2796
2797
2798
2799
2800
        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']]))
2801
        generator.functions += _parseFunctions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2802

2803
2804
2805
2806
2807
    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
2808
            raise ValueError('Illegal nonbonded method for CustomGBForce')
2809
2810
2811
2812
2813
2814
2815
2816
2817
        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])
2818
        _createFunctions(force, self.functions)
2819
        for atom in data.atoms:
2820
2821
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2822
2823
2824
2825
2826
2827
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

2828

2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
## @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'):
2867
            types = ff._findAtomTypes(donor.attrib, 3)[:generator.particlesPerDonor]
2868
2869
2870
2871
2872
2873
2874
2875
            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'):
2876
            types = ff._findAtomTypes(acceptor.attrib, 3)[:generator.particlesPerAcceptor]
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
            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]
2920
2921
2922
2923
                    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])
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
        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]
2952
                    if type1 in types1 and type2 in types2:
2953
                        force.addAcceptor(bond.atom1, bond.atom2, -1, self.acceptorParamValues[i])
2954
2955
                    elif type1 in types2 and type2 in types1:
                        force.addAcceptor(bond.atom2, bond.atom1, -1, self.acceptorParamValues[i])
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
2988
2989
2990
2991
2992
2993
        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


2994
## @private
2995
class CustomManyParticleGenerator(object):
2996
2997
2998
2999
3000
3001
3002
3003
3004
3005
3006
3007
    """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 = []
3008

3009
3010
3011
3012
3013
3014
3015
3016
3017
3018
3019
3020
    @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(',')]))
3021
3022
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
3023
3024
3025
3026
3027
3028
3029
3030
3031
3032
3033
3034
3035
3036
3037
3038
3039
3040
3041
3042
3043
3044
3045
3046
3047
3048
3049
3050
3051

    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:
3052
3053
3054
            values = self.params.getAtomParameters(atom, data)
            type = int(self.params.getExtraParameters(atom, data)['filterType'])
            force.addParticle(values, type)
3055
3056
3057
3058
3059
3060
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
        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)))
3073

3074
3075
        # 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.
3076

3077
3078
3079
3080
3081
3082
3083
3084
3085
        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.
3086

3087
3088
3089
3090
3091
        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
3092
def getAtomPrint(data, atomIndex):
3093

Peter Eastman's avatar
Peter Eastman committed
3094
3095
3096
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
3097
    else:
Peter Eastman's avatar
Peter Eastman committed
3098
        returnString = "NA"
3099
3100
3101
3102
3103

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
3104
def countConstraint(data):
3105

Peter Eastman's avatar
Peter Eastman committed
3106
    bondCount = 0
3107
3108
3109
3110
3111
3112
3113
    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
3114
        if (isConstrained):
3115
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
3116

3117
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
3118

3119
## @private
3120
class AmoebaBondGenerator(object):
3121
3122
3123

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

3124
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
3125
3126

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

3128
3129
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
3130
3131
3132
3133
3134
3135
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3136

3137
3138
3139
3140
3141
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3145
        generator = AmoebaBondGenerator(float(element.attrib['bond-cubic']), float(element.attrib['bond-quartic']))
3146
3147
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
3148
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
3149
            if None not in types:
3150
3151
3152
3153
3154
                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:
3155
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
3156
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
3157
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3158
3159
                raise ValueError(outputString)

3160
3161
    #=============================================================================================

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

Mark Friedrichs's avatar
Mark Friedrichs committed
3164
        #countConstraint(data)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3165

Peter Eastman's avatar
Peter Eastman committed
3166
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3167
        existing = [f for f in existing if type(f) == mm.AmoebaBondForce]
3168
        if len(existing) == 0:
3169
            force = mm.AmoebaBondForce()
3170
3171
3172
3173
            sys.addForce(force)
        else:
            force = existing[0]

3174
3175
        force.setAmoebaGlobalBondCubic(self.cubic)
        force.setAmoebaGlobalBondQuartic(self.quartic)
3176
3177
3178
3179
3180
3181
3182
3183
3184
3185

        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:
3186
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
3187
3188
3189
                    if self.k[i] != 0:
                        if not bond.isConstrained or args.get('flexibleConstraints', False):
                            force.addBond(bond.atom1, bond.atom2, self.length[i], self.k[i])
3190
3191
                    break

3192
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
3193
3194
3195
3196

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

Peter Eastman's avatar
Peter Eastman committed
3198
def addAngleConstraint(angle, idealAngle, data, sys):
3199
3200

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

Peter Eastman's avatar
Peter Eastman committed
3202
3203
    bond1 = None
    bond2 = None
3204
3205
3206
3207
3208
3209
3210
    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
3211

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

3214
3215
3216
3217
        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
3218
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
3219
                data.addConstraint(sys, angle[0], angle[2], length)
3220
3221
3222
                return

#=============================================================================================
3223
## @private
3224
class AmoebaAngleGenerator(object):
3225
3226

    #=============================================================================================
3227
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
3228
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3229

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

Peter Eastman's avatar
Peter Eastman committed
3232
3233
3234
3235
3236
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3237

Peter Eastman's avatar
Peter Eastman committed
3238
3239
3240
        self.types1 = []
        self.types2 = []
        self.types3 = []
3241

Peter Eastman's avatar
Peter Eastman committed
3242
3243
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3244

3245
3246
3247
3248
3249
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3253
        generator = AmoebaAngleGenerator(forceField, float(element.attrib['angle-cubic']), float(element.attrib['angle-quartic']),  float(element.attrib['angle-pentic']), float(element.attrib['angle-sextic']))
3254
3255
        forceField._forces.append(generator)
        for angle in element.findall('Angle'):
3256
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
3257
            if None not in types:
3258
3259
3260
3261
3262
3263

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
3264
                angleList.append(float(angle.attrib['angle1']))
3265
3266

                try:
Peter Eastman's avatar
Peter Eastman committed
3267
                    angleList.append(float(angle.attrib['angle2']))
3268
                    try:
Peter Eastman's avatar
Peter Eastman committed
3269
                        angleList.append(float(angle.attrib['angle3']))
3270
3271
3272
3273
3274
3275
3276
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
            else:
3277
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
3278
3279
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
3280
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
3281
3282
                raise ValueError(outputString)

3283
3284
3285
3286
    #=============================================================================================
    # 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
3287

Peter Eastman's avatar
Peter Eastman committed
3288
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3289
3290
3291
3292
3293
3294
        pass

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

Peter Eastman's avatar
Peter Eastman committed
3296
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3297
3298
3299
3300

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3301
        existing = [f for f in existing if type(f) == mm.AmoebaAngleForce]
3302
3303

        if len(existing) == 0:
3304
            force = mm.AmoebaAngleForce()
3305
3306
3307
3308
            sys.addForce(force)
        else:
            force = existing[0]

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3309
        # set scalars
3310

3311
3312
3313
3314
        force.setAmoebaGlobalAngleCubic(self.cubic)
        force.setAmoebaGlobalAngleQuartic(self.quartic)
        force.setAmoebaGlobalAnglePentic(self.pentic)
        force.setAmoebaGlobalAngleSextic(self.sextic)
3315

3316
3317
        DEG_TO_RAD = math.pi / 180

3318
        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
3319
3320
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3321

Peter Eastman's avatar
Peter Eastman committed
3322
3323
3324
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3325
3326
3327
3328
3329
3330
3331
            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]
3332
3333
                        addAngleConstraint(angle, self.angle[i][0]*DEG_TO_RAD, data, sys)
                    if self.k[i] != 0 and (not isConstrained or args.get('flexibleConstraints', False)):
Peter Eastman's avatar
Peter Eastman committed
3334
3335
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
3336
3337
3338
3339
3340
3341
                            # 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
3342
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
3343
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3344
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
3345
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3346
                            if (numberOfHydrogens < lenAngle):
3347
3348
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
3349
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
3350
                                raise ValueError(outputString)
3351
3352
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
3353

3354
                        angleDict['idealAngle'] = angleValue
Peter Eastman's avatar
Peter Eastman committed
3355
                        force.addAngle(angle[0], angle[1], angle[2], angleValue, self.k[i])
3356
3357
3358
3359
3360
3361
                    break

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

Peter Eastman's avatar
Peter Eastman committed
3363
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3364
3365
3366
3367

        # get force

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3368
        existing = [f for f in existing if type(f) == mm.AmoebaInPlaneAngleForce]
3369
3370

        if len(existing) == 0:
3371
            force = mm.AmoebaInPlaneAngleForce()
3372
3373
3374
3375
3376
3377
            sys.addForce(force)
        else:
            force = existing[0]

        # scalars

3378
3379
3380
3381
        force.setAmoebaGlobalInPlaneAngleCubic(self.cubic)
        force.setAmoebaGlobalInPlaneAngleQuartic(self.quartic)
        force.setAmoebaGlobalInPlaneAnglePentic(self.pentic)
        force.setAmoebaGlobalInPlaneAngleSextic(self.sextic)
3382
3383

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

Peter Eastman's avatar
Peter Eastman committed
3385
3386
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3387

Peter Eastman's avatar
Peter Eastman committed
3388
3389
3390
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3391
3392
3393
3394
3395
3396
3397
3398
3399

            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
3400
                    if (isConstrained and self.k[i] != 0.0):
3401
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3402
                    if self.k[i] != 0.0 and (not isConstrained or args.get('flexibleConstraints', False)):
3403
3404
3405
                        force.addAngle(angle[0], angle[1], angle[2], angle[3], self.angle[i][0], self.k[i])
                    break

3406
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
3407
3408
3409

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
3410
3411
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
3412
3413
#=============================================================================================

3414
## @private
3415
class AmoebaOutOfPlaneBendGenerator(object):
3416
3417
3418
3419

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

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

3421
3422
3423
3424
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3425
3426
3427
3428
3429
3430
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3431

Peter Eastman's avatar
Peter Eastman committed
3432
3433
3434
3435
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3436

Peter Eastman's avatar
Peter Eastman committed
3437
        self.ks = []
3438
3439
3440
3441
3442

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

3444
3445
3446
3447
3448
3449
3450
3451
3452
3453
3454
3455
3456
3457
3458
3459
3460
3461
3462
3463
3464
3465
3466
3467
3468
    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
3469

3470
3471
        # get global scalar parameters

Peter Eastman's avatar
Peter Eastman committed
3472
        generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'],
3473
3474
3475
                                                   float(element.attrib['opbend-cubic']),
                                                   float(element.attrib['opbend-quartic']),
                                                   float(element.attrib['opbend-pentic']),
Peter Eastman's avatar
Peter Eastman committed
3476
                                                   float(element.attrib['opbend-sextic']))
3477
3478
3479
3480

        forceField._forces.append(generator)

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

Peter Eastman's avatar
Peter Eastman committed
3484
3485
3486
3487
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
3488

Peter Eastman's avatar
Peter Eastman committed
3489
                generator.ks.append(float(angle.attrib['k']))
3490
3491

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3492
3493
                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
3494
3495
                raise ValueError(outputString)

3496
3497
3498
3499
3500
3501
    #=============================================================================================
    # 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
3502

Peter Eastman's avatar
Peter Eastman committed
3503
    def getMiddleAtom(self, angle, data):
3504
3505
3506

        # find atom shared by both bonds making up the angle

Peter Eastman's avatar
Peter Eastman committed
3507
        middleAtom = -1
Justin MacCallum's avatar
Justin MacCallum committed
3508
        for atomIndex in angle:
Peter Eastman's avatar
Peter Eastman committed
3509
            isMiddle = 0
3510
3511
3512
            for bond in data.atomBonds[atomIndex]:
                atom1 = data.bonds[bond].atom1
                atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3513
                if (atom1 != atomIndex):
3514
3515
3516
                    partner = atom1
                else:
                    partner = atom2
Justin MacCallum's avatar
Justin MacCallum committed
3517
                if (partner == angle[0] or partner == angle[1] or partner == angle[2]):
3518
3519
                    isMiddle += 1

Peter Eastman's avatar
Peter Eastman committed
3520
            if (isMiddle == 2):
3521
3522
3523
3524
3525
                return atomIndex
        return -1

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

Peter Eastman's avatar
Peter Eastman committed
3526
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3527
3528
3529
3530
3531
3532
3533
3534
3535
3536
3537
3538
3539

        # 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
3540
3541
3542
3543
        force.setAmoebaGlobalOutOfPlaneBendCubic(  self.cubic)
        force.setAmoebaGlobalOutOfPlaneBendQuartic(self.quartic)
        force.setAmoebaGlobalOutOfPlaneBendPentic( self.pentic)
        force.setAmoebaGlobalOutOfPlaneBendSextic( self.sextic)
3544
3545
3546
3547

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

Peter Eastman's avatar
Peter Eastman committed
3548
        skipAtoms = dict()
3549
3550
3551
3552

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

Peter Eastman's avatar
Peter Eastman committed
3553
3554
        inPlaneAngles = []
        nonInPlaneAngles = []
3555
        nonInPlaneAnglesConstrained = []
Peter Eastman's avatar
Peter Eastman committed
3556
        idealAngles = []*len(data.angles)
3557
3558
3559

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

Peter Eastman's avatar
Peter Eastman committed
3560
3561
3562
            middleAtom = self.getMiddleAtom(angle, data)
            if (middleAtom > -1):
                middleType = data.atomType[data.atoms[middleAtom]]
3563
3564
                middleCovalency = len(data.atomBonds[middleAtom])
            else:
Peter Eastman's avatar
Peter Eastman committed
3565
                middleType = -1
3566
3567
                middleCovalency = -1

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

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

Peter Eastman's avatar
Peter Eastman committed
3576
3577
3578
3579
                partners = []
                partnerSet = set()
                partnerTypes = []
                partnerK = []
3580
3581
3582
3583

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
Peter Eastman's avatar
Peter Eastman committed
3584
                    if (atom1 != middleAtom):
3585
3586
3587
3588
3589
3590
3591
3592
                        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
3593
3594
3595
3596
3597
                        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
3598

Peter Eastman's avatar
Peter Eastman committed
3599
                if (len(partners) == 3):
3600

Peter Eastman's avatar
Peter Eastman committed
3601
3602
3603
                    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])
3604

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

3607
                    skipAtoms[middleAtom] = set()
Peter Eastman's avatar
Peter Eastman committed
3608
3609
3610
3611
                    skipAtoms[middleAtom].add(partners[0])
                    skipAtoms[middleAtom].add(partners[1])
                    skipAtoms[middleAtom].add(partners[2])

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3612
3613
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
3614
3615
3616
3617
3618
                    angleDict = {}
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
3619
3620
                    angleDict['angle'] = angleList

Peter Eastman's avatar
Peter Eastman committed
3621
                    angleDict['isConstrained'] = 0
3622

Peter Eastman's avatar
Peter Eastman committed
3623
3624
3625
3626
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3627
3628

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3629
3630
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3631

Peter Eastman's avatar
Peter Eastman committed
3632
                    inPlaneAngles.append(angleDict)
3633
3634

                else:
Peter Eastman's avatar
Peter Eastman committed
3635
3636
3637
3638
                    angleDict = {}
                    angleDict['angle'] = angle
                    angleDict['isConstrained'] = isConstrained
                    nonInPlaneAngles.append(angleDict)
3639
            else:
Peter Eastman's avatar
Peter Eastman committed
3640
                if (middleAtom > -1 and middleCovalency == 3 and middleAtom in skipAtoms):
3641

Peter Eastman's avatar
Peter Eastman committed
3642
                    partnerSet = skipAtoms[middleAtom]
Justin MacCallum's avatar
Justin MacCallum committed
3643

Peter Eastman's avatar
Peter Eastman committed
3644
                    angleDict = {}
3645

Peter Eastman's avatar
Peter Eastman committed
3646
3647
3648
3649
3650
                    angleList = []
                    angleList.append(angle[0])
                    angleList.append(angle[1])
                    angleList.append(angle[2])
                    angleDict['angle'] = angleList
3651

Peter Eastman's avatar
Peter Eastman committed
3652
                    angleDict['isConstrained'] = isConstrained
3653

Peter Eastman's avatar
Peter Eastman committed
3654
3655
3656
3657
                    angleSet = set()
                    angleSet.add(angle[0])
                    angleSet.add(angle[1])
                    angleSet.add(angle[2])
3658
3659

                    for atomIndex in partnerSet:
Peter Eastman's avatar
Peter Eastman committed
3660
3661
                        if (atomIndex not in angleSet):
                            angleList.append(atomIndex)
3662

Peter Eastman's avatar
Peter Eastman committed
3663
                    inPlaneAngles.append(angleDict)
3664
3665

                else:
Peter Eastman's avatar
Peter Eastman committed
3666
3667
                    angleDict = {}
                    angleDict['angle'] = angle
3668
                    angleDict['isConstrained'] = isConstrained
Peter Eastman's avatar
Peter Eastman committed
3669
                    nonInPlaneAngles.append(angleDict)
3670

3671
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
3672
3673

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3674
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
3675
3676
                force.createForcePostOpBendAngle(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, nonbondedMethod, nonbondedCutoff, inPlaneAngles, args)
3677
3678

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3679
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
3680
                for angleDict in inPlaneAngles:
Peter Eastman's avatar
Peter Eastman committed
3681
                    nonInPlaneAngles.append(angleDict)
3682
                force.createForcePostAmoebaBondForce(sys, data, nonbondedMethod, nonbondedCutoff, nonInPlaneAngles, args)
3683
3684
3685
3686
3687

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

3688
## @private
3689
class AmoebaTorsionGenerator(object):
3690
3691
3692
3693
3694
3695
3696

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
3697
        self.torsionUnit = torsionUnit
3698

Peter Eastman's avatar
Peter Eastman committed
3699
3700
3701
3702
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3703

Peter Eastman's avatar
Peter Eastman committed
3704
3705
3706
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
3707

3708
3709
3710
3711
3712
3713
3714
3715
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3717
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
3718
3719
3720
3721
3722
3723
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
3724
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
3725
            if None not in types:
3726
3727
3728
3729
3730
3731
3732

                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
3733
3734
                    tInfo = []
                    suffix = str(ii)
3735
3736
3737
3738
3739
3740
                    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
3741
3742
3743
3744
3745
3746
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
3747
3748
3749

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
3750
3751
3752
3753
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
3754
3755
                raise ValueError(outputString)

3756
3757
3758
3759
3760
    #=============================================================================================

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

        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
3761
        existing = [f for f in existing if type(f) == mm.PeriodicTorsionForce]
3762
        if len(existing) == 0:
3763
            force = mm.PeriodicTorsionForce()
3764
3765
3766
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3767

3768
3769
3770
3771
3772
3773
3774
3775
3776
3777
3778
3779
3780
3781
3782
3783
        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
3784
                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):
3785
3786
3787
3788
3789
3790
                    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])
3791
3792
3793
3794
3795
3796
                    break

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

3797
## @private
3798
class AmoebaPiTorsionGenerator(object):
3799
3800
3801
3802
3803
3804

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

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

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

3806
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
3807
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
3808
3809
3810
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3811

3812
3813
3814
3815
3816
3817
3818
3819
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
3820
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
3821
3822
3823
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
3824
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
3825
            if None not in types:
3826
3827
3828
3829
3830
3831
                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
3832
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3833
3834
                raise ValueError(outputString)

3835
3836
3837
3838
3839
3840
3841
3842
3843
3844
3845
3846
3847
3848
3849
3850
    #=============================================================================================

    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
3851

3852
3853
            atom1 = bond.atom1
            atom2 = bond.atom2
Justin MacCallum's avatar
Justin MacCallum committed
3854

3855
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
3856
3857
3858
3859
3860
3861
3862
3863
3864
3865
3866

                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
3867
3868
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
3869
3870
3871
3872
3873
3874
3875
3876
3877
3878
3879
3880

                       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
3881
                           if (bondedAtom1 != atom1):
3882
3883
3884
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
3885
3886
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3887
                                   piTorsionAtom1 = b1
3888
3889
3890
3891
3892
3893
                               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
3894
                           if (bondedAtom1 != atom2):
3895
3896
3897
3898
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
3899
3900
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
3901
                                   piTorsionAtom5 = b1
3902
3903
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
3904

Peter Eastman's avatar
Peter Eastman committed
3905
                       force.addPiTorsion(piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6, self.k[i])
3906
3907
3908
3909
3910

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

3911
## @private
3912
class AmoebaTorsionTorsionGenerator(object):
3913
3914
3915
3916
3917

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

Peter Eastman's avatar
Peter Eastman committed
3918
    def __init__(self):
3919

Peter Eastman's avatar
Peter Eastman committed
3920
3921
3922
3923
3924
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
3925

Peter Eastman's avatar
Peter Eastman committed
3926
        self.gridIndex = []
3927

Peter Eastman's avatar
Peter Eastman committed
3928
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
3929

3930
3931
3932
3933
3934
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
3935
        generator = AmoebaTorsionTorsionGenerator()
3936
3937
3938
3939
3940
3941
3942
        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'):
3943
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
3944
            if None not in types:
3945
3946
3947
3948
3949
3950
3951

                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
3952
3953
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
3954
3955
3956
3957
3958
3959
3960
3961
3962
                    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
3963
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
3964
3965
                raise ValueError(outputString)

3966
3967
3968
3969
3970
3971
3972
3973
3974
3975
        # 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
3976
3977
3978
3979
3980
3981
        #     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
3982
3983

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

Peter Eastman's avatar
Peter Eastman committed
3987
3988
3989
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
3990

Peter Eastman's avatar
Peter Eastman committed
3991
3992
            grid = []
            gridCol = []
3993
3994
3995
3996
3997

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
3998
3999
4000
4001
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
4002
                if 'fx' in gridEntry.attrib:
4003
4004
4005
                    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
4006
                gridCol.append(gridRow)
4007
4008

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
4009
4010
4011
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
4012
4013
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
4014

Peter Eastman's avatar
Peter Eastman committed
4015
4016
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
4017
            else:
Peter Eastman's avatar
Peter Eastman committed
4018
4019
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
4020
4021
4022
4023
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
4024
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
4025
4026
4027
4028
4029
4030
4031
4032

        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
4033
        if (len(data.atomBonds[atomC]) == 4):
4034
4035
4036
4037
4038
            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
4039
4040
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
4041
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
4042
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
4043
4044
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
4045
4046
                if (hit > -1):
                    if (atomE == -1):
4047
4048
4049
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
4050

4051
4052
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
4053
            if (atomE == -1 or atomF == -1):
4054
                outputString = "getChiralAtomIndex: error getting bonded partners of atomC=%s %d %s" % (atomC.name, atomC.residue.index, atomC.residue.name,)
Justin MacCallum's avatar
Justin MacCallum committed
4055
                raise ValueError(outputString)
4056
4057
4058
4059
4060

            # 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
4061
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
4062
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
4063
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
4064
                chiralAtomIndex = atomF
4065

Peter Eastman's avatar
Peter Eastman committed
4066
4067
4068
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
4069
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
4070
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
4071
                chiralAtomIndex = massF
4072
4073
4074
4075

        return chiralAtomIndex

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

4077
4078
4079
4080
4081
4082
4083
4084
4085
4086
4087
4088
4089
    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
4090
4091
            # search for bitorsions; based on TINKER subroutine bitors()

4092
4093
4094
4095
4096
4097
4098
            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
4099
                if (bondedAtom1 != ib):
4100
4101
4102
4103
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4104
                if (ia != ic and ia != id):
4105
4106
4107
                    for bondIndex in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex].atom1
                        bondedAtom2 = data.bonds[bondIndex].atom2
Peter Eastman's avatar
Peter Eastman committed
4108
                        if (bondedAtom1 != id):
4109
4110
4111
4112
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4113
                        if (ie != ic and ie != ib and ie != ia):
4114
4115
4116
4117
4118
4119
4120
4121
4122
4123
4124
4125
4126
4127
4128
4129
4130
4131
4132
4133

                            # 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
4134
4135
4136
                                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])
4137
4138
4139

                                # match in reverse order

4140
                                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
4141
4142
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
4143
4144
4145
4146

        # set grids

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

4149
4150
4151
4152
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

4153
## @private
4154
class AmoebaStretchBendGenerator(object):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4155
4156

    #=============================================================================================
4157
4158
4159
4160
4161
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

    def __init__(self):

Peter Eastman's avatar
Peter Eastman committed
4162
4163
4164
        self.types1 = []
        self.types2 = []
        self.types3 = []
4165

Peter Eastman's avatar
Peter Eastman committed
4166
4167
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
4168

4169
4170
4171
4172
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
Peter Eastman's avatar
Peter Eastman committed
4173
        generator = AmoebaStretchBendGenerator()
4174
4175
4176
4177
4178
4179
4180
        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'):
4181
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
4182
            if None not in types:
4183
4184
4185
4186
4187
4188
4189
4190
4191
4192
4193
4194

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

4198
4199
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
4200
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
4201
4202
    # 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
4203
4204
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
4205
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
4206
4207
4208

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

Peter Eastman's avatar
Peter Eastman committed
4209
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4210
4211
4212
4213
4214
4215
4216
4217
        pass

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

    # Note: request for constrained bonds is ignored.

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

4218
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
4219
4220
4221
4222
4223
4224
4225
4226
4227
4228
4229
4230

        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
4231
            if ('isConstrained' in angleDict):
4232
4233
4234
4235
4236
4237
4238
4239
                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
4240
            radian = 57.2957795130
4241
4242
4243
4244
4245
4246
4247
4248
4249
4250
            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
4251
                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
4252
4253
4254
4255
4256
4257
4258
4259
4260
4261
4262
4263
4264
4265
4266
                    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
4267

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4270
                    if ('idealAngle' not in angleDict):
4271

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4272
4273
4274
4275
4276
                       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
4277
                       raise ValueError(outputString)
4278

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4279
4280
4281
4282
4283
4284
4285
                    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
4286
                       raise ValueError(outputString)
4287

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4291
                    break
4292
4293
4294
4295
4296

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

4297
## @private
4298
class AmoebaVdwGenerator(object):
4299
4300

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

4302
4303
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4306
        self.type = type
4307

Peter Eastman's avatar
Peter Eastman committed
4308
4309
4310
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
4311

Peter Eastman's avatar
Peter Eastman committed
4312
        self.epsilonrule = epsilonrule
4313

Peter Eastman's avatar
Peter Eastman committed
4314
4315
4316
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
4317
4318
4319
4320
4321
4322
4323

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

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

4327
4328
4329
        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
4330
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
4331
4332
4333
4334
4335
4336
4337
4338
4339
4340
4341
4342
            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')
4343
        generator.params.parseDefinitions(element)
4344
4345
4346
4347
4348
        two_six = 1.122462048309372

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

    @staticmethod
4349
    def getBondedParticleSets(sys, data):
4350

4351
4352
4353
4354
4355
4356
        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
4357

4358
4359
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
4362
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
4363
4364
        epsilonMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'HARMONIC':1, 'HHG':1}

4365
4366
        force = mm.AmoebaVdwForce()
        sys.addForce(force)
4367

4368
        # sigma and epsilon combining rules
4369

4370
4371
4372
4373
        if ('sigmaCombiningRule' in args):
            sigmaRule = args['sigmaCombiningRule'].upper()
            if (sigmaRule.upper() in sigmaMap):
                force.setSigmaCombiningRule(sigmaRule.upper())
4374
            else:
4375
4376
4377
4378
                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)
4379

4380
4381
4382
4383
        if ('epsilonCombiningRule' in args):
            epsilonRule = args['epsilonCombiningRule'].upper()
            if (epsilonRule.upper() in epsilonMap):
                force.setEpsilonCombiningRule(epsilonRule.upper())
4384
            else:
4385
4386
4387
4388
                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
4389

4390
        # cutoff
4391

4392
4393
4394
4395
        if ('vdwCutoff' in args):
            force.setCutoff(args['vdwCutoff'])
        else:
            force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4396

4397
        # dispersion correction
4398

4399
4400
        if ('useDispersionCorrection' in args):
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
4401

4402
4403
        if (nonbondedMethod == PME):
            force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
4404
4405
4406

        # add particles to force

4407
4408
4409
4410
4411
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
4412
        for (i, atom) in enumerate(data.atoms):
4413
4414
4415
4416
4417
4418
4419
4420
4421
4422
            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
4423

4424
            force.addParticle(ivIndex, values[0]*sigmaScale, values[1], values[2])
4425
4426
4427
4428
4429
4430
4431

        # 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

4432
        bondedParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4433
4434

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

4436
4437
4438
4439
4440
4441
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
4442
            if (self.vdw13Scale == 0.0):
4443
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
4444
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
4445
4446
4447
4448
4449

            # self

            exclusionSet.add(i)

4450
            force.setParticleExclusions(i, tuple(exclusionSet))
4451
4452
4453
4454
4455

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

4456
## @private
4457
class AmoebaMultipoleGenerator(object):
4458
4459
4460
4461

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

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

4463
4464
    #=============================================================================================

4465
    def __init__(self, forceField):
Peter Eastman's avatar
Peter Eastman committed
4466
4467
        self.forceField = forceField
        self.typeMap = {}
4468
4469
4470
4471
4472
4473

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
4474
    def setAxisType(kIndices):
4475
4476
4477

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
4478
4479
4480
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
4481
                else:
Peter Eastman's avatar
Peter Eastman committed
4482
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
4483

Peter Eastman's avatar
Peter Eastman committed
4484
4485
                if (kIndicesLen > 2):
                    kx = kIndices[2]
4486
                else:
Peter Eastman's avatar
Peter Eastman committed
4487
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
4488

Peter Eastman's avatar
Peter Eastman committed
4489
4490
                if (kIndicesLen > 1):
                    kz = kIndices[1]
4491
                else:
Peter Eastman's avatar
Peter Eastman committed
4492
                    kz = 0
4493

Peter Eastman's avatar
Peter Eastman committed
4494
4495
                while(len(kIndices) < 4):
                    kIndices.append(0)
4496
4497

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
4498
                if (kz == 0):
4499
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
4500
                if (kz != 0 and kx == 0):
4501
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
4502
                if (kz < 0 or kx < 0):
4503
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
4504
                if (kx < 0 and ky < 0):
4505
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
4506
                if (kz < 0 and kx < 0 and ky  < 0):
4507
4508
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
4509
4510
4511
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
4512
4513
4514
4515
4516
4517
4518
4519

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
4520
        #   <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"  >
4521
4522
4523
        # <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"  />

4524
4525
4526
4527
4528
4529
4530
        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]
4531
4532
4533
4534

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

        for atom in element.findall('Multipole'):
4535
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4536
            if None not in types:
4537
4538
4539
4540
4541
4542
4543
4544

                # 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
4545
4546
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
4547
                    except:
4548
4549
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
4550
                # set axis type based on k-Indices
4551

Peter Eastman's avatar
Peter Eastman committed
4552
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
4553
4554
4555

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
4558
4559
4560
4561
4562
4563
4564
4565
4566
4567
4568
4569
4570
                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']))
4571
4572

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4573
                    if (t not in generator.typeMap):
4574
4575
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
4576
4577
4578
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
4579
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
4580
4581
4582
4583
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
4584

4585
            else:
Peter Eastman's avatar
Peter Eastman committed
4586
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4587
4588
                raise ValueError(outputString)

4589
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
4590

4591
        for atom in element.findall('Polarize'):
4592
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4593
            if None not in types:
4594

Peter Eastman's avatar
Peter Eastman committed
4595
4596
4597
4598
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
4599
4600
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
4601
                    pdamp = pow(polarizability, 1.0/6.0)
4602

Peter Eastman's avatar
Peter Eastman committed
4603
4604
                pgrpMap = dict()
                for index in range(1, 7):
4605
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
4606
                    if (pgrp in atom.attrib):
4607
4608
4609
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4610
4611
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
4612
                        raise ValueError(outputString)
4613
4614
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
4615
4616
4617
4618
4619
4620
4621
                        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
4622
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
4623
4624
4625
4626
4627
                                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
4628
4629
                            raise ValueError(outputString)

4630
            else:
Peter Eastman's avatar
Peter Eastman committed
4631
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
4632
4633
                raise ValueError(outputString)

4634
4635
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
4636
    def setPolarGroups(self, data, bonded12ParticleSets, force):
4637
4638
4639
4640
4641

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
4642
4643
4644
4645
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
4646
4647
4648
4649
            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
4650
4651
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
4652
4653
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
4654

4655
4656
4657
4658
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
4659
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
4660
4661
                continue

Peter Eastman's avatar
Peter Eastman committed
4662
4663
            group = set()
            visited = set()
4664
4665
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4666
4667
4668
4669
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
4670
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
4671
4672
                if (nextAtom not in visited):
                   visited.add(nextAtom)
4673
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
4674
4675
4676
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
4677
4678
4679

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
4680
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
4681
4682

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4683
4684
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
4685
4686
4687
4688
4689

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
4690
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
4691
4692
                continue

Peter Eastman's avatar
Peter Eastman committed
4693
            pgrp11 = set(atom.polarizationGroupSet[0])
4694
4695
4696
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4697
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
4698
4699
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4700
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
4701

4702
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4703
4704
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
4705
4706
4707
4708
4709

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
4710
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
4711
4712
                continue

Peter Eastman's avatar
Peter Eastman committed
4713
4714
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
4715
4716
4717
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4718
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
4719
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4720
            pgrp13 = pgrp13 - set(pgrp11)
4721
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
4722
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
4723

4724
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4725
4726
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
4727
4728
4729
4730
4731

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
4732
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
4733
4734
                continue

Peter Eastman's avatar
Peter Eastman committed
4735
4736
4737
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
4738
4739
4740
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
4741
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
4742
4743
4744

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
4745
            pgrp14 = pgrp14 - set(pgrp11)
4746
4747

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

4750
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
4751
4752
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
4753
4754
4755

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

Peter Eastman's avatar
Peter Eastman committed
4756
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4757
4758
4759
4760

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

4761
4762
4763
4764
4765
4766
4767
        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)
4768

4769
4770
        if ('ewaldErrorTolerance' in args):
            force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))
4771

4772
4773
4774
4775
4776
4777
4778
4779
        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)
4780

4781
4782
        if ('aEwald' in args):
            force.setAEwald(float(args['aEwald']))
4783

4784
4785
        if ('pmeGridDimensions' in args):
            force.setPmeGridDimensions(args['pmeGridDimensions'])
4786

4787
4788
        if ('mutualInducedMaxIterations' in args):
            force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
4789

4790
4791
        if ('mutualInducedTargetEpsilon' in args):
            force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))
4792
4793

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
4794
        # throw error if particle type not available
4795
4796
4797
4798
4799

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

        # 1-2

4800
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4801
4802
4803
4804
4805

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4806
            bonded13Set = set()
4807
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4808
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4809
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
4810
4811
4812
4813

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4814
            selfSet = set()
4815
4816
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4817
4818
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
4819
4820
4821
4822
4823

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4824
4825
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4826
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4827
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
4828

4829
4830
4831
4832
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
4833
            selfSet = set()
4834
4835
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4836
4837
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
4838
4839
4840
4841
4842

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
4843
4844
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
4845
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
4846
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
4847
4848
4849
4850
4851
4852

            # 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
4853
            selfSet = set()
4854
4855
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
4856
4857
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
4858
4859
4860
4861
4862

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

Peter Eastman's avatar
Peter Eastman committed
4863
4864
                multipoleList = self.typeMap[t]
                hit = 0
4865
4866
4867
4868
4869
4870
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4871
                    if (hit != 0):
4872
4873
                        break

Peter Eastman's avatar
Peter Eastman committed
4874
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4875
4876

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4877
4878
                    kx = kIndices[2]
                    ky = kIndices[3]
4879
4880
4881
4882

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

4884
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
4885
4886
4887
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4888
4889
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
4890
                       if (hit != 0):
4891
4892
4893
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
4894
4895
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
4896
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4897
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4898
4899
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4900
4901
4902
4903
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4904
4905
4906
4907
4908
4909
4910
4911
4912
4913
                                      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

4914
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4915
                                      hit = 1
4916
4917
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
4918
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4919
4920
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4921
4922
4923
4924
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4925
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4926
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
4927

4928
4929
4930
4931
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4932
                    if (hit != 0):
4933
4934
                        break

Peter Eastman's avatar
Peter Eastman committed
4935
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
4936
4937

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
4938
4939
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
4940

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

4945
4946
4947
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
4948
4949
4950
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
4951
4952
4953

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
4954
                       if (hit != 0):
4955
4956
4957
                           break

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

Peter Eastman's avatar
Peter Eastman committed
4960
                       if (bondedAtomZType == kz):
4961
4962
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
4963
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
4964
4965
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
4966
4967
4968
4969
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
4970
4971
4972
4973
4974
4975
4976
4977

                                      # 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

4978
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4979
                                      hit = 3
4980
4981
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
4982
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
4983
4984
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
4985
4986
4987
4988
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
4989
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
4990
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
4991

4992
4993
4994
4995
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
4996
                    if (hit != 0):
4997
4998
                        break

Peter Eastman's avatar
Peter Eastman committed
4999
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5000
5001
5002
5003

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

Peter Eastman's avatar
Peter Eastman committed
5004
5005
5006
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5007
5008
5009

                    for bondedAtomZIndex in bondedAtom12Indices:

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

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

Peter Eastman's avatar
Peter Eastman committed
5016
                        if (kx == 0 and kz == bondedAtomZType):
5017
                            zaxis = bondedAtomZIndex
5018
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5019
                            hit = 5
5020
5021
5022
5023
5024

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5025
                    if (hit != 0):
5026
5027
                        break

Peter Eastman's avatar
Peter Eastman committed
5028
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5029
5030
5031

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
5032
5033
5034
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5035

Peter Eastman's avatar
Peter Eastman committed
5036
                    if (kz == 0):
5037
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5038
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
5039

5040
5041
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
5042
                if (hit != 0):
5043

Peter Eastman's avatar
Peter Eastman committed
5044
                    atom.multipoleDict = savedMultipoleDict
5045
                    atom.polarizationGroups = dict()
5046
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
5047
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
5048
                    if (atomIndex == newIndex):
5049
5050
5051
5052
                        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]))
5053
                    else:
5054
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
5055
                else:
Peter Eastman's avatar
Peter Eastman committed
5056
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
5057
            else:
Peter Eastman's avatar
Peter Eastman committed
5058
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
5059
5060
5061

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
5062
        self.setPolarGroups(data, bonded12ParticleSets, force)
5063
5064
5065
5066
5067

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

5068
## @private
5069
class AmoebaWcaDispersionGenerator(object):
5070
5071

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

5073
5074
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
5077
5078
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
5079
5080
5081
5082
5083
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
5084
        self.shctd = shctd
5085
5086
5087
5088
5089
5090
5091
5092
5093

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

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

Peter Eastman's avatar
Peter Eastman committed
5095
        generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
5096
5097
5098
                                                  element.attrib['epsh'],
                                                  element.attrib['rmino'],
                                                  element.attrib['rminh'],
Justin MacCallum's avatar
Justin MacCallum committed
5099
                                                  element.attrib['awater'],
5100
5101
                                                  element.attrib['slevy'],
                                                  element.attrib['dispoff'],
Justin MacCallum's avatar
Justin MacCallum committed
5102
                                                  element.attrib['shctd'])
5103
        forceField._forces.append(generator)
5104
5105
        generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaWcaDispersionForce', 'WcaDispersion', ('radius', 'epsilon'))
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
5106

5107
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
5108

Peter Eastman's avatar
Peter Eastman committed
5109
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
5110
5111
5112
5113
5114
5115
5116
5117
5118
5119
5120
5121

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

Peter Eastman's avatar
Peter Eastman committed
5124
5125
5126
5127
5128
5129
5130
5131
        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  ))
5132

5133
5134
5135
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
5136
5137
5138
5139
5140

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

5141
## @private
5142
class AmoebaGeneralizedKirkwoodGenerator(object):
5143
5144

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

5146
5147
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5148
5149
5150
5151
5152
5153
5154
5155
5156
5157
5158
    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
5159
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5160
5161
5162
5163
5164
5165
5166
5167
5168
5169
5170
5171
5172
5173
5174
5175
5176
5177
5178
5179
5180
5181
5182
5183
5184
        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
5185
5186
5187

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

Peter Eastman's avatar
Peter Eastman committed
5188
    def getObcShct(self, data, atomIndex):
5189

Peter Eastman's avatar
Peter Eastman committed
5190
        atom = data.atoms[atomIndex]
5191
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5192
        shct = -1.0
5193
5194

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
5195

Peter Eastman's avatar
Peter Eastman committed
5196
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
5197
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5198
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5199
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
5200
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
5201
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
5202
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
5203
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5204
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
5205
5206
5207
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
5208
        elif (atomicNumber == 16):              # S(16)
5209
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
5210
        elif (atomicNumber == 26):              # Fe(26)
5211
5212
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
5213
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5214
            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
5215
5216

        return shct
5217
5218
5219

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

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

Peter Eastman's avatar
Peter Eastman committed
5222
        atom = data.atoms[atomIndex]
5223
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5224
        radius = -1.0
5225

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

Peter Eastman's avatar
Peter Eastman committed
5228
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
5229

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

5233
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5234
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
5235

Peter Eastman's avatar
Peter Eastman committed
5236
            if (bondedAtomAtomicNumber == 7):
5237
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
5238
            if (bondedAtomAtomicNumber == 8):
5239
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
5240

Peter Eastman's avatar
Peter Eastman committed
5241
        elif (atomicNumber == 3):               # Li(3)
5242
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
5243
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5244

5245
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
5246
            if (len(bondedAtomIndices) == 3):
5247
5248
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
5249
            elif (len(bondedAtomIndices) == 4):
5250
5251
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5252
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
5253
5254
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
5255
        elif (atomicNumber == 7):               # N(7)
5256
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
5257
        elif (atomicNumber == 8):               # O(8)
5258
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
5259
            if (len(bondedAtomIndices) == 2):
5260
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
5261
        elif (atomicNumber == 9):               # F(9)
5262
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
5263
        elif (atomicNumber == 10):
5264
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
5265
        elif (atomicNumber == 11):
5266
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
5267
        elif (atomicNumber == 12):
5268
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5269
        elif (atomicNumber == 14):
5270
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
5271
        elif (atomicNumber == 15):              # P(15)
5272
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
5273
        elif (atomicNumber == 16):              # S(16)
5274
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
5275
        elif (atomicNumber == 17):
5276
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
5277
        elif (atomicNumber == 18):
5278
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5279
        elif (atomicNumber == 19):
5280
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
5281
        elif (atomicNumber == 20):
5282
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
5283
        elif (atomicNumber == 35):
5284
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
5285
        elif (atomicNumber == 36):
5286
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
5287
        elif (atomicNumber == 37):
5288
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
5289
        elif (atomicNumber == 53):
5290
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
5291
        elif (atomicNumber == 54):
5292
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
5293
        elif (atomicNumber == 55):
5294
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
5295
        elif (atomicNumber == 56):
5296
5297
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
5298
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5299
5300
            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
5301

5302
5303
5304
5305
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
5308
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5309
        atom = data.atoms[atomIndex]
5310
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
5311
        if (atomicNumber in bondiMap):
5312
5313
            radius = bondiMap[atomicNumber]
        else:
peastman's avatar
peastman committed
5314
            outputString = "Warning no Bondi radius for atom %s of %s %d using default value" % (atom.name, atom.residue.name, atom.residue.index)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5315
            raise ValueError( outputString )
Justin MacCallum's avatar
Justin MacCallum committed
5316

5317
5318
5319
5320
5321
5322
5323
5324
5325
5326
        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
5327

Peter Eastman's avatar
Peter Eastman committed
5328
        generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'], element.attrib['soluteDielectric'],
Justin MacCallum's avatar
Justin MacCallum committed
5329
5330
                                                        element.attrib['includeCavityTerm'],
                                                        element.attrib['probeRadius'], element.attrib['surfaceAreaFactor'])
5331
5332
5333
        forceField._forces.append(generator)

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

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

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

5340
5341
5342
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

Peter Eastman's avatar
Peter Eastman committed
5343
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5344
        amoebaMultipoleForceList = [f for f in existing if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
5345
        if (len(amoebaMultipoleForceList) > 0):
5346
5347
5348
5349
5350
            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
5351
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
5352
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
5353

5354
5355
5356
5357
5358
5359
5360
        # 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
5361

Peter Eastman's avatar
Peter Eastman committed
5362
5363
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
5364
            else:
Peter Eastman's avatar
Peter Eastman committed
5365
                force.setSolventDielectric(   float(self.solventDielectric))
5366

Peter Eastman's avatar
Peter Eastman committed
5367
5368
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
5369
            else:
Peter Eastman's avatar
Peter Eastman committed
5370
                force.setSoluteDielectric(    float(self.soluteDielectric))
5371

Peter Eastman's avatar
Peter Eastman committed
5372
5373
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
5374
            else:
Peter Eastman's avatar
Peter Eastman committed
5375
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
5376
5377
5378
5379
5380

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
5383
5384
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
5385
5386
5387

        # 1-2

5388
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5389
5390

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
5391
5392
5393
5394
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5395
            else:
Peter Eastman's avatar
Peter Eastman committed
5396
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5397
5398
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
5399
            force.addParticle(multipoleParameters[0], radius, shct)
5400
5401
5402
5403
5404

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

5405
## @private
5406
class AmoebaUreyBradleyGenerator(object):
5407
5408
5409
5410

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

5412
    def __init__(self):
5413

peastman's avatar
peastman committed
5414
        self.anglesForAtom2Type = defaultdict(list)
Peter Eastman's avatar
Peter Eastman committed
5415
5416
5417
        self.types1 = []
        self.types2 = []
        self.types3 = []
5418

Peter Eastman's avatar
Peter Eastman committed
5419
5420
        self.length = []
        self.k = []
5421
5422
5423
5424
5425
5426

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

    @staticmethod
    def parseElement(element, forceField):

5427
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
5428
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
5429

5430
        generator = AmoebaUreyBradleyGenerator()
5431
5432
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
5433
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
5434
            if None not in types:
peastman's avatar
peastman committed
5435
                index = len(generator.types1)
5436
5437
5438
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
peastman's avatar
peastman committed
5439
5440
                for t in types[1]:
                    generator.anglesForAtom2Type[t].append(index)
5441
5442
5443
5444
5445
                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
5446
                                    bond.attrib['class1'], bond.attrib['class2'], bond.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
5447
5448
                raise ValueError(outputString)

5449
5450
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
5453
        existing = [sys.getForce(i) for i in range(sys.getNumForces())]
5454
        existing = [f for f in existing if type(f) == mm.HarmonicBondForce]
5455
5456

        if len(existing) == 0:
5457
            force = mm.HarmonicBondForce()
5458
5459
5460
5461
5462
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
5463
            if (isConstrained and not args.get('flexibleConstraints', False)):
5464
5465
5466
5467
                continue
            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
5468
            for i in self.anglesForAtom2Type[type2]:
5469
5470
5471
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
Peter Eastman's avatar
Peter Eastman committed
5472
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
5473
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
5474
5475
5476
5477
5478
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
5479
5480


5481
5482
5483
5484
5485
5486
5487
5488
5489
5490
5491
5492
5493
5494
5495
5496
5497
5498
5499
5500
5501
5502
5503
5504
5505
5506
5507
5508
5509
5510
5511
5512
5513
5514
5515
5516
5517
5518
5519
5520
5521
5522
5523
5524
5525
5526
5527
5528
5529
5530
5531
5532
5533
5534
5535
5536
5537
5538
5539
## @private
class HippoNonbondedGenerator(object):
    """A HippoNonbondedGenerator constructs a HippoNonbondedForce."""

    def __init__(self, forcefield, extrapCoeff):
        self.ff = forcefield
        self.extrapCoeff = extrapCoeff
        self.exceptions = {}

    @staticmethod
    def parseElement(element, ff):
        extrapCoeff = [float(c) for c in element.attrib['extrapolationCoefficients'].split(',')]
        generator = HippoNonbondedGenerator(ff, extrapCoeff)
        ff.registerGenerator(generator)
        scaleNames = ('mmScale', 'dmScale', 'ddScale', 'dispScale', 'repScale', 'ctScale')
        paramNames = ('charge', 'coreCharge', 'alpha', 'epsilon', 'damping', 'c6', 'pauliK', 'pauliQ', 'pauliAlpha', 'polarizability', 'axisType', 'd0', 'd1', 'd2', 'q11', 'q12', 'q13', 'q21', 'q22', 'q23', 'q31', 'q32', 'q33')
        for ex in element.findall('Exception'):
            separation = int(ex.attrib['separation'])
            ingroup = ex.attrib['ingroup'].lower() == 'true'
            key = (separation, ingroup)
            if key in generator.exceptions:
                raise ValueError('HippoNonbondedForce: multiple exceptions with separation=%d ingroup=%s' % (separation, ingroup))
            generator.exceptions[key] = [float(ex.attrib[s]) for s in scaleNames]
        generator.params = ForceField._AtomTypeParameters(ff, 'HippoNonbondedForce', 'Atom', paramNames)
        generator.params.parseDefinitions(element)

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.HippoNonbondedForce.NoCutoff,
                     PME:mm.HippoNonbondedForce.PME}
        if nonbondedMethod not in methodMap:
            raise ValueError('Illegal nonbonded method for HippoNonbondedForce')

        # Build data structures we'll need for building local coordinate frames.

        bondIndices = _findBondsForExclusions(data, sys)
        pairs = _findExclusions(bondIndices, 2, len(data.atoms))
        bonded12 = [set() for i in range(len(data.atoms))]
        bonded13 = [set() for i in range(len(data.atoms))]
        for atom1, atom2, sep in pairs:
            if sep == 1:
                bonded12[atom1].add(data.atoms[atom2])
                bonded12[atom2].add(data.atoms[atom1])
            else:
                bonded13[atom1].add(data.atoms[atom2])
                bonded13[atom2].add(data.atoms[atom1])

        # Create the force.

        force = mm.HippoNonbondedForce()
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            params = [float(v) for v in values[:10]]
            axisType = int(values[10])
            dipole = [float(v) for v in values[11:14]]
            quadrupole = [float(v) for v in values[14:23]]
            extra = self.params.getExtraParameters(atom, data)
            zAtom = self._findAxisAtom('zAtomType', extra, bonded12[atom.index], None, data, [])
            xAtom = self._findAxisAtom('xAtomType', extra, bonded12[atom.index], bonded13[atom.index], data, [zAtom])
            yAtom = self._findAxisAtom('yAtomType', extra, bonded12[atom.index], bonded13[atom.index], data, [zAtom, xAtom])
Peter Eastman's avatar
Peter Eastman committed
5540
            force.addParticle(params[0], dipole, quadrupole, *params[1:], axisType=axisType, multipoleAtomZ=zAtom, multipoleAtomX=xAtom, multipoleAtomY=yAtom)
5541
5542
5543
5544
5545
5546
5547
5548
5549
5550
5551
5552
5553
5554
5555
5556
5557
5558
5559
5560
5561
5562
5563
5564
5565
5566
5567
5568
5569
5570
5571
5572
5573
5574
5575
5576
5577
5578
5579
5580
5581
5582
5583
5584
5585
5586
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setExtrapolationCoefficients(self.extrapCoeff)
        force.setCutoffDistance(nonbondedCutoff)
        if args['switchDistance'] is not None:
            force.setSwitchingDistance(args['switchDistance'])
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
        sys.addForce(force)

    def _findAxisAtom(self, paramName, params, bonded12, bonded13, data, exclude):
        if paramName not in params:
            return -1
        atomType = params[paramName]
        for atom in bonded12:
            if data.atomType[atom] == atomType and atom.index not in exclude:
                return atom.index
        if bonded13 is not None:
            for atom in bonded13:
                if data.atomType[atom] == atomType and atom.index not in exclude:
                    return atom.index
        raise ValueError('No bonded atom of type %s' % atomType)

    def postprocessSystem(self, sys, data, args):
        # Identify polarization groups.

        bondIndices = _findBondsForExclusions(data, sys)
        groupBondTypes = [self.params.getExtraParameters(atom, data)['groupTypes'].split(',') for atom in data.atoms]
        groupBonds = [[] for i in range(len(data.atoms))]
        for i,j in bondIndices:
            if data.atomType[data.atoms[i]] in groupBondTypes[j]:
                groupBonds[i].append(j)
                groupBonds[j].append(i)
        polarizationGroup = _findGroups(groupBonds)

        # Create the exclusions.

        maxSeparation = max(e[0] for e in self.exceptions)
        hippo = [f for f in sys.getForces() if isinstance(f, mm.HippoNonbondedForce)][0]
        pairs = _findExclusions(bondIndices, maxSeparation, hippo.getNumParticles())
        for atom1, atom2, sep in pairs:
            params = self.exceptions[(sep, polarizationGroup[atom1] == polarizationGroup[atom2])]
            hippo.addException(atom1, atom2, *params)

parsers["HippoNonbondedForce"] = HippoNonbondedGenerator.parseElement


peastman's avatar
peastman committed
5587
## @private
5588
class DrudeGenerator(object):
peastman's avatar
peastman committed
5589
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
5590

5591
5592
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
5593
5594
5595
5596
5597
5598
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
5599
5600
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
5601
5602
5603
5604
        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'):
5605
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
5606
5607
5608
5609
5610
5611
5612
5613
5614
5615
            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
5616

peastman's avatar
peastman committed
5617
5618
5619
5620
    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
5621

peastman's avatar
peastman committed
5622
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
5623

peastman's avatar
peastman committed
5624
5625
5626
5627
5628
5629
5630
5631
5632
5633
5634
5635
5636
5637
5638
5639
        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
5640
5641
                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
5642
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
5643

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

peastman's avatar
peastman committed
5647
5648
5649
5650
5651
5652
5653
        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)
5654
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
5655
5656
5657
5658
5659
                # 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]
5660
5661
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
5662
5663
5664
5665
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

Justin MacCallum's avatar
Justin MacCallum committed
5666
parsers["DrudeForce"] = DrudeGenerator.parseElement
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
5667
5668

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