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

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
            from importlib_metadata import entry_points
61
62
63
64
65
66
        except:
            try:
                from importlib.metadata import entry_points
            except:
                pass
        try:
67
            for entry in entry_points().select(group='openmm.forcefielddir'):
68
69
                _dataDirectories.append(entry.load()())
        except:
70
            pass # importlib_metadata is not installed
71
    return _dataDirectories
72

73
74
def _convertParameterToNumber(param):
    if unit.is_quantity(param):
75
76
77
        if param.unit.is_compatible(unit.bar):
            return param / unit.bar
        return param.value_in_unit_system(unit.md_unit_system)
78
79
    return float(param)

80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
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])
95
        if functionType.startswith('Continuous'):
96
97
98
99
100
            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')
101
102
103
104
105
106
107
        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':
108
109
110
111
            force.addTabulatedFunction(
                name,
                mm.Continuous1DFunction(values, params['min'], params['max'], params['periodic']),
            )
112
        elif type == 'Continuous2D':
113
114
115
116
117
118
119
120
121
122
            force.addTabulatedFunction(
                name,
                mm.Continuous2DFunction(
                    params['xsize'], params['ysize'],
                    values,
                    params['xmin'], params['xmax'],
                    params['ymin'], params['ymax'],
                    params['periodic'],
                ),
            )
123
        elif type == 'Continuous3D':
124
125
            force.addTabulatedFunction(
                name,
126
                mm.Continuous3DFunction(
127
128
129
130
131
132
133
134
                    params['xsize'], params['ysize'], params['zsize'],
                    values,
                    params['xmin'], params['xmax'],
                    params['ymin'], params['ymax'],
                    params['zmin'], params['zmax'],
                    params['periodic'],
                ),
            )
135
136
137
138
139
        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':
140
            force.addTabulatedFunction(name, mm.Discrete3DFunction(params['xsize'], params['ysize'], params['zsize'], values))
141

142
143
# Enumerated values for nonbonded method

144
class NoCutoff(Singleton):
145
146
147
148
    def __repr__(self):
        return 'NoCutoff'
NoCutoff = NoCutoff()

149
class CutoffNonPeriodic(Singleton):
150
151
152
153
    def __repr__(self):
        return 'CutoffNonPeriodic'
CutoffNonPeriodic = CutoffNonPeriodic()

154
class CutoffPeriodic(Singleton):
155
156
157
158
    def __repr__(self):
        return 'CutoffPeriodic'
CutoffPeriodic = CutoffPeriodic()

159
class Ewald(Singleton):
160
161
162
163
    def __repr__(self):
        return 'Ewald'
Ewald = Ewald()

164
class PME(Singleton):
165
166
167
    def __repr__(self):
        return 'PME'
PME = PME()
168

Peter Eastman's avatar
Peter Eastman committed
169
170
171
172
173
class LJPME(Singleton):
    def __repr__(self):
        return 'LJPME'
LJPME = LJPME()

174
175
# Enumerated values for constraint type

176
class HBonds(Singleton):
177
178
179
180
    def __repr__(self):
        return 'HBonds'
HBonds = HBonds()

181
class AllBonds(Singleton):
182
183
184
185
    def __repr__(self):
        return 'AllBonds'
AllBonds = AllBonds()

186
class HAngles(Singleton):
187
188
189
    def __repr__(self):
        return 'HAngles'
HAngles = HAngles()
190
191
192
193
194
195
196
197
198
199

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

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

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

Robert McGibbon's avatar
Robert McGibbon committed
225
226
        Parameters
        ----------
227
228
229
        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
230
231
232
            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.
233
234
235
        prefix : string
            An optional string to be prepended to each residue name found in the
            loaded files.
peastman's avatar
peastman committed
236
        """
237

238
239
240
241
        if isinstance(files, tuple):
            files = list(files)
        else:
            files = [files]
242
243
244

        trees = []

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

            trees.append(tree)
273
            i += 1
274

275
            # Process includes in this file.
276

277
278
            if isinstance(file, str):
                parentDir = os.path.dirname(file)
peastman's avatar
peastman committed
279
280
            else:
                parentDir = ''
peastman's avatar
peastman committed
281
282
            for included in tree.getroot().findall('Include'):
                includeFile = included.attrib['file']
peastman's avatar
peastman committed
283
284
285
                joined = os.path.join(parentDir, includeFile)
                if os.path.isfile(joined):
                    includeFile = joined
286
287
                if includeFile not in files:
                    files.append(includeFile)
peastman's avatar
peastman committed
288
289
290

        # Load the atom types.

291
292
293
294
        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
295
296
297

        # Load the residue templates.

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

342
        # Load the patch definitions.
343
344
345
346
347
348
349
350
351
352

        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)
353
354
                    for key in patch.attrib:
                        patchData.attributes[key] = patch.attrib[key]
355
356
357
358
359
360
                    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
361
                        if atomName in patchData.allAtomNames:
362
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
363
                        patchData.allAtomNames.add(atomName)
364
365
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
366
                        patchData.addedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
367
368
369
370
371
372
                    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
373
                        if atomName in patchData.allAtomNames:
374
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
375
                        patchData.allAtomNames.add(atomName)
376
377
                        atomDescription = ForceField._PatchAtomData(atomName)
                        typeName = atom.attrib['type']
378
                        patchData.changedAtoms[atomDescription.residue].append(ForceField._TemplateAtomData(atomDescription.name, typeName, self._atomTypes[typeName].element, params))
379
380
                    for atom in patch.findall('RemoveAtom'):
                        atomName = atom.attrib['name']
peastman's avatar
peastman committed
381
                        if atomName in patchData.allAtomNames:
382
                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
peastman's avatar
peastman committed
383
                        patchData.allAtomNames.add(atomName)
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
                        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)
400
401
402
                    # The following three lines are only correct for single residue patches.  Multi-residue patches with
                    # virtual sites currently don't work correctly.  See issue #2848.
                    atomIndices = dict((atom.name, i) for i, atom in enumerate(patchData.addedAtoms[0]+patchData.changedAtoms[0]))
peastman's avatar
peastman committed
403
                    for site in patch.findall('VirtualSite'):
404
                        patchData.virtualSites[0].append(ForceField._VirtualSiteData(site, atomIndices))
405
406
407
408
409
410
411
412
413
                    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
414
415
        # Load force definitions

416
417
418
419
        for tree in trees:
            for child in tree.getroot():
                if child.tag in parsers:
                    parsers[child.tag](child, self)
peastman's avatar
peastman committed
420
421
422

        # Load scripts

423
424
425
        for tree in trees:
            for node in tree.getroot().findall('Script'):
                self.registerScript(node.text)
426

427
428
429
430
431
432
        # Execute initialization scripts.

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

433
434
435
    def getGenerators(self):
        """Get the list of all registered generators."""
        return self._forces
436

437
438
439
    def registerGenerator(self, generator):
        """Register a new generator."""
        self._forces.append(generator)
440

441
442
443
444
    def registerAtomType(self, parameters):
        """Register a new atom type."""
        name = parameters['name']
        if name in self._atomTypes:
445
446
            #  allow multiple registrations of the same atom type provided the definitions are identical
            existing = self._atomTypes[name]
447
448
            elementsMatch = ((existing.element is None and 'element' not in parameters) or (existing.element is not None and 'element' in parameters and existing.element.symbol == parameters['element']))
            if existing.atomClass == parameters['class'] and existing.mass == float(parameters['mass']) and elementsMatch:
449
                return
450
451
452
453
454
455
456
457
            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)
458
        self._atomTypes[name] = ForceField._AtomType(name, atomClass, mass, element)
459
460
461
462
463
464
465
        if atomClass in self._atomClasses:
            typeSet = self._atomClasses[atomClass]
        else:
            typeSet = set()
            self._atomClasses[atomClass] = typeSet
        typeSet.add(name)
        self._atomClasses[''].add(name)
466

467
468
    def registerResidueTemplate(self, template):
        """Register a new residue template."""
469
470
        if template.name in self._templates:
            # There is already a template with this name, so check the override levels.
471

472
473
474
475
476
477
478
479
480
481
482
            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))
483

484
        # Register the template.
485

486
487
488
        self._templates[template.name] = template
        signature = _createResidueSignature([atom.element for atom in template.atoms])
        if signature in self._templateSignatures:
489
            self._templateSignatures[signature].append(template)
490
491
        else:
            self._templateSignatures[signature] = [template]
492

493
494
    def registerPatch(self, patch):
        """Register a new patch that can be applied to templates."""
495
        patch.index = len(self._patches)
496
        self._patches[patch.name] = patch
497

498
499
500
    def registerTemplatePatch(self, residue, patch, patchResidueIndex):
        """Register that a particular patch can be used with a particular residue."""
        if residue not in self._templatePatches:
501
502
            self._templatePatches[residue] = set()
        self._templatePatches[residue].add((patch, patchResidueIndex))
503

504
505
506
    def registerScript(self, script):
        """Register a new script to be executed after building the System."""
        self._scripts.append(script)
507

508
509
510
511
    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::
512

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

Peter Eastman's avatar
Peter Eastman committed
515
516
517
518
519
        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.

520
521
522
        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.

523
        .. CAUTION:: This method is experimental, and its API is subject to change.
524
525
        """
        self._templateMatchers.append(matcher)
526

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
527
    def registerTemplateGenerator(self, generator):
528
529
530
531
        """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
532
533
        .. CAUTION:: This method is experimental, and its API is subject to change.

534
535
        Parameters
        ----------
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
536
        generator : function
537
538
            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
539
        When a residue without a template is encountered, the ``generator`` function is called with:
540

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
541
542
        ::
           success = generator(forcefield, residue)
543

544
        where ``forcefield`` is the calling ``ForceField`` object and ``residue`` is a openmm.app.topology.Residue object.
John Chodera's avatar
John Chodera committed
545
546

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

        ::
John Chodera's avatar
John Chodera committed
549
550
551
           generator API

           Parameters
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
552
           ----------
553
           forcefield : openmm.app.ForceField
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
554
               The ForceField object to which residue templates and/or parameters are to be added.
555
           residue : openmm.app.Topology.Residue
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
556
               The residue topology for which a template is to be generated.
557

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
558
559
560
561
562
           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`.
563

John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
564
565
566
567
568
           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.
569
570

        """
571
        self._templateGenerators.append(generator)
572

573
    def _findAtomTypes(self, attrib, num):
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
        """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.

        """
589
590
591
592
593
594
595
        types = []
        for i in range(num):
            if num == 1:
                suffix = ''
            else:
                suffix = str(i+1)
            classAttrib = 'class'+suffix
596
            typeAttrib = 'type'+suffix
597
            if classAttrib in attrib:
598
                if typeAttrib in attrib:
599
                    raise ValueError('Specified both a type and a class for the same atom: '+str(attrib))
600
                if attrib[classAttrib] not in self._atomClasses:
peastman's avatar
peastman committed
601
602
603
                    types.append(None) # Unknown atom class
                else:
                    types.append(self._atomClasses[attrib[classAttrib]])
604
605
            elif typeAttrib in attrib:
                if attrib[typeAttrib] == '':
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
606
                    types.append(self._atomClasses[''])
607
                elif attrib[typeAttrib] not in self._atomTypes:
peastman's avatar
peastman committed
608
609
610
                    types.append(None) # Unknown atom type
                else:
                    types.append([attrib[typeAttrib]])
611
612
            else:
                types.append(None) # Unknown atom type
613
614
        return types

615
    def _parseTorsion(self, attrib):
616
        """Parse the node defining a torsion."""
617
        types = self._findAtomTypes(attrib, 4)
peastman's avatar
peastman committed
618
        if None in types:
619
620
621
622
623
            return None
        torsion = PeriodicTorsion(types)
        index = 1
        while 'phase%d'%index in attrib:
            torsion.periodicity.append(int(attrib['periodicity%d'%index]))
624
625
            torsion.phase.append(_convertParameterToNumber(attrib['phase%d'%index]))
            torsion.k.append(_convertParameterToNumber(attrib['k%d'%index]))
626
            index += 1
Justin MacCallum's avatar
Justin MacCallum committed
627
628
        return torsion

629
    class _SystemData(object):
630
        """Inner class used to encapsulate data about the system being created."""
631
        def __init__(self, topology):
632
            self.atomType = {}
633
            self.atomParameters = {}
634
            self.atomTemplateIndexes = {}
635
            self.atoms = list(topology.atoms())
Peter Eastman's avatar
Peter Eastman committed
636
            self.excludeAtomWith = [[] for _ in self.atoms]
637
            self.virtualSites = {}
638
            self.bonds = [ForceField._BondData(bond[0].index, bond[1].index) for bond in topology.bonds()]
639
640
641
            self.angles = []
            self.propers = []
            self.impropers = []
Peter Eastman's avatar
Peter Eastman committed
642
            self.atomBonds = [[] for _ in self.atoms]
643
            self.isAngleConstrained = []
644
            self.constraints = {}
Peter Eastman's avatar
Peter Eastman committed
645
            self.bondedToAtom = [set() for _ in self.atoms]
646

647
            # Record which atoms are bonded to each other atom
648

649
650
651
652
653
654
            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)
655
            self.bondedToAtom = [sorted(b) for b in self.bondedToAtom]
656
657
658
659
660

        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:
661
                if self.constraints[key] != distance:
662
663
664
665
                    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)
666

667
668
669
670
671
672
        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
673
                self.atomTemplateIndexes[atom] = match
674
675
676
                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)
677

678
    class _TemplateData(object):
679
680
681
682
        """Inner class used to encapsulate data about a residue template definition."""
        def __init__(self, name):
            self.name = name
            self.atoms = []
683
            self.atomIndices = {}
684
            self.virtualSites = []
685
686
            self.bonds = []
            self.externalBonds = []
687
            self.overrideLevel = 0
688
            self.rigidWater = True
689
            self.attributes = {}
690

691
692
        def getAtomIndexByName(self, atom_name):
            """Look up an atom index by atom name, providing a helpful error message if not found."""
693
694
695
            index = self.atomIndices.get(atom_name, None)
            if index is not None:
                return index
696
697

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

702
703
        def addAtom(self, atom):
            self.atoms.append(atom)
704
            self.atomIndices[atom.name] = len(self.atoms)-1
705

706
        def addBond(self, atom1, atom2):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
707
            """Add a bond between two atoms in a template given their indices in the template."""
708
709
710
            self.bonds.append((atom1, atom2))
            self.atoms[atom1].bondedTo.append(atom2)
            self.atoms[atom2].bondedTo.append(atom1)
711

712
        def addBondByName(self, atom1_name, atom2_name):
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
713
            """Add a bond between two atoms in a template given their atom names."""
714
715
716
717
718
            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
719
            """Designate that an atom in a residue template has an external bond, using atom index within template."""
720
721
722
723
            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
724
            """Designate that an atom in a residue template has an external bond, using atom name within template."""
725
726
727
            atom = self.getAtomIndexByName(atom_name)
            self.addExternalBond(atom)

peastman's avatar
peastman committed
728
729
        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
730

peastman's avatar
peastman committed
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
            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

751
    class _TemplateAtomData(object):
752
        """Inner class used to encapsulate data about an atom in a residue template definition."""
753
        def __init__(self, name, type, element, parameters={}):
754
755
756
            self.name = name
            self.type = type
            self.element = element
757
            self.parameters = parameters
758
759
760
            self.bondedTo = []
            self.externalBonds = 0

761
    class _BondData(object):
762
763
764
765
766
767
        """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
768

769
    class _VirtualSiteData(object):
770
        """Inner class used to encapsulate data about a virtual site."""
771
        def __init__(self, node, atomIndices):
772
773
774
            attrib = node.attrib
            self.type = attrib['type']
            if self.type == 'average2':
775
                numAtoms = 2
776
777
                self.weights = [float(attrib['weight1']), float(attrib['weight2'])]
            elif self.type == 'average3':
778
                numAtoms = 3
779
780
                self.weights = [float(attrib['weight1']), float(attrib['weight2']), float(attrib['weight3'])]
            elif self.type == 'outOfPlane':
781
                numAtoms = 3
782
                self.weights = [float(attrib['weight12']), float(attrib['weight13']), float(attrib['weightCross'])]
783
            elif self.type == 'localCoords':
784
785
786
787
788
789
790
791
792
                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]))
793
                self.localPos = [float(attrib['p1']), float(attrib['p2']), float(attrib['p3'])]
794
795
            else:
                raise ValueError('Unknown virtual site type: %s' % self.type)
796
797
798
799
800
801
            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)]
802
803
804
805
            if 'excludeWith' in attrib:
                self.excludeWith = int(attrib['excludeWith'])
            else:
                self.excludeWith = self.atoms[0]
806

807
808
809
810
811
812
813
814
815
816
817
        def __eq__(self, other):
            if not isinstance(other, ForceField._VirtualSiteData):
                return False
            if self.type != other.type or self.index != other.index or self.atoms != other.atoms or self.excludeWith != other.excludeWith:
                return False
            if self.type in ('average2', 'average3', 'outOfPlane'):
                return self.weights == other.weights
            elif self.type == 'localCoords':
                return self.originWeights == other.originWeights and self.xWeights == other.xWeights and self.yWeights == other.yWeights and self.localPos == other.localPos
            return False

818
819
820
821
822
    class _PatchData(object):
        """Inner class used to encapsulate data about a patch definition."""
        def __init__(self, name, numResidues):
            self.name = name
            self.numResidues = numResidues
823
824
            self.addedAtoms = [[] for i in range(numResidues)]
            self.changedAtoms = [[] for i in range(numResidues)]
825
826
827
828
829
            self.deletedAtoms = []
            self.addedBonds = []
            self.deletedBonds = []
            self.addedExternalBonds = []
            self.deletedExternalBonds = []
peastman's avatar
peastman committed
830
            self.allAtomNames = set()
peastman's avatar
peastman committed
831
            self.virtualSites = [[] for i in range(numResidues)]
832
            self.attributes = {}
833
834
835
836
            self.index = None

        def __lt__(self, other):
            return self.index < other.index
837

838
839
840
841
        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)))
842

843
            # Construct a new version of each template.
844

845
846
847
848
            newTemplates = []
            for index, template in enumerate(templates):
                newTemplate = ForceField._TemplateData("%s-%s" % (template.name, self.name))
                newTemplates.append(newTemplate)
849

850
                # Build the list of atoms in it.
851

852
853
                for atom in template.atoms:
                    if not any(deleted.name == atom.name and deleted.residue == index for deleted in self.deletedAtoms):
854
                        newTemplate.addAtom(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
855
                for atom in self.addedAtoms[index]:
856
857
                    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))
858
                    newTemplate.addAtom(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
859
860
861
862
863
                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))
864
                    newTemplate.atoms[newAtomIndex[atom.name]] = ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters)
865

866
                # Copy over the virtual sites, translating the atom indices.
867

868
869
870
871
872
873
                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]
874
                        newSite.excludeWith = indexMap[site.excludeWith]
875
                        newTemplate.virtualSites.append(newSite)
876

877
                # Build the lists of bonds and external bonds.
878

879
880
881
882
883
                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]
884
                    if a1 in atomMap and a2 in atomMap and (a1.name, a2.name) not in deletedBonds and (a2.name, a1.name) not in deletedBonds:
885
886
887
888
                        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:
889
                        newTemplate.addExternalBond(indexMap[atom])
890
891
892
893
894
895
896
897
898
                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
899
900
901
902
903
904
905
906

                # 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]
907
                    newSite.excludeWith = indexMap[site.excludeWith]
peastman's avatar
peastman committed
908
909
                    newTemplate.virtualSites = [site for site in newTemplate.virtualSites if site.index != newSite.index]
                    newTemplate.virtualSites.append(newSite)
910
            return newTemplates
911

912
913
914
915
916
917
918
919
920
921
922
    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

923
    class _AtomType(object):
924
925
926
927
928
929
930
        """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

