forcefield.py 279 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-2022 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
import warnings
41
from math import sqrt, cos
42
from copy import deepcopy
peastman's avatar
peastman committed
43
from collections import defaultdict
44
45
import openmm as mm
import openmm.unit as unit
46
from . import element as elem
47
48
from openmm.app.internal.singleton import Singleton
from openmm.app.internal import compiled
49
from openmm.app.internal.argtracker import ArgTracker
50

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

53
54
55
56
57
58
59
_dataDirectories = None

def _getDataDirectories():
    global _dataDirectories
    if _dataDirectories is None:
        _dataDirectories = [os.path.join(os.path.dirname(__file__), 'data')]
        try:
60
61
            from importlib_metadata import entry_points
            for entry in entry_points().select(group='openmm.forcefielddir'):
62
63
                _dataDirectories.append(entry.load()())
        except:
64
            pass # importlib_metadata is not installed
65
    return _dataDirectories
66

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

74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
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])
89
        if functionType.startswith('Continuous'):
90
91
92
93
94
            periodicStr = function.attrib.get('periodic', 'false').lower()
            if periodicStr in ['true', 'false', 'yes', 'no', '1', '0']:
                params['periodic'] = periodicStr in ['true', 'yes', '1']
            else:
                raise ValueError('ForceField: non-boolean value for periodic attribute in tabulated function definition')
95
96
97
98
99
100
101
        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':
102
103
104
105
            force.addTabulatedFunction(
                name,
                mm.Continuous1DFunction(values, params['min'], params['max'], params['periodic']),
            )
106
        elif type == 'Continuous2D':
107
108
109
110
111
112
113
114
115
116
            force.addTabulatedFunction(
                name,
                mm.Continuous2DFunction(
                    params['xsize'], params['ysize'],
                    values,
                    params['xmin'], params['xmax'],
                    params['ymin'], params['ymax'],
                    params['periodic'],
                ),
            )
117
        elif type == 'Continuous3D':
118
119
            force.addTabulatedFunction(
                name,
120
                mm.Continuous3DFunction(
121
122
123
124
125
126
127
128
                    params['xsize'], params['ysize'], params['zsize'],
                    values,
                    params['xmin'], params['xmax'],
                    params['ymin'], params['ymax'],
                    params['zmin'], params['zmax'],
                    params['periodic'],
                ),
            )
129
130
131
132
133
        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':
134
            force.addTabulatedFunction(name, mm.Discrete3DFunction(params['xsize'], params['ysize'], params['zsize'], values))
135

136
137
# Enumerated values for nonbonded method

138
class NoCutoff(Singleton):
139
140
141
142
    def __repr__(self):
        return 'NoCutoff'
NoCutoff = NoCutoff()

143
class CutoffNonPeriodic(Singleton):
144
145
146
147
    def __repr__(self):
        return 'CutoffNonPeriodic'
CutoffNonPeriodic = CutoffNonPeriodic()

148
class CutoffPeriodic(Singleton):
149
150
151
152
    def __repr__(self):
        return 'CutoffPeriodic'
CutoffPeriodic = CutoffPeriodic()

153
class Ewald(Singleton):
154
155
156
157
    def __repr__(self):
        return 'Ewald'
Ewald = Ewald()

158
class PME(Singleton):
159
160
161
    def __repr__(self):
        return 'PME'
PME = PME()
162

Peter Eastman's avatar
Peter Eastman committed
163
164
165
166
167
class LJPME(Singleton):
    def __repr__(self):
        return 'LJPME'
LJPME = LJPME()

168
169
# Enumerated values for constraint type

170
class HBonds(Singleton):
171
172
173
174
    def __repr__(self):
        return 'HBonds'
HBonds = HBonds()

175
class AllBonds(Singleton):
176
177
178
179
    def __repr__(self):
        return 'AllBonds'
AllBonds = AllBonds()

180
class HAngles(Singleton):
181
182
183
    def __repr__(self):
        return 'HAngles'
HAngles = HAngles()
184
185
186
187
188
189
190
191
192
193

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

Robert McGibbon's avatar
Robert McGibbon committed
195
196
197
198
199
200
201
202
        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.
203
204
205
        """
        self._atomTypes = {}
        self._templates = {}
206
207
        self._patches = {}
        self._templatePatches = {}
208
        self._templateSignatures = {None:[]}
209
        self._atomClasses = {'':set()}
210
        self._forces = []
211
        self._scripts = []
212
        self._templateMatchers = []
213
        self._templateGenerators = []
214
        self.loadFile(files)
215

216
    def loadFile(self, files, resname_prefix=''):
217
        """Load an XML file and add the definitions from it to this ForceField.
218

Robert McGibbon's avatar
Robert McGibbon committed
219
220
        Parameters
        ----------
221
222
223
        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
224
225
226
            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.
227
228
229
        prefix : string
            An optional string to be prepended to each residue name found in the
            loaded files.
peastman's avatar
peastman committed
230
        """
231

232
233
234
235
        if isinstance(files, tuple):
            files = list(files)
        else:
            files = [files]
236
237
238

        trees = []

239
240
241
        i = 0
        while i < len(files):
            file = files[i]
242
            tree = None
243
244
245
246
            try:
                # this handles either filenames or open file-like objects
                tree = etree.parse(file)
            except IOError:
247
                for dataDir in _getDataDirectories():
248
249
250
251
                    f = os.path.join(dataDir, file)
                    if os.path.isfile(f):
                        tree = etree.parse(f)
                        break
252
253
254
255
256
257
258
259
260
261
262
            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)
263
264
            if tree is None:
                raise ValueError('Could not locate file "%s"' % file)
265
266

            trees.append(tree)
267
            i += 1
268

269
            # Process includes in this file.
270

271
272
            if isinstance(file, str):
                parentDir = os.path.dirname(file)
peastman's avatar
peastman committed
273
274
            else:
                parentDir = ''
peastman's avatar
peastman committed
275
276
            for included in tree.getroot().findall('Include'):
                includeFile = included.attrib['file']
peastman's avatar
peastman committed
277
278
279
                joined = os.path.join(parentDir, includeFile)
                if os.path.isfile(joined):
                    includeFile = joined
280
281
                if includeFile not in files:
                    files.append(includeFile)
peastman's avatar
peastman committed
282
283
284

        # Load the atom types.

285
286
287
288
        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
289
290
291

        # Load the residue templates.

292
293
294
        for tree in trees:
            if tree.getroot().find('Residues') is not None:
                for residue in tree.getroot().find('Residues').findall('Residue'):
295
                    resName = resname_prefix+residue.attrib['name']
296
                    template = ForceField._TemplateData(resName)
297
298
                    if 'override' in residue.attrib:
                        template.overrideLevel = int(residue.attrib['override'])
299
300
                    if 'rigidWater' in residue.attrib:
                        template.rigidWater = (residue.attrib['rigidWater'].lower() == 'true')
301
302
                    for key in residue.attrib:
                        template.attributes[key] = residue.attrib[key]
303
304
                    atomIndices = template.atomIndices
                    for ia, atom in enumerate(residue.findall('Atom')):
305
                        params = {}
306
307
308
309
                        for key in atom.attrib:
                            if key not in ('name', 'type'):
                                params[key] = _convertParameterToNumber(atom.attrib[key])
                        atomName = atom.attrib['name']
310
311
                        if atomName in atomIndices:
                            raise ValueError('Residue '+resName+' contains multiple atoms named '+atomName)
312
                        typeName = atom.attrib['type']
313
                        atomIndices[atomName] = ia
314
315
316
317
318
319
320
321
322
323
324
325
326
                        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']))
327
                    for patch in residue.findall('AllowPatch'):
328
                        patchName = patch.attrib['name']
329
330
                        if ':' in patchName:
                            colonIndex = patchName.find(':')
331
332
333
                            self.registerTemplatePatch(resName, patchName[:colonIndex], int(patchName[colonIndex+1:])-1)
                        else:
                            self.registerTemplatePatch(resName, patchName, 0)
334
                    self.registerResidueTemplate(template)
peastman's avatar
peastman committed
335

336
        # Load the patch definitions.
337
338
339
340
341
342
343
344
345
346

        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)
347
348
                    for key in patch.attrib:
                        patchData.attributes[key] = patch.attrib[key]
349
350
351
352
353
354
                    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
355
                        if atomName in patchData.allAtomNames:
356
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
357
                        patchData.allAtomNames.add(atomName)
358
359
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
360
                        patchData.addedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
361
362
363
364
365
366
                    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
367
                        if atomName in patchData.allAtomNames:
368
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
369
                        patchData.allAtomNames.add(atomName)
370
371
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
372
                        patchData.changedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
373
374
                    for atom in patch.findall('RemoveAtom'):
                        atomName = atom.attrib['name']
peastman's avatar
peastman committed
375
                        if atomName in patchData.allAtomNames:
376
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
377
                        patchData.allAtomNames.add(atomName)
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
                        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)
peastman's avatar
peastman committed
394
395
396
                    atomIndices = dict((atom.name, i) for i, atom in enumerate(patchData.addedAtoms[atomDescription.residue]+patchData.changedAtoms[atomDescription.residue]))
                    for site in patch.findall('VirtualSite'):
                        patchData.virtualSites[atomDescription.residue].append(ForceField._VirtualSiteData(site, atomIndices))
397
398
399
400
401
402
403
404
405
                    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
406
407
        # Load force definitions

408
409
410
411
        for tree in trees:
            for child in tree.getroot():
                if child.tag in parsers:
                    parsers[child.tag](child, self)
peastman's avatar
peastman committed
412
413
414

        # Load scripts

415
416
417
        for tree in trees:
            for node in tree.getroot().findall('Script'):
                self.registerScript(node.text)
418

419
420
421
422
423
424
        # Execute initialization scripts.

        for tree in trees:
            for node in tree.getroot().findall('InitializationScript'):
                exec(node.text, locals())

425
426
427
    def getGenerators(self):
        """Get the list of all registered generators."""
        return self._forces
428

429
430
431
    def registerGenerator(self, generator):
        """Register a new generator."""
        self._forces.append(generator)
432

433
434
435
436
437
438
439
440
441
442
443
444
    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)
445
        self._atomTypes[name] = ForceField._AtomType(name, atomClass, mass, element)
446
447
448
449
450
451
452
        if atomClass in self._atomClasses:
            typeSet = self._atomClasses[atomClass]
        else:
            typeSet = set()
            self._atomClasses[atomClass] = typeSet
        typeSet.add(name)
        self._atomClasses[''].add(name)
453

454
455
    def registerResidueTemplate(self, template):
        """Register a new residue template."""
456
457
        if template.name in self._templates:
            # There is already a template with this name, so check the override levels.
458

459
460
461
462
463
464
465
466
467
468
469
            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))
470

471
        # Register the template.
472

473
474
475
        self._templates[template.name] = template
        signature = _createResidueSignature([atom.element for atom in template.atoms])
        if signature in self._templateSignatures:
476
            self._templateSignatures[signature].append(template)
477
478
        else:
            self._templateSignatures[signature] = [template]
479

480
481
482
    def registerPatch(self, patch):
        """Register a new patch that can be applied to templates."""
        self._patches[patch.name] = patch
483

484
485
486
    def registerTemplatePatch(self, residue, patch, patchResidueIndex):
        """Register that a particular patch can be used with a particular residue."""
        if residue not in self._templatePatches:
487
488
            self._templatePatches[residue] = set()
        self._templatePatches[residue].add((patch, patchResidueIndex))
489

490
491
492
    def registerScript(self, script):
        """Register a new script to be executed after building the System."""
        self._scripts.append(script)
493

494
495
496
497
    def registerTemplateMatcher(self, matcher):
        """Register an object that can override the default logic for matching templates to residues.

        A template matcher is a callable object that can be invoked as::
498

Peter Eastman's avatar
Peter Eastman committed
499
            template = f(forcefield, residue, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
500

Peter Eastman's avatar
Peter Eastman committed
501
502
503
504
505
        where ``forcefield`` is the ForceField invoking it, ``residue`` is an openmm.app.Residue object,
        ``bondedToAtom[i]`` is the set of atoms bonded to atom index i, and ``ignoreExternalBonds`` and
        ``ignoreExtraParticles`` indicate whether external bonds and extra particules should be considered
        in matching templates.

506
507
508
        It should return a _TemplateData object that matches the residue.  Alternatively it may return
        None, in which case the standard logic will be used to find a template for the residue.

509
        .. CAUTION:: This method is experimental, and its API is subject to change.
510
511
        """
        self._templateMatchers.append(matcher)
512

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
513
    def registerTemplateGenerator(self, generator):
514
515
516
517
        """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
518
519
        .. CAUTION:: This method is experimental, and its API is subject to change.

520
521
        Parameters
        ----------
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
522
        generator : function
523
524
            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
525
        When a residue without a template is encountered, the ``generator`` function is called with:
526

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
527
528
        ::
           success = generator(forcefield, residue)
529

530
        where ``forcefield`` is the calling ``ForceField`` object and ``residue`` is a openmm.app.topology.Residue object.
John Chodera's avatar
John Chodera committed
531
532

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

        ::
John Chodera's avatar
John Chodera committed
535
536
537
           generator API

           Parameters
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
538
           ----------
539
           forcefield : openmm.app.ForceField
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
540
               The ForceField object to which residue templates and/or parameters are to be added.
541
           residue : openmm.app.Topology.Residue
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
542
               The residue topology for which a template is to be generated.
543

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
544
545
546
547
548
           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`.
549

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
550
551
552
553
554
           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.
555
556

        """
557
        self._templateGenerators.append(generator)
558

559
    def _findAtomTypes(self, attrib, num):
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
        """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.

        """
575
576
577
578
579
580
581
        types = []
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
582
            typeAttrib = 'type'+suffix
583
            if classAttrib in attrib:
584
                if typeAttrib in attrib:
585
                    raise ValueError('Specified both a type and a class for the same atom: '+str(attrib))
586
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
587
588
589
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
590
591
            elif typeAttrib in attrib:
                if attrib[typeAttrib] == '':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
592
                    types.append(self._atomClasses[''])
593
                elif attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
594
595
596
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
597
598
            else:
                types.append(None) # Unknown atom type
599
600
        return types

601
    def _parseTorsion(self, attrib):
602
        """Parse the node defining a torsion."""
603
        types = self._findAtomTypes(attrib, 4)
peastman's avatar
peastman committed
604
        if None in types:
605
606
607
608
609
            return None
        torsion = PeriodicTorsion(types)
        index = 1
        while 'phase%d'%index in attrib:
            torsion.periodicity.append(int(attrib['periodicity%d'%index]))
610
611
            torsion.phase.append(_convertParameterToNumber(attrib['phase%d'%index]))
            torsion.k.append(_convertParameterToNumber(attrib['k%d'%index]))
612
            index += 1
Justin MacCallum's avatar
Justin MacCallum committed
613
614
        return torsion

615
    class _SystemData(object):
616
        """Inner class used to encapsulate data about the system being created."""
617
        def __init__(self, topology):
618
            self.atomType = {}
619
            self.atomParameters = {}
620
            self.atomTemplateIndexes = {}
621
            self.atoms = list(topology.atoms())
Peter Eastman's avatar
Peter Eastman committed
622
            self.excludeAtomWith = [[] for _ in self.atoms]
623
            self.virtualSites = {}
624
            self.bonds = [ForceField._BondData(bond[0].index, bond[1].index) for bond in topology.bonds()]
625
626
627
            self.angles = []
            self.propers = []
            self.impropers = []
Peter Eastman's avatar
Peter Eastman committed
628
            self.atomBonds = [[] for _ in self.atoms]
629
            self.isAngleConstrained = []
630
            self.constraints = {}
Peter Eastman's avatar
Peter Eastman committed
631
            self.bondedToAtom = [set() for _ in self.atoms]
632

633
            # Record which atoms are bonded to each other atom
634

635
636
637
638
639
640
            for i in range(len(self.bonds)):
                bond = self.bonds[i]
                self.bondedToAtom[bond.atom1].add(bond.atom2)
                self.bondedToAtom[bond.atom2].add(bond.atom1)
                self.atomBonds[bond.atom1].append(i)
                self.atomBonds[bond.atom2].append(i)
641
            self.bondedToAtom = [sorted(b) for b in self.bondedToAtom]
642
643
644
645
646

        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:
647
                if self.constraints[key] != distance:
648
649
650
651
                    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)
652

653
654
655
656
657
658
        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
659
                self.atomTemplateIndexes[atom] = match
660
661
662
                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)
663

664
    class _TemplateData(object):
665
666
667
668
        """Inner class used to encapsulate data about a residue template definition."""
        def __init__(self, name):
            self.name = name
            self.atoms = []
669
            self.atomIndices = {}
670
            self.virtualSites = []
671
672
            self.bonds = []
            self.externalBonds = []
673
            self.overrideLevel = 0
674
            self.rigidWater = True
675
            self.attributes = {}
676

677
678
        def getAtomIndexByName(self, atom_name):
            """Look up an atom index by atom name, providing a helpful error message if not found."""
679
680
681
            index = self.atomIndices.get(atom_name, None)
            if index is not None:
                return index
682
683

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

688
689
        def addAtom(self, atom):
            self.atoms.append(atom)
690
            self.atomIndices[atom.name] = len(self.atoms)-1
691

692
        def addBond(self, atom1, atom2):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
693
            """Add a bond between two atoms in a template given their indices in the template."""
694
695
696
            self.bonds.append((atom1, atom2))
            self.atoms[atom1].bondedTo.append(atom2)
            self.atoms[atom2].bondedTo.append(atom1)
697

698
        def addBondByName(self, atom1_name, atom2_name):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
699
            """Add a bond between two atoms in a template given their atom names."""
700
701
702
703
704
            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
705
            """Designate that an atom in a residue template has an external bond, using atom index within template."""
706
707
708
709
            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
710
            """Designate that an atom in a residue template has an external bond, using atom name within template."""
711
712
713
            atom = self.getAtomIndexByName(atom_name)
            self.addExternalBond(atom)

peastman's avatar
peastman committed
714
715
        def areParametersIdentical(self, template2, matchingAtoms, matchingAtoms2):
            """Get whether this template and another one both assign identical atom types and parameters to all atoms.
John Chodera's avatar
John Chodera committed
716

peastman's avatar
peastman committed
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
            Parameters
            ----------
            template2: _TemplateData
                the template to compare this one to
            matchingAtoms: list
                the indices of atoms in this template that atoms of the residue are matched to
            matchingAtoms2: list
                the indices of atoms in template2 that atoms of the residue are matched to
            """
            atoms1 = [self.atoms[m] for m in matchingAtoms]
            atoms2 = [template2.atoms[m] for m in matchingAtoms2]
            if any(a1.type != a2.type or a1.parameters != a2.parameters for a1,a2 in zip(atoms1, atoms2)):
                return False
            # Properly comparing virtual sites really needs a much more complicated analysis.  This simple check
            # could easily fail for templates containing vsites, even if they're actually identical.  Since we
            # currently have no force fields that include both patches and vsites, I'm not going to worry about it now.
            if self.virtualSites != template2.virtualSites:
                return False
            return True

737
    class _TemplateAtomData(object):
738
        """Inner class used to encapsulate data about an atom in a residue template definition."""
739
        def __init__(self, name, type, element, parameters={}):
740
741
742
            self.name = name
            self.type = type
            self.element = element
743
            self.parameters = parameters
744
745
746
            self.bondedTo = []
            self.externalBonds = 0

747
    class _BondData(object):
748
749
750
751
752
753
        """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
754

755
    class _VirtualSiteData(object):
756
        """Inner class used to encapsulate data about a virtual site."""
757
        def __init__(self, node, atomIndices):
758
759
760
            attrib = node.attrib
            self.type = attrib['type']
            if self.type == 'average2':
761
                numAtoms = 2
762
763
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
764
                numAtoms = 3
765
766
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
767
                numAtoms = 3
768
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
769
            elif self.type == 'localCoords':
770
771
772
773
774
775
776
777
778
                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]))
779
                self.localPos = [float(attrib['p1']), float(attrib['p2']), float(attrib['p3'])]
780
781
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
782
783
784
785
786
787
            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)]
788
789
790
791
            if 'excludeWith' in attrib:
                self.excludeWith = int(attrib['excludeWith'])
            else:
                self.excludeWith = self.atoms[0]
792

793
794
795
796
797
    class _PatchData(object):
        """Inner class used to encapsulate data about a patch definition."""
        def __init__(self, name, numResidues):
            self.name = name
            self.numResidues = numResidues
798
799
            self.addedAtoms = [[] for i in range(numResidues)]
            self.changedAtoms = [[] for i in range(numResidues)]
800
801
802
803
804
            self.deletedAtoms = []
            self.addedBonds = []
            self.deletedBonds = []
            self.addedExternalBonds = []
            self.deletedExternalBonds = []
peastman's avatar
peastman committed
805
            self.allAtomNames = set()
peastman's avatar
peastman committed
806
            self.virtualSites = [[] for i in range(numResidues)]
807
            self.attributes = {}
808

809
810
811
812
        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)))
813

814
            # Construct a new version of each template.
815

816
817
818
819
            newTemplates = []
            for index, template in enumerate(templates):
                newTemplate = ForceField._TemplateData("%s-%s" % (template.name, self.name))
                newTemplates.append(newTemplate)
820

821
                # Build the list of atoms in it.
822

823
824
                for atom in template.atoms:
                    if not any(deleted.name == atom.name and deleted.residue == index for deleted in self.deletedAtoms):
825
                        newTemplate.addAtom(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
826
                for atom in self.addedAtoms[index]:
827
828
                    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))
829
                    newTemplate.addAtom(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
830
831
832
833
834
                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))
835
                    newTemplate.atoms[newAtomIndex[atom.name]] = ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters)
836

837
                # Copy over the virtual sites, translating the atom indices.
838

839
840
841
842
843
844
845
                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)
846

847
                # Build the lists of bonds and external bonds.
848

849
850
851
852
853
                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]
854
                    if a1 in atomMap and a2 in atomMap and (a1.name, a2.name) not in deletedBonds and (a2.name, a1.name) not in deletedBonds:
855
856
857
858
                        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:
859
                        newTemplate.addExternalBond(indexMap[atom])
860
861
862
863
864
865
866
867
868
                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)
peastman's avatar
peastman committed
869
870
871
872
873
874
875
876

                # Add new virtual sites.

                indexMap = dict((i, newAtomIndex[atom.name]) for i, atom in enumerate(self.addedAtoms[index]+self.changedAtoms[index]))
                for site in self.virtualSites[index]:
                    newSite = deepcopy(site)
                    newSite.index = indexMap[site.index]
                    newSite.atoms = [indexMap[i] for i in site.atoms]
877
                    newSite.excludeWith = indexMap[site.excludeWith]
peastman's avatar
peastman committed
878
879
                    newTemplate.virtualSites = [site for site in newTemplate.virtualSites if site.index != newSite.index]
                    newTemplate.virtualSites.append(newSite)
880
            return newTemplates
881

882
883
884
885
886
887
888
889
890
891
892
    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

893
    class _AtomType(object):
894
895
896
897
898
899
900
        """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

901
    class _AtomTypeParameters(object):
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
        """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))


972
973
    def _getResidueTemplateMatches(self, res, bondedToAtom, templateSignatures=None, ignoreExternalBonds=False, ignoreExtraParticles=False):
        """Return the templates that match a residue, or None if none are found.
974
975
976
977
978

        Parameters
        ----------
        res : Topology.Residue
            The residue for which template matches are to be retrieved.
979
980
        bondedToAtom : list of set of int
            bondedToAtom[i] is the set of atoms bonded to atom index i
981
982
983

        Returns
        -------
peastman's avatar
peastman committed
984
        template : _TemplateData
985
986
987
988
989
990
991
992
            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
993
        for matcher in self._templateMatchers:
Peter Eastman's avatar
Peter Eastman committed
994
            template = matcher(self, res, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
995
996
997
            if template is not None:
                match = compiled.matchResidueToTemplate(res, template, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
                if match is None:
Peter Eastman's avatar
Peter Eastman committed
998
                    raise ValueError('A custom template matcher returned a template for residue %d (%s), but it does not match the residue.' % (res.index, res.name))
999
                return [template, match]
1000
1001
        if templateSignatures is None:
            templateSignatures = self._templateSignatures
1002
        signature = _createResidueSignature([atom.element for atom in res.atoms()])
1003
        if signature in templateSignatures:
1004
            allMatches = []
1005
            for t in templateSignatures[signature]:
1006
                match = compiled.matchResidueToTemplate(res, t, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1007
1008
1009
1010
1011
1012
                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:
1013
1014
1015
                # We found multiple matches.  This is OK if and only if they assign identical types and parameters to all atoms.
                t1, m1 = allMatches[0]
                for t2, m2 in allMatches[1:]:
peastman's avatar
peastman committed
1016
1017
                    if not t1.areParametersIdentical(t2, m1, m2):
                        raise Exception('Multiple non-identical matching templates found for residue %d (%s): %s.' % (res.index+1, res.name, ', '.join(match[0].name for match in allMatches)))
1018
1019
                template = allMatches[0][0]
                matches = allMatches[0][1]
1020
1021
        return [template, matches]

1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
    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
        -------
1032
1033
        bondedToAtom : list of list of int
            bondedToAtom[index] is the list of atom indices bonded to atom `index`
1034
1035

        """
Peter Eastman's avatar
Peter Eastman committed
1036
        bondedToAtom = [set() for _ in topology.atoms()]
1037
1038
1039
        for (atom1, atom2) in topology.bonds():
            bondedToAtom[atom1.index].add(atom2.index)
            bondedToAtom[atom2.index].add(atom1.index)
1040
        bondedToAtom = [sorted(b) for b in bondedToAtom]
1041
        return bondedToAtom
1042

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

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

1048
1049
1050
        Parameters
        ----------
        topology : Topology
1051
            The Topology whose residues are to be checked against the forcefield residue templates.
1052
1053
1054
1055

        Returns
        -------
        unmatched_residues : list of Residue
1056
            List of Residue objects from `topology` for which no forcefield residue templates are available.
1057
1058
1059
1060
1061
            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.
1062
        bondedToAtom = self._buildBondedToAtomList(topology)
1063
        unmatched_residues = list() # list of unmatched residues
1064
1065
1066
1067
1068
1069
        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)
1070
1071
1072

        return unmatched_residues

1073
    def getMatchingTemplates(self, topology, ignoreExternalBonds=False):
1074
1075
1076
1077
1078
1079
1080
1081
        """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.
1082
1083
        ignoreExternalBonds : bool=False
            If true, ignore external bonds when matching residues to templates.
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
        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
1095
        for residue in topology.residues():
1096
            # Attempt to match one of the existing templates.
1097
            [template, matches] = self._getResidueTemplateMatches(residue, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1098
1099
            # Raise an exception if we have found no templates that match.
            if matches is None:
1100
                raise ValueError('No template found for chainid <%s> resid <%s> resname <%s> (residue index within topology %d).\n%s' % (residue.chain.id, residue.id, residue.name, residue.index, _findMatchErrors(self, residue)))
1101
1102
1103
1104
1105
            else:
                templates.append(template)

        return templates

1106
1107
    def generateTemplatesForUnmatchedResidues(self, topology):
        """Generate forcefield residue templates for residues in specified topology for which no forcefield templates are available.
1108

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

1111
1112
1113
        Parameters
        ----------
        topology : Topology
1114
            The Topology whose residues are to be checked against the forcefield residue templates.
1115
1116
1117

        Returns
        -------
1118
1119
1120
1121
1122
1123
        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]`
1124
1125
1126
1127
1128

        """
        # Get a non-unique list of unmatched residues.
        unmatched_residues = self.getUnmatchedResidues(topology)
        # Generate a unique list of unmatched residues by comparing fingerprints.
1129
        bondedToAtom = self._buildBondedToAtomList(topology)
1130
1131
        unique_unmatched_residues = list() # list of unique unmatched Residue objects from topology
        templates = list() # corresponding _TemplateData templates
1132
1133
1134
        signatures = set()
        for residue in unmatched_residues:
            signature = _createResidueSignature([ atom.element for atom in residue.atoms() ])
1135
            template = _createResidueTemplate(residue)
1136
1137
1138
1139
            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
1140
                    matches = compiled.matchResidueToTemplate(check_residue, template, bondedToAtom, False)
1141
1142
1143
1144
1145
1146
                    if matches is not None:
                        is_unique = False
            if is_unique:
                # Residue is unique.
                unique_unmatched_residues.append(residue)
                signatures.add(signature)
1147
                templates.append(template)
1148

1149
        return [templates, unique_unmatched_residues]
1150

1151
    def createSystem(self, topology, nonbondedMethod=NoCutoff, nonbondedCutoff=1.0*unit.nanometer,
1152
                     constraints=None, rigidWater=None, removeCMMotion=True, hydrogenMass=None, residueTemplates=dict(),
1153
                     ignoreExternalBonds=False, switchDistance=None, flexibleConstraints=False, drudeMass=0.4*unit.amu, **args):
1154
        """Construct an OpenMM System representing a Topology with this force field.
Justin MacCallum's avatar
Justin MacCallum committed
1155

Robert McGibbon's avatar
Robert McGibbon committed
1156
1157
1158
1159
1160
1161
        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
1162
            NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME, or LJPME.
Robert McGibbon's avatar
Robert McGibbon committed
1163
1164
1165
1166
1167
        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.
1168
        rigidWater : boolean=None
Robert McGibbon's avatar
Robert McGibbon committed
1169
            If true, water molecules will be fully rigid regardless of the value
1170
1171
            passed for the constraints argument.  If None (the default), it uses the
            default behavior for this force field's water model.
Robert McGibbon's avatar
Robert McGibbon committed
1172
1173
1174
1175
1176
        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
1177
1178
            their total mass the same.  If rigidWater is used to make water molecules
            rigid, then water hydrogens are not altered.
1179
        residueTemplates : dict=dict()
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
            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.
1191
1192
1193
        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.
1194
1195
        flexibleConstraints : boolean=False
            If True, parameters for constrained degrees of freedom will be added to the System
1196
        drudeMass : mass=0.4*amu
1197
1198
            The mass to use for Drude particles.  Any mass added to a Drude particle is
            subtracted from its parent atom to keep their total mass the same.
Robert McGibbon's avatar
Robert McGibbon committed
1199
        args
1200
1201
1202
            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
1203
1204
1205
1206
1207

        Returns
        -------
        system
            the newly created System
1208
        """
1209
1210
        args['switchDistance'] = switchDistance
        args['flexibleConstraints'] = flexibleConstraints
1211
        args['drudeMass'] = drudeMass
1212
        args = ArgTracker(args)
1213
        data = ForceField._SystemData(topology)
1214
        rigidResidue = [False]*topology.getNumResidues()
1215
1216

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

1218
1219
1220
1221
        templateForResidue = self._matchAllResiduesToTemplates(data, topology, residueTemplates, ignoreExternalBonds)
        for res in topology.residues():
            if res.name == 'HOH':
                # Determine whether this should be a rigid water.
1222

1223
1224
1225
1226
                if rigidWater is None:
                    rigidResidue[res.index] = templateForResidue[res.index].rigidWater
                elif rigidWater:
                    rigidResidue[res.index] = True
1227
1228

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

1230
1231
        sys = mm.System()
        for atom in topology.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1232
            # Look up the atom type name, returning a helpful error message if it cannot be found.
1233
1234
1235
1236
            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
1237
            # Look up the type name in the list of registered atom types, returning a helpful error message if it cannot be found.
1238
1239
1240
1241
            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
1242
1243

            # Add the particle to the OpenMM system.
1244
            mass = self._atomTypes[typename].mass
1245
            sys.addParticle(mass)
1246

1247
        # Adjust hydrogen masses if requested.
1248

1249
        if hydrogenMass is not None:
1250
1251
            if not unit.is_quantity(hydrogenMass):
                hydrogenMass *= unit.dalton
1252
            for atom1, atom2 in topology.bonds():
1253
                if atom1.element is elem.hydrogen:
1254
                    (atom1, atom2) = (atom2, atom1)
1255
                if atom2.element is elem.hydrogen and atom1.element not in (elem.hydrogen, None) and not rigidResidue[atom2.residue.index]:
1256
1257
1258
                    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
1259

1260
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
1261

1262
1263
1264
        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
1265
        elif nonbondedMethod not in [NoCutoff, CutoffNonPeriodic]:
1266
1267
1268
            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
1269

1270
1271
        uniqueAngles = set()
        for bond in data.bonds:
1272
            for atom in data.bondedToAtom[bond.atom1]:
1273
1274
1275
1276
1277
                if atom != bond.atom2:
                    if atom < bond.atom2:
                        uniqueAngles.add((atom, bond.atom1, bond.atom2))
                    else:
                        uniqueAngles.add((bond.atom2, bond.atom1, atom))
1278
            for atom in data.bondedToAtom[bond.atom2]:
1279
1280
1281
1282
1283
1284
                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
1285

1286
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1287

1288
1289
        uniquePropers = set()
        for angle in data.angles:
1290
            for atom in data.bondedToAtom[angle[0]]:
pgrinaway's avatar
pgrinaway committed
1291
                if atom not in angle:
1292
1293
1294
1295
                    if atom < angle[2]:
                        uniquePropers.add((atom, angle[0], angle[1], angle[2]))
                    else:
                        uniquePropers.add((angle[2], angle[1], angle[0], atom))
1296
            for atom in data.bondedToAtom[angle[2]]:
pgrinaway's avatar
pgrinaway committed
1297
                if atom not in angle:
1298
1299
1300
1301
1302
                    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
1303

1304
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1305

1306
1307
        for atom in range(len(data.bondedToAtom)):
            bondedTo = data.bondedToAtom[atom]
1308
1309
1310
            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
1311

1312
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1313

1314
1315
1316
1317
1318
1319
1320
        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]
1321
                bond.isConstrained = atom1.element is elem.hydrogen or atom2.element is elem.hydrogen
1322
1323
1324
1325
1326
        for bond in data.bonds:
            atom1 = data.atoms[bond.atom1]
            atom2 = data.atoms[bond.atom2]
            if rigidResidue[atom1.residue.index] and rigidResidue[atom2.residue.index]:
                bond.isConstrained = True
Justin MacCallum's avatar
Justin MacCallum committed
1327

1328
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1329

1330
1331
1332
1333
1334
1335
        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
1336
                if atom1.element is elem.hydrogen:
1337
                    numH += 1
1338
                if atom3.element is elem.hydrogen:
1339
                    numH += 1
1340
                data.isAngleConstrained.append(numH == 2 or (numH == 1 and atom2.element is elem.oxygen))
1341
1342
        else:
            data.isAngleConstrained = len(data.angles)*[False]
1343
1344
1345
1346
1347
1348
1349
        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 rigidResidue[atom1.residue.index] and rigidResidue[atom2.residue.index] and rigidResidue[atom3.residue.index]:
                data.isAngleConstrained[i] = True
Justin MacCallum's avatar
Justin MacCallum committed
1350

1351
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
1352

1353
        for atom in data.virtualSites:
1354
            (site, atoms, excludeWith) = data.virtualSites[atom]
1355
            index = atom.index
1356
            data.excludeAtomWith[excludeWith].append(index)
1357
            if site.type == 'average2':
1358
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(atoms[0], atoms[1], site.weights[0], site.weights[1]))
1359
            elif site.type == 'average3':
1360
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
1361
            elif site.type == 'outOfPlane':
1362
1363
                sys.setVirtualSite(index, mm.OutOfPlaneSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
            elif site.type == 'localCoords':
1364
                sys.setVirtualSite(index, mm.LocalCoordinatesSite(atoms, site.originWeights, site.xWeights, site.yWeights, site.localPos))
Justin MacCallum's avatar
Justin MacCallum committed
1365

1366
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
1367

1368
1369
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
1370
1371
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
1372

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

peastman's avatar
peastman committed
1375
1376
1377
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
1378

1379
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
1380

1381
        for script in self._scripts:
1382
            exec(script, locals())
1383
        args.checkArgs(self.createSystem)
1384
1385
1386
        return sys


1387
    def _matchAllResiduesToTemplates(self, data, topology, residueTemplates, ignoreExternalBonds, ignoreExtraParticles=False, recordParameters=True):
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
        """Return a list of which template matches each residue in the topology, and assign atom types."""
        templateForResidue = [None]*topology.getNumResidues()
        unmatchedResidues = []
        for chain in topology.chains():
            for res in chain.residues():
                if res in residueTemplates:
                    tname = residueTemplates[res]
                    template = self._templates[tname]
                    matches = compiled.matchResidueToTemplate(res, template, data.bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
                    if matches is None:
                        raise Exception('User-supplied template %s does not match the residue %d (%s)' % (tname, res.index+1, res.name))
                else:
                    # Attempt to match one of the existing templates.
                    [template, matches] = self._getResidueTemplateMatches(res, data.bondedToAtom, ignoreExternalBonds=ignoreExternalBonds, ignoreExtraParticles=ignoreExtraParticles)
                if matches is None:
                    unmatchedResidues.append(res)
                else:
1405
1406
                    if recordParameters:
                        data.recordMatchedAtomParameters(res, template, matches)
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
                    templateForResidue[res.index] = template

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

        if len(unmatchedResidues) > 0:
            unmatchedResidues = _applyPatchesToMatchResidues(self, data, unmatchedResidues, templateForResidue, data.bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)

        # 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).

        for res in unmatchedResidues:
            # A template might have been generated on an earlier iteration of this loop.
            [template, matches] = self._getResidueTemplateMatches(res, data.bondedToAtom, ignoreExternalBonds=ignoreExternalBonds, ignoreExtraParticles=ignoreExtraParticles)
            if matches is None:
                # Try all generators.
                for generator in self._templateGenerators:
                    if generator(self, res):
                        # This generator has registered a new residue template that should match.
                        [template, matches] = self._getResidueTemplateMatches(res, data.bondedToAtom, ignoreExternalBonds=ignoreExternalBonds, ignoreExtraParticles=ignoreExtraParticles)
                        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:
1433
                raise ValueError('No template found for residue %d (%s).  %s  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#template' % (res.index+1, res.name, _findMatchErrors(self, res)))
1434
            else:
1435
1436
                if recordParameters:
                    data.recordMatchedAtomParameters(res, template, matches)
1437
1438
1439
1440
                templateForResidue[res.index] = template
        return templateForResidue


1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
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))
1465
1466
1467
    for atom in data.atoms:
        for child in data.excludeAtomWith[atom.index]:
            bondIndices.append((child, atom.index))
1468
1469
    return bondIndices

1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
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

1534
1535
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
1536
1537
    counts = {}
    for element in elements:
1538
        if element in counts:
1539
1540
1541
            counts[element] += 1
        else:
            counts[element] = 1
1542
1543
1544
1545
1546
1547
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
1548
1549
    sig = []
    for c in counts:
1550
1551
        if c is not None:
            sig.append((c, counts[c]))
1552
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
1553

1554
    # Convert it to a string.
1555
1556

    s = ''
1557
    for element, count in sig:
1558
1559
1560
1561
        s += element.symbol+str(count)
    return s


1562
def _applyPatchesToMatchResidues(forcefield, data, residues, templateForResidue, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles):
1563
1564
1565
1566
    """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.
1567

1568
1569
1570
1571
1572
1573
1574
1575
    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
1576
                _generatePatchedSingleResidueTemplates(template, patches, 0, newTemplates, set())
1577
1578
1579
1580
1581
1582
                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]
1583

1584
    # Now see if any of those templates matches any of the residues.
1585

1586
1587
    unmatchedResidues = []
    for res in residues:
1588
        [template, matches] = forcefield._getResidueTemplateMatches(res, bondedToAtom, patchedTemplateSignatures, ignoreExternalBonds, ignoreExtraParticles)
1589
1590
1591
1592
        if matches is None:
            unmatchedResidues.append(res)
        else:
            data.recordMatchedAtomParameters(res, template, matches)
1593
            templateForResidue[res.index] = template
1594
1595
    if len(unmatchedResidues) == 0:
        return []
1596

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

1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
    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]
1616

1617
    # Record which unmatched residues are bonded to each other.
1618

1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
    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)
1629

1630
1631
    # 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.
1632

1633
1634
1635
1636
    clusterSize = 2
    clusters = bonds
    while clusterSize <= maxPatchSize:
        # Try to apply patches to clusters of this size.
1637

1638
1639
1640
        for patchName in patches:
            patch = forcefield._patches[patchName]
            if patch.numResidues == clusterSize:
1641
                matchedClusters = _matchToMultiResiduePatchedTemplates(data, clusters, patch, patches[patchName], bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1642
1643
1644
1645
1646
1647
                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.
1648

1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
        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

1664
1665
1666
    return unmatchedResidues


peastman's avatar
peastman committed
1667
def _generatePatchedSingleResidueTemplates(template, patches, index, newTemplates, alteredAtoms):
1668
1669
    """Apply all possible combinations of a set of single-residue patches to a template."""
    try:
peastman's avatar
peastman committed
1670
1671
1672
1673
1674
1675
1676
        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)
1677
1678
1679
1680
    except:
        # This probably means the patch is inconsistent with another one that has already been applied,
        # so just ignore it.
        patchedTemplate = None
1681

1682
    # Call this function recursively to generate combinations of patches.
1683

1684
    if index+1 < len(patches):
peastman's avatar
peastman committed
1685
        _generatePatchedSingleResidueTemplates(template, patches, index+1, newTemplates, alteredAtoms)
1686
        if patchedTemplate is not None:
peastman's avatar
peastman committed
1687
1688
            newAlteredAtoms = alteredAtoms.union(patches[index].allAtomNames)
            _generatePatchedSingleResidueTemplates(patchedTemplate, patches, index+1, newTemplates, newAlteredAtoms)
1689
1690


1691
def _matchToMultiResiduePatchedTemplates(data, clusters, patch, residueTemplates, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles):
1692
1693
1694
    """Apply a multi-residue patch to templates, then try to match them against clusters of residues."""
    matchedClusters = []
    selectedTemplates = [None]*patch.numResidues
1695
    _applyMultiResiduePatch(data, clusters, patch, residueTemplates, selectedTemplates, 0, matchedClusters, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1696
1697
1698
    return matchedClusters


1699
def _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index, matchedClusters, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles):
1700
1701
1702
1703
1704
    """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
1705
            _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index+1, matchedClusters, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1706
1707
1708
    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.
1709

1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
        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):
1721
                    matches = compiled.matchResidueToTemplate(residue, template, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1722
1723
1724
1725
1726
1727
                    if matches is None:
                        residueMatches = None
                        break
                    else:
                        residueMatches.append(matches)
                if residueMatches is not None:
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
                    # 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.
1746

1747
1748
1749
1750
                        for i in range(patch.numResidues):
                            data.recordMatchedAtomParameters(residues[i], patchedTemplates[i], residueMatches[i])
                        newlyMatchedClusters.append(cluster)
                        break
1751

1752
        # Record which clusters were successfully matched.
1753

1754
1755
1756
        matchedClusters += newlyMatchedClusters
        for cluster in newlyMatchedClusters:
            clusters.remove(cluster)
1757

1758

1759
1760
1761
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()])
1762
    numResidueAtoms = sum(residueCounts.values())
1763
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
1764

1765
    # Loop over templates and see how closely each one might match.
1766

1767
1768
1769
1770
1771
1772
    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])
1773

1774
        # Does the residue have any atoms that clearly aren't in the template?
1775

1776
1777
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
1778

1779
        # If there are too many missing atoms, discard this template.
1780

1781
        numTemplateAtoms = sum(templateCounts.values())
1782
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
1783
1784
1785
1786
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
1787

1788
1789
        # 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.
1790

1791
1792
1793
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
1794

1795
        # Accept this as our new best match.
1796

1797
1798
1799
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
1800
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
1801

1802
    # Return an appropriate error message.
1803

1804
    if numBestMatchAtoms == numResidueAtoms:
1805
1806
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
1807
1808
1809
1810
            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:
1811
1812
1813
1814
1815
            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)
1816
1817
1818
        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.'

1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
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():
1837
        template.addAtom(ForceField._TemplateAtomData(atom.name, None, atom.element))
1838
1839
1840
1841
1842
1843
1844
1845
1846
    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
1847

1848
1849
1850
1851
1852
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
1853
    wildcard = generator.ff._atomClasses['']
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
    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
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
                    elif tordef.ordering == 'smirnoff':
                        # topology atom indexes
                        a1 = torsion[0]
                        a2 = torsion[t2[1]]
                        a3 = torsion[t3[1]]
                        a4 = torsion[t4[1]]
                        # enforce exact match
                        match = (a1, a2, a3, a4, tordef)
                        break

Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1958
    return match
1959

1960

1961
1962
1963
1964
1965
# 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.

1966
## @private
1967
class HarmonicBondGenerator(object):
1968
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
1969

1970
1971
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
1972
        self.bondsForAtomType = defaultdict(set)
1973
1974
1975
1976
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
1977

1978
1979
1980
    def registerBond(self, parameters):
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
peastman's avatar
peastman committed
1981
            index = len(self.types1)
1982
1983
            self.types1.append(types[0])
            self.types2.append(types[1])
peastman's avatar
peastman committed
1984
1985
1986
1987
            for t in types[0]:
                self.bondsForAtomType[t].add(index)
            for t in types[1]:
                self.bondsForAtomType[t].add(index)
1988
1989
            self.length.append(_convertParameterToNumber(parameters['length']))
            self.k.append(_convertParameterToNumber(parameters['k']))
Justin MacCallum's avatar
Justin MacCallum committed
1990

1991
1992
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
1993
1994
1995
1996
1997
1998
        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]
1999
        for bond in element.findall('Bond'):
2000
            generator.registerBond(bond.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
2001

2002
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2003
        existing = [f for f in sys.getForces() if type(f) == mm.HarmonicBondForce]
2004
2005
2006
2007
2008
2009
2010
2011
        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
2012
            for i in self.bondsForAtomType[type1]:
2013
2014
2015
2016
2017
                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:
2018
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
2019
2020
2021
2022
2023
                    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])
2024
2025
2026
2027
2028
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


2029
## @private
2030
class HarmonicAngleGenerator(object):
2031
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2032

2033
2034
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
2035
        self.anglesForAtom2Type = defaultdict(list)
2036
2037
2038
2039
2040
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2041

2042
2043
2044
    def registerAngle(self, parameters):
        types = self.ff._findAtomTypes(parameters, 3)
        if None not in types:
peastman's avatar
peastman committed
2045
            index = len(self.types1)
2046
2047
2048
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.types3.append(types[2])
peastman's avatar
peastman committed
2049
2050
            for t in types[1]:
                self.anglesForAtom2Type[t].append(index)
2051
2052
2053
            self.angle.append(_convertParameterToNumber(parameters['angle']))
            self.k.append(_convertParameterToNumber(parameters['k']))

2054
2055
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2056
2057
2058
2059
2060
2061
        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]
2062
        for angle in element.findall('Angle'):
2063
            generator.registerAngle(angle.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
2064

2065
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2066
2067
2068
2069
2070
        pass

    def postprocessSystem(self, sys, data, args):
        # We need to wait until after all bonds have been added so their lengths will be set correctly.

2071
        existing = [f for f in sys.getForces() if type(f) == mm.HarmonicAngleForce]
2072
2073
2074
2075
2076
2077
2078
2079
2080
        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
2081
            for i in self.anglesForAtom2Type[type2]:
2082
2083
2084
2085
2086
2087
                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
2088

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

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

2101
2102
2103
2104
2105
                        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]))
2106
                                data.addConstraint(sys, angle[0], angle[2], length)
2107
2108
2109
                    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])
2110
2111
2112
2113
2114
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


2115
## @private
2116
class PeriodicTorsion(object):
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
    """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 = []
2127
        self.ordering = 'default'
2128

2129
## @private
2130
class PeriodicTorsionGenerator(object):
2131
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2132

2133
2134
    def __init__(self, forcefield):
        self.ff = forcefield
2135
2136
        self.proper = []
        self.improper = []
2137
        self.propersForAtomType = defaultdict(set)
Justin MacCallum's avatar
Justin MacCallum committed
2138

2139
2140
2141
    def registerProperTorsion(self, parameters):
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
2142
            index = len(self.proper)
2143
            self.proper.append(torsion)
2144
2145
2146
2147
            for t in torsion.types2:
                self.propersForAtomType[t].add(index)
            for t in torsion.types3:
                self.propersForAtomType[t].add(index)
2148

2149
    def registerImproperTorsion(self, parameters, ordering='default'):
2150
2151
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
2152
            if ordering in ['default', 'charmm', 'amber', 'smirnoff']:
2153
2154
                torsion.ordering = ordering
            else:
2155
                raise ValueError('Illegal ordering type %s for improper torsion %s' % (ordering, torsion))
2156
2157
            self.improper.append(torsion)

2158
2159
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2160
2161
2162
2163
2164
2165
        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]
2166
        for torsion in element.findall('Proper'):
2167
            generator.registerProperTorsion(torsion.attrib)
2168
        for torsion in element.findall('Improper'):
2169
2170
2171
2172
            if 'ordering' in element.attrib:
                generator.registerImproperTorsion(torsion.attrib, element.attrib['ordering'])
            else:
                generator.registerImproperTorsion(torsion.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
2173

2174
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2175
        existing = [f for f in sys.getForces() if type(f) == mm.PeriodicTorsionForce]
2176
2177
2178
2179
2180
2181
        if len(existing) == 0:
            force = mm.PeriodicTorsionForce()
            sys.addForce(force)
        else:
            force = existing[0]
        wildcard = self.ff._atomClasses['']
tic20's avatar
tic20 committed
2182
        proper_cache = {}
2183
        for torsion in data.propers:
tic20's avatar
tic20 committed
2184
2185
            type1, type2, type3, type4 = [data.atomType[data.atoms[torsion[i]]] for i in range(4)]
            sig = (type1, type2, type3, type4)
2186
            sig = frozenset((sig, sig[::-1]))
tic20's avatar
tic20 committed
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
            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
2207
2208
2209
2210
            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])
2211
        impr_cache = {}
2212
        for torsion in data.impropers:
Peter Eastman's avatar
Peter Eastman committed
2213
            t1, t2, t3, t4 = [data.atomType[data.atoms[torsion[i]]] for i in range(4)]
2214
            sig = (t1, t2, t3, t4)
2215
2216
2217
2218
            match = impr_cache.get(sig, None)
            if match == -1:
                # Previously checked, and doesn't appear in the database
                continue
2219
2220
2221
2222
            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)
2223
2224
2225
            if match is None:
                match = _matchImproper(data, torsion, self)
                if match is not None:
2226
2227
2228
                    order = match[:4]
                    i1, i2, i3, i4 = tuple(torsion.index(a) for a in order)
                    impr_cache[sig] = (i1, i2, i3, i4, match[-1])
2229
2230
                else:
                    impr_cache[sig] = -1
2231
2232
2233
2234
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                for i in range(len(tordef.phase)):
                    if tordef.k[i] != 0:
2235
2236
2237
2238
2239
2240
2241
                        if tordef.ordering == 'smirnoff':
                            # Add all torsions in trefoil
                            force.addTorsion(a1, a2, a3, a4, tordef.periodicity[i], tordef.phase[i], tordef.k[i])
                            force.addTorsion(a1, a3, a4, a2, tordef.periodicity[i], tordef.phase[i], tordef.k[i])
                            force.addTorsion(a1, a4, a2, a3, tordef.periodicity[i], tordef.phase[i], tordef.k[i])
                        else:
                            force.addTorsion(a1, a2, a3, a4, tordef.periodicity[i], tordef.phase[i], tordef.k[i])
2242
2243
2244
parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


2245
## @private
2246
class RBTorsion(object):
2247
2248
    """An RBTorsion records the information for a Ryckaert-Bellemans torsion definition."""

2249
    def __init__(self, types, c, ordering='charmm'):
2250
2251
2252
2253
2254
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.c = c
2255
        if ordering in ['default', 'charmm', 'amber']:
2256
2257
            self.ordering = ordering
        else:
2258
            raise ValueError('Illegal ordering type %s for RBTorsion (%s,%s,%s,%s)' % (ordering, types[0], types[1], types[2], types[3]))
2259

2260
## @private
2261
class RBTorsionGenerator(object):
2262
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2263

2264
2265
    def __init__(self, forcefield):
        self.ff = forcefield
2266
2267
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
2268

2269
2270
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2271
2272
2273
2274
2275
2276
        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]
2277
        for torsion in element.findall('Proper'):
2278
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2279
            if None not in types:
2280
2281
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
2282
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2283
            if None not in types:
2284
2285
2286
2287
                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
2288

2289
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2290
        existing = [f for f in sys.getForces() if type(f) == mm.RBTorsionForce]
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
        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:
2317
            match = _matchImproper(data, torsion, self)
2318
2319
2320
            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])
2321
2322
2323
2324

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


2325
## @private
2326
class CMAPTorsion(object):
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
    """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

2337
## @private
2338
class CMAPTorsionGenerator(object):
2339
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2340

2341
2342
    def __init__(self, forcefield):
        self.ff = forcefield
2343
2344
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
2345

2346
2347
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2348
2349
2350
2351
2352
2353
        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]
2354
2355
2356
2357
2358
2359
2360
        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'):
2361
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
2362
            if None not in types:
2363
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
2364

2365
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2366
        existing = [f for f in sys.getForces() if type(f) == mm.CMAPTorsionForce]
2367
2368
2369
2370
2371
2372
2373
        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
2374

2375
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
2376

2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
        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


2420
## @private
2421
class NonbondedGenerator(object):
2422
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2423

2424
2425
    SCALETOL = 1e-5

2426
    def __init__(self, forcefield, coulomb14scale, lj14scale, useDispersionCorrection):
2427
        self.ff = forcefield
2428
2429
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
2430
        self.useDispersionCorrection = useDispersionCorrection
2431
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
2432

2433
    def registerAtom(self, parameters):
2434
        self.params.registerAtom(parameters)
2435

2436
2437
2438
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
2439
2440
2441
2442
        if 'useDispersionCorrection' in element.attrib:
            useDispersionCorrection = bool(eval(element.attrib.get('useDispersionCorrection')))
        else:
            useDispersionCorrection = None
2443
        if len(existing) == 0:
2444
2445
2446
2447
2448
2449
            generator = NonbondedGenerator(
                ff,
                float(element.attrib['coulomb14scale']),
                float(element.attrib['lj14scale']),
                useDispersionCorrection
            )
2450
            ff.registerGenerator(generator)
2451
2452
2453
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2454
2455
2456
            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
2457
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
2458
2459
2460
2461
2462
2463
            if (
                    generator.useDispersionCorrection is not None
                    and useDispersionCorrection is not None
                    and generator.useDispersionCorrection != useDispersionCorrection
            ):
                raise ValueError('Found multiple NonbondedForce tags with different useDispersionCorrection settings.')
2464
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2465

2466
2467
2468
2469
2470
    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
2471
2472
                     PME:mm.NonbondedForce.PME,
                     LJPME:mm.NonbondedForce.LJPME}
2473
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2474
            raise ValueError('Illegal nonbonded method for NonbondedForce')
2475
2476
        force = mm.NonbondedForce()
        for atom in data.atoms:
2477
2478
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
2479
2480
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2481
2482
2483
        if args['switchDistance'] is not None:
            force.setUseSwitchingFunction(True)
            force.setSwitchingDistance(args['switchDistance'])
peastman's avatar
peastman committed
2484
2485
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
2486
        if 'useDispersionCorrection' in args:
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
            lrc_keyword = bool(args['useDispersionCorrection'])
            if self.useDispersionCorrection is not None and lrc_keyword != self.useDispersionCorrection:
                warnings.warn(
                    "Conflicting settings for useDispersionCorrection in createSystem() and forcefield file. "
                    "Using the one specified in createSystem()."
                )
            force.setUseDispersionCorrection(lrc_keyword)
        elif self.useDispersionCorrection is not None:
            force.setUseDispersionCorrection(self.useDispersionCorrection)
        else:
            # by default
            force.setUseDispersionCorrection(True)
peastman's avatar
peastman committed
2499
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
2500

peastman's avatar
peastman committed
2501
    def postprocessSystem(self, sys, data, args):
2502
2503
        # Create the exceptions.

2504
        bondIndices = _findBondsForExclusions(data, sys)
peastman's avatar
peastman committed
2505
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
2506
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
2507
2508
2509

parsers["NonbondedForce"] = NonbondedGenerator.parseElement

2510
2511

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

2515
    def __init__(self, forcefield, lj14scale, useDispersionCorrection):
ChayaSt's avatar
ChayaSt committed
2516
        self.ff = forcefield
2517
        self.nbfixTypes = {}
ChayaSt's avatar
ChayaSt committed
2518
        self.lj14scale = lj14scale
2519
        self.useDispersionCorrection = useDispersionCorrection
2520
        self.ljTypes = ForceField._AtomTypeParameters(forcefield, 'LennardJonesForce', 'Atom', ('sigma', 'epsilon'))
ChayaSt's avatar
ChayaSt committed
2521

ChayaSt's avatar
ChayaSt committed
2522
    def registerNBFIX(self, parameters):
ChayaSt's avatar
ChayaSt committed
2523
2524
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
2525
2526
2527
2528
2529
2530
            for type1 in types[0]:
                for type2 in types[1]:
                    epsilon = _convertParameterToNumber(parameters['epsilon'])
                    sigma = _convertParameterToNumber(parameters['sigma'])
                    self.nbfixTypes[(type1, type2)] = [sigma, epsilon]
                    self.nbfixTypes[(type2, type1)] = [sigma, epsilon]
2531

ChayaSt's avatar
ChayaSt committed
2532
    def registerLennardJones(self, parameters):
2533
        self.ljTypes.registerAtom(parameters)
ChayaSt's avatar
ChayaSt committed
2534
2535
2536

    @staticmethod
    def parseElement(element, ff):
ChayaSt's avatar
ChayaSt committed
2537
        existing = [f for f in ff._forces if isinstance(f, LennardJonesGenerator)]
2538
2539
2540
2541
        if 'useDispersionCorrection' in element.attrib:
            useDispersionCorrection = bool(eval(element.attrib.get('useDispersionCorrection')))
        else:
            useDispersionCorrection = None
ChayaSt's avatar
ChayaSt committed
2542
        if len(existing) == 0:
2543
2544
2545
2546
2547
            generator = LennardJonesGenerator(
                ff,
                float(element.attrib['lj14scale']),
                useDispersionCorrection=useDispersionCorrection
            )
ChayaSt's avatar
ChayaSt committed
2548
2549
            ff.registerGenerator(generator)
        else:
ChayaSt's avatar
ChayaSt committed
2550
            # Multiple <LennardJonesForce> tags were found, probably in different files
ChayaSt's avatar
ChayaSt committed
2551
            generator = existing[0]
2552
2553
            if abs(generator.lj14scale - float(element.attrib['lj14scale'])) > NonbondedGenerator.SCALETOL:
                raise ValueError('Found multiple LennardJonesForce tags with different 1-4 scales')
2554
2555
2556
2557
2558
2559
            if (
                    generator.useDispersionCorrection is not None
                    and useDispersionCorrection is not None
                    and generator.useDispersionCorrection != useDispersionCorrection
            ):
                raise ValueError('Found multiple LennardJonesForce tags with different useDispersionCorrection settings.')
ChayaSt's avatar
ChayaSt committed
2560
2561
        for LJ in element.findall('Atom'):
            generator.registerLennardJones(LJ.attrib)
ChayaSt's avatar
ChayaSt committed
2562
        for Nbfix in element.findall('NBFixPair'):
ChayaSt's avatar
ChayaSt committed
2563
            generator.registerNBFIX(Nbfix.attrib)
ChayaSt's avatar
ChayaSt committed
2564

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

        nbfixTypeSet = set().union(*self.nbfixTypes)
peastman's avatar
peastman committed
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
        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]
2586
            else:
peastman's avatar
peastman committed
2587
2588
2589
2590
2591
                # This is a new type.
                typeToMergedType[t] = len(mergedTypes)
                paramsToMergedType[params] = len(mergedTypes)
                mergedTypes.append(t)
                mergedTypeParams.append(params)
2592

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

peastman's avatar
peastman committed
2595
        numLjTypes = len(mergedTypes)
2596
        acoef = [0]*(numLjTypes*numLjTypes)
ChayaSt's avatar
ChayaSt committed
2597
        bcoef = acoef[:]
2598
2599
        for m in range(numLjTypes):
            for n in range(numLjTypes):
peastman's avatar
peastman committed
2600
                pair = (mergedTypes[m], mergedTypes[n])
2601
2602
2603
2604
2605
2606
                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
2607
                    continue
ChayaSt's avatar
ChayaSt committed
2608
                else:
peastman's avatar
peastman committed
2609
                    sigma = 0.5*(mergedTypeParams[m][0]+mergedTypeParams[n][0])
2610
                    sigma6 = sigma**6
peastman's avatar
peastman committed
2611
                    epsilon = math.sqrt(mergedTypeParams[m][1]*mergedTypeParams[n][1])
2612
2613
2614
2615
2616
2617
                    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
2618
        self.force.addPerParticleParameter('type')
2619
        self.force.setName('LennardJones')
Peter Eastman's avatar
Peter Eastman committed
2620
        if nonbondedMethod in [CutoffPeriodic, Ewald, PME, LJPME]:
ChayaSt's avatar
ChayaSt committed
2621
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
ChayaSt's avatar
ChayaSt committed
2622
        elif nonbondedMethod is NoCutoff:
ChayaSt's avatar
ChayaSt committed
2623
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
ChayaSt's avatar
ChayaSt committed
2624
        elif nonbondedMethod is CutoffNonPeriodic:
ChayaSt's avatar
ChayaSt committed
2625
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic)
ChayaSt's avatar
ChayaSt committed
2626
        else:
2627
            raise AssertionError('Unrecognized nonbonded method [%s]' % nonbondedMethod)
2628
        if args['switchDistance'] is not None:
2629
2630
            self.force.setUseSwitchingFunction(True)
            self.force.setSwitchingDistance(args['switchDistance'])
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
        if 'useDispersionCorrection' in args:
            lrc_keyword = bool(args['useDispersionCorrection'])
            if self.useDispersionCorrection is not None and lrc_keyword != self.useDispersionCorrection:
                warnings.warn(
                    "Conflicting settings for useDispersionCorrection in createSystem() and forcefield file. "
                    "Using the one specified in createSystem()."
                )
            self.force.setUseLongRangeCorrection(lrc_keyword)
        elif self.useDispersionCorrection is not None:
            self.force.setUseLongRangeCorrection(self.useDispersionCorrection)
        else:
            # by default
            self.force.setUseLongRangeCorrection(True)
2644

ChayaSt's avatar
ChayaSt committed
2645
        # Add the particles
2646

peastman's avatar
peastman committed
2647
2648
        for atom in data.atoms:
            self.force.addParticle((typeToMergedType[data.atomType[atom]],))
ChayaSt's avatar
ChayaSt committed
2649
        self.force.setCutoffDistance(nonbondedCutoff)
ChayaSt's avatar
ChayaSt committed
2650
2651
2652
        sys.addForce(self.force)

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

2655
        bondIndices = _findBondsForExclusions(data, sys)
2656
2657
2658
2659
2660
2661
2662
2663
2664
        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')
2665
            bonded.setName('LennardJones14')
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
2679
            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)
2680
2681
2682
2683
2684
2685
2686
2687
                        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)
2688
                    bonded.addBond(p1, p2, (sigma, epsilon))
ChayaSt's avatar
ChayaSt committed
2689

ChayaSt's avatar
ChayaSt committed
2690
parsers["LennardJonesForce"] = LennardJonesGenerator.parseElement
2691

2692

2693
## @private
2694
class GBSAOBCGenerator(object):
2695
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2696

2697
2698
    def __init__(self, forcefield):
        self.ff = forcefield
2699
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
2700

2701
    def registerAtom(self, parameters):
2702
        self.params.registerAtom(parameters)
2703

2704
2705
    @staticmethod
    def parseElement(element, ff):
2706
2707
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
2708
2709
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
2710
2711
2712
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2713
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2714

2715
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2716
2717
2718
2719
2720
2721
        methodMap = {NoCutoff:mm.GBSAOBCForce.NoCutoff,
                     CutoffNonPeriodic:mm.GBSAOBCForce.CutoffNonPeriodic,
                     CutoffPeriodic:mm.GBSAOBCForce.CutoffPeriodic,
                     Ewald:mm.GBSAOBCForce.CutoffPeriodic,
                     PME:mm.GBSAOBCForce.CutoffPeriodic,
                     LJPME:mm.GBSAOBCForce.CutoffPeriodic}
2722
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2723
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
2724
2725
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
2726
2727
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
2728
2729
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2730
2731
2732
2733
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
2734
2735
        sys.addForce(force)

2736
2737
    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
2738

2739
2740
2741
2742
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

2743
2744
2745
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


2746
## @private
2747
class CustomBondGenerator(object):
2748
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2749

2750
2751
    def __init__(self, forcefield):
        self.ff = forcefield
2752
2753
2754
2755
2756
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2757

2758
2759
    @staticmethod
    def parseElement(element, ff):
2760
2761
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
2762
2763
2764
2765
2766
2767
        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'):
2768
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2769
            if None not in types:
2770
2771
2772
                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
2773

2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
    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


2794
## @private
2795
class CustomAngleGenerator(object):
2796
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2797

2798
2799
    def __init__(self, forcefield):
        self.ff = forcefield
2800
2801
2802
2803
2804
2805
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2806

2807
2808
    @staticmethod
    def parseElement(element, ff):
2809
2810
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
2811
2812
2813
2814
2815
2816
        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'):
2817
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2818
            if None not in types:
2819
2820
2821
2822
                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
2823

2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
    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


2846
## @private
2847
class CustomTorsion(object):
2848
2849
    """A CustomTorsion records the information for a custom torsion definition."""

2850
    def __init__(self, types, paramValues, ordering='charmm'):
2851
2852
2853
2854
2855
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.paramValues = paramValues
2856
        if ordering in ['default', 'charmm', 'amber']:
2857
2858
            self.ordering = ordering
        else:
2859
            raise ValueError('Illegal ordering type %s for CustomTorsion (%s,%s,%s,%s)' % (ordering, types[0], types[1], types[2], types[3]))
2860

2861
## @private
2862
class CustomTorsionGenerator(object):
2863
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2864

2865
2866
    def __init__(self, forcefield):
        self.ff = forcefield
2867
2868
2869
2870
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
2871

2872
2873
    @staticmethod
    def parseElement(element, ff):
2874
2875
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
2876
2877
2878
2879
2880
2881
        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'):
2882
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2883
            if None not in types:
2884
2885
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
2886
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2887
            if None not in types:
2888
2889
2890
2891
                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
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
2920
    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:
2921
            match = _matchImproper(data, torsion, self)
2922
2923
2924
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
2925
2926
2927
2928

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


2929
## @private
2930
class CustomNonbondedGenerator(object):
2931
2932
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

2933
2934
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
2935
2936
2937
2938
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
2939
        self.computedValues = {}
2940
2941
2942
2943
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
2944
2945
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
2946
2947
2948
2949
        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'])
2950
2951
        for value in element.findall('ComputedValue'):
            generator.computedValues[value.attrib['name']] = value.attrib['expression']
2952
2953
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
2954
        generator.functions += _parseFunctions(element)
2955
2956
2957
2958

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomNonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomNonbondedForce.CutoffNonPeriodic,
2959
2960
2961
2962
                     CutoffPeriodic:mm.CustomNonbondedForce.CutoffPeriodic,
                     Ewald:mm.CustomNonbondedForce.CutoffPeriodic,
                     PME:mm.CustomNonbondedForce.CutoffPeriodic,
                     LJPME:mm.CustomNonbondedForce.CutoffPeriodic}
2963
2964
2965
2966
2967
2968
2969
        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)
2970
2971
        for name in self.computedValues:
            force.addComputedValue(name, self.computedValues[name])
2972
        _createFunctions(force, self.functions)
2973
        for atom in data.atoms:
2974
2975
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
2976
2977
2978
2979
2980
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

    def postprocessSystem(self, sys, data, args):
2981
        # Create the exclusions.
2982

2983
        bondIndices = _findBondsForExclusions(data, sys)
2984
2985
2986
2987
2988
2989
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.CustomNonbondedForce)][0]
        nonbonded.createExclusionsFromBonds(bondIndices, self.bondCutoff)

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


2990
## @private
2991
class CustomGBGenerator(object):
2992
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2993

2994
2995
    def __init__(self, forcefield):
        self.ff = forcefield
2996
2997
2998
2999
3000
3001
3002
3003
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
3004
3005
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
3006
3007
3008
3009
        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'])
3010
3011
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
3012
3013
3014
3015
3016
3017
3018
        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']]))
3019
        generator.functions += _parseFunctions(element)
Justin MacCallum's avatar
Justin MacCallum committed
3020

3021
3022
3023
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomGBForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomGBForce.CutoffNonPeriodic,
3024
3025
3026
3027
                     CutoffPeriodic:mm.CustomGBForce.CutoffPeriodic,
                     Ewald:mm.CustomGBForce.CutoffPeriodic,
                     PME:mm.CustomGBForce.CutoffPeriodic,
                     LJPME:mm.CustomGBForce.CutoffPeriodic}
3028
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
3029
            raise ValueError('Illegal nonbonded method for CustomGBForce')
3030
3031
3032
3033
3034
3035
3036
3037
3038
        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])
3039
        _createFunctions(force, self.functions)
3040
        for atom in data.atoms:
3041
3042
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
3043
3044
3045
3046
3047
3048
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

3049

3050
3051
3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
3073
3074
3075
3076
3077
3078
3079
3080
3081
3082
3083
3084
3085
3086
3087
## @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'):
3088
            types = ff._findAtomTypes(donor.attrib, 3)[:generator.particlesPerDonor]
3089
3090
3091
3092
3093
3094
3095
3096
            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'):
3097
            types = ff._findAtomTypes(acceptor.attrib, 3)[:generator.particlesPerAcceptor]
3098
3099
3100
3101
3102
3103
3104
3105
3106
3107
3108
3109
            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,
3110
3111
3112
3113
                     CutoffPeriodic:mm.CustomHbondForce.CutoffPeriodic,
                     Ewald:mm.CustomHbondForce.CutoffPeriodic,
                     PME:mm.CustomHbondForce.CutoffPeriodic,
                     LJPME:mm.CustomHbondForce.CutoffPeriodic}
3114
3115
3116
3117
3118
3119
3120
3121
3122
3123
3124
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
        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]
Peter Eastman's avatar
Peter Eastman committed
3135
                    if type1 in types1:
3136
3137
3138
3139
3140
3141
3142
3143
                        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]
3144
3145
3146
3147
                    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])
3148
3149
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
        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]
Peter Eastman's avatar
Peter Eastman committed
3167
                    if type1 in types1:
3168
3169
3170
3171
3172
3173
3174
3175
                        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]
3176
                    if type1 in types1 and type2 in types2:
3177
                        force.addAcceptor(bond.atom1, bond.atom2, -1, self.acceptorParamValues[i])
3178
3179
                    elif type1 in types2 and type2 in types1:
                        force.addAcceptor(bond.atom2, bond.atom1, -1, self.acceptorParamValues[i])
3180
3181
3182
3183
3184
3185
3186
3187
3188
3189
3190
3191
3192
3193
3194
3195
3196
3197
3198
3199
3200
3201
3202
3203
3204
3205
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
        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


3218
## @private
3219
class CustomManyParticleGenerator(object):
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230
3231
    """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 = []
