Commit 6f98dd5f authored by peastman's avatar peastman
Browse files

Implemented Modeller.addExtraParticles(). Also some cleanup in ForceField.

parent 7644bc07
......@@ -114,15 +114,11 @@ class ForceField(object):
template.externalBonds.append(b)
template.atoms[b].externalBonds += 1
for template in self._templates.values():
template.signature = _createResidueSignature([atom.element for atom in template.atoms])
if template.signature is None:
sigString = None
signature = _createResidueSignature([atom.element for atom in template.atoms])
if signature in self._templateSignatures:
self._templateSignatures[signature].append(template)
else:
sigString = _signatureToString(template.signature)
if sigString in self._templateSignatures:
self._templateSignatures[sigString].append(template)
else:
self._templateSignatures[sigString] = [template]
self._templateSignatures[signature] = [template]
# Build sets of every atom type belonging to each class
......@@ -260,20 +256,13 @@ class ForceField(object):
particular force fields.
Returns: the newly created System
"""
# Record atom indices
data = ForceField._SystemData()
atomIndices = {}
for index, atom in enumerate(topology.atoms()):
data.atoms.append(atom)
atomIndices[atom] = index
data.atoms = list(topology.atoms())
# Make a list of all bonds
for bond in topology.bonds():
if bond[0] in atomIndices and bond[1] in atomIndices:
data.bonds.append(ForceField._BondData(atomIndices[bond[0]], atomIndices[bond[1]]))
data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
# Record which atoms are bonded to each other atom
......@@ -294,19 +283,10 @@ class ForceField(object):
for res in chain.residues():
template = None
matches = None
sig = _createResidueSignature([atom.element for atom in res.atoms()])
if sig is not None:
signature = _signatureToString(sig)
if signature in self._templateSignatures:
for t in self._templateSignatures[signature]:
matches = _matchResidue(res, t, bondedToAtom, atomIndices)
if matches is not None:
template = t
break
if matches is None:
# Check templates involving virtual sites
for t in self._templateSignatures[None]:
matches = _matchResidue(res, t, bondedToAtom, atomIndices)
signature = _createResidueSignature([atom.element for atom in res.atoms()])
if signature in self._templateSignatures:
for t in self._templateSignatures[signature]:
matches = _matchResidue(res, t, bondedToAtom)
if matches is not None:
template = t
break
......@@ -421,7 +401,7 @@ class ForceField(object):
for atom in data.virtualSites:
site = data.virtualSites[atom]
index = atomIndices[atom]
index = atom.index
if site.type == 'average2':
sys.setVirtualSite(index, mm.TwoParticleAverageSite(index+site.atoms[0], index+site.atoms[1], site.weights[0], site.weights[1]))
elif site.type == 'average3':
......@@ -448,8 +428,8 @@ def _createResidueSignature(elements):
counts = {}
for element in elements:
if element is None:
return None # This residue contains "atoms" (probably virtual sites) that should match any element
if element in counts:
pass # This residue contains "atoms" (probably virtual sites) that should match any element
elif element in counts:
counts[element] += 1
else:
counts[element] = 1
......@@ -457,33 +437,28 @@ def _createResidueSignature(elements):
for c in counts:
sig.append((c, counts[c]))
sig.sort(key=lambda x: -x[0].mass)
return sig
# Convert it to a string.
def _signatureToString(signature):
"""Convert the signature returned by _createResidueSignature() to a string."""
s = ''
for element, count in signature:
for element, count in sig:
s += element.symbol+str(count)
return s
def _matchResidue(res, template, bondedToAtom, atomIndices):
def _matchResidue(res, template, bondedToAtom):
"""Determine whether a residue matches a template and return a list of corresponding atoms.
Parameters:
- res (Residue) The residue to check
- template (_TemplateData) The template to compare it to
- bondedToAtom (list) Enumerates which other atoms each atom is bonded to
- atomIndices (map) Maps from atoms to their indices in the System
Returns: a list specifying which atom of the template each atom of the residue corresponds to,
or None if it does not match the template
"""
atoms = list(res.atoms())
if len(atoms) != len(template.atoms):
return None
residueAtomBonds = []
templateAtomBonds = []
matches = len(atoms)*[0]
hasMatch = len(atoms)*[False]
......@@ -491,13 +466,13 @@ def _matchResidue(res, template, bondedToAtom, atomIndices):
renumberAtoms = {}
for i in range(len(atoms)):
renumberAtoms[atomIndices[atoms[i]]] = i
renumberAtoms[atoms[i].index] = i
bondedTo = []
externalBonds = []
for atom in atoms:
bonds = [renumberAtoms[x] for x in bondedToAtom[atomIndices[atom]] if x in renumberAtoms]
bonds = [renumberAtoms[x] for x in bondedToAtom[atom.index] if x in renumberAtoms]
bondedTo.append(bonds)
externalBonds.append(len([x for x in bondedToAtom[atomIndices[atom]] if x not in renumberAtoms]))
externalBonds.append(len([x for x in bondedToAtom[atom.index] if x not in renumberAtoms]))
if _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch, 0):
return matches
return None
......
......@@ -33,11 +33,12 @@ from __future__ import division
__author__ = "Peter Eastman"
__version__ = "1.0"
from simtk.openmm.app import Topology, PDBFile
from simtk.openmm.app.forcefield import HAngles
from simtk.openmm.app import Topology, PDBFile, ForceField
from simtk.openmm.app.forcefield import HAngles, _createResidueSignature, _matchResidue
from simtk.openmm.app.topology import Residue
from simtk.openmm.vec3 import Vec3
from simtk.openmm import System, Context, NonbondedForce, VerletIntegrator, LocalEnergyMinimizer
from simtk.unit import nanometer, molar, elementary_charge, amu, gram, liter, degree, sqrt, acos, is_quantity, dot, norm
from simtk.unit import nanometer, molar, elementary_charge, amu, gram, liter, degree, sqrt, acos, is_quantity, dot, norm, sum
import element as elem
import os
import random
......@@ -516,7 +517,7 @@ class Modeller(object):
terminal = hydrogen.attrib['terminal']
data.hydrogens.append(Modeller._Hydrogen(hydrogen.attrib['name'], hydrogen.attrib['parent'], maxph, atomVariants, terminal))
def addHydrogens(self, forcefield, pH=7.0, variants=None, hydrogenDefinitions=None):
def addHydrogens(self, forcefield, pH=7.0, variants=None):
"""Add missing hydrogens to the model.
Some residues can exist in multiple forms depending on the pH and properties of the local environment. These
......@@ -763,3 +764,145 @@ class Modeller(object):
self.topology = newTopology
self.positions = context.getState(getPositions=True).getPositions()
return actualVariants
def addExtraParticles(self, forcefield):
"""Add missing extra particles to the model that are required by a force field.
Some force fields use "extra particles" that do not represent actual atoms, but still need to be included in
the System. Examples include lone pairs, Drude particles, and the virtual sites used in some water models
to adjust the charge distribution. Extra particles can be recognized by the fact that their element is None.
This method is primarily used to add extra particles, but it can also remove them. It tries to match every
residue in the Topology to a template in the force field. If there is no match, it will both add and remove
extra particles as necessary to make it match.
Parameters:
- forcefield (ForceField) the ForceField defining what extra particles should be present
"""
# Create copies of all residue templates that have had all extra points removed.
templatesNoEP = {}
for resName, template in forcefield._templates.iteritems():
if any(atom.element is None for atom in template.atoms):
index = 0
newIndex = {}
newTemplate = ForceField._TemplateData(resName)
for i, atom in enumerate(template.atoms):
if atom.element is not None:
newIndex[i] = index
index += 1
newTemplate.atoms.append(ForceField._TemplateAtomData(atom.name, atom.type, atom.element))
for b1, b2 in template.bonds:
if b1 in newIndex and b2 in newIndex:
newTemplate.bonds.append((newIndex[b1], newIndex[b2]))
newTemplate.atoms[newIndex[b1]].bondedTo.append(newIndex[b2])
newTemplate.atoms[newIndex[b2]].bondedTo.append(newIndex[b1])
for b in template.externalBonds:
if b in newIndex:
newTemplate.externalBonds.append(newIndex[b])
templatesNoEP[template] = newTemplate
# Record which atoms are bonded to each other atom, with and without extra particles.
bondedToAtom = []
bondedToAtomNoEP = []
for atom in self.topology.atoms():
bondedToAtom.append(set())
if atom.element is not None:
bondedToAtomNoEP.append(set())
for atom1, atom2 in self.topology.bonds():
bondedToAtom[atom1.index].add(atom2.index)
bondedToAtom[atom2.index].add(atom1.index)
if atom1.element is not None and atom2.element is not None:
bondedToAtomNoEP[atom1.index].add(atom2.index)
bondedToAtomNoEP[atom2.index].add(atom1.index)
# Create the new Topology.
newTopology = Topology()
newTopology.setUnitCellDimensions(deepcopy(self.topology.getUnitCellDimensions()))
newAtoms = {}
newPositions = []*nanometer
for chain in self.topology.chains():
newChain = newTopology.addChain()
for residue in chain.residues():
newResidue = newTopology.addResidue(residue.name, newChain)
# Look for a matching template.
matchFound = False
signature = _createResidueSignature([atom.element for atom in residue.atoms()])
if signature in forcefield._templateSignatures:
for t in forcefield._templateSignatures[signature]:
if _matchResidue(residue, t, bondedToAtom) is not None:
matchFound = True
if matchFound:
# Just copy the residue over.
for atom in residue.atoms():
newAtom = newTopology.addAtom(atom.name, atom.element, newResidue)
newAtoms[atom] = newAtom
newPositions.append(deepcopy(self.positions[atom.index]))
else:
# There's no matching template. Try to find one that matches based on everything except
# extra points.
template = None
residueNoEP = Residue(residue.name, residue.index, residue.chain)
residueNoEP._atoms = [atom for atom in residue.atoms() if atom.element is not None]
if signature in forcefield._templateSignatures:
for t in forcefield._templateSignatures[signature]:
if t in templatesNoEP:
matches = _matchResidue(residueNoEP, templatesNoEP[t], bondedToAtomNoEP)
if matches is not None:
template = t;
# Record the corresponding atoms.
matchingAtoms = {}
for atom, match in zip(residueNoEP.atoms(), matches):
templateAtomName = t.atoms[match].name
for templateAtom in template.atoms:
if templateAtom.name == templateAtomName:
matchingAtoms[templateAtom] = atom
break
if template is None:
raise ValueError('Residue %d (%s) does not match any template defined by the ForceField.' % (residue.index+1, residue.name))
# Add the regular atoms.
for atom in residue.atoms():
if atom.element is not None:
newAtoms[atom] = newTopology.addAtom(atom.name, atom.element, newResidue)
newPositions.append(deepcopy(self.positions[atom.index]))
# Add the extra points.
templateAtomPositions = len(template.atoms)*[None]
for index, atom in enumerate(template.atoms):
if atom in matchingAtoms:
templateAtomPositions[index] = self.positions[matchingAtoms[atom].index]
for index, atom in enumerate(template.atoms):
if atom.element is None:
newTopology.addAtom(atom.name, None, newResidue)
position = None
for site in template.virtualSites:
if site.index == index:
# This is a virtual site. Compute its position by the correct rule.
if site.type == 'average2':
position = site.weights[0]*templateAtomPositions[index+site.atoms[0]] + site.weights[1]*templateAtomPositions[index+site.atoms[1]]
elif site.type == 'average3':
position = site.weights[0]*templateAtomPositions[index+site.atoms[0]] + site.weights[1]*templateAtomPositions[index+site.atoms[1]] + site.weights[2]*templateAtomPositions[index+site.atoms[2]]
elif site.type == 'outOfPlane':
position = site.weights[0]*templateAtomPositions[index+site.atoms[0]] + site.weights[1]*templateAtomPositions[index+site.atoms[1]] + site.weights[2]*templateAtomPositions[index+site.atoms[2]]
if position is None:
# We couldn't figure out the correct position. As a wild guess, just put it at the center of the residue
# and hope that energy minimization will fix it.
knownPositions = [x for x in templateAtomPositions if x is not None]
position = sum(knownPositions)/len(knownPositions)
newPositions.append(position)
for bond in self.topology.bonds():
if bond[0] in newAtoms and bond[1] in newAtoms:
newTopology.addBond(newAtoms[bond[0]], newAtoms[bond[1]])
self.topology = newTopology
self.positions = newPositions
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment