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
874
                indexMap = dict([(oldAtomIndex[name], newAtomIndex[name]) for name in newAtomIndex if name in oldAtomIndex])
                for site in template.virtualSites:
                    if site.index in indexMap and all(i in indexMap for i in site.atoms):
                        newSite = deepcopy(site)
                        newSite.index = indexMap[site.index]
                        newSite.atoms = [indexMap[i] for i in site.atoms]
                        newTemplate.virtualSites.append(newSite)
875

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

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

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

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

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

930
    class _AtomTypeParameters(object):
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
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
        """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))


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

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

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

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

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

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

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

1077
1078
1079
        Parameters
        ----------
        topology : Topology
1080
            The Topology whose residues are to be checked against the forcefield residue templates.
1081
1082
1083
1084
1085
1086
        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).
1087
1088
1089
1090

        Returns
        -------
        unmatched_residues : list of Residue
1091
            List of Residue objects from `topology` for which no forcefield residue templates are available.
1092
1093
1094
1095
1096
            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.
1097
        bondedToAtom = self._buildBondedToAtomList(topology)
1098
        unmatched_residues = list() # list of unmatched residues
1099
        for res in topology.residues():
1100
1101
1102
1103
1104
1105
1106
            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)
1107
1108
1109
            if matches is None:
                # No existing templates match.
                unmatched_residues.append(res)
1110
1111
1112

        return unmatched_residues

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

        return templates

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

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

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

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

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

1189
        return [templates, unique_unmatched_residues]
1190

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

Robert McGibbon's avatar
Robert McGibbon committed
1196
1197
1198
1199
1200
1201
        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
1202
            NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME, or LJPME.
Robert McGibbon's avatar
Robert McGibbon committed
1203
1204
1205
1206
1207
        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.
1208
        rigidWater : boolean=None
Robert McGibbon's avatar
Robert McGibbon committed
1209
            If true, water molecules will be fully rigid regardless of the value
1210
1211
            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
1212
1213
1214
1215
1216
        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
1217
1218
            their total mass the same.  If rigidWater is used to make water molecules
            rigid, then water hydrogens are not altered.
1219
        residueTemplates : dict=dict()
1220
1221
1222
1223
1224
            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).
1225
1226
1227
1228
1229
1230
        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.
1231
1232
1233
        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.
1234
1235
        flexibleConstraints : boolean=False
            If True, parameters for constrained degrees of freedom will be added to the System
1236
        drudeMass : mass=0.4*amu
1237
1238
            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
1239
        args
1240
1241
1242
            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
1243
1244
1245
1246
1247

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

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

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

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

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

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

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

1287
        # Adjust hydrogen masses if requested.
1288

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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


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

1510
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
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

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


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

1594
    # Convert it to a string.
1595
1596

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


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

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

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

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

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

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

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

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

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

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

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

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

1704
1705
1706
    return unmatchedResidues


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

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

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


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


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

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

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

1792
        # Record which clusters were successfully matched.
1793

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

1798

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

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

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

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

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

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

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

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

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

1835
        # Accept this as our new best match.
1836

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

1842
    # Return an appropriate error message.
1843

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

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

1888
1889
1890
1891
1892
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
1893
    wildcard = generator.ff._atomClasses['']
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
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
    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
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
                    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
1998
    return match
1999

2000

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

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

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

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

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

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

parsers["HarmonicBondForce"] = HarmonicBondGenerator.parseElement


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

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

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

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

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

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

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

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

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

parsers["HarmonicAngleForce"] = HarmonicAngleGenerator.parseElement


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

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

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

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

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

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

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


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

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

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

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

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

2329
    def createForce(self, sys, data, nonbondedMethod, nonbondedCutoff, args):
2330
        existing = [f for f in sys.getForces() if type(f) == mm.RBTorsionForce]
2331
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
        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:
2357
            match = _matchImproper(data, torsion, self)
2358
2359
2360
            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])
2361
2362
2363
2364

parsers["RBTorsionForce"] = RBTorsionGenerator.parseElement


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

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

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

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

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

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

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


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

2464
2465
    SCALETOL = 1e-5

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

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

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

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

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

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

parsers["NonbondedForce"] = NonbondedGenerator.parseElement

2550
2551

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

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

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

ChayaSt's avatar
ChayaSt committed
2567
        if None not in types:
2568
2569
2570
2571
2572
2573
2574
2575
            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])
2576
            for type1 in types[0]:
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
                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
2589

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

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

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

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

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

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

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

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

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

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

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

2748

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

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

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

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

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

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

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

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


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

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

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

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


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

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

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

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


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

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

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

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

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

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
    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:
2977
            match = _matchImproper(data, torsion, self)
2978
2979
2980
            if match is not None:
                (a1, a2, a3, a4, tordef) = match
                force.addTorsion(a1, a2, a3, a4, tordef.paramValues)
2981
2982
2983
2984

parsers["CustomTorsionForce"] = CustomTorsionGenerator.parseElement


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

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

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

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

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

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

parsers["CustomNonbondedForce"] = CustomNonbondedGenerator.parseElement


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

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

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

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

parsers["CustomGBForce"] = CustomGBGenerator.parseElement

3106

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
## @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'):
3145
            types = ff._findAtomTypes(donor.attrib, 3)[:generator.particlesPerDonor]
3146
3147
3148
3149
3150
3151
3152
3153
            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'):
3154
            types = ff._findAtomTypes(acceptor.attrib, 3)[:generator.particlesPerAcceptor]
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
            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,
3167
3168
3169
3170
                     CutoffPeriodic:mm.CustomHbondForce.CutoffPeriodic,
                     Ewald:mm.CustomHbondForce.CutoffPeriodic,
                     PME:mm.CustomHbondForce.CutoffPeriodic,
                     LJPME:mm.CustomHbondForce.CutoffPeriodic}
3171
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
3183
3184
3185
3186
3187
3188
3189
3190
3191
        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
3192
                    if type1 in types1:
3193
3194
3195
3196
3197
3198
3199
3200
                        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]
3201
3202
3203
3204
                    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])
3205
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
3218
3219
3220
3221
3222
3223
        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
3224
                    if type1 in types1:
3225
3226
3227
3228
3229
3230
3231
3232
                        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]
3233
                    if type1 in types1 and type2 in types2:
3234
                        force.addAcceptor(bond.atom1, bond.atom2, -1, self.acceptorParamValues[i])
3235
3236
                    elif type1 in types2 and type2 in types1:
                        force.addAcceptor(bond.atom2, bond.atom1, -1, self.acceptorParamValues[i])
3237
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
        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


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

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

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

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

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

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

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

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

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

    return returnString

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

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

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

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

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

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

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

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

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

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

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

    @staticmethod
    def parseElement(element, forceField):

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

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

3444
3445
    #=============================================================================================

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

3528
3529
3530
3531
3532
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

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

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

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

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

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

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

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

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

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

        # get force

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

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

3603
3604
        DEG_TO_RAD = math.pi / 180

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

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

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

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

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

        # get force

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

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

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

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

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

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

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

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

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

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

3718
3719
3720
3721
    #=============================================================================================

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

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

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

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

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

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
    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
3766

3767
3768
        # get global scalar parameters

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

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

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

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

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

3798
3799
    #=============================================================================================

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

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

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

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

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

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

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

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

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

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

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

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

3870
                if len(partners) == 3:
3871

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

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

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

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

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

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

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

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

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

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

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

parsers["AmoebaOutOfPlaneBendForce"] = AmoebaOutOfPlaneBendGenerator.parseElement

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

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

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

    def __init__(self, torsionUnit):

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

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

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

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

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

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

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

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

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

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

4007
4008
4009
4010
    #=============================================================================================

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

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

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

parsers["AmoebaTorsionForce"] = AmoebaTorsionGenerator.parseElement

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

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

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

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

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

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

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

    @staticmethod
    def parseElement(element, forceField):

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

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

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

4085
4086
    #=============================================================================================

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

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

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

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

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

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

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

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

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

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

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

parsers["AmoebaPiTorsionForce"] = AmoebaPiTorsionGenerator.parseElement

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

4171
4172
4173
4174
4175
4176
4177
4178
4179
4180
4181
4182
4183
4184
4185
4186
4187
4188
4189
4190
4191
4192
4193
4194
4195
4196
4197
4198
4199
4200
4201
4202
4203
4204
4205
4206
4207
4208
4209
4210
4211
4212
4213
4214
4215
4216
4217
4218
4219
4220
4221
4222
4223
4224
4225
4226
4227
4228
4229
4230
4231
4232
4233
4234
4235
4236
4237
4238
4239
4240
4241
4242
4243
4244
4245
4246
4247
4248
4249
4250
4251
4252
4253
4254
4255
4256
4257
4258
4259
4260
4261
4262
4263
4264
4265
4266
4267
4268
4269
4270
4271
4272
4273
4274
4275
4276
4277
4278
4279
4280
4281
4282
4283
4284
4285
4286
4287
4288
4289
4290
4291
4292
4293
4294
4295
4296
4297
4298
4299
4300
4301
4302
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
## @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

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

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

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

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

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

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

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

4376
4377
4378
4379
4380
    #=============================================================================================

    @staticmethod
    def parseElement(element, forceField):

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

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

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

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

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

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

            gridColIndex = 0

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

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

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

Justin MacCallum's avatar
Justin MacCallum committed
4460

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

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

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

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

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

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

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

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

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

        return chiralAtomIndex

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

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

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

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

        for angle in data.angles:

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

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

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

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

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

                                # match in reverse order

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

        # set grids

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

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

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

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

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

4606
    def __init__(self, forcefield):
4607

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

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

4616
4617
4618
4619
    #=============================================================================================

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

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

4645
4646
    #=============================================================================================

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

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

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

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

    # Note: request for constrained bonds is ignored.

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

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

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

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

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

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

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

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

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

parsers["AmoebaStretchBendForce"] = AmoebaStretchBendGenerator.parseElement

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

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

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

4756
4757
    #=============================================================================================

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

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

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

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

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

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

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

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

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

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

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

4816
4817
    #=============================================================================================

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

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

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

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

4835
        # sigma and epsilon combining rules
4836

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

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

4857
        # cutoff
4858

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

4864
        # dispersion correction
4865

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

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

4872
        # Define types
4873

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

        # add particles to force

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

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

        # Add pairs
4902

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

        # 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

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

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

4916
4917
4918
4919
4920
4921
            # 1-2 partners

            exclusionSet = bondedParticleSets[i].copy()

            # 1-3 partners

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

            # self

            exclusionSet.add(i)

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

parsers["AmoebaVdwForce"] = AmoebaVdwGenerator.parseElement

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

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

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

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

4943
4944
    #=============================================================================================

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

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

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

                # set axis type

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

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

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

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

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

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

                return axisType

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

    @staticmethod
    def parseElement(element, forceField):

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

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

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

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

                # k-indices not provided default to 0

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

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

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

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

                # set multipole

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

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

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

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

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

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

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

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

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

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

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

5110
5111
    #=============================================================================================

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

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

            # assign multipole parameters via only 1-2 connected atoms

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

5131
5132
5133
5134
        # pgrp11

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

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

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

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

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

        # pgrp12

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

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

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

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

        # pgrp13

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

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

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

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

        # pgrp14

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        # 1-2

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

        # 1-3

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

            # remove 1-2 and self from set

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

        # 1-4

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

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

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

        # 1-5

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

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

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

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

                # assign multipole parameters via only 1-2 connected atoms

                for multipoleDict in multipoleList:

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

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

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

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

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

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

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

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

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

                for multipoleDict in multipoleList:

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

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

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

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

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

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

                    for bondedAtomZIndex in bondedAtom12Indices:

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

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

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

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

                                      # select xaxis w/ smallest index

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

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

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

                for multipoleDict in multipoleList:

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

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

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

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

                    for bondedAtomZIndex in bondedAtom12Indices:

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

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

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

                # assign multipole parameters via no connected atoms

                for multipoleDict in multipoleList:

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

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

                    kz = kIndices[1]

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

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

5516
5517
                # add particle if there was a hit

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

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

        # set polar groups

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

parsers["AmoebaMultipoleForce"] = AmoebaMultipoleGenerator.parseElement

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

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

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

5549
5550
    #=========================================================================================

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

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

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

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

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

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

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

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

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

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

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

parsers["AmoebaWcaDispersionForce"] = AmoebaWcaDispersionGenerator.parseElement

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

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

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

5626
5627
    #=========================================================================================

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

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

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

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

        # shct
Justin MacCallum's avatar
Justin MacCallum committed
5675

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

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

        return shct
5697
5698
5699

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

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

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

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

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

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

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

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

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

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

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

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

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

5782
5783
5784
5785
        return radius

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        else:
            force = existing[0]

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

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

        # 1-2

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

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

parsers["AmoebaGeneralizedKirkwoodForce"] = AmoebaGeneralizedKirkwoodGenerator.parseElement

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

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

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

5898
    def __init__(self):
5899

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

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

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

    @staticmethod
    def parseElement(element, forceField):

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

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

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

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

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

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

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

parsers["AmoebaUreyBradleyForce"] = AmoebaUreyBradleyGenerator.parseElement

#=============================================================================================
peastman's avatar
peastman committed
5959
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
## @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
6020
            force.addParticle(params[0], dipole, quadrupole, *params[1:], axisType=axisType, multipoleAtomZ=zAtom, multipoleAtomX=xAtom, multipoleAtomY=yAtom)
6021
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
        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
6067
## @private
6068
class DrudeGenerator(object):
peastman's avatar
peastman committed
6069
    """A DrudeGenerator constructs a DrudeForce."""
Justin MacCallum's avatar
Justin MacCallum committed
6070

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

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

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

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

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

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

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

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

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