3232

3233
3234
3235
3236
3237
3238
3239
3240
3241
3242
3243
3244
    @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(',')]))
3245
3246
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
3247
3248
3249
3250

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomManyParticleForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomManyParticleForce.CutoffNonPeriodic,
3251
3252
3253
3254
                     CutoffPeriodic:mm.CustomManyParticleForce.CutoffPeriodic,
                     Ewald:mm.CustomManyParticleForce.CutoffPeriodic,
                     PME:mm.CustomManyParticleForce.CutoffPeriodic,
                     LJPME:mm.CustomManyParticleForce.CutoffPeriodic}
3255
3256
3257
3258
3259
3260
3261
3262
3263
3264
3265
3266
        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':
3267
                force.addTabulatedFunction(name, mm.Continuous1DFunction(values, params['min'], params['max'], params['periodic']))
3268
3269
3270
3271
3272
3273
3274
3275
3276
3277
3278
            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:
3279
3280
3281
            values = self.params.getAtomParameters(atom, data)
            type = int(self.params.getExtraParameters(atom, data)['filterType'])
            force.addParticle(values, type)
3282
3283
3284
3285
3286
3287
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

3289
3290
3291
3292
3293
3294
3295
3296
3297
3298
3299
        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)))
3300

3301
3302
        # 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.
3303

3304
3305
3306
3307
3308
3309
3310
3311
3312
        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.
3313

3314
3315
3316
3317
3318
        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
3319
def getAtomPrint(data, atomIndex):
3320

Peter Eastman's avatar
Peter Eastman committed
3321
3322
3323
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
3324
    else:
Peter Eastman's avatar
Peter Eastman committed
3325
        returnString = "NA"
3326
3327
3328
3329
3330

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
3331
def countConstraint(data):
3332

Peter Eastman's avatar
Peter Eastman committed
3333
    bondCount = 0
3334
3335
3336
3337
3338
3339
3340
    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
3341
        if (isConstrained):
3342
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
3343

3344
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
3345

3346
## @private
3347
class AmoebaBondGenerator(object):
3348
3349
3350

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

3351
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
3352
3353

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

3355
3356
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
3357
3358
3359
3360
3361
3362
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3363

3364
3365
3366
3367
3368
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3372
        generator = AmoebaBondGenerator(element.attrib['bond-cubic'], element.attrib['bond-quartic'])
3373
3374
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
3375
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
3376
            if None not in types:
3377
3378
3379
3380
3381
                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:
3382
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
3383
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
3384
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3385
3386
                raise ValueError(outputString)

3387
3388
    #=============================================================================================

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

3391
3392
        energy = "k*(d^2 + %s*d^3 + %s*d^4); d=r-r0" % (self.cubic, self.quartic)
        existing = [f for f in sys.getForces() if type(f) == mm.CustomBondForce and f.getEnergyFunction() == energy]
3393
        if len(existing) == 0:
3394
3395
3396
            force = mm.CustomBondForce(energy)
            force.addPerBondParameter('r0')
            force.addPerBondParameter('k')
3397
            force.setName('AmoebaBond')
3398
3399
3400
3401
3402
3403
3404
3405
3406
3407
3408
3409
3410
            sys.addForce(force)
        else:
            force = existing[0]

        for bond in data.bonds:
            type1 = data.atomType[data.atoms[bond.atom1]]
            type2 = data.atomType[data.atoms[bond.atom2]]
            for i in range(len(self.types1)):
                types1 = self.types1[i]
                types2 = self.types2[i]
                if (type1 in types1 and type2 in types2) or (type1 in types2 and type2 in types1):
                    bond.length = self.length[i]
                    if bond.isConstrained:
3411
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
3412
3413
                    if self.k[i] != 0:
                        if not bond.isConstrained or args.get('flexibleConstraints', False):
3414
                            force.addBond(bond.atom1, bond.atom2, [self.length[i], self.k[i]])
3415
3416
                    break

3417
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
3418
3419
3420
3421

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

Peter Eastman's avatar
Peter Eastman committed
3423
def addAngleConstraint(angle, idealAngle, data, sys):
3424
3425

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

Peter Eastman's avatar
Peter Eastman committed
3427
3428
    bond1 = None
    bond2 = None
3429
3430
3431
3432
3433
3434
3435
    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
3436

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

3439
3440
3441
3442
        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
3443
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
3444
                data.addConstraint(sys, angle[0], angle[2], length)
3445
3446
3447
                return

#=============================================================================================
3448
## @private
3449
class AmoebaAngleGenerator(object):
3450
3451

    #=============================================================================================
3452
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
3453
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3454

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

Peter Eastman's avatar
Peter Eastman committed
3457
3458
3459
3460
3461
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3462

Peter Eastman's avatar
Peter Eastman committed
3463
3464
3465
        self.types1 = []
        self.types2 = []
        self.types3 = []
3466

Peter Eastman's avatar
Peter Eastman committed
3467
3468
        self.angle = []
        self.k = []
3469
        self.inPlane = []
Justin MacCallum's avatar
Justin MacCallum committed
3470

3471
3472
3473
3474
3475
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3479
3480
3481
3482
3483
3484
3485
3486
        existing = [f for f in forceField._forces if isinstance(f, AmoebaAngleGenerator)]
        if len(existing) == 0:
            generator = AmoebaAngleGenerator(forceField, element.attrib['angle-cubic'], element.attrib['angle-quartic'],  element.attrib['angle-pentic'], element.attrib['angle-sextic'])
            forceField.registerGenerator(generator)
        else:
            generator = existing[0]
            if tuple(element.attrib[x] for x in ('angle-cubic', 'angle-quartic', 'angle-pentic', 'angle-sextic')) != (generator.cubic, generator.quartic, generator.pentic, generator.sextic):
                raise ValueError('All <AmoebaAngleForce> tags must use identical scale factors')
3487
        for angle in element.findall('Angle'):
3488
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
3489
            if None not in types:
3490
3491
3492
3493
3494
3495

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
3496
                angleList.append(float(angle.attrib['angle1']))
3497
3498

                try:
Peter Eastman's avatar
Peter Eastman committed
3499
                    angleList.append(float(angle.attrib['angle2']))
3500
                    try:
Peter Eastman's avatar
Peter Eastman committed
3501
                        angleList.append(float(angle.attrib['angle3']))
3502
3503
3504
3505
3506
3507
                    except:
                        pass
                except:
                    pass
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
3508
3509
3510
3511
                if 'inPlane' in angle.attrib:
                    generator.inPlane.append(angle.attrib['inPlane'].lower() == 'true')
                else:
                    generator.inPlane.append(None) # for backward compatibility with older versions of AMOEBA
3512
            else:
3513
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
3514
3515
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
3516
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
3517
3518
                raise ValueError(outputString)

3519
3520
3521
3522
    #=============================================================================================
    # 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
3523

Peter Eastman's avatar
Peter Eastman committed
3524
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3525
3526
        if not any(isinstance(f, AmoebaOutOfPlaneBendGenerator) for f in self.forceField.getGenerators()):
            raise ValueError('A ForceField containing an <AmoebaAngleForce> must also contain an <AmoebaOutOfPlaneBendForce>')
3527
3528
3529
3530
3531

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

Peter Eastman's avatar
Peter Eastman committed
3533
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3534
3535
3536

        # get force

3537
3538
        energy = "k*(d^2 + %s*d^3 + %s*d^4 + %s*d^5 + %s*d^6); d=%.15g*theta-theta0" % (self.cubic, self.quartic, self.pentic, self.sextic, 180/math.pi)
        existing = [f for f in sys.getForces() if type(f) == mm.CustomAngleForce and f.getEnergyFunction() == energy]
3539
3540

        if len(existing) == 0:
3541
3542
3543
            force = mm.CustomAngleForce(energy)
            force.addPerAngleParameter('theta0')
            force.addPerAngleParameter('k')
3544
            force.setName('AmoebaAngle')
3545
3546
3547
3548
            sys.addForce(force)
        else:
            force = existing[0]

3549
3550
        DEG_TO_RAD = math.pi / 180

3551
        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
3552
3553
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3554
            inPlane = angleDict['inPlane']
3555

Peter Eastman's avatar
Peter Eastman committed
3556
3557
3558
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3559
            for i in range(len(self.types1)):
3560
3561
3562
                # self.inPlane is used for modern force fields.  inPlane is used for legacy ones that don't specify it.
                if self.inPlane[i] or (self.inPlane[i] is None and inPlane):
                    continue
3563
3564
3565
3566
3567
3568
                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]
3569
3570
                        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
3571
3572
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
3573
3574
3575
3576
3577
3578
                            # 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
3579
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
3580
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3581
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
3582
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3583
                            if (numberOfHydrogens < lenAngle):
3584
3585
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
3586
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
3587
                                raise ValueError(outputString)
3588
3589
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
3590

3591
                        angleDict['idealAngle'] = angleValue
3592
                        force.addAngle(angle[0], angle[1], angle[2], [angleValue, self.k[i]])
3593
3594
3595
3596
3597
3598
                    break

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

Peter Eastman's avatar
Peter Eastman committed
3600
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3601
3602
3603

        # get force

3604
3605
3606
3607
3608
3609
3610
3611
3612
3613
        energy = """k*(d^2 + %s*d^3 + %s*d^4 + %s*d^5 + %s*d^6); d=theta-theta0;
                    theta = %.15g*pointangle(x1, y1, z1, projx, projy, projz, x3, y3, z3);
                    projx = x2-nx*dot; projy = y2-ny*dot; projz = z2-nz*dot;
                    dot = nx*(x2-x3) + ny*(y2-y3) + nz*(z2-z3);
                    nx = px/norm; ny = py/norm; nz = pz/norm;
                    norm = sqrt(px*px + py*py + pz*pz);
                    px = (d1y*d2z-d1z*d2y); py = (d1z*d2x-d1x*d2z); pz = (d1x*d2y-d1y*d2x);
                    d1x = x1-x4; d1y = y1-y4; d1z = z1-z4;
                    d2x = x3-x4; d2y = y3-y4; d2z = z3-z4""" % (self.cubic, self.quartic, self.pentic, self.sextic, 180/math.pi)
        existing = [f for f in sys.getForces() if type(f) == mm.CustomCompoundBondForce and f.getEnergyFunction() == energy]
3614
        if len(existing) == 0:
3615
3616
3617
            force = mm.CustomCompoundBondForce(4, energy)
            force.addPerBondParameter("theta0")
            force.addPerBondParameter("k")
3618
            force.setName('AmoebaInPlaneAngle')
3619
3620
3621
3622
3623
            sys.addForce(force)
        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
3625
3626
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3627
            inPlane = angleDict['inPlane']
3628

Peter Eastman's avatar
Peter Eastman committed
3629
3630
3631
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3632
3633

            for i in range(len(self.types1)):
3634
3635
3636
                # self.inPlane is used for modern force fields.  inPlane is used for legacy ones that don't specify it.
                if self.inPlane[i] == False or (self.inPlane[i] is None and not inPlane):
                    continue
3637
3638
3639
3640
3641
3642
                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
3643
                    if (isConstrained and self.k[i] != 0.0):
3644
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3645
                    if self.k[i] != 0.0 and (not isConstrained or args.get('flexibleConstraints', False)):
3646
                        force.addBond((angle[0], angle[1], angle[2], angle[3]), (self.angle[i][0], self.k[i]))
3647
3648
                    break

3649
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
3650
3651
3652

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
3653
3654
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
3655
3656
#=============================================================================================

3657
## @private
3658
class AmoebaOutOfPlaneBendGenerator(object):
3659
3660
3661
3662

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

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

3664
3665
3666
3667
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3668
3669
3670
3671
3672
3673
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3674

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

Peter Eastman's avatar
Peter Eastman committed
3680
        self.ks = []
3681
3682
3683
3684
3685

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

3687
3688
3689
3690
3691
3692
3693
3694
3695
3696
3697
3698
3699
3700
3701
3702
3703
3704
3705
3706
3707
3708
3709
3710
3711
    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
3712

3713
3714
        # get global scalar parameters

3715
3716
3717
3718
3719
3720
3721
3722
        existing = [f for f in forceField._forces if isinstance(f, AmoebaOutOfPlaneBendGenerator)]
        if len(existing) == 0:
            generator = AmoebaOutOfPlaneBendGenerator(forceField, element.attrib['type'], element.attrib['opbend-cubic'], element.attrib['opbend-quartic'],  element.attrib['opbend-pentic'], element.attrib['opbend-sextic'])
            forceField.registerGenerator(generator)
        else:
            generator = existing[0]
            if tuple(element.attrib[x] for x in ('type', 'opbend-cubic', 'opbend-quartic', 'opbend-pentic', 'opbend-sextic')) != (generator.type, generator.cubic, generator.quartic, generator.pentic, generator.sextic):
                raise ValueError('All <AmoebaOutOfPlaneBendForce> tags must use identical scale factors')
3723
3724

        for angle in element.findall('Angle'):
3725
3726
3727
3728
            if 'class3' in angle.attrib and 'class4' in angle.attrib and angle.attrib['class3'] == '0' and angle.attrib['class4'] == '0':
                # This is needed for backward compatibility with old AMOEBA force fields that specified wildcards in a nonstandard way.
                angle.attrib['class3'] = ''
                angle.attrib['class4'] = ''
Peter Eastman's avatar
Peter Eastman committed
3729
            types = generator.findAtomTypes(forceField, angle, 4)
3730
3731
            if types is not None:

Peter Eastman's avatar
Peter Eastman committed
3732
3733
3734
3735
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
3736

Peter Eastman's avatar
Peter Eastman committed
3737
                generator.ks.append(float(angle.attrib['k']))
3738
3739

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3740
3741
                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
3742
3743
                raise ValueError(outputString)

3744
3745
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3746
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3747
3748
        self._nonbondedMethod = nonbondedMethod
        self._nonbondedCutoff = nonbondedCutoff
3749

3750
3751
    def postprocessSystem(self, sys, data, args):
        # We need to wait until after all bonds have been added so their lengths will be set correctly.
3752

3753
3754
3755
3756
3757
3758
3759
3760
3761
3762
        energy = """k*(theta^2 + %s*theta^3 + %s*theta^4 + %s*theta^5 + %s*theta^6);
                    theta = %.15g*pointangle(x2, y2, z2, x4, y4, z4, projx, projy, projz);
                    projx = x2-nx*dot; projy = y2-ny*dot; projz = z2-nz*dot;
                    dot = nx*(x2-x3) + ny*(y2-y3) + nz*(z2-z3);
                    nx = px/norm; ny = py/norm; nz = pz/norm;
                    norm = sqrt(px*px + py*py + pz*pz);
                    px = (d1y*d2z-d1z*d2y); py = (d1z*d2x-d1x*d2z); pz = (d1x*d2y-d1y*d2x);
                    d1x = x1-x4; d1y = y1-y4; d1z = z1-z4;
                    d2x = x3-x4; d2y = y3-y4; d2z = z3-z4""" % (self.cubic, self.quartic, self.pentic, self.sextic, 180/math.pi)
        existing = [f for f in sys.getForces() if type(f) == mm.CustomCompoundBondForce and f.getEnergyFunction() == energy]
3763
        if len(existing) == 0:
3764
3765
            force = mm.CustomCompoundBondForce(4, energy)
            force.addPerBondParameter("k")
3766
            force.setName('AmoebaOutOfPlaneBend')
3767
3768
3769
        else:
            force = existing[0]

3770
        # this hash is used to ensure the out-of-plane-bend bonds
3771
3772
        # are only added once

Peter Eastman's avatar
Peter Eastman committed
3773
        skipAtoms = dict()
3774
        angles = []
3775