931
    class _AtomTypeParameters(object):
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
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
        """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))


1002
1003
    def _getResidueTemplateMatches(self, res, bondedToAtom, templateSignatures=None, ignoreExternalBonds=False, ignoreExtraParticles=False):
        """Return the templates that match a residue, or None if none are found.
1004
1005
1006
1007
1008

        Parameters
        ----------
        res : Topology.Residue
            The residue for which template matches are to be retrieved.
1009
1010
        bondedToAtom : list of set of int
            bondedToAtom[i] is the set of atoms bonded to atom index i
1011
1012
1013

        Returns
        -------
peastman's avatar
peastman committed
1014
        template : _TemplateData
1015
1016
1017
1018
1019
1020
1021
1022
            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
1023
        for matcher in self._templateMatchers:
Peter Eastman's avatar
Peter Eastman committed
1024
            template = matcher(self, res, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1025
1026
1027
            if template is not None:
                match = compiled.matchResidueToTemplate(res, template, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
                if match is None:
Peter Eastman's avatar
Peter Eastman committed
1028
                    raise ValueError('A custom template matcher returned a template for residue %d (%s), but it does not match the residue.' % (res.index, res.name))
1029
                return [template, match]
1030
1031
        if templateSignatures is None:
            templateSignatures = self._templateSignatures
1032
        signature = _createResidueSignature([atom.element for atom in res.atoms()])
1033
        if signature in templateSignatures:
1034
            allMatches = []
1035
            for t in templateSignatures[signature]:
1036
                match = compiled.matchResidueToTemplate(res, t, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1037
1038
1039
1040
1041
1042
                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:
1043
1044
1045
                # 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
1046
1047
                    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)))
1048
1049
                template = allMatches[0][0]
                matches = allMatches[0][1]
1050
1051
        return [template, matches]

1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
    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
        -------
1062
1063
        bondedToAtom : list of list of int
            bondedToAtom[index] is the list of atom indices bonded to atom `index`
1064
1065

        """
Peter Eastman's avatar
Peter Eastman committed
1066
        bondedToAtom = [set() for _ in topology.atoms()]
1067
1068
1069
        for (atom1, atom2) in topology.bonds():
            bondedToAtom[atom1.index].add(atom2.index)
            bondedToAtom[atom2.index].add(atom1.index)
1070
        bondedToAtom = [sorted(b) for b in bondedToAtom]
1071
        return bondedToAtom
1072

1073
    def getUnmatchedResidues(self, topology, residueTemplates=dict()):
1074
1075
        """Return a list of Residue objects from specified topology for which no forcefield templates are available.

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

1078
1079
1080
        Parameters
        ----------
        topology : Topology
1081
            The Topology whose residues are to be checked against the forcefield residue templates.
1082
1083
1084
1085
1086
1087
        residueTemplates : dict=dict()
            Specifies which template to use for particular residues.  The keys should be Residue
            objects from the Topology, and the values should be the names of the templates to
            use for them.  This is useful when a ForceField contains multiple templates that
            can match the same residue (e.g Fe2+ and Fe3+ templates in the ForceField for a
            monoatomic iron ion in the Topology).
1088
1089
1090
1091

        Returns
        -------
        unmatched_residues : list of Residue
1092
            List of Residue objects from `topology` for which no forcefield residue templates are available.
1093
1094
1095
1096
1097
            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.
1098
        bondedToAtom = self._buildBondedToAtomList(topology)
1099
        unmatched_residues = list() # list of unmatched residues
1100
        for res in topology.residues():