3776
3777
3778
3779
3780
3781
3782
        def addBond(particles):
            types = [data.atomType[data.atoms[p]] for p in particles]
            for i in range(len(self.types1)):
                if types[1] in self.types2[i] and types[3] in self.types1[i]:
                    if (types[0] in self.types3[i] and types[2] in self.types4[i]) or (types[2] in self.types3[i] and types[0] in self.types4[i]):
                        force.addBond(particles, [self.ks[i]])
                        return
3783
3784
3785

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

3786
3787
3788
            middleAtom = angle[1]
            middleType = data.atomType[data.atoms[middleAtom]]
            middleCovalency = len(data.atomBonds[middleAtom])
3789

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

3796
            if middleCovalency == 3 and middleAtom not in skipAtoms:
3797

Peter Eastman's avatar
Peter Eastman committed
3798
                partners = []
3799
3800
3801
3802

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
3803
                    if atom1 != middleAtom:
3804
3805
3806
3807
3808
3809
3810
3811
                        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]
3812
                        if middleType in types2 and partnerType in types1:
Peter Eastman's avatar
Peter Eastman committed
3813
                            partners.append(partner)
3814
                            break
Justin MacCallum's avatar
Justin MacCallum committed
3815

3816
                if len(partners) == 3:
3817

3818
3819
3820
                    addBond([partners[0], middleAtom, partners[1], partners[2]])
                    addBond([partners[2], middleAtom, partners[0], partners[1]])
                    addBond([partners[1], middleAtom, partners[2], partners[0]])
3821

3822
                    # skipAtoms is used to ensure angles are only included once
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3823

3824
                    skipAtoms[middleAtom] = set(partners[:3])
Peter Eastman's avatar
Peter Eastman committed
3825

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3826
3827
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
3828
                    angleDict = {}
3829
3830
3831
3832
                    angleList = list(angle[:3])
                    for atomIndex in partners:
                        if atomIndex not in angleList:
                            angleList.append(atomIndex)
3833
                    angleDict['angle'] = angleList
Peter Eastman's avatar
Peter Eastman committed
3834
                    angleDict['isConstrained'] = 0
3835
3836
                    angleDict['inPlane'] = True
                    angles.append(angleDict)
3837
3838

                else:
Peter Eastman's avatar
Peter Eastman committed
3839
                    angleDict = {}
3840
3841
3842
3843
                    angleList = list(angle[:3])
                    for atomIndex in partners:
                        if atomIndex not in angleList:
                            angleList.append(atomIndex)
Peter Eastman's avatar
Peter Eastman committed
3844
3845
                    angleDict['angle'] = angleList
                    angleDict['isConstrained'] = isConstrained
3846
3847
3848
3849
3850
3851
3852
3853
3854
3855
3856
3857
3858
                    angleDict['inPlane'] = False
                    angles.append(angleDict)
            elif middleCovalency == 3 and middleAtom in skipAtoms:

                angleDict = {}
                angleList = list(angle[:3])
                for atomIndex in skipAtoms[middleAtom]:
                    if atomIndex not in angleList:
                        angleList.append(atomIndex)
                angleDict['angle'] = angleList
                angleDict['isConstrained'] = isConstrained
                angleDict['inPlane'] = True
                angles.append(angleDict)
3859

3860
3861
3862
3863
3864
3865
            else:
                angleDict = {}
                angleDict['angle'] = angle
                angleDict['isConstrained'] = isConstrained
                angleDict['inPlane'] = False
                angles.append(angleDict)
3866

3867
3868
        if len(existing) == 0 and force.getNumBonds() > 0:
            sys.addForce(force)
3869

3870
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
3871
3872

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3873
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
3874
3875
                force.createForcePostOpBendAngle(sys, data, self._nonbondedMethod, self._nonbondedCutoff, angles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, self._nonbondedMethod, self._nonbondedCutoff, angles, args)
3876
3877

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3878
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
3879
                force.createForcePostAmoebaBondForce(sys, data, self._nonbondedMethod, self._nonbondedCutoff, angles, args)
3880
3881
3882
3883
3884

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

3885
## @private
3886
class AmoebaTorsionGenerator(object):
3887
3888
3889
3890
3891
3892
3893

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
3894
        self.torsionUnit = torsionUnit
3895

Peter Eastman's avatar
Peter Eastman committed
3896
3897
3898
3899
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3900

Peter Eastman's avatar
Peter Eastman committed
3901
3902
3903
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
3904

3905
3906
3907
3908
3909
3910
3911
3912
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3914
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
3915
3916
3917
3918
3919
3920
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
3921
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
3922
            if None not in types:
3923
3924
3925
3926
3927
3928
3929

                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
3930
3931
                    tInfo = []
                    suffix = str(ii)
3932
3933
3934
3935
3936
3937
                    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
3938
3939
3940
3941
3942
3943
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
3944
3945
3946

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
3947
3948
3949
3950
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
3951
3952
                raise ValueError(outputString)

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

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

3957
        existing = [f for f in sys.getForces() if type(f) == mm.PeriodicTorsionForce]
3958
        if len(existing) == 0:
3959
            force = mm.PeriodicTorsionForce()
3960
3961
3962
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3963

3964
3965
3966
3967
3968
3969
3970
3971
3972
3973
3974
3975
3976
3977
3978
3979
        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
3980
                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):
3981
3982
3983
3984
3985
3986
                    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])
3987
3988
3989
3990
3991
3992
                    break

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

3993
## @private
3994
class AmoebaPiTorsionGenerator(object):
3995
3996
3997
3998
3999
4000

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

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

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

4002
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
4003
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
4004
4005
4006
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
4007

4008
4009
4010
4011
4012
4013
4014
4015
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
4016
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
4017
4018
4019
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
4020
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
4021
            if None not in types:
4022
4023
4024
4025
4026
4027
                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
4028
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
4029
4030
                raise ValueError(outputString)

4031
4032
    #=============================================================================================

4033
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4034

4035
4036
4037
4038
4039
4040
4041
4042
4043
        energy = """2*k*sin(phi)^2;
                    phi = pointdihedral(x3+c1x, y3+c1y, z3+c1z, x3, y3, z3, x4, y4, z4, x4+c2x, y4+c2y, z4+c2z);
                    c1x = (d14y*d24z-d14z*d24y); c1y = (d14z*d24x-d14x*d24z); c1z = (d14x*d24y-d14y*d24x);
                    c2x = (d53y*d63z-d53z*d63y); c2y = (d53z*d63x-d53x*d63z); c2z = (d53x*d63y-d53y*d63x);
                    d14x = x1-x4; d14y = y1-y4; d14z = z1-z4;
                    d24x = x2-x4; d24y = y2-y4; d24z = z2-z4;
                    d53x = x5-x3; d53y = y5-y3; d53z = z5-z3;
                    d63x = x6-x3; d63y = y6-y3; d63z = z6-z3"""
        existing = [f for f in sys.getForces() if type(f) == mm.CustomCompoundBondForce and f.getEnergyFunction() == energy]
4044
4045

        if len(existing) == 0:
4046
4047
            force = mm.CustomCompoundBondForce(6, energy)
            force.addPerBondParameter('k')
4048
            force.setName('AmoebaPiTorsion')
4049
4050
4051
        else:
            force = existing[0]

Peter Eastman's avatar
Peter Eastman committed
4052
        for bond1 in data.bonds:
4053
4054

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

Peter Eastman's avatar
Peter Eastman committed
4056
4057
            atom1 = bond1.atom1
            atom2 = bond1.atom2
Justin MacCallum's avatar
Justin MacCallum committed
4058

4059
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
4060
4061
4062
4063
4064
4065
4066
4067
4068
4069
4070

                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
4071
4072
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
4073
4074
4075
4076
4077
4078
4079
4080
4081
4082
4083
4084

                       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
4085
                           if (bondedAtom1 != atom1):
4086
4087
4088
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
4089
4090
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
4091
                                   piTorsionAtom1 = b1
4092
4093
4094
4095
4096
4097
                               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
4098
                           if (bondedAtom1 != atom2):
4099
4100
4101
4102
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4103
4104
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
4105
                                   piTorsionAtom5 = b1
4106
4107
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
4108

4109
                       force.addBond([piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6], [self.k[i]])
4110
4111
        if len(existing) == 0 and force.getNumBonds() > 0:
            sys.addForce(force)
4112
4113
4114
4115
4116

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

4117
4118
4119
4120
4121
4122
4123
4124
4125
4126
4127
4128
4129
4130
4131
4132
4133
4134
4135
4136
4137
4138
4139
4140
4141
4142
4143
4144
4145
4146
4147
4148
4149
4150
4151
4152
4153
4154
4155
4156
4157
4158
4159
4160
4161
4162
4163
4164
4165
4166
4167
4168
4169
4170
4171
4172
4173
4174
4175
4176
4177
4178
4179
4180
4181
4182
4183
4184
4185
4186
4187
4188
4189
4190
4191
4192
4193
4194
4195
4196
4197
4198
4199
4200
4201
4202
4203
4204
4205
4206
4207
4208
4209
4210
4211
4212
4213
4214
4215
4216
4217
4218
4219
4220
4221
4222
4223
4224
4225
4226
4227
4228
4229
4230
4231
4232
4233
4234
4235
4236
4237
4238
4239
4240
4241
4242
4243
4244
4245
4246
4247
4248
4249
4250
4251
4252
4253
4254
4255
4256
4257
4258
4259
4260
4261
4262
4263
4264
4265
4266
4267
4268
4269
4270
4271
4272
4273
4274
4275
4276
4277
4278
4279
4280
4281
4282
4283
4284
4285
4286
4287
4288
4289
4290
4291
4292
4293
4294
4295
4296
4297
4298
4299
4300
4301
4302
## @private
class AmoebaStretchTorsionGenerator(object):
    """An AmoebaStretchTorsionGenerator constructs a AmoebaStretchTorsionForce."""

    def __init__(self):
        self.torsions = []

    @staticmethod
    def parseElement(element, forceField):
        generator = AmoebaStretchTorsionGenerator()
        forceField._forces.append(generator)
        params = ('v11', 'v12', 'v13', 'v21', 'v22', 'v23', 'v31', 'v32', 'v33')
        for torsion in element.findall('Torsion'):
            types = forceField._findAtomTypes(torsion.attrib, 4)
            if None not in types:
                v = [float(torsion.attrib[param]) for param in params]
                generator.torsions.append((types, v))

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        pass

    def postprocessSystem(self, sys, data, args):
        # We need to wait until after all bonds and torsions have been added before adding the stretch-torsions,
        # since it needs parameters from them.

        energy = """v11*(distance(p1,p2)-length1)*phi1 +
                    v12*(distance(p1,p2)-length1)*phi2 +
                    v13*(distance(p1,p2)-length1)*phi3 +
                    v21*(distance(p2,p3)-length2)*phi1 +
                    v22*(distance(p2,p3)-length2)*phi2 +
                    v23*(distance(p2,p3)-length2)*phi3 +
                    v31*(distance(p3,p4)-length3)*phi1 +
                    v32*(distance(p3,p4)-length3)*phi2 +
                    v33*(distance(p3,p4)-length3)*phi3;
                    phi1=1+cos(phi+phase1); phi2=1+cos(2*phi+phase2); phi3=1+cos(3*phi+phase3);
                    phi=dihedral(p1,p2,p3,p4)"""
        existing = [f for f in sys.getForces() if type(f) == mm.CustomCompoundBondForce and f.getEnergyFunction() == energy]
        if len(existing) == 0:
            force = mm.CustomCompoundBondForce(4, energy)
            for param in ('v11', 'v12', 'v13', 'v21', 'v22', 'v23', 'v31', 'v32', 'v33'):
                force.addPerBondParameter(param)
            for i in range(3):
                force.addPerBondParameter(f'length{i+1}')
            for i in range(3):
                force.addPerBondParameter(f'phase{i+1}')
            force.setName('AmoebaStretchTorsion')
        else:
            force = existing[0]

        # Record parameters for bonds and torsions so we can look them up quickly.

        bondForce = [f for f in sys.getForces() if type(f) == mm.CustomBondForce and f.getName() == 'AmoebaBond'][0]
        torsionForce = [f for f in sys.getForces() if type(f) == mm.PeriodicTorsionForce][0]
        bondLength = {}
        torsionPhase = defaultdict(lambda: [0.0, math.pi, 0.0])
        for i in range(bondForce.getNumBonds()):
            p1, p2, params = bondForce.getBondParameters(i)
            bondLength[(p1, p2)] = params[0]
            bondLength[(p2, p1)] = params[0]
        for i in range(torsionForce.getNumTorsions()):
            p1, p2, p3, p4, periodicity, phase, k = torsionForce.getTorsionParameters(i)
            if periodicity < 4:
                phase = phase.value_in_unit(unit.radian)
                torsionPhase[(p1, p2, p3, p4)][periodicity-1] = phase
                torsionPhase[(p4, p3, p2, p1)][periodicity-1] = phase

        # Add stretch-torsions.

        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 types, v in self.torsions:
                if (type1 in types[3] and type2 in types[2] and type3 in types[1] and type4 in types[0]):
                    type1, type2, type3, type4 = type4, type3, type2, type1
                    torsion = tuple(reversed(torsion))
                if (type1 in types[0] and type2 in types[1] and type3 in types[2] and type4 in types[3]):
                    params = list(v)
                    params.append(bondLength[(torsion[0], torsion[1])])
                    params.append(bondLength[(torsion[1], torsion[2])])
                    params.append(bondLength[(torsion[2], torsion[3])])
                    params += torsionPhase[torsion]
                    force.addBond(torsion, params)
                    break
        if len(existing) == 0 and force.getNumBonds() > 0:
            sys.addForce(force)

parsers["AmoebaStretchTorsionForce"] = AmoebaStretchTorsionGenerator.parseElement

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

## @private
class AmoebaAngleTorsionGenerator(object):
    """An AmoebaAngleTorsionGenerator constructs a AmoebaAngleTorsionForce."""

    def __init__(self):
        self.torsions = []

    @staticmethod
    def parseElement(element, forceField):
        generator = AmoebaAngleTorsionGenerator()
        forceField._forces.append(generator)
        params = ('v11', 'v12', 'v13', 'v21', 'v22', 'v23')
        for torsion in element.findall('Torsion'):
            types = forceField._findAtomTypes(torsion.attrib, 4)
            if None not in types:
                v = [float(torsion.attrib[param]) for param in params]
                generator.torsions.append((types, v))

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        pass

    def postprocessSystem(self, sys, data, args):
        # We need to wait until after all angles and torsions have been added before adding the angle-torsions,
        # since it needs parameters from them.

        energy = """v11*(angle(p1,p2,p3)-angle1)*phi1 +
                    v12*(angle(p1,p2,p3)-angle1)*phi2 +
                    v13*(angle(p1,p2,p3)-angle1)*phi3 +
                    v21*(angle(p2,p3,p4)-angle2)*phi1 +
                    v22*(angle(p2,p3,p4)-angle2)*phi2 +
                    v23*(angle(p2,p3,p4)-angle2)*phi3;
                    phi1=1+cos(phi+phase1); phi2=1+cos(2*phi+phase2); phi3=1+cos(3*phi+phase3);
                    phi=dihedral(p1,p2,p3,p4)"""
        existing = [f for f in sys.getForces() if type(f) == mm.CustomCompoundBondForce and f.getEnergyFunction() == energy]
        if len(existing) == 0:
            force = mm.CustomCompoundBondForce(4, energy)
            for param in ('v11', 'v12', 'v13', 'v21', 'v22', 'v23'):
                force.addPerBondParameter(param)
            for i in range(2):
                force.addPerBondParameter(f'angle{i+1}')
            for i in range(3):
                force.addPerBondParameter(f'phase{i+1}')
            force.setName('AmoebaAngleTorsion')
        else:
            force = existing[0]

        # Record parameters for angles and torsions so we can look them up quickly.

        angleForce = [f for f in sys.getForces() if type(f) == mm.CustomAngleForce and f.getName() == 'AmoebaAngle'][0]
        inPlaneAngleForce = [f for f in sys.getForces() if type(f) == mm.CustomCompoundBondForce and f.getName() == 'AmoebaInPlaneAngle'][0]
        torsionForce = [f for f in sys.getForces() if type(f) == mm.PeriodicTorsionForce][0]
        equilAngle = {}
        torsionPhase = defaultdict(lambda: [0.0, math.pi, 0.0])
        angleScale = math.pi/180
        for i in range(angleForce.getNumAngles()):
            p1, p2, p3, params = angleForce.getAngleParameters(i)
            equilAngle[(p1, p2, p3)] = params[0]*angleScale
            equilAngle[(p3, p2, p1)] = params[0]*angleScale
        for i in range(inPlaneAngleForce.getNumBonds()):
            particles, params = inPlaneAngleForce.getBondParameters(i)
            equilAngle[tuple(particles[:3])] = params[0]*angleScale
            equilAngle[tuple(reversed(particles[:3]))] = params[0]*angleScale
        for i in range(torsionForce.getNumTorsions()):
            p1, p2, p3, p4, periodicity, phase, k = torsionForce.getTorsionParameters(i)
            if periodicity < 4:
                phase = phase.value_in_unit(unit.radian)
                torsionPhase[(p1, p2, p3, p4)][periodicity-1] = phase
                torsionPhase[(p4, p3, p2, p1)][periodicity-1] = phase

        # Add stretch-torsions.

        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 types, v in self.torsions:
                if (type1 in types[3] and type2 in types[2] and type3 in types[1] and type4 in types[0]):
                    type1, type2, type3, type4 = type4, type3, type2, type1
                    torsion = tuple(reversed(torsion))
                if (type1 in types[0] and type2 in types[1] and type3 in types[2] and type4 in types[3]):
                    params = list(v)
                    params.append(equilAngle[(torsion[0], torsion[1], torsion[2])])
                    params.append(equilAngle[(torsion[1], torsion[2], torsion[3])])
                    params += torsionPhase[torsion]
                    force.addBond(torsion, params)
                    break
        if len(existing) == 0 and force.getNumBonds() > 0:
            sys.addForce(force)

parsers["AmoebaAngleTorsionForce"] = AmoebaAngleTorsionGenerator.parseElement

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

4303
## @private
4304
class AmoebaTorsionTorsionGenerator(object):
4305
4306
4307
4308
4309

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

Peter Eastman's avatar
Peter Eastman committed
4310
    def __init__(self):
4311

Peter Eastman's avatar
Peter Eastman committed
4312
4313
4314
4315
4316
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
4317

Peter Eastman's avatar
Peter Eastman committed
4318
        self.gridIndex = []
4319

Peter Eastman's avatar
Peter Eastman committed
4320
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
4321

4322
4323
4324
4325
4326
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
4327
        generator = AmoebaTorsionTorsionGenerator()
4328
4329
4330
4331
4332
4333
4334
        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'):
4335
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
4336
            if None not in types:
4337
4338
4339
4340
4341
4342
4343

                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
4344
4345
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
4346
4347
4348
4349
4350
4351
4352
4353
4354
                    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
4355
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
4356
4357
                raise ValueError(outputString)

4358
4359
4360
4361
4362
4363
4364
4365
4366
4367
        # 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
4368
4369
4370
4371
4372
4373
        #     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
4374
4375

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

Peter Eastman's avatar
Peter Eastman committed
4379
4380
4381
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
4382

Peter Eastman's avatar
Peter Eastman committed
4383
4384
            grid = []
            gridCol = []
4385
4386
4387
4388
4389

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
4390
4391
4392
4393
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
4394
                if 'fx' in gridEntry.attrib:
4395
4396
4397
                    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
4398
                gridCol.append(gridRow)
4399
4400

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
4401
4402
4403
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
4404
4405
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
4406

Peter Eastman's avatar
Peter Eastman committed
4407
4408
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
4409
            else:
Peter Eastman's avatar
Peter Eastman committed
4410
4411
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
4412
4413
4414
4415
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
4416
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
4417
4418
4419
4420
4421
4422
4423
4424

        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
4425
        if (len(data.atomBonds[atomC]) == 4):
4426
4427
4428
4429
4430
            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
4431
4432
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
4433
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
4434
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
4435
4436
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
4437
4438
                if (hit > -1):
                    if (atomE == -1):
4439
4440
4441
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
4442

4443
4444
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
4445
            if (atomE == -1 or atomF == -1):
4446
                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
4447
                raise ValueError(outputString)
4448
4449
4450
4451
4452

            # 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
4453
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
4454
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
4455
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
4456
                chiralAtomIndex = atomF
4457

Peter Eastman's avatar
Peter Eastman committed
4458
4459
4460
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
4461
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
4462
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
4463
                chiralAtomIndex = massF
4464
4465
4466
4467

        return chiralAtomIndex

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

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

4471
        existing = [f for f in sys.getForces() if type(f) == mm.AmoebaTorsionTorsionForce]
4472
4473
4474
4475
4476
4477
4478
4479

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

        for angle in data.angles:

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

4482
4483
4484
4485
4486
4487
4488
            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
4489
                if (bondedAtom1 != ib):
4490
4491
4492
4493
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4494
                if (ia != ic and ia != id):
Peter Eastman's avatar
Peter Eastman committed
4495
4496
4497
                    for bondIndex2 in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex2].atom1
                        bondedAtom2 = data.bonds[bondIndex2].atom2
Peter Eastman's avatar
Peter Eastman committed
4498
                        if (bondedAtom1 != id):
4499
4500
4501
4502
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4503
                        if (ie != ic and ie != ib and ie != ia):
4504
4505
4506
4507
4508
4509
4510
4511
4512
4513
4514
4515
4516
4517
4518
4519
4520
4521
4522
4523

                            # 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
4524
4525
4526
                                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])
4527
4528
4529

                                # match in reverse order

4530
                                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
4531
4532
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
4533
4534
4535
4536

        # set grids

        for (index, grid) in enumerate(self.grids):
Peter Eastman's avatar
Peter Eastman committed
4537
            force.setTorsionTorsionGrid(index, grid)
4538
4539
        if len(existing) == 0 and force.getNumTorsionTorsions() > 0:
            sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
4540

4541
4542
4543
4544
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

4545
## @private
4546
class AmoebaStretchBendGenerator(object):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4547
4548

    #=============================================================================================
4549
4550
4551
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

4552
    def __init__(self, forcefield):
4553

4554
        self.forcefield = forcefield
Peter Eastman's avatar
Peter Eastman committed
4555
4556
4557
        self.types1 = []
        self.types2 = []
        self.types3 = []
4558

Peter Eastman's avatar
Peter Eastman committed
4559
4560
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
4561

4562
4563
4564
4565
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
4566
        generator = AmoebaStretchBendGenerator(forceField)
4567
4568
4569
4570
4571
4572
4573
        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'):
4574
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
4575
            if None not in types:
4576
4577
4578
4579
4580
4581
4582
4583
4584
4585
4586
4587

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

4591
4592
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
4593
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
4594
4595
    # 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
4596
4597
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
4598
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
4599
4600
4601

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

Peter Eastman's avatar
Peter Eastman committed
4602
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4603
4604
        if not any(isinstance(f, AmoebaOutOfPlaneBendGenerator) for f in self.forcefield.getGenerators()):
            raise ValueError('A ForceField containing an <AmoebaStretchBendForce> must also contain an <AmoebaOutOfPlaneBendForce>')
4605
4606
4607
4608
4609
4610
4611

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

    # Note: request for constrained bonds is ignored.

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

4612
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
4613

4614
4615
        energy = "(k1*(distance(p1,p2)-r12) + k2*(distance(p2,p3)-r23))*(%.15g*(angle(p1,p2,p3)-theta0))" % (180/math.pi)
        existing = [f for f in sys.getForces() if type(f) == mm.CustomCompoundBondForce and f.getEnergyFunction() == energy]
4616
        if len(existing) == 0:
4617
4618
4619
4620
4621
4622
            force = mm.CustomCompoundBondForce(3, energy)
            force.addPerBondParameter("r12")
            force.addPerBondParameter("r23")
            force.addPerBondParameter("theta0")
            force.addPerBondParameter("k1")
            force.addPerBondParameter("k2")
4623
            force.setName('AmoebaStretchBend')
4624
4625
4626
4627
4628
4629
4630
4631
4632
4633
4634
4635
            sys.addForce(force)
        else:
            force = existing[0]

        for angleDict in angleList:

            angle = angleDict['angle']

            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
4636
            radian = 57.2957795130
4637
4638
4639
4640
4641
4642
4643
4644
4645
4646
            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
4647
                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
4648
4649
4650
4651
4652
4653
4654
4655
4656
4657
4658
4659
4660
4661
4662
                    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
4663

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4666
                    if ('idealAngle' not in angleDict):
4667

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4668
4669
4670
4671
4672
                       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
4673
                       raise ValueError(outputString)
4674

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4675
4676
4677
4678
4679
4680
4681
                    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
4682
                       raise ValueError(outputString)
4683

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4684
                    else:
4685
4686
4687
4688
4689
                        if type1 in types1 and type3 in types3:
                            k1, k2 = self.k1[i], self.k2[i]
                        else:
                            k1, k2 = self.k2[i], self.k1[i]
                        force.addBond((angle[0], angle[1], angle[2]), (bondAB, bondCB, angleDict['idealAngle']/radian, k1, k2))
4690

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4691
                    break
4692
4693
4694
4695
4696

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

4697
## @private
4698
class AmoebaVdwGenerator(object):
4699
4700

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

4702
4703
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4706
        self.type = type
4707

Peter Eastman's avatar
Peter Eastman committed
4708
4709
4710
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
4711

Peter Eastman's avatar
Peter Eastman committed
4712
        self.epsilonrule = epsilonrule
4713

Peter Eastman's avatar
Peter Eastman committed
4714
4715
4716
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
4717
4718
4719
4720
4721
4722
4723

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

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

4727
4728
4729
        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
4730
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
4731
            forceField.registerGenerator(generator)
4732
4733
            generator.params = {}
            generator.pairs = []
4734
4735
4736
4737
4738
4739
4740
4741
4742
4743
        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')
4744
4745
4746
4747
4748
        for vdw in element.findall('Vdw'):
            generator.params[vdw.attrib['class']] = tuple(float(vdw.attrib[name]) for name in ('sigma', 'epsilon', 'reduction'))
        for pair in element.findall('Pair'):
            generator.pairs.append((pair.attrib['class1'], pair.attrib['class2'], float(pair.attrib['sigma']), float(pair.attrib['epsilon'])))
        generator.classNameForType = dict((t.name, t.atomClass) for t in forceField._atomTypes.values())
4749
4750
4751
4752

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

    @staticmethod
4753
    def getBondedParticleSets(sys, data):
4754

4755
4756
4757
4758
4759
4760
        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
4761

4762
4763
    #=============================================================================================

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

4766
        potentialMap = {'BUFFERED-14-7':0, 'LENNARD-JONES':1}
Peter Eastman's avatar
Peter Eastman committed
4767
        sigmaMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'CUBIC-MEAN':1}
Chengwen Liu's avatar
Chengwen Liu committed
4768
        epsilonMap = {'ARITHMETIC':1, 'GEOMETRIC':1, 'HARMONIC':1, 'W-H':1, 'HHG':1}
4769

4770
4771
        force = mm.AmoebaVdwForce()
        sys.addForce(force)
4772

4773
4774
4775
4776
4777
4778
4779
4780
        # Potential function

        if (self.type.upper() in potentialMap):
            force.setPotentialFunction(potentialMap[self.type.upper()])
        else:
            stringList = ' '.join(str(x) for x in potentialMap.keys())
            raise ValueError("AmoebaVdwGenerator: potential type %s not recognized; valid values are %s; using default." % (self.type, stringList))

4781
        # sigma and epsilon combining rules
4782

4783
4784
4785
4786
        if ('sigmaCombiningRule' in args):
            sigmaRule = args['sigmaCombiningRule'].upper()
            if (sigmaRule.upper() in sigmaMap):
                force.setSigmaCombiningRule(sigmaRule.upper())
4787
            else:
4788
4789
4790
4791
                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)
4792

4793
4794
4795
4796
        if ('epsilonCombiningRule' in args):
            epsilonRule = args['epsilonCombiningRule'].upper()
            if (epsilonRule.upper() in epsilonMap):
                force.setEpsilonCombiningRule(epsilonRule.upper())
4797
            else:
4798
4799
4800
4801
                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
4802

4803
        # cutoff
4804

4805
4806
4807
4808
        if ('vdwCutoff' in args):
            force.setCutoff(args['vdwCutoff'])
        else:
            force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4809

4810
        # dispersion correction
4811

4812
4813
        if ('useDispersionCorrection' in args):
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
4814

4815
4816
        if (nonbondedMethod == PME):
            force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
4817

4818
        # Define types
4819

4820
4821
4822
4823
4824
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
4825
4826
4827
4828
4829
4830
4831
        classToTypeMap = {}
        for className in self.params:
            sigma, epsilon, _ = self.params[className]
            classToTypeMap[className] = force.addParticleType(sigma*sigmaScale, epsilon)

        # add particles to force

4832
        for (i, atom) in enumerate(data.atoms):
4833
4834
            className = self.classNameForType[data.atomType[atom]]
            _, _, reduction = self.params[className]
4835
4836
4837
4838
4839
4840
4841
4842
4843
            # 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
4844

4845
4846
4847
            force.addParticle(ivIndex, classToTypeMap[className], reduction)

        # Add pairs
4848

4849
4850
        for c1, c2, sigma, epsilon in self.pairs:
            force.addTypePair(classToTypeMap[c1], classToTypeMap[c2], sigma, epsilon)
4851
4852
4853
4854
4855
4856
4857

        # 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

4858
        bondedParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4859
4860

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

4862
4863
4864
4865
4866
4867
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
4868
            if (self.vdw13Scale == 0.0):
4869
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
4870
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
4871
4872
4873
4874
4875

            # self

            exclusionSet.add(i)

4876
            force.setParticleExclusions(i, tuple(exclusionSet))
4877
4878
4879
4880
4881

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

4882
## @private
4883
class AmoebaMultipoleGenerator(object):
4884
4885
4886
4887

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

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

4889
4890
    #=============================================================================================

4891
    def __init__(self, forceField):
Peter Eastman's avatar
Peter Eastman committed
4892
4893
        self.forceField = forceField
        self.typeMap = {}
4894
4895
4896
4897
4898
4899

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
4900
    def setAxisType(kIndices):
4901
4902
4903

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
4904
4905
4906
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
4907
                else:
Peter Eastman's avatar
Peter Eastman committed
4908
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
4909

Peter Eastman's avatar
Peter Eastman committed
4910
4911
                if (kIndicesLen > 2):
                    kx = kIndices[2]
4912
                else:
Peter Eastman's avatar
Peter Eastman committed
4913
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
4914

Peter Eastman's avatar
Peter Eastman committed
4915
4916
                if (kIndicesLen > 1):
                    kz = kIndices[1]
4917
                else:
Peter Eastman's avatar
Peter Eastman committed
4918
                    kz = 0
4919

Peter Eastman's avatar
Peter Eastman committed
4920
4921
                while(len(kIndices) < 4):
                    kIndices.append(0)
4922
4923

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
4924
                if (kz == 0):
4925
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
4926
                if (kz != 0 and kx == 0):
4927
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
4928
                if (kz < 0 or kx < 0):
4929
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
4930
                if (kx < 0 and ky < 0):
4931
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
4932
                if (kz < 0 and kx < 0 and ky  < 0):
4933
4934
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
4935
4936
4937
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
4938
4939
4940
4941
4942
4943
4944
4945

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
4946
        #   <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"  >
4947
4948
4949
        # <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"  />

4950
4951
4952
4953
4954
4955
4956
        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]
4957
4958
4959
4960

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

        for atom in element.findall('Multipole'):
4961
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
4962
            if None not in types:
4963
4964
4965
4966
4967
4968
4969
4970

                # 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
4971
4972
                        if (atom.attrib[kString]):
                             kIndices.append(int(atom.attrib[kString]))
Justin MacCallum's avatar
Justin MacCallum committed
4973
                    except:
4974
4975
                        pass

Justin MacCallum's avatar
Justin MacCallum committed
4976
                # set axis type based on k-Indices
4977

Peter Eastman's avatar
Peter Eastman committed
4978
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
4979
4980
4981

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
4984
4985
4986
4987
4988
4989
4990
4991
4992
4993
4994
4995
4996
                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']))
4997
4998

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
4999
                    if (t not in generator.typeMap):
5000
5001
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
5002
5003
5004
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
5005
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
5006
5007
5008
5009
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
5010

5011
            else:
Peter Eastman's avatar
Peter Eastman committed
5012
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
5013
5014
                raise ValueError(outputString)

5015
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
5016

5017
        for atom in element.findall('Polarize'):
5018
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
5019
            if None not in types:
5020

Peter Eastman's avatar
Peter Eastman committed
5021
5022
5023
5024
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
5025
5026
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
5027
                    pdamp = pow(polarizability, 1.0/6.0)
5028

Peter Eastman's avatar
Peter Eastman committed
5029
5030
                pgrpMap = dict()
                for index in range(1, 7):
5031
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
5032
                    if (pgrp in atom.attrib):
5033
5034
5035
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
5036
5037
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
5038
                        raise ValueError(outputString)
5039
5040
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
5041
5042
5043
5044
5045
5046
5047
                        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
5048
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
5049
5050
5051
5052
5053
                                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
5054
5055
                            raise ValueError(outputString)

5056
            else:
Peter Eastman's avatar
Peter Eastman committed
5057
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
5058
5059
                raise ValueError(outputString)

5060
5061
    #=============================================================================================

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

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
5068
5069
5070
5071
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
5072
5073
5074
5075
            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
5076
5077
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
5078
5079
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
5080

5081
5082
5083
5084
        # pgrp11

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

Peter Eastman's avatar
Peter Eastman committed
5085
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 0):
5086
5087
                continue

Peter Eastman's avatar
Peter Eastman committed
5088
5089
            group = set()
            visited = set()
5090
5091
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
5092
5093
5094
5095
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
5096
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
5097
5098
                if (nextAtom not in visited):
                   visited.add(nextAtom)
5099
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
5100
5101
5102
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
5103
5104
5105

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
5106
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
5107
5108

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
5109
5110
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
5111
5112
5113
5114
5115

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
5116
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
5117
5118
                continue

Peter Eastman's avatar
Peter Eastman committed
5119
            pgrp11 = set(atom.polarizationGroupSet[0])
5120
5121
5122
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
5123
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
5124
5125
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
5126
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
5127

5128
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
5129
5130
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
5131
5132
5133
5134
5135

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
5136
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
5137
5138
                continue

Peter Eastman's avatar
Peter Eastman committed
5139
5140
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
5141
5142
5143
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
5144
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
5145
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
5146
            pgrp13 = pgrp13 - set(pgrp11)
5147
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
5148
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
5149

5150
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
5151
5152
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
5153
5154
5155
5156
5157

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
5158
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
5159
5160
                continue

Peter Eastman's avatar
Peter Eastman committed
5161
5162
5163
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
5164
5165
5166
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
5167
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
5168
5169
5170

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
5171
            pgrp14 = pgrp14 - set(pgrp11)
5172
5173

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

5176
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
5177
5178
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
5179
5180
5181

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