1101
1102
1103
1104
1105
1106
1107
            if res in residueTemplates:
                # Make sure the specified template matches.
                template = self._templates[residueTemplates[res]]
                matches = compiled.matchResidueToTemplate(res, template, bondedToAtom, False, False)
            else:
                # Attempt to match one of the existing templates.
                [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
1108
1109
1110
            if matches is None:
                # No existing templates match.
                unmatched_residues.append(res)
1111
1112
1113

        return unmatched_residues

1114
    def getMatchingTemplates(self, topology, ignoreExternalBonds=False):
1115
1116
1117
1118
1119
1120
1121
1122
        """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.
1123
1124
        ignoreExternalBonds : bool=False
            If true, ignore external bonds when matching residues to templates.
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
        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
1136
        for residue in topology.residues():
1137
            # Attempt to match one of the existing templates.
1138
            [template, matches] = self._getResidueTemplateMatches(residue, bondedToAtom, ignoreExternalBonds=ignoreExternalBonds)
1139
1140
            # Raise an exception if we have found no templates that match.
            if matches is None:
1141
                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)))
1142
1143
1144
1145
1146
            else:
                templates.append(template)

        return templates

1147
1148
    def generateTemplatesForUnmatchedResidues(self, topology):
        """Generate forcefield residue templates for residues in specified topology for which no forcefield templates are available.
1149

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

1152
1153
1154
        Parameters
        ----------
        topology : Topology
1155
            The Topology whose residues are to be checked against the forcefield residue templates.
1156
1157
1158

        Returns
        -------
1159
1160
1161
1162
1163
1164
        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]`
1165
1166
1167
1168
1169

        """
        # Get a non-unique list of unmatched residues.
        unmatched_residues = self.getUnmatchedResidues(topology)
        # Generate a unique list of unmatched residues by comparing fingerprints.
1170
        bondedToAtom = self._buildBondedToAtomList(topology)
1171
1172
        unique_unmatched_residues = list() # list of unique unmatched Residue objects from topology
        templates = list() # corresponding _TemplateData templates
1173
1174
1175
        signatures = set()
        for residue in unmatched_residues:
            signature = _createResidueSignature([ atom.element for atom in residue.atoms() ])
1176
            template = _createResidueTemplate(residue)
1177
1178
1179
1180
            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
1181
                    matches = compiled.matchResidueToTemplate(check_residue, template, bondedToAtom, False)
1182
1183
1184
1185
1186
1187
                    if matches is not None:
                        is_unique = False
            if is_unique:
                # Residue is unique.
                unique_unmatched_residues.append(residue)
                signatures.add(signature)
1188
                templates.append(template)
1189

1190
        return [templates, unique_unmatched_residues]
1191

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

Robert McGibbon's avatar
Robert McGibbon committed
1197
1198
1199
1200
1201
1202
        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
1203
            NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME, or LJPME.
Robert McGibbon's avatar
Robert McGibbon committed
1204
1205
1206
1207
1208
        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.
1209
        rigidWater : boolean=None
Robert McGibbon's avatar
Robert McGibbon committed
1210
            If true, water molecules will be fully rigid regardless of the value
1211
1212
            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
1213
1214
1215
1216
1217
        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
1218
1219
            their total mass the same.  If rigidWater is used to make water molecules
            rigid, then water hydrogens are not altered.
1220
        residueTemplates : dict=dict()
1221
1222
1223
1224
1225
            Specifies which template to use for particular residues.  The keys should be Residue
            objects from the Topology, and the values should be the names of the templates to
            use for them.  This is useful when a ForceField contains multiple templates that
            can match the same residue (e.g Fe2+ and Fe3+ templates in the ForceField for a
            monoatomic iron ion in the Topology).
1226
1227
1228
1229
1230
1231
        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.
1232
1233
1234
        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.
1235
1236
        flexibleConstraints : boolean=False
            If True, parameters for constrained degrees of freedom will be added to the System
1237
        drudeMass : mass=0.4*amu
1238
1239
            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
1240
        args
1241
1242
1243
            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
1244
1245
1246
1247
1248

        Returns
        -------
        system
            the newly created System
1249
        """
1250
1251
        args['switchDistance'] = switchDistance
        args['flexibleConstraints'] = flexibleConstraints
1252
        args['drudeMass'] = drudeMass
1253
        args = ArgTracker(args)
1254
        data = ForceField._SystemData(topology)
1255
        rigidResidue = [False]*topology.getNumResidues()
1256
1257

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

1259
1260
1261
1262
        templateForResidue = self._matchAllResiduesToTemplates(data, topology, residueTemplates, ignoreExternalBonds)
        for res in topology.residues():
            if res.name == 'HOH':
                # Determine whether this should be a rigid water.
1263

1264
1265
1266
1267
                if rigidWater is None:
                    rigidResidue[res.index] = templateForResidue[res.index].rigidWater
                elif rigidWater:
                    rigidResidue[res.index] = True
1268
1269

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

1271
1272
        sys = mm.System()
        for atom in topology.atoms():
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
1273
            # Look up the atom type name, returning a helpful error message if it cannot be found.
1274
1275
1276
1277
            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
1278
            # Look up the type name in the list of registered atom types, returning a helpful error message if it cannot be found.
1279
1280
1281
1282
            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
1283
1284

            # Add the particle to the OpenMM system.
1285
            mass = self._atomTypes[typename].mass
1286
            sys.addParticle(mass)
1287

1288
        # Adjust hydrogen masses if requested.
1289

1290
        if hydrogenMass is not None:
1291
1292
            if not unit.is_quantity(hydrogenMass):
                hydrogenMass *= unit.dalton
1293
            for atom1, atom2 in topology.bonds():
1294
                if atom1.element is elem.hydrogen:
1295
                    (atom1, atom2) = (atom2, atom1)
1296
                if atom2.element is elem.hydrogen and atom1.element not in (elem.hydrogen, None) and not rigidResidue[atom2.residue.index]:
1297
1298
1299
                    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
1300

1301
        # Set periodic boundary conditions.
Justin MacCallum's avatar
Justin MacCallum committed
1302

1303
1304
1305
        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
1306
        elif nonbondedMethod not in [NoCutoff, CutoffNonPeriodic]:
1307
1308
1309
            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
1310

1311
1312
        uniqueAngles = set()
        for bond in data.bonds:
1313
            for atom in data.bondedToAtom[bond.atom1]:
1314
1315
1316
1317
1318
                if atom != bond.atom2:
                    if atom < bond.atom2:
                        uniqueAngles.add((atom, bond.atom1, bond.atom2))
                    else:
                        uniqueAngles.add((bond.atom2, bond.atom1, atom))
1319
            for atom in data.bondedToAtom[bond.atom2]:
1320
1321
1322
1323
1324
1325
                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
1326

1327
        # Make a list of all unique proper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1328

1329
1330
        uniquePropers = set()
        for angle in data.angles:
1331
            for atom in data.bondedToAtom[angle[0]]:
pgrinaway's avatar
pgrinaway committed
1332
                if atom not in angle:
1333
1334
1335
1336
                    if atom < angle[2]:
                        uniquePropers.add((atom, angle[0], angle[1], angle[2]))
                    else:
                        uniquePropers.add((angle[2], angle[1], angle[0], atom))
1337
            for atom in data.bondedToAtom[angle[2]]:
pgrinaway's avatar
pgrinaway committed
1338
                if atom not in angle:
1339
1340
1341
1342
1343
                    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
1344

1345
        # Make a list of all unique improper torsions
Justin MacCallum's avatar
Justin MacCallum committed
1346

1347
1348
        for atom in range(len(data.bondedToAtom)):
            bondedTo = data.bondedToAtom[atom]
1349
1350
1351
            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
1352

1353
        # Identify bonds that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1354

1355
1356
1357
1358
1359
1360
1361
        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]
1362
                bond.isConstrained = atom1.element is elem.hydrogen or atom2.element is elem.hydrogen
1363
1364
1365
1366
1367
        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
1368

1369
        # Identify angles that should be implemented with constraints
Justin MacCallum's avatar
Justin MacCallum committed
1370

1371
1372
1373
1374
1375
1376
        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
1377
                if atom1.element is elem.hydrogen:
1378
                    numH += 1
1379
                if atom3.element is elem.hydrogen:
1380
                    numH += 1
1381
                data.isAngleConstrained.append(numH == 2 or (numH == 1 and atom2.element is elem.oxygen))
1382
1383
        else:
            data.isAngleConstrained = len(data.angles)*[False]
1384
1385
1386
1387
1388
1389
1390
        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
1391

1392
        # Add virtual sites
Justin MacCallum's avatar
Justin MacCallum committed
1393

1394
        for atom in data.virtualSites:
1395
            (site, atoms, excludeWith) = data.virtualSites[atom]
1396
            index = atom.index
1397
            data.excludeAtomWith[excludeWith].append(index)
1398
            if site.type == 'average2':
1399
                sys.setVirtualSite(index, mm.TwoParticleAverageSite(atoms[0], atoms[1], site.weights[0], site.weights[1]))
1400
            elif site.type == 'average3':
1401
                sys.setVirtualSite(index, mm.ThreeParticleAverageSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
1402
            elif site.type == 'outOfPlane':
1403
1404
                sys.setVirtualSite(index, mm.OutOfPlaneSite(atoms[0], atoms[1], atoms[2], site.weights[0], site.weights[1], site.weights[2]))
            elif site.type == 'localCoords':
1405
                sys.setVirtualSite(index, mm.LocalCoordinatesSite(atoms, site.originWeights, site.xWeights, site.yWeights, site.localPos))
Justin MacCallum's avatar
Justin MacCallum committed
1406

1407
        # Add forces to the System
Justin MacCallum's avatar
Justin MacCallum committed
1408

1409
1410
        for force in self._forces:
            force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
1411
1412
        if removeCMMotion:
            sys.addForce(mm.CMMotionRemover())
Justin MacCallum's avatar
Justin MacCallum committed
1413

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

peastman's avatar
peastman committed
1416
1417
1418
        for force in self._forces:
            if 'postprocessSystem' in dir(force):
                force.postprocessSystem(sys, data, args)
Justin MacCallum's avatar
Justin MacCallum committed
1419

1420
        # Execute scripts found in the XML files.
Justin MacCallum's avatar
Justin MacCallum committed
1421

1422
        for script in self._scripts:
1423
            exec(script, locals())
1424
        args.checkArgs(self.createSystem)
1425
1426
1427
        return sys


1428
    def _matchAllResiduesToTemplates(self, data, topology, residueTemplates, ignoreExternalBonds, ignoreExtraParticles=False, recordParameters=True):
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
        """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:
1446
1447
                    if recordParameters:
                        data.recordMatchedAtomParameters(res, template, matches)
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
                    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:
1474
                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)))
1475
            else:
1476
1477
                if recordParameters:
                    data.recordMatchedAtomParameters(res, template, matches)
1478
1479
1480
1481
                templateForResidue[res.index] = template
        return templateForResidue


1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
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))
1506
1507
1508
    for atom in data.atoms:
        for child in data.excludeAtomWith[atom.index]:
            bondIndices.append((child, atom.index))
1509
1510
    return bondIndices

1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
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

1575
1576
def _countResidueAtoms(elements):
    """Count the number of atoms of each element in a residue."""
1577
1578
    counts = {}
    for element in elements:
1579
        if element in counts:
1580
1581
1582
            counts[element] += 1
        else:
            counts[element] = 1
1583
1584
1585
1586
1587
1588
    return counts


def _createResidueSignature(elements):
    """Create a signature for a residue based on the elements of the atoms it contains."""
    counts = _countResidueAtoms(elements)
1589
1590
    sig = []
    for c in counts:
1591
1592
        if c is not None:
            sig.append((c, counts[c]))
1593
    sig.sort(key=lambda x: -x[0].mass)
Justin MacCallum's avatar
Justin MacCallum committed
1594

1595
    # Convert it to a string.
1596
1597

    s = ''
1598
    for element, count in sig:
1599
1600
1601
1602
        s += element.symbol+str(count)
    return s


1603
def _applyPatchesToMatchResidues(forcefield, data, residues, templateForResidue, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles):
1604
1605
1606
1607
    """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.
1608

1609
1610
1611
1612
    patchedTemplateSignatures = {}
    patchedTemplates = {}
    for name, template in forcefield._templates.items():
        if name in forcefield._templatePatches:
1613
            patches = sorted([forcefield._patches[patchName] for patchName, patchResidueIndex in forcefield._templatePatches[name] if forcefield._patches[patchName].numResidues == 1])
1614
1615
1616
            if len(patches) > 0:
                newTemplates = []
                patchedTemplates[name] = newTemplates
peastman's avatar
peastman committed
1617
                _generatePatchedSingleResidueTemplates(template, patches, 0, newTemplates, set())
1618
1619
1620
1621
1622
1623
                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]
1624

1625
    # Now see if any of those templates matches any of the residues.
1626

1627
1628
    unmatchedResidues = []
    for res in residues:
1629
        [template, matches] = forcefield._getResidueTemplateMatches(res, bondedToAtom, patchedTemplateSignatures, ignoreExternalBonds, ignoreExtraParticles)
1630
1631
1632
1633
        if matches is None:
            unmatchedResidues.append(res)
        else:
            data.recordMatchedAtomParameters(res, template, matches)
1634
            templateForResidue[res.index] = template
1635
1636
    if len(unmatchedResidues) == 0:
        return []
1637

1638
1639
1640
    # 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.
1641

1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
    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]
1657

1658
    # Record which unmatched residues are bonded to each other.
1659

1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
    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)
1670

1671
1672
    # 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.
1673

1674
1675
1676
1677
    clusterSize = 2
    clusters = bonds
    while clusterSize <= maxPatchSize:
        # Try to apply patches to clusters of this size.
1678

1679
1680
1681
        for patchName in patches:
            patch = forcefield._patches[patchName]
            if patch.numResidues == clusterSize:
1682
                matchedClusters = _matchToMultiResiduePatchedTemplates(data, clusters, patch, patches[patchName], bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1683
1684
1685
1686
1687
1688
                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.
1689

1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
        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

1705
1706
1707
    return unmatchedResidues


peastman's avatar
peastman committed
1708
def _generatePatchedSingleResidueTemplates(template, patches, index, newTemplates, alteredAtoms):
1709
1710
    """Apply all possible combinations of a set of single-residue patches to a template."""
    try:
peastman's avatar
peastman committed
1711
1712
1713
1714
1715
1716
1717
        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)
1718
1719
1720
1721
    except:
        # This probably means the patch is inconsistent with another one that has already been applied,
        # so just ignore it.
        patchedTemplate = None
1722

1723
    # Call this function recursively to generate combinations of patches.
1724

1725
    if index+1 < len(patches):
peastman's avatar
peastman committed
1726
        _generatePatchedSingleResidueTemplates(template, patches, index+1, newTemplates, alteredAtoms)
1727
        if patchedTemplate is not None:
peastman's avatar
peastman committed
1728
1729
            newAlteredAtoms = alteredAtoms.union(patches[index].allAtomNames)
            _generatePatchedSingleResidueTemplates(patchedTemplate, patches, index+1, newTemplates, newAlteredAtoms)
1730
1731


1732
def _matchToMultiResiduePatchedTemplates(data, clusters, patch, residueTemplates, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles):
1733
1734
1735
    """Apply a multi-residue patch to templates, then try to match them against clusters of residues."""
    matchedClusters = []
    selectedTemplates = [None]*patch.numResidues
1736
    _applyMultiResiduePatch(data, clusters, patch, residueTemplates, selectedTemplates, 0, matchedClusters, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1737
1738
1739
    return matchedClusters


1740
def _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index, matchedClusters, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles):
1741
1742
1743
1744
1745
    """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
1746
            _applyMultiResiduePatch(data, clusters, patch, candidateTemplates, selectedTemplates, index+1, matchedClusters, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1747
1748
1749
    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.
1750

1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
        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):
1762
                    matches = compiled.matchResidueToTemplate(residue, template, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
1763
1764
1765
1766
1767
1768
                    if matches is None:
                        residueMatches = None
                        break
                    else:
                        residueMatches.append(matches)
                if residueMatches is not None:
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
                    # 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.
1787

1788
1789
1790
1791
                        for i in range(patch.numResidues):
                            data.recordMatchedAtomParameters(residues[i], patchedTemplates[i], residueMatches[i])
                        newlyMatchedClusters.append(cluster)
                        break
1792

1793
        # Record which clusters were successfully matched.
1794

1795
1796
1797
        matchedClusters += newlyMatchedClusters
        for cluster in newlyMatchedClusters:
            clusters.remove(cluster)
1798

1799

1800
1801
1802
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()])
1803
    numResidueAtoms = sum(residueCounts.values())
1804
    numResidueHeavyAtoms = sum(residueCounts[element] for element in residueCounts if element not in (None, elem.hydrogen))
1805

1806
    # Loop over templates and see how closely each one might match.
1807

1808
1809
1810
1811
1812
1813
    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])
1814

1815
        # Does the residue have any atoms that clearly aren't in the template?
1816

1817
1818
        if any(element not in templateCounts or templateCounts[element] < residueCounts[element] for element in residueCounts):
            continue
1819

1820
        # If there are too many missing atoms, discard this template.
1821

1822
        numTemplateAtoms = sum(templateCounts.values())
1823
        numTemplateHeavyAtoms = sum(templateCounts[element] for element in templateCounts if element not in (None, elem.hydrogen))
1824
1825
1826
1827
        if numTemplateAtoms > numBestMatchAtoms:
            continue
        if numTemplateHeavyAtoms > numBestMatchHeavyAtoms:
            continue
1828

1829
1830
        # 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.
1831

1832
1833
1834
        if numTemplateAtoms == numBestMatchAtoms:
            if bestMatchName == res.name or res.name not in templateName:
                continue
1835

1836
        # Accept this as our new best match.
1837

1838
1839
1840
        bestMatchName = templateName
        numBestMatchAtoms = numTemplateAtoms
        numBestMatchHeavyAtoms = numTemplateHeavyAtoms
1841
        numBestMatchExtraParticles = len([atom for atom in template.atoms if atom.element is None])
1842

1843
    # Return an appropriate error message.
1844

1845
    if numBestMatchAtoms == numResidueAtoms:
1846
1847
        chainResidues = list(res.chain.residues())
        if len(chainResidues) > 1 and (res == chainResidues[0] or res == chainResidues[-1]):
1848
1849
1850
1851
            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:
1852
1853
1854
1855
1856
            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)
1857
1858
1859
        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.'

1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
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():
1878
        template.addAtom(ForceField._TemplateAtomData(atom.name, None, atom.element))
1879
1880
1881
1882
1883
1884
1885
1886
1887
    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
1888

1889
1890
1891
1892
1893
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
1894
    wildcard = generator.ff._atomClasses['']
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
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
    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
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
                    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
1999
    return match
2000

2001

2002
2003
2004
2005
2006
# 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.

2007
## @private
2008
class HarmonicBondGenerator(object):
2009
    """A HarmonicBondGenerator constructs a HarmonicBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2010

2011
2012
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
2013
        self.bondsForAtomType = defaultdict(set)
2014
2015
2016
2017
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
2018

2019
2020
2021
    def registerBond(self, parameters):
        types = self.ff._findAtomTypes(parameters, 2)
        if None not in types:
peastman's avatar
peastman committed
2022
            index = len(self.types1)
2023
2024
            self.types1.append(types[0])
            self.types2.append(types[1])
peastman's avatar
peastman committed
2025
2026
2027
2028
            for t in types[0]:
                self.bondsForAtomType[t].add(index)
            for t in types[1]:
                self.bondsForAtomType[t].add(index)
2029
2030
            self.length.append(_convertParameterToNumber(parameters['length']))
            self.k.append(_convertParameterToNumber(parameters['k']))
Justin MacCallum's avatar
Justin MacCallum committed
2031

2032
2033
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2034
2035
2036
2037
2038
2039
        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]
2040
        for bond in element.findall('Bond'):
2041
            generator.registerBond(bond.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
2042

2043
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2044
        existing = [f for f in sys.getForces() if type(f) == mm.HarmonicBondForce]
2045
2046
2047
2048
2049
2050
2051
2052
        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
2053
            for i in self.bondsForAtomType[type1]:
2054
2055
2056
2057
2058
                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:
2059
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
2060
2061
2062
2063
2064
                    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])
2065
2066
2067
2068
2069
                    break

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


2070
## @private
2071
class HarmonicAngleGenerator(object):
2072
    """A HarmonicAngleGenerator constructs a HarmonicAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2073

2074
2075
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
2076
        self.anglesForAtom2Type = defaultdict(list)
2077
2078
2079
2080
2081
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.angle = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
2082

2083
2084
2085
    def registerAngle(self, parameters):
        types = self.ff._findAtomTypes(parameters, 3)
        if None not in types:
peastman's avatar
peastman committed
2086
            index = len(self.types1)
2087
2088
2089
            self.types1.append(types[0])
            self.types2.append(types[1])
            self.types3.append(types[2])
peastman's avatar
peastman committed
2090
2091
            for t in types[1]:
                self.anglesForAtom2Type[t].append(index)
2092
2093
2094
            self.angle.append(_convertParameterToNumber(parameters['angle']))
            self.k.append(_convertParameterToNumber(parameters['k']))

2095
2096
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2097
2098
2099
2100
2101
2102
        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]
2103
        for angle in element.findall('Angle'):
2104
            generator.registerAngle(angle.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
2105

2106
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2107
2108
2109
2110
2111
        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.

2112
        existing = [f for f in sys.getForces() if type(f) == mm.HarmonicAngleForce]
2113
2114
2115
2116
2117
2118
2119
2120
2121
        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
2122
            for i in self.anglesForAtom2Type[type2]:
2123
2124
2125
2126
2127
2128
                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
2129

2130
2131
2132
2133
2134
2135
2136
2137
2138
                        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
2139

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

2142
2143
2144
2145
2146
                        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]))
2147
                                data.addConstraint(sys, angle[0], angle[2], length)
2148
2149
2150
                    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])
2151
2152
2153
2154
2155
                    break

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


2156
## @private
2157
class PeriodicTorsion(object):
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
    """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 = []
2168
        self.ordering = 'default'
2169

2170
## @private
2171
class PeriodicTorsionGenerator(object):
2172
    """A PeriodicTorsionGenerator constructs a PeriodicTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2173

2174
2175
    def __init__(self, forcefield):
        self.ff = forcefield
2176
2177
        self.proper = []
        self.improper = []
2178
        self.propersForAtomType = defaultdict(set)
Justin MacCallum's avatar
Justin MacCallum committed
2179

2180
2181
2182
    def registerProperTorsion(self, parameters):
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
2183
            index = len(self.proper)
2184
            self.proper.append(torsion)
2185
2186
2187
2188
            for t in torsion.types2:
                self.propersForAtomType[t].add(index)
            for t in torsion.types3:
                self.propersForAtomType[t].add(index)
2189

2190
    def registerImproperTorsion(self, parameters, ordering='default'):
2191
2192
        torsion = self.ff._parseTorsion(parameters)
        if torsion is not None:
2193
            if ordering in ['default', 'charmm', 'amber', 'smirnoff']:
2194
2195
                torsion.ordering = ordering
            else:
2196
                raise ValueError('Illegal ordering type %s for improper torsion %s' % (ordering, torsion))
2197
2198
            self.improper.append(torsion)

2199
2200
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2201
2202
2203
2204
2205
2206
        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]
2207
        for torsion in element.findall('Proper'):
2208
            generator.registerProperTorsion(torsion.attrib)
2209
        for torsion in element.findall('Improper'):
2210
2211
2212
2213
            if 'ordering' in element.attrib:
                generator.registerImproperTorsion(torsion.attrib, element.attrib['ordering'])
            else:
                generator.registerImproperTorsion(torsion.attrib)
Justin MacCallum's avatar
Justin MacCallum committed
2214

2215
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2216
        existing = [f for f in sys.getForces() if type(f) == mm.PeriodicTorsionForce]
2217
2218
2219
2220
2221
2222
        if len(existing) == 0:
            force = mm.PeriodicTorsionForce()
            sys.addForce(force)
        else:
            force = existing[0]
        wildcard = self.ff._atomClasses['']
tic20's avatar
tic20 committed
2223
        proper_cache = {}
2224
        for torsion in data.propers:
tic20's avatar
tic20 committed
2225
2226
            type1, type2, type3, type4 = [data.atomType[data.atoms[torsion[i]]] for i in range(4)]
            sig = (type1, type2, type3, type4)
2227
            sig = frozenset((sig, sig[::-1]))
tic20's avatar
tic20 committed
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
            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
2248
2249
2250
2251
            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])
2252
        impr_cache = {}
2253
        for torsion in data.impropers:
Peter Eastman's avatar
Peter Eastman committed
2254
            t1, t2, t3, t4 = [data.atomType[data.atoms[torsion[i]]] for i in range(4)]
2255
            sig = (t1, t2, t3, t4)
2256
2257
2258
2259
            match = impr_cache.get(sig, None)
            if match == -1:
                # Previously checked, and doesn't appear in the database
                continue
2260
2261
2262
2263
            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)
2264
2265
2266
            if match is None:
                match = _matchImproper(data, torsion, self)
                if match is not None:
2267
2268
2269
                    order = match[:4]
                    i1, i2, i3, i4 = tuple(torsion.index(a) for a in order)
                    impr_cache[sig] = (i1, i2, i3, i4, match[-1])
2270
2271
                else:
                    impr_cache[sig] = -1
2272
2273
2274
2275
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                for i in range(len(tordef.phase)):
                    if tordef.k[i] != 0:
2276
2277
2278
2279
2280
2281
2282
                        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])
2283
2284
2285
parsers["PeriodicTorsionForce"] = PeriodicTorsionGenerator.parseElement


2286
## @private
2287
class RBTorsion(object):
2288
2289
    """An RBTorsion records the information for a Ryckaert-Bellemans torsion definition."""

2290
    def __init__(self, types, c, ordering='charmm'):
2291
2292
2293
2294
2295
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.c = c
2296
        if ordering in ['default', 'charmm', 'amber']:
2297
2298
            self.ordering = ordering
        else:
2299
            raise ValueError('Illegal ordering type %s for RBTorsion (%s,%s,%s,%s)' % (ordering, types[0], types[1], types[2], types[3]))
2300

2301
## @private
2302
class RBTorsionGenerator(object):
2303
    """An RBTorsionGenerator constructs an RBTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2304

2305
2306
    def __init__(self, forcefield):
        self.ff = forcefield
2307
2308
        self.proper = []
        self.improper = []
Justin MacCallum's avatar
Justin MacCallum committed
2309

2310
2311
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2312
2313
2314
2315
2316
2317
        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]
2318
        for torsion in element.findall('Proper'):
2319
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2320
            if None not in types:
2321
2322
                generator.proper.append(RBTorsion(types, [float(torsion.attrib['c'+str(i)]) for i in range(6)]))
        for torsion in element.findall('Improper'):
2323
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2324
            if None not in types:
2325
2326
2327
2328
                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
2329

2330
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2331
        existing = [f for f in sys.getForces() if type(f) == mm.RBTorsionForce]
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
        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:
2358
            match = _matchImproper(data, torsion, self)
2359
2360
2361
            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])
2362
2363
2364
2365

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


2366
## @private
2367
class CMAPTorsion(object):
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
    """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

2378
## @private
2379
class CMAPTorsionGenerator(object):
2380
    """A CMAPTorsionGenerator constructs a CMAPTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2381

2382
2383
    def __init__(self, forcefield):
        self.ff = forcefield
2384
2385
        self.torsions = []
        self.maps = []
Justin MacCallum's avatar
Justin MacCallum committed
2386

2387
2388
    @staticmethod
    def parseElement(element, ff):
Rafal P. Wiewiora's avatar
Rafal P. Wiewiora committed
2389
2390
2391
2392
2393
2394
        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]
2395
2396
2397
2398
2399
2400
2401
        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'):
2402
            types = ff._findAtomTypes(torsion.attrib, 5)
peastman's avatar
peastman committed
2403
            if None not in types:
2404
                generator.torsions.append(CMAPTorsion(types, int(torsion.attrib['map'])))
Justin MacCallum's avatar
Justin MacCallum committed
2405

2406
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2407
        existing = [f for f in sys.getForces() if type(f) == mm.CMAPTorsionForce]
2408
2409
2410
2411
2412
2413
2414
        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
2415

2416
        # Find all chains of length 5
Justin MacCallum's avatar
Justin MacCallum committed
2417

2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
        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


2461
## @private
2462
class NonbondedGenerator(object):
2463
    """A NonbondedGenerator constructs a NonbondedForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2464

2465
2466
    SCALETOL = 1e-5

2467
    def __init__(self, forcefield, coulomb14scale, lj14scale, useDispersionCorrection):
2468
        self.ff = forcefield
2469
2470
        self.coulomb14scale = coulomb14scale
        self.lj14scale = lj14scale
2471
        self.useDispersionCorrection = useDispersionCorrection
2472
        self.params = ForceField._AtomTypeParameters(forcefield, 'NonbondedForce', 'Atom', ('charge', 'sigma', 'epsilon'))
2473

2474
    def registerAtom(self, parameters):
2475
        self.params.registerAtom(parameters)
2476

2477
2478
2479
    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, NonbondedGenerator)]
2480
2481
2482
2483
        if 'useDispersionCorrection' in element.attrib:
            useDispersionCorrection = bool(eval(element.attrib.get('useDispersionCorrection')))
        else:
            useDispersionCorrection = None
2484
        if len(existing) == 0:
2485
2486
2487
2488
2489
2490
            generator = NonbondedGenerator(
                ff,
                float(element.attrib['coulomb14scale']),
                float(element.attrib['lj14scale']),
                useDispersionCorrection
            )
2491
            ff.registerGenerator(generator)
2492
2493
2494
        else:
            # Multiple <NonbondedForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2495
2496
2497
            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
2498
                raise ValueError('Found multiple NonbondedForce tags with different 1-4 scales')
2499
2500
2501
2502
2503
2504
            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.')
2505
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2506

2507
2508
2509
2510
2511
    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
2512
2513
                     PME:mm.NonbondedForce.PME,
                     LJPME:mm.NonbondedForce.LJPME}
2514
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2515
            raise ValueError('Illegal nonbonded method for NonbondedForce')
2516
2517
        force = mm.NonbondedForce()
        for atom in data.atoms:
2518
2519
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
peastman's avatar
peastman committed
2520
2521
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2522
2523
2524
        if args['switchDistance'] is not None:
            force.setUseSwitchingFunction(True)
            force.setSwitchingDistance(args['switchDistance'])
peastman's avatar
peastman committed
2525
2526
        if 'ewaldErrorTolerance' in args:
            force.setEwaldErrorTolerance(args['ewaldErrorTolerance'])
2527
        if 'useDispersionCorrection' in args:
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
            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
2540
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
2541

peastman's avatar
peastman committed
2542
    def postprocessSystem(self, sys, data, args):
2543
2544
        # Create the exceptions.

2545
        bondIndices = _findBondsForExclusions(data, sys)
peastman's avatar
peastman committed
2546
        nonbonded = [f for f in sys.getForces() if isinstance(f, mm.NonbondedForce)][0]
Justin MacCallum's avatar
Justin MacCallum committed
2547
        nonbonded.createExceptionsFromBonds(bondIndices, self.coulomb14scale, self.lj14scale)
2548
2549
2550

parsers["NonbondedForce"] = NonbondedGenerator.parseElement

2551
2552

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

2556
    def __init__(self, forcefield, lj14scale, useDispersionCorrection):
ChayaSt's avatar
ChayaSt committed
2557
        self.ff = forcefield
2558
2559
2560
        self.nbfixParameters = []
        self.nbfixTypes1 = defaultdict(set)
        self.nbfixTypes2 = defaultdict(set)
ChayaSt's avatar
ChayaSt committed
2561
        self.lj14scale = lj14scale
2562
        self.useDispersionCorrection = useDispersionCorrection
2563
        self.ljTypes = ForceField._AtomTypeParameters(forcefield, 'LennardJonesForce', 'Atom', ('sigma', 'epsilon'))
ChayaSt's avatar
ChayaSt committed
2564

ChayaSt's avatar
ChayaSt committed
2565
    def registerNBFIX(self, parameters):
ChayaSt's avatar
ChayaSt committed
2566
        types = self.ff._findAtomTypes(parameters, 2)
2567

ChayaSt's avatar
ChayaSt committed
2568
        if None not in types:
2569
2570
2571
2572
2573
2574
2575
2576
            sigma = _convertParameterToNumber(parameters['sigma'])
            epsilon = _convertParameterToNumber(parameters['epsilon'])

            # Retrieve the index of nbfixParameters into which this sigma and
            # epsilon will be stored, then register this index with the atom
            # types that should have this sigma and epsilon applied.
            nbfixIndex = len(self.nbfixParameters)
            self.nbfixParameters.append([sigma, epsilon])
2577
            for type1 in types[0]:
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
                self.nbfixTypes1[type1].add(nbfixIndex)
            for type2 in types[1]:
                self.nbfixTypes2[type2].add(nbfixIndex)

    def getNBFIX(self, type1, type2):
        nbfixIndices = (self.nbfixTypes1[type1] & self.nbfixTypes2[type2]) | (self.nbfixTypes2[type1] & self.nbfixTypes1[type2])
        if nbfixIndices:
            if len(nbfixIndices) > 1:
                raise ValueError('Multiple NBFixPair entries match atom types %s-%s.' % (type1, type2))
            return self.nbfixParameters[nbfixIndices.pop()]
        else:
            return None
2590

ChayaSt's avatar
ChayaSt committed
2591
    def registerLennardJones(self, parameters):
2592
        self.ljTypes.registerAtom(parameters)
ChayaSt's avatar
ChayaSt committed
2593
2594
2595

    @staticmethod
    def parseElement(element, ff):
ChayaSt's avatar
ChayaSt committed
2596
        existing = [f for f in ff._forces if isinstance(f, LennardJonesGenerator)]
2597
2598
2599
2600
        if 'useDispersionCorrection' in element.attrib:
            useDispersionCorrection = bool(eval(element.attrib.get('useDispersionCorrection')))
        else:
            useDispersionCorrection = None
ChayaSt's avatar
ChayaSt committed
2601
        if len(existing) == 0:
2602
2603
2604
2605
2606
            generator = LennardJonesGenerator(
                ff,
                float(element.attrib['lj14scale']),
                useDispersionCorrection=useDispersionCorrection
            )
ChayaSt's avatar
ChayaSt committed
2607
2608
            ff.registerGenerator(generator)
        else:
ChayaSt's avatar
ChayaSt committed
2609
            # Multiple <LennardJonesForce> tags were found, probably in different files
ChayaSt's avatar
ChayaSt committed
2610
            generator = existing[0]
2611
2612
            if abs(generator.lj14scale - float(element.attrib['lj14scale'])) > NonbondedGenerator.SCALETOL:
                raise ValueError('Found multiple LennardJonesForce tags with different 1-4 scales')
2613
2614
2615
2616
2617
2618
            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
2619
2620
        for LJ in element.findall('Atom'):
            generator.registerLennardJones(LJ.attrib)
ChayaSt's avatar
ChayaSt committed
2621
        for Nbfix in element.findall('NBFixPair'):
ChayaSt's avatar
ChayaSt committed
2622
            generator.registerNBFIX(Nbfix.attrib)
ChayaSt's avatar
ChayaSt committed
2623

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

2628
        nbfixTypeSet = {t for nbfixTypes in (self.nbfixTypes1, self.nbfixTypes2) for t in nbfixTypes if nbfixTypes[t]}
peastman's avatar
peastman committed
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
        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]
2645
            else:
peastman's avatar
peastman committed
2646
2647
2648
2649
2650
                # This is a new type.
                typeToMergedType[t] = len(mergedTypes)
                paramsToMergedType[params] = len(mergedTypes)
                mergedTypes.append(t)
                mergedTypeParams.append(params)
2651

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

peastman's avatar
peastman committed
2654
        numLjTypes = len(mergedTypes)
2655
        acoef = [0]*(numLjTypes*numLjTypes)
ChayaSt's avatar
ChayaSt committed
2656
        bcoef = acoef[:]
2657
2658
        for m in range(numLjTypes):
            for n in range(numLjTypes):
2659
2660
2661
                nbfix = self.getNBFIX(mergedTypes[m], mergedTypes[n])
                if nbfix is not None:
                    sigma, epsilon = nbfix
2662
2663
2664
                    sigma6 = sigma**6
                    acoef[m+numLjTypes*n] = 4*epsilon*sigma6*sigma6
                    bcoef[m+numLjTypes*n] = 4*epsilon*sigma6
ChayaSt's avatar
cleanup  
ChayaSt committed
2665
                    continue
ChayaSt's avatar
ChayaSt committed
2666
                else:
peastman's avatar
peastman committed
2667
                    sigma = 0.5*(mergedTypeParams[m][0]+mergedTypeParams[n][0])
2668
                    sigma6 = sigma**6
peastman's avatar
peastman committed
2669
                    epsilon = math.sqrt(mergedTypeParams[m][1]*mergedTypeParams[n][1])
2670
2671
2672
2673
2674
2675
                    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
2676
        self.force.addPerParticleParameter('type')
2677
        self.force.setName('LennardJones')
Peter Eastman's avatar
Peter Eastman committed
2678
        if nonbondedMethod in [CutoffPeriodic, Ewald, PME, LJPME]:
ChayaSt's avatar
ChayaSt committed
2679
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
ChayaSt's avatar
ChayaSt committed
2680
        elif nonbondedMethod is NoCutoff:
ChayaSt's avatar
ChayaSt committed
2681
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
ChayaSt's avatar
ChayaSt committed
2682
        elif nonbondedMethod is CutoffNonPeriodic:
ChayaSt's avatar
ChayaSt committed
2683
            self.force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffNonPeriodic)
ChayaSt's avatar
ChayaSt committed
2684
        else:
2685
            raise AssertionError('Unrecognized nonbonded method [%s]' % nonbondedMethod)
2686
        if args['switchDistance'] is not None:
2687
2688
            self.force.setUseSwitchingFunction(True)
            self.force.setSwitchingDistance(args['switchDistance'])
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
        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)
2702

ChayaSt's avatar
ChayaSt committed
2703
        # Add the particles
2704

peastman's avatar
peastman committed
2705
2706
        for atom in data.atoms:
            self.force.addParticle((typeToMergedType[data.atomType[atom]],))
ChayaSt's avatar
ChayaSt committed
2707
        self.force.setCutoffDistance(nonbondedCutoff)
ChayaSt's avatar
ChayaSt committed
2708
2709
2710
        sys.addForce(self.force)

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

2713
        bondIndices = _findBondsForExclusions(data, sys)
2714
2715
2716
2717
2718
2719
2720
2721
2722
        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')
2723
            bonded.setName('LennardJones14')
2724
2725
2726
2727
2728
2729
2730
            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:
2731
2732
2733
                    nbfix = self.getNBFIX(data.atomType[a1], data.atomType[a2])
                    if nbfix is not None:
                        sigma, epsilon = nbfix
2734
2735
2736
                    else:
                        values1 = self.ljTypes.getAtomParameters(a1, data)
                        values2 = self.ljTypes.getAtomParameters(a2, data)
2737
2738
2739
2740
2741
2742
2743
2744
                        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)
2745
                    bonded.addBond(p1, p2, (sigma, epsilon))
ChayaSt's avatar
ChayaSt committed
2746

ChayaSt's avatar
ChayaSt committed
2747
parsers["LennardJonesForce"] = LennardJonesGenerator.parseElement
2748

2749

2750
## @private
2751
class GBSAOBCGenerator(object):
2752
    """A GBSAOBCGenerator constructs a GBSAOBCForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2753

2754
2755
    def __init__(self, forcefield):
        self.ff = forcefield
2756
        self.params = ForceField._AtomTypeParameters(forcefield, 'GBSAOBCForce', 'Atom', ('charge', 'radius', 'scale'))
2757

2758
    def registerAtom(self, parameters):
2759
        self.params.registerAtom(parameters)
2760

2761
2762
    @staticmethod
    def parseElement(element, ff):
2763
2764
        existing = [f for f in ff._forces if isinstance(f, GBSAOBCGenerator)]
        if len(existing) == 0:
2765
2766
            generator = GBSAOBCGenerator(ff)
            ff.registerGenerator(generator)
2767
2768
2769
        else:
            # Multiple <GBSAOBCForce> tags were found, probably in different files.  Simply add more types to the existing one.
            generator = existing[0]
2770
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
2771

2772
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2773
2774
2775
2776
2777
2778
        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}
2779
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
2780
            raise ValueError('Illegal nonbonded method for GBSAOBCForce')
2781
2782
        force = mm.GBSAOBCForce()
        for atom in data.atoms:
2783
2784
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1], values[2])
2785
2786
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
2787
2788
2789
2790
        if 'soluteDielectric' in args:
            force.setSoluteDielectric(float(args['soluteDielectric']))
        if 'solventDielectric' in args:
            force.setSolventDielectric(float(args['solventDielectric']))
2791
2792
        sys.addForce(force)

2793
2794
    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
2795

2796
2797
2798
2799
        for force in sys.getForces():
            if isinstance(force, mm.NonbondedForce):
                force.setReactionFieldDielectric(1.0)

2800
2801
2802
parsers["GBSAOBCForce"] = GBSAOBCGenerator.parseElement


2803
## @private
2804
class CustomBondGenerator(object):
2805
    """A CustomBondGenerator constructs a CustomBondForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2806

2807
2808
    def __init__(self, forcefield):
        self.ff = forcefield
2809
2810
2811
2812
2813
        self.types1 = []
        self.types2 = []
        self.globalParams = {}
        self.perBondParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2814

2815
2816
    @staticmethod
    def parseElement(element, ff):
2817
2818
        generator = CustomBondGenerator(ff)
        ff.registerGenerator(generator)
2819
2820
2821
2822
2823
2824
        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'):
2825
            types = ff._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
2826
            if None not in types:
2827
2828
2829
                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
2830

2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
    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


2851
## @private
2852
class CustomAngleGenerator(object):
2853
    """A CustomAngleGenerator constructs a CustomAngleForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2854

2855
2856
    def __init__(self, forcefield):
        self.ff = forcefield
2857
2858
2859
2860
2861
2862
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.globalParams = {}
        self.perAngleParams = []
        self.paramValues = []
Justin MacCallum's avatar
Justin MacCallum committed
2863

2864
2865
    @staticmethod
    def parseElement(element, ff):
2866
2867
        generator = CustomAngleGenerator(ff)
        ff.registerGenerator(generator)
2868
2869
2870
2871
2872
2873
        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'):
2874
            types = ff._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
2875
            if None not in types:
2876
2877
2878
2879
                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
2880

2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
    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


2903
## @private
2904
class CustomTorsion(object):
2905
2906
    """A CustomTorsion records the information for a custom torsion definition."""

2907
    def __init__(self, types, paramValues, ordering='charmm'):
2908
2909
2910
2911
2912
        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.paramValues = paramValues
2913
        if ordering in ['default', 'charmm', 'amber']:
2914
2915
            self.ordering = ordering
        else:
2916
            raise ValueError('Illegal ordering type %s for CustomTorsion (%s,%s,%s,%s)' % (ordering, types[0], types[1], types[2], types[3]))
2917

2918
## @private
2919
class CustomTorsionGenerator(object):
2920
    """A CustomTorsionGenerator constructs a CustomTorsionForce."""
Justin MacCallum's avatar
Justin MacCallum committed
2921

2922
2923
    def __init__(self, forcefield):
        self.ff = forcefield
2924
2925
2926
2927
        self.proper = []
        self.improper = []
        self.globalParams = {}
        self.perTorsionParams = []
Justin MacCallum's avatar
Justin MacCallum committed
2928

2929
2930
    @staticmethod
    def parseElement(element, ff):
2931
2932
        generator = CustomTorsionGenerator(ff)
        ff.registerGenerator(generator)
2933
2934
2935
2936
2937
2938
        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'):
2939
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2940
            if None not in types:
2941
2942
                generator.proper.append(CustomTorsion(types, [float(torsion.attrib[param]) for param in generator.perTorsionParams]))
        for torsion in element.findall('Improper'):
2943
            types = ff._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
2944
            if None not in types:
2945
2946
2947
2948
                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
2949

2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
    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:
2978
            match = _matchImproper(data, torsion, self)
2979
2980
2981
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
2982
2983
2984
2985

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


2986
## @private
2987
class CustomNonbondedGenerator(object):
2988
2989
    """A CustomNonbondedGenerator constructs a CustomNonbondedForce."""

2990
2991
    def __init__(self, forcefield, energy, bondCutoff):
        self.ff = forcefield
2992
2993
2994
2995
        self.energy = energy
        self.bondCutoff = bondCutoff
        self.globalParams = {}
        self.perParticleParams = []
2996
        self.computedValues = {}
2997
2998
2999
3000
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
3001
3002
        generator = CustomNonbondedGenerator(ff, element.attrib['energy'], int(element.attrib['bondCutoff']))
        ff.registerGenerator(generator)
3003
3004
3005
3006
        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'])
3007
3008
        for value in element.findall('ComputedValue'):
            generator.computedValues[value.attrib['name']] = value.attrib['expression']
3009
3010
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomNonbondedForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
3011
        generator.functions += _parseFunctions(element)
3012
3013
3014
3015

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomNonbondedForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomNonbondedForce.CutoffNonPeriodic,
3016
3017
3018
3019
                     CutoffPeriodic:mm.CustomNonbondedForce.CutoffPeriodic,
                     Ewald:mm.CustomNonbondedForce.CutoffPeriodic,
                     PME:mm.CustomNonbondedForce.CutoffPeriodic,
                     LJPME:mm.CustomNonbondedForce.CutoffPeriodic}
3020
3021
3022
3023
3024
3025
3026
        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)
3027
3028
        for name in self.computedValues:
            force.addComputedValue(name, self.computedValues[name])
3029
        _createFunctions(force, self.functions)
3030
        for atom in data.atoms:
3031
3032
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
3033
3034
3035
3036
3037
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

    def postprocessSystem(self, sys, data, args):
3038
        # Create the exclusions.
3039

3040
        bondIndices = _findBondsForExclusions(data, sys)
3041
3042
3043
        for f in sys.getForces():
            if isinstance(f, mm.CustomNonbondedForce) and f.getEnergyFunction() == self.energy:
                f.createExclusionsFromBonds(bondIndices, self.bondCutoff)
3044
3045
3046
3047

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


3048
## @private
3049
class CustomGBGenerator(object):
3050
    """A CustomGBGenerator constructs a CustomGBForce."""
Justin MacCallum's avatar
Justin MacCallum committed
3051

3052
3053
    def __init__(self, forcefield):
        self.ff = forcefield
3054
3055
3056
3057
3058
3059
3060
3061
        self.globalParams = {}
        self.perParticleParams = []
        self.computedValues = []
        self.energyTerms = []
        self.functions = []

    @staticmethod
    def parseElement(element, ff):
3062
3063
        generator = CustomGBGenerator(ff)
        ff.registerGenerator(generator)
3064
3065
3066
3067
        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'])
3068
3069
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomGBForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
3070
3071
3072
3073
3074
3075
3076
        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']]))
3077
        generator.functions += _parseFunctions(element)
Justin MacCallum's avatar
Justin MacCallum committed
3078

3079
3080
3081
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomGBForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomGBForce.CutoffNonPeriodic,
3082
3083
3084
3085
                     CutoffPeriodic:mm.CustomGBForce.CutoffPeriodic,
                     Ewald:mm.CustomGBForce.CutoffPeriodic,
                     PME:mm.CustomGBForce.CutoffPeriodic,
                     LJPME:mm.CustomGBForce.CutoffPeriodic}
3086
        if nonbondedMethod not in methodMap:
Justin MacCallum's avatar
Justin MacCallum committed
3087
            raise ValueError('Illegal nonbonded method for CustomGBForce')
3088
3089
3090
3091
3092
3093
3094
3095
3096
        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])
3097
        _createFunctions(force, self.functions)
3098
        for atom in data.atoms:
3099
3100
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values)
3101
3102
3103
3104
3105
3106
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

3107

3108
3109
3110
3111
3112
3113
3114
3115
3116
3117
3118
3119
3120
3121
3122
3123
3124
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140
3141
3142
3143
3144
3145
## @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'):
3146
            types = ff._findAtomTypes(donor.attrib, 3)[:generator.particlesPerDonor]
3147
3148
3149
3150
3151
3152
3153
3154
            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'):
3155
            types = ff._findAtomTypes(acceptor.attrib, 3)[:generator.particlesPerAcceptor]
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
            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,
3168
3169
3170
3171
                     CutoffPeriodic:mm.CustomHbondForce.CutoffPeriodic,
                     Ewald:mm.CustomHbondForce.CutoffPeriodic,
                     PME:mm.CustomHbondForce.CutoffPeriodic,
                     LJPME:mm.CustomHbondForce.CutoffPeriodic}
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
3183
3184
3185
3186
3187
3188
3189
3190
3191
3192
        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
3193
                    if type1 in types1:
3194
3195
3196
3197
3198
3199
3200
3201
                        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]
3202
3203
3204
3205
                    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])
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
3218
3219
3220
3221
3222
3223
3224
        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
3225
                    if type1 in types1:
3226
3227
3228
3229
3230
3231
3232
3233
                        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]
3234
                    if type1 in types1 and type2 in types2:
3235
                        force.addAcceptor(bond.atom1, bond.atom2, -1, self.acceptorParamValues[i])
3236
3237
                    elif type1 in types2 and type2 in types1:
                        force.addAcceptor(bond.atom2, bond.atom1, -1, self.acceptorParamValues[i])
3238
3239
3240
3241
3242
3243
3244
3245
3246
3247
3248
3249
3250
3251
3252
3253
3254
3255
3256
3257
3258
3259
3260
3261
3262
3263
3264
3265
3266
3267
3268
3269
3270
3271
3272
3273
3274
3275
        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


3276
## @private
3277
class CustomManyParticleGenerator(object):
3278
3279
3280
3281
3282
3283
3284
3285
3286
3287
3288
3289
    """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 = []
3290

3291
3292
3293
3294
3295
3296
3297
3298
3299
3300
3301
3302
    @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(',')]))
3303
3304
        generator.params = ForceField._AtomTypeParameters(ff, 'CustomManyParticleForce', 'Atom', generator.perParticleParams)
        generator.params.parseDefinitions(element)
3305
3306
3307
3308

    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
        methodMap = {NoCutoff:mm.CustomManyParticleForce.NoCutoff,
                     CutoffNonPeriodic:mm.CustomManyParticleForce.CutoffNonPeriodic,
3309
3310
3311
3312
                     CutoffPeriodic:mm.CustomManyParticleForce.CutoffPeriodic,
                     Ewald:mm.CustomManyParticleForce.CutoffPeriodic,
                     PME:mm.CustomManyParticleForce.CutoffPeriodic,
                     LJPME:mm.CustomManyParticleForce.CutoffPeriodic}
3313
3314
3315
3316
3317
3318
3319
3320
3321
3322
3323
3324
        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':
3325
                force.addTabulatedFunction(name, mm.Continuous1DFunction(values, params['min'], params['max'], params['periodic']))
3326
3327
3328
3329
3330
3331
3332
3333
3334
3335
3336
            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:
3337
3338
3339
            values = self.params.getAtomParameters(atom, data)
            type = int(self.params.getExtraParameters(atom, data)['filterType'])
            force.addParticle(values, type)
3340
3341
3342
3343
3344
3345
        force.setNonbondedMethod(methodMap[nonbondedMethod])
        force.setCutoffDistance(nonbondedCutoff)
        sys.addForce(force)

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

3347
3348
3349
3350
3351
3352
3353
3354
3355
3356
3357
        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)))
3358

3359
3360
        # 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.
3361

3362
3363
3364
3365
3366
3367
3368
3369
3370
        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.
3371

3372
3373
3374
3375
3376
        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
3377
def getAtomPrint(data, atomIndex):
3378

Peter Eastman's avatar
Peter Eastman committed
3379
3380
3381
    if (atomIndex < len(data.atoms)):
        atom = data.atoms[atomIndex]
        returnString = "%4s %4s %5d" % (atom.name, atom.residue.name, atom.residue.index)
3382
    else:
Peter Eastman's avatar
Peter Eastman committed
3383
        returnString = "NA"
3384
3385
3386
3387
3388

    return returnString

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

Peter Eastman's avatar
Peter Eastman committed
3389
def countConstraint(data):
3390

Peter Eastman's avatar
Peter Eastman committed
3391
    bondCount = 0
3392
3393
3394
3395
3396
3397
3398
    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
3399
        if (isConstrained):
3400
            angleCount += 1
Justin MacCallum's avatar
Justin MacCallum committed
3401

3402
    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
3403

3404
## @private
3405
class AmoebaBondGenerator(object):
3406
3407
3408

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

3409
    """An AmoebaBondGenerator constructs a AmoebaBondForce."""
3410
3411

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

3413
3414
    def __init__(self, cubic, quartic):

Peter Eastman's avatar
Peter Eastman committed
3415
3416
3417
3418
3419
3420
        self.cubic = cubic
        self.quartic = quartic
        self.types1 = []
        self.types2 = []
        self.length = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
3421

3422
3423
3424
3425
3426
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3430
        generator = AmoebaBondGenerator(element.attrib['bond-cubic'], element.attrib['bond-quartic'])
3431
3432
        forceField._forces.append(generator)
        for bond in element.findall('Bond'):
3433
            types = forceField._findAtomTypes(bond.attrib, 2)
peastman's avatar
peastman committed
3434
            if None not in types:
3435
3436
3437
3438
3439
                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:
3440
                outputString = "AmoebaBondGenerator: error getting types: %s %s" % (
3441
                                    bond.attrib['class1'],
Peter Eastman's avatar
Peter Eastman committed
3442
                                    bond.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
3443
3444
                raise ValueError(outputString)

3445
3446
    #=============================================================================================

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

3449
3450
        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]
3451
        if len(existing) == 0:
3452
3453
3454
            force = mm.CustomBondForce(energy)
            force.addPerBondParameter('r0')
            force.addPerBondParameter('k')
3455
            force.setName('AmoebaBond')
3456
3457
3458
3459
3460
3461
3462
3463
3464
3465
3466
3467
3468
            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:
3469
                        data.addConstraint(sys, bond.atom1, bond.atom2, self.length[i])
3470
3471
                    if self.k[i] != 0:
                        if not bond.isConstrained or args.get('flexibleConstraints', False):
3472
                            force.addBond(bond.atom1, bond.atom2, [self.length[i], self.k[i]])
3473
3474
                    break

3475
parsers["AmoebaBondForce"] = AmoebaBondGenerator.parseElement
3476
3477
3478
3479

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

Peter Eastman's avatar
Peter Eastman committed
3481
def addAngleConstraint(angle, idealAngle, data, sys):
3482
3483

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

Peter Eastman's avatar
Peter Eastman committed
3485
3486
    bond1 = None
    bond2 = None
3487
3488
3489
3490
3491
3492
3493
    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
3494

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

3497
3498
3499
3500
        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
3501
                length = sqrt(l1*l1 + l2*l2 - 2*l1*l2*cos(idealAngle))
3502
                data.addConstraint(sys, angle[0], angle[2], length)
3503
3504
3505
                return

#=============================================================================================
3506
## @private
3507
class AmoebaAngleGenerator(object):
3508
3509

    #=============================================================================================
3510
    """An AmoebaAngleGenerator constructs a AmoebaAngleForce."""
3511
    #=============================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
3512

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

Peter Eastman's avatar
Peter Eastman committed
3515
3516
3517
3518
3519
        self.forceField = forceField
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3520

Peter Eastman's avatar
Peter Eastman committed
3521
3522
3523
        self.types1 = []
        self.types2 = []
        self.types3 = []
3524

Peter Eastman's avatar
Peter Eastman committed
3525
3526
        self.angle = []
        self.k = []
3527
        self.inPlane = []
Justin MacCallum's avatar
Justin MacCallum committed
3528

3529
3530
3531
3532
3533
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

3537
3538
3539
3540
3541
3542
3543
3544
        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')
3545
        for angle in element.findall('Angle'):
3546
            types = forceField._findAtomTypes(angle.attrib, 3)
peastman's avatar
peastman committed
3547
            if None not in types:
3548
3549
3550
3551
3552
3553

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

                angleList = []
Peter Eastman's avatar
Peter Eastman committed
3554
                angleList.append(float(angle.attrib['angle1']))
3555

3556
                if 'angle2' in angle.attrib:
Peter Eastman's avatar
Peter Eastman committed
3557
                    angleList.append(float(angle.attrib['angle2']))
3558
3559
3560
                if 'angle3' in angle.attrib:
                    angleList.append(float(angle.attrib['angle3']))

3561
3562
                generator.angle.append(angleList)
                generator.k.append(float(angle.attrib['k']))
3563
3564
3565
3566
                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
3567
            else:
3568
                outputString = "AmoebaAngleGenerator: error getting types: %s %s %s" % (
3569
3570
                                    angle.attrib['class1'],
                                    angle.attrib['class2'],
Peter Eastman's avatar
Peter Eastman committed
3571
                                    angle.attrib['class3'])
Justin MacCallum's avatar
Justin MacCallum committed
3572
3573
                raise ValueError(outputString)

3574
3575
3576
3577
    #=============================================================================================
    # 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
3578

Peter Eastman's avatar
Peter Eastman committed
3579
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3580
3581
        if not any(isinstance(f, AmoebaOutOfPlaneBendGenerator) for f in self.forceField.getGenerators()):
            raise ValueError('A ForceField containing an <AmoebaAngleForce> must also contain an <AmoebaOutOfPlaneBendForce>')
3582
3583
3584
3585
3586

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

Peter Eastman's avatar
Peter Eastman committed
3588
    def createForcePostOpBendAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3589
3590
3591

        # get force

3592
3593
        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]
3594
3595

        if len(existing) == 0:
3596
3597
3598
            force = mm.CustomAngleForce(energy)
            force.addPerAngleParameter('theta0')
            force.addPerAngleParameter('k')
3599
            force.setName('AmoebaAngle')
3600
3601
3602
3603
            sys.addForce(force)
        else:
            force = existing[0]

3604
3605
        DEG_TO_RAD = math.pi / 180

3606
        for angleDict in angleList:
Peter Eastman's avatar
Peter Eastman committed
3607
3608
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3609
            inPlane = angleDict['inPlane']
3610

Peter Eastman's avatar
Peter Eastman committed
3611
3612
3613
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3614
            for i in range(len(self.types1)):
3615
3616
3617
                # 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
3618
3619
3620
3621
3622
3623
                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]
3624
3625
                        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
3626
3627
                        lenAngle = len(self.angle[i])
                        if (lenAngle > 1):
3628
3629
3630
3631
3632
3633
                            # 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
3634
                                if (atom1 == angle[1] and atom2 != angle[0] and atom2 != angle[2] and (sys.getParticleMass(atom2)/unit.dalton) < 1.90):
3635
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3636
                                if (atom2 == angle[1] and atom1 != angle[0] and atom1 != angle[2] and (sys.getParticleMass(atom1)/unit.dalton) < 1.90):
3637
                                    numberOfHydrogens += 1
Peter Eastman's avatar
Peter Eastman committed
3638
                            if (numberOfHydrogens < lenAngle):
3639
3640
                                angleValue =  self.angle[i][numberOfHydrogens]
                            else:
3641
                                outputString = "AmoebaAngleGenerator angle index=%d is out of range: [0, %5d] " % (numberOfHydrogens, lenAngle)
Justin MacCallum's avatar
Justin MacCallum committed
3642
                                raise ValueError(outputString)
3643
3644
                        else:
                            angleValue =  self.angle[i][0]
Justin MacCallum's avatar
Justin MacCallum committed
3645

3646
                        angleDict['idealAngle'] = angleValue
3647
                        force.addAngle(angle[0], angle[1], angle[2], [angleValue, self.k[i]])
3648
3649
3650
3651
3652
3653
                    break

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

Peter Eastman's avatar
Peter Eastman committed
3655
    def createForcePostOpBendInPlaneAngle(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
3656
3657
3658

        # get force

3659
3660
3661
3662
3663
3664
3665
3666
3667
3668
        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]
3669
        if len(existing) == 0:
3670
3671
3672
            force = mm.CustomCompoundBondForce(4, energy)
            force.addPerBondParameter("theta0")
            force.addPerBondParameter("k")
3673
            force.setName('AmoebaInPlaneAngle')
3674
3675
3676
3677
3678
            sys.addForce(force)
        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
3680
3681
            angle = angleDict['angle']
            isConstrained = angleDict['isConstrained']
3682
            inPlane = angleDict['inPlane']
3683

Peter Eastman's avatar
Peter Eastman committed
3684
3685
3686
            type1 = data.atomType[data.atoms[angle[0]]]
            type2 = data.atomType[data.atoms[angle[1]]]
            type3 = data.atomType[data.atoms[angle[2]]]
3687
3688

            for i in range(len(self.types1)):
3689
3690
3691
                # 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
3692
3693
3694
3695
3696
3697
                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
3698
                    if (isConstrained and self.k[i] != 0.0):
3699
                        addAngleConstraint(angle, self.angle[i][0]*math.pi/180.0, data, sys)
3700
                    if self.k[i] != 0.0 and (not isConstrained or args.get('flexibleConstraints', False)):
3701
                        force.addBond((angle[0], angle[1], angle[2], angle[3]), (self.angle[i][0], self.k[i]))
3702
3703
                    break

3704
parsers["AmoebaAngleForce"] = AmoebaAngleGenerator.parseElement
3705
3706
3707

#=============================================================================================
# Generator for the AmoebaOutOfPlaneBend covalent force; also calls methods in the
3708
3709
# AmoebaAngleGenerator to generate the AmoebaAngleForce and
# AmoebaInPlaneAngleForce
3710
3711
#=============================================================================================

3712
## @private
3713
class AmoebaOutOfPlaneBendGenerator(object):
3714
3715
3716
3717

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

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

3719
3720
3721
3722
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3723
3724
3725
3726
3727
3728
        self.forceField = forceField
        self.type = type
        self.cubic = cubic
        self.quartic = quartic
        self.pentic = pentic
        self.sextic = sextic
3729

Peter Eastman's avatar
Peter Eastman committed
3730
3731
3732
3733
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3734

Peter Eastman's avatar
Peter Eastman committed
3735
        self.ks = []
3736
3737
3738
3739
3740

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

3742
3743
3744
3745
3746
3747
3748
3749
3750
3751
3752
3753
3754
3755
3756
3757
3758
3759
3760
3761
3762
3763
3764
3765
3766
    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
3767

3768
3769
        # get global scalar parameters

3770
3771
3772
3773
3774
3775
3776
3777
        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')
3778
3779

        for angle in element.findall('Angle'):
3780
3781
3782
3783
            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
3784
            types = generator.findAtomTypes(forceField, angle, 4)
3785
3786
            if types is not None:

Peter Eastman's avatar
Peter Eastman committed
3787
3788
3789
3790
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
                generator.types4.append(types[3])
3791

Peter Eastman's avatar
Peter Eastman committed
3792
                generator.ks.append(float(angle.attrib['k']))
3793
3794

            else:
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3795
3796
                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
3797
3798
                raise ValueError(outputString)

3799
3800
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
3801
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
3802
3803
        self._nonbondedMethod = nonbondedMethod
        self._nonbondedCutoff = nonbondedCutoff
3804

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

3808
3809
3810
3811
3812
3813
3814
3815
3816
3817
        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]
3818
        if len(existing) == 0:
3819
3820
            force = mm.CustomCompoundBondForce(4, energy)
            force.addPerBondParameter("k")
3821
            force.setName('AmoebaOutOfPlaneBend')
3822
3823
3824
        else:
            force = existing[0]

3825
        # this hash is used to ensure the out-of-plane-bend bonds
3826
3827
        # are only added once

Peter Eastman's avatar
Peter Eastman committed
3828
        skipAtoms = dict()
3829
        angles = []
3830

3831
3832
3833
3834
3835
3836
3837
        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
3838
3839
3840

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

3841
3842
3843
            middleAtom = angle[1]
            middleType = data.atomType[data.atoms[middleAtom]]
            middleCovalency = len(data.atomBonds[middleAtom])
3844

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

3851
            if middleCovalency == 3 and middleAtom not in skipAtoms:
3852

Peter Eastman's avatar
Peter Eastman committed
3853
                partners = []
3854
3855
3856
3857

                for bond in data.atomBonds[middleAtom]:
                    atom1 = data.bonds[bond].atom1
                    atom2 = data.bonds[bond].atom2
3858
                    if atom1 != middleAtom:
3859
3860
3861
3862
3863
3864
3865
3866
                        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]
3867
                        if middleType in types2 and partnerType in types1:
Peter Eastman's avatar
Peter Eastman committed
3868
                            partners.append(partner)
3869
                            break
Justin MacCallum's avatar
Justin MacCallum committed
3870

3871
                if len(partners) == 3:
3872

3873
3874
3875
                    addBond([partners[0], middleAtom, partners[1], partners[2]])
                    addBond([partners[2], middleAtom, partners[0], partners[1]])
                    addBond([partners[1], middleAtom, partners[2], partners[0]])
3876

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

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
3881
3882
                    # in-plane angle

Peter Eastman's avatar
Peter Eastman committed
3883
                    angleDict = {}
3884
3885
3886
3887
                    angleList = list(angle[:3])
                    for atomIndex in partners:
                        if atomIndex not in angleList:
                            angleList.append(atomIndex)
3888
                    angleDict['angle'] = angleList
Peter Eastman's avatar
Peter Eastman committed
3889
                    angleDict['isConstrained'] = 0
3890
3891
                    angleDict['inPlane'] = True
                    angles.append(angleDict)
3892
3893

                else:
Peter Eastman's avatar
Peter Eastman committed
3894
                    angleDict = {}
3895
3896
3897
3898
                    angleList = list(angle[:3])
                    for atomIndex in partners:
                        if atomIndex not in angleList:
                            angleList.append(atomIndex)
Peter Eastman's avatar
Peter Eastman committed
3899
3900
                    angleDict['angle'] = angleList
                    angleDict['isConstrained'] = isConstrained
3901
3902
3903
3904
3905
3906
3907
3908
3909
3910
3911
3912
3913
                    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)
3914

3915
3916
3917
3918
3919
3920
            else:
                angleDict = {}
                angleDict['angle'] = angle
                angleDict['isConstrained'] = isConstrained
                angleDict['inPlane'] = False
                angles.append(angleDict)
3921

3922
3923
        if len(existing) == 0 and force.getNumBonds() > 0:
            sys.addForce(force)
3924

3925
        # get AmoebaAngleGenerator and add AmoebaAngle and AmoebaInPlaneAngle forces
3926
3927

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3928
            if (force.__class__.__name__ == 'AmoebaAngleGenerator'):
3929
3930
                force.createForcePostOpBendAngle(sys, data, self._nonbondedMethod, self._nonbondedCutoff, angles, args)
                force.createForcePostOpBendInPlaneAngle(sys, data, self._nonbondedMethod, self._nonbondedCutoff, angles, args)
3931
3932

        for force in self.forceField._forces:
Justin MacCallum's avatar
Justin MacCallum committed
3933
            if (force.__class__.__name__ == 'AmoebaStretchBendGenerator'):
3934
                force.createForcePostAmoebaBondForce(sys, data, self._nonbondedMethod, self._nonbondedCutoff, angles, args)
3935
3936
3937
3938
3939

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

3940
## @private
3941
class AmoebaTorsionGenerator(object):
3942
3943
3944
3945
3946
3947
3948

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

    def __init__(self, torsionUnit):

Peter Eastman's avatar
Peter Eastman committed
3949
        self.torsionUnit = torsionUnit
3950

Peter Eastman's avatar
Peter Eastman committed
3951
3952
3953
3954
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
3955

Peter Eastman's avatar
Peter Eastman committed
3956
3957
3958
        self.t1 = []
        self.t2 = []
        self.t3 = []
Justin MacCallum's avatar
Justin MacCallum committed
3959

3960
3961
3962
3963
3964
3965
3966
3967
    #=============================================================================================

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

Peter Eastman's avatar
Peter Eastman committed
3969
        generator = AmoebaTorsionGenerator(float(element.attrib['torsionUnit']))
3970
3971
3972
3973
3974
3975
        forceField._forces.append(generator)

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

        for torsion in element.findall('Torsion'):
3976
            types = forceField._findAtomTypes(torsion.attrib, 4)
peastman's avatar
peastman committed
3977
            if None not in types:
3978
3979
3980
3981
3982
3983
3984

                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
3985
3986
                    tInfo = []
                    suffix = str(ii)
3987
3988
3989
3990
3991
3992
                    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
3993
3994
3995
3996
3997
3998
                    if (ii == 1):
                        generator.t1.append(tInfo)
                    elif (ii == 2):
                        generator.t2.append(tInfo)
                    elif (ii == 3):
                        generator.t3.append(tInfo)
3999
4000
4001

            else:
                outputString = "AmoebaTorsionGenerator: error getting types: %s %s %s %s" % (
4002
4003
4004
4005
                                    torsion.attrib['class1'],
                                    torsion.attrib['class2'],
                                    torsion.attrib['class3'],
                                    torsion.attrib['class4'])
Justin MacCallum's avatar
Justin MacCallum committed
4006
4007
                raise ValueError(outputString)

4008
4009
4010
4011
    #=============================================================================================

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

4012
        existing = [f for f in sys.getForces() if type(f) == mm.PeriodicTorsionForce]
4013
        if len(existing) == 0:
4014
            force = mm.PeriodicTorsionForce()
4015
4016
4017
            sys.addForce(force)
        else:
            force = existing[0]
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4018

4019
4020
4021
4022
4023
4024
4025
4026
4027
4028
4029
4030
4031
4032
4033
4034
        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
4035
                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):
4036
4037
4038
4039
4040
4041
                    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])
4042
4043
4044
4045
4046
4047
                    break

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

4048
## @private
4049
class AmoebaPiTorsionGenerator(object):
4050
4051
4052
4053
4054
4055

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

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

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

4057
    def __init__(self, piTorsionUnit):
Justin MacCallum's avatar
Justin MacCallum committed
4058
        self.piTorsionUnit = piTorsionUnit
Peter Eastman's avatar
Peter Eastman committed
4059
4060
4061
        self.types1 = []
        self.types2 = []
        self.k = []
Justin MacCallum's avatar
Justin MacCallum committed
4062

4063
4064
4065
4066
4067
4068
4069
4070
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

Peter Eastman's avatar
Peter Eastman committed
4071
        generator = AmoebaPiTorsionGenerator(float(element.attrib['piTorsionUnit']))
4072
4073
4074
        forceField._forces.append(generator)

        for piTorsion in element.findall('PiTorsion'):
4075
            types = forceField._findAtomTypes(piTorsion.attrib, 2)
peastman's avatar
peastman committed
4076
            if None not in types:
4077
4078
4079
4080
4081
4082
                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
4083
                                    piTorsion.attrib['class2'])
Justin MacCallum's avatar
Justin MacCallum committed
4084
4085
                raise ValueError(outputString)

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

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

4090
4091
4092
4093
4094
4095
4096
4097
4098
        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]
4099
4100

        if len(existing) == 0:
4101
4102
            force = mm.CustomCompoundBondForce(6, energy)
            force.addPerBondParameter('k')
4103
            force.setName('AmoebaPiTorsion')
4104
4105
4106
        else:
            force = existing[0]

Peter Eastman's avatar
Peter Eastman committed
4107
        for bond1 in data.bonds:
4108
4109

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

Peter Eastman's avatar
Peter Eastman committed
4111
4112
            atom1 = bond1.atom1
            atom2 = bond1.atom2
Justin MacCallum's avatar
Justin MacCallum committed
4113

4114
            if (len(data.atomBonds[atom1]) == 3 and len(data.atomBonds[atom2]) == 3):
4115
4116
4117
4118
4119
4120
4121
4122
4123
4124
4125

                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
4126
4127
                       # piTorsionAtom1, piTorsionAtom2 are the atoms bonded to atom1, excluding atom2
                       # piTorsionAtom5, piTorsionAtom6 are the atoms bonded to atom2, excluding atom1
4128
4129
4130
4131
4132
4133
4134
4135
4136
4137
4138
4139

                       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
4140
                           if (bondedAtom1 != atom1):
4141
4142
4143
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
4144
4145
                           if (b1 != atom2):
                               if (piTorsionAtom1 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
4146
                                   piTorsionAtom1 = b1
4147
4148
4149
4150
4151
4152
                               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
4153
                           if (bondedAtom1 != atom2):
4154
4155
4156
4157
                               b1 = bondedAtom1
                           else:
                               b1 = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4158
4159
                           if (b1 != atom1):
                               if (piTorsionAtom5 == -1):
Justin MacCallum's avatar
Justin MacCallum committed
4160
                                   piTorsionAtom5 = b1
4161
4162
                               else:
                                   piTorsionAtom6 = b1
Justin MacCallum's avatar
Justin MacCallum committed
4163

4164
                       force.addBond([piTorsionAtom1, piTorsionAtom2, piTorsionAtom3, piTorsionAtom4, piTorsionAtom5, piTorsionAtom6], [self.k[i]])
4165
4166
        if len(existing) == 0 and force.getNumBonds() > 0:
            sys.addForce(force)
4167
4168
4169
4170
4171

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

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
4303
4304
4305
4306
4307
4308
4309
4310
4311
4312
4313
4314
4315
4316
4317
4318
4319
4320
4321
4322
4323
4324
4325
4326
4327
4328
4329
4330
4331
4332
4333
4334
4335
4336
4337
4338
4339
4340
4341
4342
4343
4344
4345
4346
4347
4348
4349
4350
4351
4352
4353
4354
4355
4356
4357
## @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

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

4358
## @private
4359
class AmoebaTorsionTorsionGenerator(object):
4360
4361
4362
4363
4364

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

Peter Eastman's avatar
Peter Eastman committed
4365
    def __init__(self):
4366

Peter Eastman's avatar
Peter Eastman committed
4367
4368
4369
4370
4371
        self.types1 = []
        self.types2 = []
        self.types3 = []
        self.types4 = []
        self.types5 = []
4372

Peter Eastman's avatar
Peter Eastman committed
4373
        self.gridIndex = []
4374

Peter Eastman's avatar
Peter Eastman committed
4375
        self.grids = []
Justin MacCallum's avatar
Justin MacCallum committed
4376

4377
4378
4379
4380
4381
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

Peter Eastman's avatar
Peter Eastman committed
4382
        generator = AmoebaTorsionTorsionGenerator()
4383
4384
4385
4386
4387
4388
4389
        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'):
4390
            types = forceField._findAtomTypes(torsionTorsion.attrib, 5)
peastman's avatar
peastman committed
4391
            if None not in types:
4392
4393
4394
4395
4396
4397
4398

                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
4399
4400
                gridIndex = int(torsionTorsion.attrib['grid'])
                if (gridIndex > maxGridIndex):
4401
4402
4403
4404
4405
4406
4407
4408
4409
                    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
4410
                                    torsionTorsion.attrib['class5'] )
Justin MacCallum's avatar
Justin MacCallum committed
4411
4412
                raise ValueError(outputString)

4413
4414
4415
4416
4417
4418
4419
4420
4421
4422
        # 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
4423
4424
4425
4426
4427
4428
        #     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
4429
4430

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

Peter Eastman's avatar
Peter Eastman committed
4434
4435
4436
            gridIndex = int(torsionTorsionGrid.attrib[ "grid"])
            nx = int(torsionTorsionGrid.attrib[ "nx"])
            ny = int(torsionTorsionGrid.attrib[ "ny"])
4437

Peter Eastman's avatar
Peter Eastman committed
4438
4439
            grid = []
            gridCol = []
4440
4441
4442
4443
4444

            gridColIndex = 0

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

Peter Eastman's avatar
Peter Eastman committed
4445
4446
4447
4448
                gridRow = []
                gridRow.append(float(gridEntry.attrib['angle1']))
                gridRow.append(float(gridEntry.attrib['angle2']))
                gridRow.append(float(gridEntry.attrib['f']))
4449
                if 'fx' in gridEntry.attrib:
4450
4451
4452
                    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
4453
                gridCol.append(gridRow)
4454
4455

                gridColIndex  += 1
Peter Eastman's avatar
Peter Eastman committed
4456
4457
4458
                if (gridColIndex == nx):
                    grid.append(gridCol)
                    gridCol = []
4459
4460
                    gridColIndex = 0

Justin MacCallum's avatar
Justin MacCallum committed
4461

Peter Eastman's avatar
Peter Eastman committed
4462
4463
            if (gridIndex == len(generator.grids)):
                generator.grids.append(grid)
4464
            else:
Peter Eastman's avatar
Peter Eastman committed
4465
4466
                while(len(generator.grids) < gridIndex):
                    generator.grids.append([])
4467
4468
4469
4470
                generator.grids[gridIndex] = grid

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

Peter Eastman's avatar
Peter Eastman committed
4471
    def getChiralAtomIndex(self, data, sys, atomB, atomC, atomD):
4472
4473
4474
4475
4476
4477
4478
4479

        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
4480
        if (len(data.atomBonds[atomC]) == 4):
4481
4482
4483
4484
4485
            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
4486
4487
                hit = -1
                if (  bondedAtom1 == atomC and bondedAtom2 != atomB and bondedAtom2 != atomD):
4488
                    hit = bondedAtom2
Peter Eastman's avatar
Peter Eastman committed
4489
                elif (bondedAtom2 == atomC and bondedAtom1 != atomB and bondedAtom1 != atomD):
4490
4491
                    hit = bondedAtom1

Peter Eastman's avatar
Peter Eastman committed
4492
4493
                if (hit > -1):
                    if (atomE == -1):
4494
4495
4496
                        atomE = hit
                    else:
                        atomF = hit
Justin MacCallum's avatar
Justin MacCallum committed
4497

4498
4499
            # raise error if atoms E or F not found

Peter Eastman's avatar
Peter Eastman committed
4500
            if (atomE == -1 or atomF == -1):
4501
                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
4502
                raise ValueError(outputString)
4503
4504
4505
4506
4507

            # 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
4508
            if (typeE > typeF):
Justin MacCallum's avatar
Justin MacCallum committed
4509
                chiralAtomIndex = atomE
Peter Eastman's avatar
Peter Eastman committed
4510
            if (typeF > typeE):
Justin MacCallum's avatar
Justin MacCallum committed
4511
                chiralAtomIndex = atomF
4512

Peter Eastman's avatar
Peter Eastman committed
4513
4514
4515
            massE = sys.getParticleMass(atomE)/unit.dalton
            massF = sys.getParticleMass(atomE)/unit.dalton
            if (massE > massF):
Justin MacCallum's avatar
Justin MacCallum committed
4516
                chiralAtomIndex = massE
Peter Eastman's avatar
Peter Eastman committed
4517
            if (massF > massE):
Justin MacCallum's avatar
Justin MacCallum committed
4518
                chiralAtomIndex = massF
4519
4520
4521
4522

        return chiralAtomIndex

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

4524
4525
    def createForce(self, sys, data, nonpiTorsionedMethod, nonpiTorsionedCutoff, args):

4526
        existing = [f for f in sys.getForces() if type(f) == mm.AmoebaTorsionTorsionForce]
4527
4528
4529
4530
4531
4532
4533
4534

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

        for angle in data.angles:

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

4537
4538
4539
4540
4541
4542
4543
            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
4544
                if (bondedAtom1 != ib):
4545
4546
4547
4548
                    ia = bondedAtom1
                else:
                    ia = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4549
                if (ia != ic and ia != id):
Peter Eastman's avatar
Peter Eastman committed
4550
4551
4552
                    for bondIndex2 in data.atomBonds[id]:
                        bondedAtom1 = data.bonds[bondIndex2].atom1
                        bondedAtom2 = data.bonds[bondIndex2].atom2
Peter Eastman's avatar
Peter Eastman committed
4553
                        if (bondedAtom1 != id):
4554
4555
4556
4557
                            ie = bondedAtom1
                        else:
                            ie = bondedAtom2

Peter Eastman's avatar
Peter Eastman committed
4558
                        if (ie != ic and ie != ib and ie != ia):
4559
4560
4561
4562
4563
4564
4565
4566
4567
4568
4569
4570
4571
4572
4573
4574
4575
4576
4577
4578

                            # 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
4579
4580
4581
                                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])
4582
4583
4584

                                # match in reverse order

4585
                                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
4586
4587
                                    chiralAtomIndex = self.getChiralAtomIndex(data, sys, ib, ic, id)
                                    force.addTorsionTorsion(ie, id, ic, ib, ia, chiralAtomIndex, self.gridIndex[i])
4588
4589
4590
4591

        # set grids

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

4596
4597
4598
4599
parsers["AmoebaTorsionTorsionForce"] = AmoebaTorsionTorsionGenerator.parseElement

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

4600
## @private
4601
class AmoebaStretchBendGenerator(object):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4602
4603

    #=============================================================================================
4604
4605
4606
    """An AmoebaStretchBendGenerator constructs a AmoebaStretchBendForce."""
    #=============================================================================================

4607
    def __init__(self, forcefield):
4608

4609
        self.forcefield = forcefield
Peter Eastman's avatar
Peter Eastman committed
4610
4611
4612
        self.types1 = []
        self.types2 = []
        self.types3 = []
4613

Peter Eastman's avatar
Peter Eastman committed
4614
4615
        self.k1 = []
        self.k2 = []
Justin MacCallum's avatar
Justin MacCallum committed
4616

4617
4618
4619
4620
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):
4621
        generator = AmoebaStretchBendGenerator(forceField)
4622
4623
4624
4625
4626
4627
4628
        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'):
4629
            types = forceField._findAtomTypes(stretchBend.attrib, 3)
peastman's avatar
peastman committed
4630
            if None not in types:
4631
4632
4633
4634
4635
4636
4637
4638
4639
4640
4641
4642

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

4646
4647
    #=============================================================================================

Justin MacCallum's avatar
Justin MacCallum committed
4648
    # The setup of this force is dependent on AmoebaBondForce and AmoebaAngleForce
4649
4650
    # 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
4651
4652
    # AmoebaBondForce and AmoebaAngleForce have been called prior to AmoebaStretchBendGenerator().
    # Instead, createForcePostAmoebaBondForce() is called
4653
    # after the generators for AmoebaBondForce and AmoebaAngleForce have been called
4654
4655
4656

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

Peter Eastman's avatar
Peter Eastman committed
4657
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
4658
4659
        if not any(isinstance(f, AmoebaOutOfPlaneBendGenerator) for f in self.forcefield.getGenerators()):
            raise ValueError('A ForceField containing an <AmoebaStretchBendForce> must also contain an <AmoebaOutOfPlaneBendForce>')
4660
4661
4662
4663
4664
4665
4666

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

    # Note: request for constrained bonds is ignored.

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

4667
    def createForcePostAmoebaBondForce(self, sys, data, nonbondedMethod, nonbondedCutoff, angleList, args):
4668

4669
4670
        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]
4671
        if len(existing) == 0:
4672
4673
4674
4675
4676
4677
            force = mm.CustomCompoundBondForce(3, energy)
            force.addPerBondParameter("r12")
            force.addPerBondParameter("r23")
            force.addPerBondParameter("theta0")
            force.addPerBondParameter("k1")
            force.addPerBondParameter("k2")
4678
            force.setName('AmoebaStretchBend')
4679
4680
4681
4682
4683
4684
4685
4686
4687
4688
4689
4690
            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
4691
            radian = 57.2957795130
4692
4693
4694
4695
4696
4697
4698
4699
4700
4701
            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
4702
                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
4703
4704
4705
4706
4707
4708
4709
4710
4711
4712
4713
4714
4715
4716
4717
                    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
4718

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

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4721
                    if ('idealAngle' not in angleDict):
4722

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4723
4724
4725
4726
4727
                       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
4728
                       raise ValueError(outputString)
4729

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4730
4731
4732
4733
4734
4735
4736
                    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
4737
                       raise ValueError(outputString)
4738

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4739
                    else:
4740
4741
4742
4743
4744
                        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))
4745

Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4746
                    break
4747
4748
4749
4750
4751

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

4752
## @private
4753
class AmoebaVdwGenerator(object):
4754
4755

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

4757
4758
    #=============================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
4761
        self.type = type
4762

Peter Eastman's avatar
Peter Eastman committed
4763
4764
4765
        self.radiusrule = radiusrule
        self.radiustype = radiustype
        self.radiussize = radiussize
4766

Peter Eastman's avatar
Peter Eastman committed
4767
        self.epsilonrule = epsilonrule
4768

Peter Eastman's avatar
Peter Eastman committed
4769
4770
4771
        self.vdw13Scale = vdw13Scale
        self.vdw14Scale = vdw14Scale
        self.vdw15Scale = vdw15Scale
4772
4773
4774
4775
4776
4777
4778

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

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

4782
4783
4784
        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
4785
                                        float(element.attrib['vdw-13-scale']), float(element.attrib['vdw-14-scale']), float(element.attrib['vdw-15-scale']))
4786
            forceField.registerGenerator(generator)
4787
4788
            generator.params = {}
            generator.pairs = []
4789
4790
4791
4792
4793
4794
4795
4796
4797
4798
        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')
4799
4800
4801
4802
4803
        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())
4804
4805
4806
4807

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

    @staticmethod
4808
    def getBondedParticleSets(sys, data):
4809

4810
4811
4812
4813
4814
4815
        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
4816

4817
4818
    #=============================================================================================

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

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

4825
4826
        force = mm.AmoebaVdwForce()
        sys.addForce(force)
4827

4828
4829
4830
4831
4832
4833
4834
4835
        # 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))

4836
        # sigma and epsilon combining rules
4837

4838
4839
4840
4841
        if ('sigmaCombiningRule' in args):
            sigmaRule = args['sigmaCombiningRule'].upper()
            if (sigmaRule.upper() in sigmaMap):
                force.setSigmaCombiningRule(sigmaRule.upper())
4842
            else:
4843
4844
4845
4846
                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)
4847

4848
4849
4850
4851
        if ('epsilonCombiningRule' in args):
            epsilonRule = args['epsilonCombiningRule'].upper()
            if (epsilonRule.upper() in epsilonMap):
                force.setEpsilonCombiningRule(epsilonRule.upper())
4852
            else:
4853
4854
4855
4856
                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
4857

4858
        # cutoff
4859

4860
4861
4862
4863
        if ('vdwCutoff' in args):
            force.setCutoff(args['vdwCutoff'])
        else:
            force.setCutoff(nonbondedCutoff)
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
4864

4865
        # dispersion correction
4866

4867
4868
        if ('useDispersionCorrection' in args):
            force.setUseDispersionCorrection(bool(args['useDispersionCorrection']))
4869

4870
4871
        if (nonbondedMethod == PME):
            force.setNonbondedMethod(mm.AmoebaVdwForce.CutoffPeriodic)
4872

4873
        # Define types
4874

4875
4876
4877
4878
4879
        sigmaScale = 1
        if self.radiustype == 'SIGMA':
            sigmaScale = 1.122462048309372
        if self.radiussize == 'DIAMETER':
            sigmaScale = 0.5
4880
4881
4882
4883
4884
4885
4886
        classToTypeMap = {}
        for className in self.params:
            sigma, epsilon, _ = self.params[className]
            classToTypeMap[className] = force.addParticleType(sigma*sigmaScale, epsilon)

        # add particles to force

4887
        for (i, atom) in enumerate(data.atoms):
4888
4889
            className = self.classNameForType[data.atomType[atom]]
            _, _, reduction = self.params[className]
4890
4891
4892
4893
4894
4895
4896
4897
4898
            # 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
4899

4900
4901
4902
            force.addParticle(ivIndex, classToTypeMap[className], reduction)

        # Add pairs
4903

4904
4905
        for c1, c2, sigma, epsilon in self.pairs:
            force.addTypePair(classToTypeMap[c1], classToTypeMap[c2], sigma, epsilon)
4906
4907
4908
4909
4910
4911
4912

        # 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

4913
        bondedParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
4914
4915

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

4917
4918
4919
4920
4921
4922
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

Peter Eastman's avatar
Peter Eastman committed
4923
            if (self.vdw13Scale == 0.0):
4924
                for bondedParticle in bondedParticleSets[i]:
Peter Eastman's avatar
Peter Eastman committed
4925
                    exclusionSet = exclusionSet.union(bondedParticleSets[bondedParticle])
4926
4927
4928
4929
4930

            # self

            exclusionSet.add(i)

4931
            force.setParticleExclusions(i, tuple(exclusionSet))
4932
4933
4934
4935
4936

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

4937
## @private
4938
class AmoebaMultipoleGenerator(object):
4939
4940
4941
4942

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

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

4944
4945
    #=============================================================================================

4946
    def __init__(self, forceField):
Peter Eastman's avatar
Peter Eastman committed
4947
4948
        self.forceField = forceField
        self.typeMap = {}
4949
4950
4951
4952
4953
4954

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

    @staticmethod
Peter Eastman's avatar
Peter Eastman committed
4955
    def setAxisType(kIndices):
4956
4957
4958

                # set axis type

Peter Eastman's avatar
Peter Eastman committed
4959
4960
4961
                kIndicesLen = len(kIndices)
                if (kIndicesLen > 3):
                    ky = kIndices[3]
4962
                else:
Peter Eastman's avatar
Peter Eastman committed
4963
                    ky = 0
Justin MacCallum's avatar
Justin MacCallum committed
4964

Peter Eastman's avatar
Peter Eastman committed
4965
4966
                if (kIndicesLen > 2):
                    kx = kIndices[2]
4967
                else:
Peter Eastman's avatar
Peter Eastman committed
4968
                    kx = 0
Justin MacCallum's avatar
Justin MacCallum committed
4969

Peter Eastman's avatar
Peter Eastman committed
4970
4971
                if (kIndicesLen > 1):
                    kz = kIndices[1]
4972
                else:
Peter Eastman's avatar
Peter Eastman committed
4973
                    kz = 0
4974

Peter Eastman's avatar
Peter Eastman committed
4975
4976
                while(len(kIndices) < 4):
                    kIndices.append(0)
4977
4978

                axisType = mm.AmoebaMultipoleForce.ZThenX
Peter Eastman's avatar
Peter Eastman committed
4979
                if (kz == 0):
4980
                    axisType = mm.AmoebaMultipoleForce.NoAxisType
Peter Eastman's avatar
Peter Eastman committed
4981
                if (kz != 0 and kx == 0):
4982
                    axisType = mm.AmoebaMultipoleForce.ZOnly
Peter Eastman's avatar
Peter Eastman committed
4983
                if (kz < 0 or kx < 0):
4984
                    axisType = mm.AmoebaMultipoleForce.Bisector
Peter Eastman's avatar
Peter Eastman committed
4985
                if (kx < 0 and ky < 0):
4986
                    axisType = mm.AmoebaMultipoleForce.ZBisect
Peter Eastman's avatar
Peter Eastman committed
4987
                if (kz < 0 and kx < 0 and ky  < 0):
4988
4989
                    axisType = mm.AmoebaMultipoleForce.ThreeFold

Justin MacCallum's avatar
Justin MacCallum committed
4990
4991
4992
                kIndices[1] = abs(kz)
                kIndices[2] = abs(kx)
                kIndices[3] = abs(ky)
4993
4994
4995
4996
4997
4998
4999
5000

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

Justin MacCallum's avatar
Justin MacCallum committed
5001
        #   <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"  >
5002
5003
5004
        # <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"  />

5005
5006
5007
5008
5009
5010
5011
        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]
5012
5013
5014
5015

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

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

                # k-indices not provided default to 0

                kIndices = [int(atom.attrib['type'])]

                kStrings = [ 'kz', 'kx', 'ky' ]
                for kString in kStrings:
5025
                    kIndices.append(int(atom.attrib.get(kString,0)))
5026

Justin MacCallum's avatar
Justin MacCallum committed
5027
                # set axis type based on k-Indices
5028

Peter Eastman's avatar
Peter Eastman committed
5029
                axisType = AmoebaMultipoleGenerator.setAxisType(kIndices)
5030
5031
5032

                # set multipole

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

Peter Eastman's avatar
Peter Eastman committed
5035
5036
5037
5038
5039
5040
5041
5042
5043
5044
5045
5046
5047
                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']))
5048
5049

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
5050
                    if (t not in generator.typeMap):
5051
5052
                        generator.typeMap[t] = []

Peter Eastman's avatar
Peter Eastman committed
5053
5054
5055
                    valueMap = dict()
                    valueMap['classIndex'] = atom.attrib['type']
                    valueMap['kIndices'] = kIndices
Justin MacCallum's avatar
Justin MacCallum committed
5056
                    valueMap['charge'] = charge
Peter Eastman's avatar
Peter Eastman committed
5057
5058
5059
5060
                    valueMap['dipole'] = dipole
                    valueMap['quadrupole'] = quadrupole
                    valueMap['axisType'] = axisType
                    generator.typeMap[t].append(valueMap)
Justin MacCallum's avatar
Justin MacCallum committed
5061

5062
            else:
Peter Eastman's avatar
Peter Eastman committed
5063
                outputString = "AmoebaMultipoleGenerator: error getting type for multipole: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
5064
5065
                raise ValueError(outputString)

5066
        # polarization parameters
Justin MacCallum's avatar
Justin MacCallum committed
5067

5068
        for atom in element.findall('Polarize'):
5069
            types = forceField._findAtomTypes(atom.attrib, 1)
peastman's avatar
peastman committed
5070
            if None not in types:
5071

Peter Eastman's avatar
Peter Eastman committed
5072
5073
5074
5075
                classIndex = atom.attrib['type']
                polarizability = float(atom.attrib['polarizability'])
                thole = float(atom.attrib['thole'])
                if (thole == 0):
5076
5077
                    pdamp = 0
                else:
Peter Eastman's avatar
Peter Eastman committed
5078
                    pdamp = pow(polarizability, 1.0/6.0)
5079

Peter Eastman's avatar
Peter Eastman committed
5080
5081
                pgrpMap = dict()
                for index in range(1, 7):
5082
                    pgrp = 'pgrp' + str(index)
Peter Eastman's avatar
Peter Eastman committed
5083
                    if (pgrp in atom.attrib):
5084
5085
5086
                        pgrpMap[int(atom.attrib[pgrp])] = -1

                for t in types[0]:
Peter Eastman's avatar
Peter Eastman committed
5087
5088
                    if (t not in generator.typeMap):
                        outputString = "AmoebaMultipoleGenerator: polarize type not present: %s" % (atom.attrib['type'])
Justin MacCallum's avatar
Justin MacCallum committed
5089
                        raise ValueError(outputString)
5090
5091
                    else:
                        typeMapList = generator.typeMap[t]
Peter Eastman's avatar
Peter Eastman committed
5092
5093
5094
5095
5096
5097
5098
                        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
5099
                                typeMap['pgrpMap'] = pgrpMap
Peter Eastman's avatar
Peter Eastman committed
5100
5101
5102
5103
5104
                                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
5105
5106
                            raise ValueError(outputString)

5107
            else:
Peter Eastman's avatar
Peter Eastman committed
5108
                outputString = "AmoebaMultipoleGenerator: error getting type for polarize: %s" % (atom.attrib['class'])
Justin MacCallum's avatar
Justin MacCallum committed
5109
5110
                raise ValueError(outputString)

5111
5112
    #=============================================================================================

Peter Eastman's avatar
Peter Eastman committed
5113
    def setPolarGroups(self, data, bonded12ParticleSets, force):
5114
5115
5116
5117
5118

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

            # assign multipole parameters via only 1-2 connected atoms

Peter Eastman's avatar
Peter Eastman committed
5119
5120
5121
5122
            multipoleDict = atom.multipoleDict
            pgrpMap = multipoleDict['pgrpMap']
            bondedAtomIndices = bonded12ParticleSets[atomIndex]
            atom.stage = -1
5123
5124
5125
5126
            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
5127
5128
                bondedAtom = data.atoms[bondedAtomIndex]
                if (bondedAtomType in pgrpMap):
5129
5130
                    atom.polarizationGroups[bondedAtomIndex] = 1
                    bondedAtom.polarizationGroups[atomIndex] = 1
Justin MacCallum's avatar
Justin MacCallum committed
5131

5132
5133
5134
5135
        # pgrp11

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

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

Peter Eastman's avatar
Peter Eastman committed
5139
5140
            group = set()
            visited = set()
5141
5142
            notVisited = set()
            for pgrpAtomIndex in atom.polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
5143
5144
5145
5146
                group.add(pgrpAtomIndex)
                notVisited.add(pgrpAtomIndex)
            visited.add(atomIndex)
            while(len(notVisited) > 0):
5147
                nextAtom = notVisited.pop()
Peter Eastman's avatar
Peter Eastman committed
5148
5149
                if (nextAtom not in visited):
                   visited.add(nextAtom)
5150
                   for ii in data.atoms[nextAtom].polarizationGroups:
Peter Eastman's avatar
Peter Eastman committed
5151
5152
5153
                       group.add(ii)
                       if (ii not in visited):
                           notVisited.add(ii)
5154
5155
5156

            pGroup = group
            for pgrpAtomIndex in group:
Peter Eastman's avatar
Peter Eastman committed
5157
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pGroup)
5158
5159

        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
5160
5161
            atom.polarizationGroupSet[0] = sorted(atom.polarizationGroupSet[0])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent11, atom.polarizationGroupSet[0])
5162
5163
5164
5165
5166

        # pgrp12

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

Peter Eastman's avatar
Peter Eastman committed
5167
            if (len( data.atoms[atomIndex].polarizationGroupSet) > 1):
5168
5169
                continue

Peter Eastman's avatar
Peter Eastman committed
5170
            pgrp11 = set(atom.polarizationGroupSet[0])
5171
5172
5173
            pgrp12 = set()
            for pgrpAtomIndex in pgrp11:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
5174
                    pgrp12 = pgrp12.union(data.atoms[bonded12].polarizationGroupSet[0])
5175
5176
            pgrp12 = pgrp12 - pgrp11
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
5177
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp12)
Justin MacCallum's avatar
Justin MacCallum committed
5178

5179
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
5180
5181
            atom.polarizationGroupSet[1] = sorted(atom.polarizationGroupSet[1])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent12, atom.polarizationGroupSet[1])
5182
5183
5184
5185
5186

        # pgrp13

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

Peter Eastman's avatar
Peter Eastman committed
5187
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 2):
5188
5189
                continue

Peter Eastman's avatar
Peter Eastman committed
5190
5191
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
5192
5193
5194
            pgrp13 = set()
            for pgrpAtomIndex in pgrp12:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
5195
                    pgrp13 = pgrp13.union(data.atoms[bonded12].polarizationGroupSet[0])
5196
            pgrp13 = pgrp13 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
5197
            pgrp13 = pgrp13 - set(pgrp11)
5198
            for pgrpAtomIndex in pgrp11:
Peter Eastman's avatar
Peter Eastman committed
5199
                data.atoms[pgrpAtomIndex].polarizationGroupSet.append(pgrp13)
Justin MacCallum's avatar
Justin MacCallum committed
5200

5201
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
5202
5203
            atom.polarizationGroupSet[2] = sorted(atom.polarizationGroupSet[2])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent13, atom.polarizationGroupSet[2])
5204
5205
5206
5207
5208

        # pgrp14

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

Peter Eastman's avatar
Peter Eastman committed
5209
            if (len(data.atoms[atomIndex].polarizationGroupSet) > 3):
5210
5211
                continue

Peter Eastman's avatar
Peter Eastman committed
5212
5213
5214
            pgrp11 = set(atom.polarizationGroupSet[0])
            pgrp12 = set(atom.polarizationGroupSet[1])
            pgrp13 = set(atom.polarizationGroupSet[2])
5215
5216
5217
            pgrp14 = set()
            for pgrpAtomIndex in pgrp13:
                for bonded12 in bonded12ParticleSets[pgrpAtomIndex]:
Peter Eastman's avatar
Peter Eastman committed
5218
                    pgrp14 = pgrp14.union(data.atoms[bonded12].polarizationGroupSet[0])
5219
5220
5221

            pgrp14 = pgrp14 - pgrp13
            pgrp14 = pgrp14 - pgrp12
Peter Eastman's avatar
Peter Eastman committed
5222
            pgrp14 = pgrp14 - set(pgrp11)
5223
5224

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

5227
        for (atomIndex, atom) in enumerate(data.atoms):
Peter Eastman's avatar
Peter Eastman committed
5228
5229
            atom.polarizationGroupSet[3] = sorted(atom.polarizationGroupSet[3])
            force.setCovalentMap(atomIndex, mm.AmoebaMultipoleForce.PolarizationCovalent14, atom.polarizationGroupSet[3])
5230
5231
5232

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

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

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

5238
5239
5240
5241
5242
5243
5244
        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)
5245

5246
5247
        if ('ewaldErrorTolerance' in args):
            force.setEwaldErrorTolerance(float(args['ewaldErrorTolerance']))
5248

5249
5250
5251
5252
5253
5254
5255
5256
        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)
5257

5258
5259
        if ('aEwald' in args):
            force.setAEwald(float(args['aEwald']))
5260

5261
5262
        if ('pmeGridDimensions' in args):
            force.setPmeGridDimensions(args['pmeGridDimensions'])
5263

5264
5265
        if ('mutualInducedMaxIterations' in args):
            force.setMutualInducedMaxIterations(int(args['mutualInducedMaxIterations']))
5266

5267
5268
        if ('mutualInducedTargetEpsilon' in args):
            force.setMutualInducedTargetEpsilon(float(args['mutualInducedTargetEpsilon']))
5269
5270

        # add particles to force
Justin MacCallum's avatar
Justin MacCallum committed
5271
        # throw error if particle type not available
5272
5273
5274
5275
5276

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

        # 1-2

5277
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5278
5279
5280
5281
5282

        # 1-3

        bonded13ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
5283
            bonded13Set = set()
5284
            bonded12ParticleSet = bonded12ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
5285
            for j in bonded12ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
5286
                bonded13Set = bonded13Set.union(bonded12ParticleSets[j])
5287
5288
5289
5290

            # remove 1-2 and self from set

            bonded13Set = bonded13Set - bonded12ParticleSet
Peter Eastman's avatar
Peter Eastman committed
5291
            selfSet = set()
5292
5293
            selfSet.add(i)
            bonded13Set = bonded13Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
5294
5295
            bonded13Set = set(sorted(bonded13Set))
            bonded13ParticleSets.append(bonded13Set)
5296
5297
5298
5299
5300

        # 1-4

        bonded14ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
5301
5302
            bonded14Set = set()
            bonded13ParticleSet = bonded13ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
5303
            for j in bonded13ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
5304
                bonded14Set = bonded14Set.union(bonded12ParticleSets[j])
Justin MacCallum's avatar
Justin MacCallum committed
5305

5306
5307
5308
5309
            # remove 1-3, 1-2 and self from set

            bonded14Set = bonded14Set - bonded12ParticleSets[i]
            bonded14Set = bonded14Set - bonded13ParticleSet
Peter Eastman's avatar
Peter Eastman committed
5310
            selfSet = set()
5311
5312
            selfSet.add(i)
            bonded14Set = bonded14Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
5313
5314
            bonded14Set = set(sorted(bonded14Set))
            bonded14ParticleSets.append(bonded14Set)
5315
5316
5317
5318
5319

        # 1-5

        bonded15ParticleSets = []
        for i in range(len(data.atoms)):
Peter Eastman's avatar
Peter Eastman committed
5320
5321
            bonded15Set = set()
            bonded14ParticleSet = bonded14ParticleSets[i]
Justin MacCallum's avatar
Justin MacCallum committed
5322
            for j in bonded14ParticleSet:
Peter Eastman's avatar
Peter Eastman committed
5323
                bonded15Set = bonded15Set.union(bonded12ParticleSets[j])
5324
5325
5326
5327
5328
5329

            # 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
5330
            selfSet = set()
5331
5332
            selfSet.add(i)
            bonded15Set = bonded15Set - selfSet
Peter Eastman's avatar
Peter Eastman committed
5333
5334
            bonded15Set = set(sorted(bonded15Set))
            bonded15ParticleSets.append(bonded15Set)
5335
5336
5337
5338
5339

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

Peter Eastman's avatar
Peter Eastman committed
5340
5341
                multipoleList = self.typeMap[t]
                hit = 0
5342
5343
5344
5345
5346
5347
                savedMultipoleDict = 0

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5348
                    if (hit != 0):
5349
5350
                        break

Peter Eastman's avatar
Peter Eastman committed
5351
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5352
5353

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
5354
5355
                    kx = kIndices[2]
                    ky = kIndices[3]
5356
5357
5358
5359

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

5361
                    bondedAtomIndices = bonded12ParticleSets[atomIndex]
Peter Eastman's avatar
Peter Eastman committed
5362
5363
5364
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5365
5366
                    for bondedAtomZIndex in bondedAtomIndices:

Peter Eastman's avatar
Peter Eastman committed
5367
                       if (hit != 0):
5368
5369
5370
                           break

                       bondedAtomZType = int(data.atomType[data.atoms[bondedAtomZIndex]])
Peter Eastman's avatar
Peter Eastman committed
5371
5372
                       bondedAtomZ = data.atoms[bondedAtomZIndex]
                       if (bondedAtomZType == kz):
5373
                          for bondedAtomXIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5374
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
5375
5376
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
5377
5378
5379
5380
                              if (bondedAtomXType == kx):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
5381
5382
5383
5384
5385
                                      if( bondedAtomXType == bondedAtomZType and xaxis < zaxis ):
                                          swapI = zaxis
                                          zaxis = xaxis
                                          xaxis = swapI
                                      else:
Peter Eastman's avatar
Peter Eastman committed
5386
5387
5388
5389
                                          for bondedAtomXIndex2 in bondedAtomIndices:
                                              bondedAtomX1Type = int(data.atomType[data.atoms[bondedAtomXIndex2]])
                                              if( bondedAtomX1Type == kx and bondedAtomXIndex2 != bondedAtomZIndex and bondedAtomXIndex2 < xaxis ):
                                                  xaxis = bondedAtomXIndex2
5390

5391
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5392
                                      hit = 1
5393
5394
                                  else:
                                      for bondedAtomYIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5395
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
5396
5397
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
5398
5399
5400
5401
                                          if (bondedAtomYType == ky):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
5402
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5403
                                              hit = 2
Justin MacCallum's avatar
Justin MacCallum committed
5404

5405
5406
5407
5408
                # assign multipole parameters via 1-2 and 1-3 connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5409
                    if (hit != 0):
5410
5411
                        break

Peter Eastman's avatar
Peter Eastman committed
5412
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5413
5414

                    kz = kIndices[1]
Peter Eastman's avatar
Peter Eastman committed
5415
5416
                    kx = kIndices[2]
                    ky = kIndices[3]
Justin MacCallum's avatar
Justin MacCallum committed
5417

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

5422
5423
5424
                    bondedAtom12Indices = bonded12ParticleSets[atomIndex]
                    bondedAtom13Indices = bonded13ParticleSets[atomIndex]

Peter Eastman's avatar
Peter Eastman committed
5425
5426
5427
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5428
5429
5430

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
5431
                       if (hit != 0):
5432
5433
5434
                           break

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

Peter Eastman's avatar
Peter Eastman committed
5437
                       if (bondedAtomZType == kz):
5438
5439
                          for bondedAtomXIndex in bondedAtom13Indices:

Peter Eastman's avatar
Peter Eastman committed
5440
                              if (bondedAtomXIndex == bondedAtomZIndex or hit != 0):
5441
5442
                                  continue
                              bondedAtomXType = int(data.atomType[data.atoms[bondedAtomXIndex]])
Peter Eastman's avatar
Peter Eastman committed
5443
5444
5445
5446
                              if (bondedAtomXType == kx and bondedAtomZIndex in bonded12ParticleSets[bondedAtomXIndex]):
                                  if (ky == 0):
                                      zaxis = bondedAtomZIndex
                                      xaxis = bondedAtomXIndex
5447
5448
5449

                                      # select xaxis w/ smallest index

Peter Eastman's avatar
Peter Eastman committed
5450
5451
5452
5453
                                      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
5454

5455
                                      savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5456
                                      hit = 3
5457
5458
                                  else:
                                      for bondedAtomYIndex in bondedAtom13Indices:
Peter Eastman's avatar
Peter Eastman committed
5459
                                          if (bondedAtomYIndex == bondedAtomZIndex or bondedAtomYIndex == bondedAtomXIndex or hit != 0):
5460
5461
                                              continue
                                          bondedAtomYType = int(data.atomType[data.atoms[bondedAtomYIndex]])
Peter Eastman's avatar
Peter Eastman committed
5462
5463
5464
5465
                                          if (bondedAtomYType == ky and bondedAtomZIndex in bonded12ParticleSets[bondedAtomYIndex]):
                                              zaxis = bondedAtomZIndex
                                              xaxis = bondedAtomXIndex
                                              yaxis = bondedAtomYIndex
5466
                                              savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5467
                                              hit = 4
Justin MacCallum's avatar
Justin MacCallum committed
5468

5469
5470
5471
5472
                # assign multipole parameters via only a z-defining atom

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5473
                    if (hit != 0):
5474
5475
                        break

Peter Eastman's avatar
Peter Eastman committed
5476
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5477
5478
5479
5480

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

Peter Eastman's avatar
Peter Eastman committed
5481
5482
5483
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5484
5485
5486

                    for bondedAtomZIndex in bondedAtom12Indices:

Peter Eastman's avatar
Peter Eastman committed
5487
                        if (hit != 0):
5488
5489
5490
                            break

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

Peter Eastman's avatar
Peter Eastman committed
5493
                        if (kx == 0 and kz == bondedAtomZType):
5494
                            zaxis = bondedAtomZIndex
5495
                            savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5496
                            hit = 5
5497
5498
5499
5500
5501

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

Peter Eastman's avatar
Peter Eastman committed
5502
                    if (hit != 0):
5503
5504
                        break

Peter Eastman's avatar
Peter Eastman committed
5505
                    kIndices = multipoleDict['kIndices']
Justin MacCallum's avatar
Justin MacCallum committed
5506
5507
5508

                    kz = kIndices[1]

Peter Eastman's avatar
Peter Eastman committed
5509
5510
5511
                    zaxis = -1
                    xaxis = -1
                    yaxis = -1
5512

Peter Eastman's avatar
Peter Eastman committed
5513
                    if (kz == 0):
5514
                        savedMultipoleDict = multipoleDict
Peter Eastman's avatar
Peter Eastman committed
5515
                        hit = 6
Justin MacCallum's avatar
Justin MacCallum committed
5516

5517
5518
                # add particle if there was a hit

Peter Eastman's avatar
Peter Eastman committed
5519
                if (hit != 0):
5520

Peter Eastman's avatar
Peter Eastman committed
5521
                    atom.multipoleDict = savedMultipoleDict
5522
                    atom.polarizationGroups = dict()
5523
                    newIndex = force.addMultipole(savedMultipoleDict['charge'], savedMultipoleDict['dipole'], savedMultipoleDict['quadrupole'], savedMultipoleDict['axisType'],
5524
                                                                 zaxis, xaxis, yaxis, savedMultipoleDict['thole'], savedMultipoleDict['pdamp'], savedMultipoleDict['polarizability'])
Peter Eastman's avatar
Peter Eastman committed
5525
                    if (atomIndex == newIndex):
5526
5527
5528
5529
                        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]))
5530
                    else:
5531
                        raise ValueError("Atom %s of %s %d is out of sync!." %(atom.name, atom.residue.name, atom.residue.index))
5532
                else:
Peter Eastman's avatar
Peter Eastman committed
5533
                    raise ValueError("Atom %s of %s %d was not assigned." %(atom.name, atom.residue.name, atom.residue.index))
5534
            else:
Peter Eastman's avatar
Peter Eastman committed
5535
                raise ValueError('No multipole type for atom %s %s %d' % (atom.name, atom.residue.name, atom.residue.index))
5536
5537
5538

        # set polar groups

Peter Eastman's avatar
Peter Eastman committed
5539
        self.setPolarGroups(data, bonded12ParticleSets, force)
5540
5541
5542
5543
5544

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

5545
## @private
5546
class AmoebaWcaDispersionGenerator(object):
5547
5548

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

5550
5551
    #=========================================================================================

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

Justin MacCallum's avatar
Justin MacCallum committed
5554
5555
        self.epso = epso
        self.epsh = epsh
Peter Eastman's avatar
Peter Eastman committed
5556
5557
5558
5559
5560
        self.rmino = rmino
        self.rminh = rminh
        self.awater = awater
        self.slevy = slevy
        self.dispoff = dispoff
Justin MacCallum's avatar
Justin MacCallum committed
5561
        self.shctd = shctd
5562
5563
5564
5565
5566
5567
5568
5569
5570

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

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

5572
5573
5574
5575
5576
5577
5578
5579
5580
5581
5582
5583
5584
5585
5586
        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]
5587
        generator.params.parseDefinitions(element)
Justin MacCallum's avatar
Justin MacCallum committed
5588

5589
    #=========================================================================================
Justin MacCallum's avatar
Justin MacCallum committed
5590

Peter Eastman's avatar
Peter Eastman committed
5591
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
5592
5593
5594

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

5595
        existing = [f for f in sys.getForces() if type(f) == mm.AmoebaWcaDispersionForce]
5596
5597
5598
5599
5600
5601
5602
        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
5603
        # throw error if particle type not available
5604

Peter Eastman's avatar
Peter Eastman committed
5605
5606
5607
5608
5609
5610
5611
5612
        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  ))
5613

5614
5615
5616
        for atom in data.atoms:
            values = self.params.getAtomParameters(atom, data)
            force.addParticle(values[0], values[1])
5617
5618
5619
5620
5621

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

5622
## @private
5623
class AmoebaGeneralizedKirkwoodGenerator(object):
5624
5625

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

5627
5628
    #=========================================================================================

Peter Eastman's avatar
Peter Eastman committed
5629
5630
5631
5632
5633
5634
5635
5636
5637
5638
5639
    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
5640
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5641
5642
5643
5644
5645
5646
5647
5648
5649
5650
5651
5652
5653
5654
5655
5656
5657
5658
5659
5660
5661
5662
5663
5664
5665
        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
5666
5667
5668

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

Peter Eastman's avatar
Peter Eastman committed
5669
    def getObcShct(self, data, atomIndex):
5670

Peter Eastman's avatar
Peter Eastman committed
5671
        atom = data.atoms[atomIndex]
5672
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5673
        shct = -1.0
5674
5675

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
5676

Peter Eastman's avatar
Peter Eastman committed
5677
        if (atomicNumber == 1):                 # H(1)
Justin MacCallum's avatar
Justin MacCallum committed
5678
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5679
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5680
            shct = 0.72
Peter Eastman's avatar
Peter Eastman committed
5681
        elif (atomicNumber == 7):               # N(7)
Justin MacCallum's avatar
Justin MacCallum committed
5682
            shct = 0.79
Peter Eastman's avatar
Peter Eastman committed
5683
        elif (atomicNumber == 8):               # O(8)
Justin MacCallum's avatar
Justin MacCallum committed
5684
            shct = 0.85
Peter Eastman's avatar
Peter Eastman committed
5685
        elif (atomicNumber == 9):               # F(9)
Justin MacCallum's avatar
Justin MacCallum committed
5686
5687
5688
            shct = 0.88
        elif (atomicNumber == 15):              # P(15)
            shct = 0.86
Peter Eastman's avatar
Peter Eastman committed
5689
        elif (atomicNumber == 16):              # S(16)
5690
            shct = 0.96
Peter Eastman's avatar
Peter Eastman committed
5691
        elif (atomicNumber == 26):              # Fe(26)
5692
5693
            shct = 0.88

Justin MacCallum's avatar
Justin MacCallum committed
5694
        if (shct < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5695
            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
5696
5697

        return shct
5698
5699
5700

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

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

Peter Eastman's avatar
Peter Eastman committed
5703
        atom = data.atoms[atomIndex]
5704
        atomicNumber = atom.element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5705
        radius = -1.0
5706

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

Peter Eastman's avatar
Peter Eastman committed
5709
            radius = 0.132
Justin MacCallum's avatar
Justin MacCallum committed
5710

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

5714
            for bondedAtomIndex in bondedAtomIndices:
Peter Eastman's avatar
Peter Eastman committed
5715
                bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
5716

Peter Eastman's avatar
Peter Eastman committed
5717
            if (bondedAtomAtomicNumber == 7):
5718
                radius = 0.11
Peter Eastman's avatar
Peter Eastman committed
5719
            if (bondedAtomAtomicNumber == 8):
5720
                radius = 0.105
Justin MacCallum's avatar
Justin MacCallum committed
5721

Peter Eastman's avatar
Peter Eastman committed
5722
        elif (atomicNumber == 3):               # Li(3)
5723
            radius = 0.15
Peter Eastman's avatar
Peter Eastman committed
5724
        elif (atomicNumber == 6):               # C(6)
Justin MacCallum's avatar
Justin MacCallum committed
5725

5726
            radius = 0.20
Peter Eastman's avatar
Peter Eastman committed
5727
            if (len(bondedAtomIndices) == 3):
5728
5729
                radius = 0.205

Peter Eastman's avatar
Peter Eastman committed
5730
            elif (len(bondedAtomIndices) == 4):
5731
5732
                for bondedAtomIndex in bondedAtomIndices:
                   bondedAtomAtomicNumber = data.atoms[bondedAtomIndex].element.atomic_number
Peter Eastman's avatar
Peter Eastman committed
5733
                   if (bondedAtomAtomicNumber == 7 or bondedAtomAtomicNumber == 8):
5734
5735
                       radius = 0.175

Peter Eastman's avatar
Peter Eastman committed
5736
        elif (atomicNumber == 7):               # N(7)
5737
            radius = 0.16
Peter Eastman's avatar
Peter Eastman committed
5738
        elif (atomicNumber == 8):               # O(8)
5739
            radius = 0.155
Peter Eastman's avatar
Peter Eastman committed
5740
            if (len(bondedAtomIndices) == 2):
5741
                radius = 0.145
Peter Eastman's avatar
Peter Eastman committed
5742
        elif (atomicNumber == 9):               # F(9)
5743
            radius = 0.154
Justin MacCallum's avatar
Justin MacCallum committed
5744
        elif (atomicNumber == 10):
5745
            radius = 0.146
Justin MacCallum's avatar
Justin MacCallum committed
5746
        elif (atomicNumber == 11):
5747
            radius = 0.209
Justin MacCallum's avatar
Justin MacCallum committed
5748
        elif (atomicNumber == 12):
5749
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5750
        elif (atomicNumber == 14):
5751
            radius = 0.189
Justin MacCallum's avatar
Justin MacCallum committed
5752
        elif (atomicNumber == 15):              # P(15)
5753
            radius = 0.196
Peter Eastman's avatar
Peter Eastman committed
5754
        elif (atomicNumber == 16):              # S(16)
5755
            radius = 0.186
Justin MacCallum's avatar
Justin MacCallum committed
5756
        elif (atomicNumber == 17):
5757
            radius = 0.182
Justin MacCallum's avatar
Justin MacCallum committed
5758
        elif (atomicNumber == 18):
5759
            radius = 0.179
Justin MacCallum's avatar
Justin MacCallum committed
5760
        elif (atomicNumber == 19):
5761
            radius = 0.223
Justin MacCallum's avatar
Justin MacCallum committed
5762
        elif (atomicNumber == 20):
5763
            radius = 0.191
Justin MacCallum's avatar
Justin MacCallum committed
5764
        elif (atomicNumber == 35):
5765
            radius = 2.00
Justin MacCallum's avatar
Justin MacCallum committed
5766
        elif (atomicNumber == 36):
5767
            radius = 0.190
Justin MacCallum's avatar
Justin MacCallum committed
5768
        elif (atomicNumber == 37):
5769
            radius = 0.226
Justin MacCallum's avatar
Justin MacCallum committed
5770
        elif (atomicNumber == 53):
5771
            radius = 0.237
Justin MacCallum's avatar
Justin MacCallum committed
5772
        elif (atomicNumber == 54):
5773
            radius = 0.207
Justin MacCallum's avatar
Justin MacCallum committed
5774
        elif (atomicNumber == 55):
5775
            radius = 0.263
Justin MacCallum's avatar
Justin MacCallum committed
5776
        elif (atomicNumber == 56):
5777
5778
            radius = 0.230

Justin MacCallum's avatar
Justin MacCallum committed
5779
        if (radius < 0.0):
Mark Friedrichs's avatar
Cleanup  
Mark Friedrichs committed
5780
5781
            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
5782

5783
5784
5785
5786
        return radius

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

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

Justin MacCallum's avatar
Justin MacCallum committed
5789
        bondiMap = self.radiusTypeMap['Bondi']
Peter Eastman's avatar
Peter Eastman committed
5790
        atom = data.atoms[atomIndex]
5791
        atomicNumber = atom.element.atomic_number
Justin MacCallum's avatar
Justin MacCallum committed
5792
        if (atomicNumber in bondiMap):
5793
5794
            radius = bondiMap[atomicNumber]
        else:
peastman's avatar
peastman committed
5795
            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
5796
            raise ValueError( outputString )
Justin MacCallum's avatar
Justin MacCallum committed
5797

5798
5799
5800
5801
5802
5803
5804
5805
5806
5807
        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
5808

5809
5810
5811
5812
5813
5814
5815
5816
5817
5818
5819
        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]
5820
5821

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

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

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

5828
5829
5830
        # check if AmoebaMultipoleForce exists since charges needed
        # if it has not been created, raise an error

5831
        amoebaMultipoleForceList = [f for f in sys.getForces() if type(f) == mm.AmoebaMultipoleForce]
Peter Eastman's avatar
Peter Eastman committed
5832
        if (len(amoebaMultipoleForceList) > 0):
5833
5834
5835
5836
5837
            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
5838
                if (force.__class__.__name__ == 'AmoebaMultipoleGenerator'):
Peter Eastman's avatar
Peter Eastman committed
5839
                    force.createForce(sys, data, nonbondedMethod, nonbondedCutoff, args)
Justin MacCallum's avatar
Justin MacCallum committed
5840

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

5843
        existing = [f for f in sys.getForces() if type(f) == mm.AmoebaGeneralizedKirkwoodForce]
5844
5845
5846
5847
        if len(existing) == 0:

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

Peter Eastman's avatar
Peter Eastman committed
5849
5850
            if ('solventDielectric' in args):
                force.setSolventDielectric(float(args['solventDielectric']))
5851
            else:
Peter Eastman's avatar
Peter Eastman committed
5852
                force.setSolventDielectric(   float(self.solventDielectric))
5853

Peter Eastman's avatar
Peter Eastman committed
5854
5855
            if ('soluteDielectric' in args):
                force.setSoluteDielectric(float(args['soluteDielectric']))
5856
            else:
Peter Eastman's avatar
Peter Eastman committed
5857
                force.setSoluteDielectric(    float(self.soluteDielectric))
5858

Peter Eastman's avatar
Peter Eastman committed
5859
5860
            if ('includeCavityTerm' in args):
                force.setIncludeCavityTerm(int(args['includeCavityTerm']))
5861
            else:
Peter Eastman's avatar
Peter Eastman committed
5862
               force.setIncludeCavityTerm(   int(self.includeCavityTerm))
5863
5864
5865
5866
5867

        else:
            force = existing[0]

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

Peter Eastman's avatar
Peter Eastman committed
5870
5871
        force.setProbeRadius(         float(self.probeRadius))
        force.setSurfaceAreaFactor(   float(self.surfaceAreaFactor))
5872
5873
5874

        # 1-2

5875
        bonded12ParticleSets = AmoebaVdwGenerator.getBondedParticleSets(sys, data)
5876
5877

        radiusType = 'Bondi'
Peter Eastman's avatar
Peter Eastman committed
5878
5879
5880
5881
        for atomIndex in range(0, amoebaMultipoleForce.getNumMultipoles()):
            multipoleParameters = amoebaMultipoleForce.getMultipoleParameters(atomIndex)
            if (radiusType == 'Amoeba'):
                radius = self.getAmoebaTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5882
            else:
Peter Eastman's avatar
Peter Eastman committed
5883
                radius = self.getBondiTypeRadius(data, bonded12ParticleSets[atomIndex], atomIndex)
5884
5885
            #shct = self.getObcShct(data, atomIndex)
            shct = 0.69
Peter Eastman's avatar
Peter Eastman committed
5886
            force.addParticle(multipoleParameters[0], radius, shct)
5887
5888
5889
5890
5891

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

5892
## @private
5893
class AmoebaUreyBradleyGenerator(object):
5894
5895
5896
5897

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

5899
    def __init__(self):
5900

peastman's avatar
peastman committed
5901
        self.anglesForAtom2Type = defaultdict(list)
Peter Eastman's avatar
Peter Eastman committed
5902
5903
5904
        self.types1 = []
        self.types2 = []
        self.types3 = []
5905

Peter Eastman's avatar
Peter Eastman committed
5906
5907
        self.length = []
        self.k = []
5908
5909
5910
5911
5912
5913

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

    @staticmethod
    def parseElement(element, forceField):

5914
        #  <AmoebaUreyBradleyForce>
Justin MacCallum's avatar
Justin MacCallum committed
5915
        #   <UreyBradley class1="74" class2="73" class3="74" k="16003.8" d="0.15537" />
5916

5917
        generator = AmoebaUreyBradleyGenerator()
5918
5919
        forceField._forces.append(generator)
        for bond in element.findall('UreyBradley'):
5920
            types = forceField._findAtomTypes(bond.attrib, 3)
peastman's avatar
peastman committed
5921
            if None not in types:
peastman's avatar
peastman committed
5922
                index = len(generator.types1)
5923
5924
5925
                generator.types1.append(types[0])
                generator.types2.append(types[1])
                generator.types3.append(types[2])
peastman's avatar
peastman committed
5926
5927
                for t in types[1]:
                    generator.anglesForAtom2Type[t].append(index)
5928
5929
5930
5931
5932
                generator.length.append(float(bond.attrib['d']))
                generator.k.append(float(bond.attrib['k']))

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

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

5935
        existing = [f for f in sys.getForces() if type(f) == mm.HarmonicBondForce]
5936
5937

        if len(existing) == 0:
5938
            force = mm.HarmonicBondForce()
5939
5940
5941
5942
5943
            sys.addForce(force)
        else:
            force = existing[0]

        for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
5944
            if (isConstrained and not args.get('flexibleConstraints', False)):
5945
5946
5947
5948
                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
5949
            for i in self.anglesForAtom2Type[type2]:
5950
5951
5952
                types1 = self.types1[i]
                types2 = self.types2[i]
                types3 = self.types3[i]
Peter Eastman's avatar
Peter Eastman committed
5953
                if ((type1 in types1 and type2 in types2 and type3 in types3) or (type3 in types1 and type2 in types2 and type1 in types3)):
5954
                    force.addBond(angle[0], angle[2], self.length[i], 2*self.k[i])
5955
5956
5957
5958
5959
                    break

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
5960
5961


5962
5963
5964
5965
5966
5967
5968
5969
5970
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
6017
6018
6019
6020
## @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
6021
            force.addParticle(params[0], dipole, quadrupole, *params[1:], axisType=axisType, multipoleAtomZ=zAtom, multipoleAtomX=xAtom, multipoleAtomY=yAtom)
6022
6023
6024
6025
6026
6027
6028
6029
6030
6031
6032
6033
6034
6035
6036
6037
6038
6039
6040
6041
6042
6043
6044
6045
6046
6047
6048
6049
6050
6051
6052
6053
6054
6055
6056
6057
6058
6059
6060
6061
6062
6063
6064
6065
6066
6067
        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
6068
## @private
6069
class DrudeGenerator(object):
peastman's avatar
peastman committed
6070
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
6071

6072
6073
    def __init__(self, forcefield):
        self.ff = forcefield
peastman's avatar
peastman committed
6074
6075
6076
6077
6078
6079
        self.typeMap = {}

    @staticmethod
    def parseElement(element, ff):
        existing = [f for f in ff._forces if isinstance(f, DrudeGenerator)]
        if len(existing) == 0:
6080
6081
            generator = DrudeGenerator(ff)
            ff.registerGenerator(generator)
peastman's avatar
peastman committed
6082
6083
6084
6085
        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'):
6086
            types = ff._findAtomTypes(particle.attrib, 5)
peastman's avatar
peastman committed
6087
6088
6089
6090
6091
6092
6093
6094
6095
6096
            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
6097

peastman's avatar
peastman committed
6098
6099
6100
6101
    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
6102

peastman's avatar
peastman committed
6103
        # Add Drude particles.
Justin MacCallum's avatar
Justin MacCallum committed
6104

peastman's avatar
peastman committed
6105
6106
6107
6108
6109
6110
6111
6112
6113
6114
6115
6116
6117
6118
6119
6120
        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
6121
6122
                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
6123
        sys.addForce(force)
Justin MacCallum's avatar
Justin MacCallum committed
6124

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

peastman's avatar
peastman committed
6128
6129
6130
6131
6132
6133
6134
        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)
6135
            if charge._value == 0 and epsilon._value == 0:
peastman's avatar
peastman committed
6136
6137
6138
6139
6140
                # 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]
6141
6142
                    type1 = data.atomType[data.atoms[particle1]]
                    type2 = data.atomType[data.atoms[particle2]]
peastman's avatar
peastman committed
6143
6144
6145
6146
                    thole1 = self.typeMap[type1][8]
                    thole2 = self.typeMap[type2][8]
                    drude.addScreenedPair(drude1, drude2, thole1+thole2)

6147
6148
6149
6150
6151
6152
6153
6154
6155
6156
6157
6158
6159
        # 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
6160
parsers["DrudeForce"] = DrudeGenerator.parseElement
John Chodera (MSKCC)'s avatar
John Chodera (MSKCC) committed
6161
6162

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