Peter Eastman's avatar
Peter Eastman committed
5182
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
5183
5184
5185
5186

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

5187
5188
5189
5190
5191
5192
5193
        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)
5194

5195
5196
        if ('ewaldErrorTolerance' in args):
            force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))
5197

5198
5199
5200
5201
5202
5203
5204
5205
        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)
5206

5207
5208
        if ('aEwald' in args):
            force.setAEwald(float(args['aEwald']))
5209

5210
5211
        if ('pmeGridDimensions' in args):
            force.setPmeGridDimensions(args['pmeGridDimensions'])
5212

5213
5214
        if ('mutualInducedMaxIterations' in args):
            force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
5215

5216
5217
        if ('mutualInducedTargetEpsilon' in args):
            force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))
5218
5219

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
5220
        # throw error if particle type not available
5221
5222
5223
5224
5225

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

        # 1-2

5226
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5227
5228
5229
5230
5231

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
5232
            bonded13Set = set()
5233
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
5234
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
5235
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
5236
5237
5238
5239

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
5240
            selfSet = set()
5241
5242
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
5243
5244
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
5245
5246
5247
5248
5249

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
5250
5251
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
5252
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
5253
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
5254

5255
5256
5257
5258
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
5259
            selfSet = set()
5260
5261
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
5262
5263
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
5264
5265
5266
5267
5268

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
5269
5270
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
5271
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
5272
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
5273
5274
5275
5276
5277
5278

            # 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
5279
            selfSet = set()
5280
5281
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
5282
5283
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
5284
5285
5286
5287
5288

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

Peter Eastman's avatar
Peter Eastman committed
5289
5290
                multipoleList = self.typeMap[t]
                hit = 0
5291
5292
5293
5294
5295
5296
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5297
                    if (hit != 0):
5298
5299
                        break

Peter Eastman's avatar
Peter Eastman committed
5300
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5301
5302

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
5303
5304
                    kx = kIndices[2]
                    ky = kIndices[3]
5305
5306
5307
5308

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

5310
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
5311
5312
5313
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5314
5315
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
5316
                       if (hit != 0):
5317
5318
5319
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
5320
5321
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
5322
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5323
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
5324
5325
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
5326
5327
5328
5329
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
5330
5331
5332
5333
5334
                                      if( bondedAtomXType == bondedAtomZType and xaxis < zaxis ):
                                          swapI = zaxis
                                          zaxis = xaxis
                                          xaxis = swapI
                                      else:
Peter Eastman's avatar
Peter Eastman committed
5335
5336
5337
5338
                                          for bondedAtomXIndex2 in bondedAtomIndices:
                                              bondedAtomX1Type = int(data.atomType[data.atoms[bondedAtomXIndex2]])
                                              if( bondedAtomX1Type == kx and bondedAtomXIndex2 != bondedAtomZIndex and bondedAtomXIndex2 < xaxis ):
                                                  xaxis = bondedAtomXIndex2
5339

5340
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5341
                                      hit = 1
5342
5343
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5344
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
5345
5346
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
5347
5348
5349
5350
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
5351
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5352
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
5353

5354
5355
5356
5357
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5358
                    if (hit != 0):
5359
5360
                        break

Peter Eastman's avatar
Peter Eastman committed
5361
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5362
5363

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
5364
5365
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
5366

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

5371
5372
5373
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
5374
5375
5376
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5377
5378
5379

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
5380
                       if (hit != 0):
5381
5382
5383
                           break

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

Peter Eastman's avatar
Peter Eastman committed
5386
                       if (bondedAtomZType == kz):
5387
5388
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
5389
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
5390
5391
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
5392
5393
5394
5395
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
5396
5397
5398

                                      # select xaxis w/ smallest index

Peter Eastman's avatar
Peter Eastman committed
5399
5400
5401
5402
                                      for bondedAtomXIndex2 in bondedAtom13Indices:
                                          bondedAtomX1Type = int(data.atomType[data.atoms[bondedAtomXIndex2]])
                                          if( bondedAtomX1Type == kx and bondedAtomXIndex2 != bondedAtomZIndex and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex2] and bondedAtomXIndex2 < xaxis ):
                                              xaxis = bondedAtomXIndex2
5403

5404
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5405
                                      hit = 3
5406
5407
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
5408
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
5409
5410
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
5411
5412
5413
5414
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
5415
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5416
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
5417

5418
5419
5420
5421
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5422
                    if (hit != 0):
5423
5424
                        break

Peter Eastman's avatar
Peter Eastman committed
5425
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5426
5427
5428
5429

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

Peter Eastman's avatar
Peter Eastman committed
5430
5431
5432
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5433
5434
5435

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
5436
                        if (hit != 0):
5437
5438
5439
                            break

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

Peter Eastman's avatar
Peter Eastman committed
5442
                        if (kx == 0 and kz == bondedAtomZType):
5443
                            zaxis = bondedAtomZIndex
5444
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5445
                            hit = 5
5446
5447
5448
5449
5450

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5451
                    if (hit != 0):
5452
5453
                        break

Peter Eastman's avatar
Peter Eastman committed
5454
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5455
5456
5457

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
5458
5459
5460
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5461

Peter Eastman's avatar
Peter Eastman committed
5462
                    if (kz == 0):
5463
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5464
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
5465

5466
5467
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
5468
                if (hit != 0):
5469

Peter Eastman's avatar
Peter Eastman committed
5470
                    atom.multipoleDict = savedMultipoleDict
5471
                    atom.polarizationGroups = dict()
5472
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
5473
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
5474
                    if (atomIndex == newIndex):
5475
5476
5477
5478
                        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]))
5479
                    else:
5480
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
5481
                else:
Peter Eastman's avatar
Peter Eastman committed
5482
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
5483
            else:
Peter Eastman's avatar
Peter Eastman committed
5484
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
5485
5486
5487

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
5488
        self.setPolarGroups(data, bonded12ParticleSets, force)
5489
5490
5491
5492
5493

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

5494
## @private
5495
class AmoebaWcaDispersionGenerator(object):
5496
5497

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

5499
5500
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
5503
5504
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
5505
5506
5507
5508
5509
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
5510
        self.shctd = shctd
5511
5512
5513
5514
5515
5516
5517
5518
5519

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

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

5521
5522
5523
5524
5525
5526
5527
5528
5529
5530
5531
5532
5533
5534
5535
        existing = [f for f in forceField._forces if isinstance(f, AmoebaWcaDispersionGenerator)]
        if len(existing) == 0:
            generator = AmoebaWcaDispersionGenerator(element.attrib['epso'],
                                                     element.attrib['epsh'],
                                                     element.attrib['rmino'],
                                                     element.attrib['rminh'],
                                                     element.attrib['awater'],
                                                     element.attrib['slevy'],
                                                     element.attrib['dispoff'],
                                                     element.attrib['shctd'])
            forceField.registerGenerator(generator)
            generator.params = ForceField._AtomTypeParameters(forceField, 'AmoebaWcaDispersionForce', 'WcaDispersion', ('radius', 'epsilon'))
        else:
            # Multiple <AmoebaWcaDispersionForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
5536
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
5537

5538
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
5539

Peter Eastman's avatar
Peter Eastman committed
5540
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
5541
5542
5543

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

5544
        existing = [f for f in sys.getForces() if type(f) == mm.AmoebaWcaDispersionForce]
5545
5546
5547
5548
5549
5550
5551
        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
5552
        # throw error if particle type not available
5553

Peter Eastman's avatar
Peter Eastman committed
5554
5555
5556
5557
5558
5559
5560
5561
        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  ))
5562

5563
5564
5565
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
5566
5567
5568
5569
5570

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

5571
## @private
5572
class AmoebaGeneralizedKirkwoodGenerator(object):
5573
5574

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

5576
5577
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5578
5579
5580
5581
5582
5583
5584
5585
5586
5587
5588
    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
5589
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5590
5591
5592
5593
5594
5595
5596
5597
5598
5599
5600
5601
5602
5603
5604
5605
5606
5607
5608
5609
5610
5611
5612
5613
5614
        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
5615
5616
5617

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

Peter Eastman's avatar
Peter Eastman committed
5618
    def getObcShct(self, data, atomIndex):
5619

Peter Eastman's avatar
Peter Eastman committed
5620
        atom = data.atoms[atomIndex]
5621
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5622
        shct = -1.0
5623
5624

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
5625

Peter Eastman's avatar
Peter Eastman committed
5626
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
5627
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5628
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5629
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
5630
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
5631
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
5632
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
5633
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5634
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
5635
5636
5637
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
5638
        elif (atomicNumber == 16):              # S(16)
5639
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
5640
        elif (atomicNumber == 26):              # Fe(26)
5641
5642
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
5643
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5644
            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
5645
5646

        return shct
5647
5648
5649

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

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

Peter Eastman's avatar
Peter Eastman committed
5652
        atom = data.atoms[atomIndex]
5653
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5654
        radius = -1.0
5655

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

Peter Eastman's avatar
Peter Eastman committed
5658
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
5659

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

5663
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5664
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
5665

Peter Eastman's avatar
Peter Eastman committed
5666
            if (bondedAtomAtomicNumber == 7):
5667
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
5668
            if (bondedAtomAtomicNumber == 8):
5669
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
5670

Peter Eastman's avatar
Peter Eastman committed
5671
        elif (atomicNumber == 3):               # Li(3)
5672
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
5673
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5674

5675
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
5676
            if (len(bondedAtomIndices) == 3):
5677
5678
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
5679
            elif (len(bondedAtomIndices) == 4):
5680
5681
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5682
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
5683
5684
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
5685
        elif (atomicNumber == 7):               # N(7)
5686
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
5687
        elif (atomicNumber == 8):               # O(8)
5688
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
5689
            if (len(bondedAtomIndices) == 2):
5690
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
5691
        elif (atomicNumber == 9):               # F(9)
5692
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
5693
        elif (atomicNumber == 10):
5694
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
5695
        elif (atomicNumber == 11):
5696
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
5697
        elif (atomicNumber == 12):
5698
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5699
        elif (atomicNumber == 14):
5700
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
5701
        elif (atomicNumber == 15):              # P(15)
5702
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
5703
        elif (atomicNumber == 16):              # S(16)
5704
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
5705
        elif (atomicNumber == 17):
5706
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
5707
        elif (atomicNumber == 18):
5708
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5709
        elif (atomicNumber == 19):
5710
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
5711
        elif (atomicNumber == 20):
5712
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
5713
        elif (atomicNumber == 35):
5714
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
5715
        elif (atomicNumber == 36):
5716
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
5717
        elif (atomicNumber == 37):
5718
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
5719
        elif (atomicNumber == 53):
5720
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
5721
        elif (atomicNumber == 54):
5722
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
5723
        elif (atomicNumber == 55):
5724
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
5725
        elif (atomicNumber == 56):
5726
5727
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
5728
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5729
5730
            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
5731

5732
5733
5734
5735
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
5738
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5739
        atom = data.atoms[atomIndex]
5740
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
5741
        if (atomicNumber in bondiMap):
5742
5743
            radius = bondiMap[atomicNumber]
        else:
peastman's avatar
peastman committed
5744
            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
5745
            raise ValueError( outputString )
Justin MacCallum's avatar
Justin MacCallum committed
5746

5747
5748
5749
5750
5751
5752
5753
5754
5755
5756
        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
5757

5758
5759
5760
5761
5762
5763
5764
5765
5766
5767
5768
        existing = [f for f in forceField._forces if isinstance(f, AmoebaGeneralizedKirkwoodGenerator)]
        if len(existing) == 0:
            generator = AmoebaGeneralizedKirkwoodGenerator(forceField, element.attrib['solventDielectric'],
                                                           element.attrib['soluteDielectric'],
                                                           element.attrib['includeCavityTerm'],
                                                           element.attrib['probeRadius'],
                                                           element.attrib['surfaceAreaFactor'])
            forceField.registerGenerator(generator)
        else:
            # Multiple <AmoebaGeneralizedKirkwoodFprce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
5769
5770

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

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

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

5777
5778
5779
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

5780
        amoebaMultipoleForceList = [f for f in sys.getForces() if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
5781
        if (len(amoebaMultipoleForceList) > 0):
5782
5783
5784
5785
5786
            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
5787
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
5788
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
5789

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

5792
        existing = [f for f in sys.getForces() if type(f) == mm.AmoebaGeneralizedKirkwoodForce]
5793
5794
5795
5796
        if len(existing) == 0:

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

Peter Eastman's avatar
Peter Eastman committed
5798
5799
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
5800
            else:
Peter Eastman's avatar
Peter Eastman committed
5801
                force.setSolventDielectric(   float(self.solventDielectric))
5802

Peter Eastman's avatar
Peter Eastman committed
5803
5804
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
5805
            else:
Peter Eastman's avatar
Peter Eastman committed
5806
                force.setSoluteDielectric(    float(self.soluteDielectric))
5807

Peter Eastman's avatar
Peter Eastman committed
5808
5809
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
5810
            else:
Peter Eastman's avatar
Peter Eastman committed
5811
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
5812
5813
5814
5815
5816

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
5819
5820
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
5821
5822
5823

        # 1-2

5824
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5825
5826

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
5827
5828
5829
5830
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5831
            else:
Peter Eastman's avatar
Peter Eastman committed
5832
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5833
5834
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
5835
            force.addParticle(multipoleParameters[0], radius, shct)
5836
5837
5838
5839
5840

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

5841
## @private
5842
class AmoebaUreyBradleyGenerator(object):
5843
5844
5845
5846

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

5848
    def __init__(self):
5849

peastman's avatar
peastman committed
5850
        self.anglesForAtom2Type = defaultdict(list)
Peter Eastman's avatar
Peter Eastman committed
5851
5852
5853
        self.types1 = []
        self.types2 = []
        self.types3 = []
5854

Peter Eastman's avatar
Peter Eastman committed
5855
5856
        self.length = []
        self.k = []
5857
5858
5859
5860
5861
5862

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

    @staticmethod
    def parseElement(element, forceField):

5863
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
5864
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
5865

5866
        generator = AmoebaUreyBradleyGenerator()
5867
5868
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
5869
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
5870
            if None not in types:
peastman's avatar
peastman committed
5871
                index = len(generator.types1)
5872
5873
5874
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
peastman's avatar
peastman committed
5875
5876
                for t in types[1]:
                    generator.anglesForAtom2Type[t].append(index)
5877
5878
5879
5880
5881
                generator.length.append(float(bond.attrib['d']))
                generator.k.append(float(bond.attrib['k']))

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

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

5884
        existing = [f for f in sys.getForces() if type(f) == mm.HarmonicBondForce]
5885
5886

        if len(existing) == 0:
5887
            force = mm.HarmonicBondForce()
5888
5889
5890
5891
5892
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
5893
            if (isConstrained and not args.get('flexibleConstraints', False)):
5894
5895
5896
5897
                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
5898
            for i in self.anglesForAtom2Type[type2]:
5899
5900
5901
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
Peter Eastman's avatar
Peter Eastman committed
5902
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
5903
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
5904
5905
5906
5907
5908
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
5909
5910


5911
5912
5913
5914
5915
5916
5917
5918
5919
5920
5921
5922
5923
5924
5925
5926
5927
5928
5929
5930
5931
5932
5933
5934
5935
5936
5937
5938
5939
5940
5941
5942
5943
5944
5945
5946
5947
5948
5949
5950
5951
5952
5953
5954
5955
5956
5957
5958
5959
5960
5961
5962
5963
5964
5965
5966
5967
5968
5969
## @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
5970
            force.addParticle(params[0], dipole, quadrupole, *params[1:], axisType=axisType, multipoleAtomZ=zAtom, multipoleAtomX=xAtom, multipoleAtomY=yAtom)
5971
5972
5973
5974
5975
5976
5977
5978
5979
5980
5981
5982
5983
5984
5985
5986
5987
5988
5989
5990
5991
5992
5993
5994
5995
5996
5997
5998
5999
6000
6001
6002
6003
6004
6005
6006
6007
6008
6009
6010
6011
6012
6013
6014
6015
6016
        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
6017
## @private
6018
class DrudeGenerator(object):
peastman's avatar
peastman committed
6019
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
6020

6021
6022
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
6023
6024
6025
6026
6027
6028
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
6029
6030
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
6031
6032
6033
6034
        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'):
6035
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
6036
6037
6038
6039
6040
6041
6042
6043
6044
6045
            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
6046

peastman's avatar
peastman committed
6047
6048
6049
6050
    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
6051

peastman's avatar
peastman committed
6052
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
6053

peastman's avatar
peastman committed
6054
6055
6056
6057
6058
6059
6060
6061
6062
6063
6064
6065
6066
6067
6068
6069
        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
6070
6071
                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
6072
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
6073

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

peastman's avatar
peastman committed
6077
6078
6079
6080
6081
6082
6083
        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)
6084
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
6085
6086
6087
6088
6089
                # 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]
6090
6091
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
6092
6093
6094
6095
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

6096
6097
6098
6099
6100
6101
6102
6103
6104
6105
6106
6107
6108
        # Set the masses of Drude particles.

        drudeMass = args['drudeMass']
        if not unit.is_quantity(drudeMass):
            drudeMass *= unit.dalton
        for i in range(drude.getNumParticles()):
            params = drude.getParticleParameters(i)
            particle = params[0]
            parent = params[1]
            transferMass = drudeMass-sys.getParticleMass(particle)
            sys.setParticleMass(particle, drudeMass)
            sys.setParticleMass(parent, sys.getParticleMass(parent)-transferMass)

Justin MacCallum's avatar
Justin MacCallum committed
6109
parsers["DrudeForce"] = DrudeGenerator.parseElement
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
6110
6111

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