Source code for MMTK.ForceFields.Restraints

# Harmonic restraint terms that can be added to a force field.
#
# Written by Konrad Hinsen
#

"""
Harmonic restraint terms that can be added to any force field

Example::

 from MMTK import *
 from MMTK.ForceFields import Amber99ForceField
 from MMTK.ForceFields.Restraints import HarmonicDistanceRestraint
 
 universe = InfiniteUniverse()
 universe.protein = Protein('bala1')
 force_field = Amber99ForceField() + \
               HarmonicDistanceRestraint(universe.protein[0][1].peptide.N,
                                         universe.protein[0][1].peptide.O,
                                         0.5, 10.)
 universe.setForceField(force_field)
"""

__docformat__ = 'restructuredtext'

from MMTK.ChemicalObjects import isChemicalObject
from MMTK.Collections import isCollection
from MMTK.ForceFields.ForceField import ForceField
from MMTK import Utility
from MMTK_forcefield import HarmonicDistanceTerm, HarmonicAngleTerm, \
                            CosineDihedralTerm
from MMTK_restraints import HarmonicCMTrapTerm, HarmonicCMDistanceTerm
from Scientific import N

[docs]class HarmonicDistanceRestraint(ForceField): """ Harmonic distance restraint between two atoms """ def __init__(self, obj1, obj2, distance, force_constant, nb_exclusion=False): """ :param obj1: the object defining center-of-mass 1 :type obj1: :class:`~MMTK.Collections.GroupOfAtoms` :param obj2: the object defining center-of-mass 2 :type obj2: :class:`~MMTK.Collections.GroupOfAtoms` :param distance: the distance between cm 1 and cm2 at which the restraint is zero :type distance: float :param force_constant: the force constant of the restraint term. The functional form of the restraint is force_constant*((cm1-cm2).length()-distance)**2, where cm1 and cm2 are the centrer-of-mass positions of the two objects. :type force_constant: float :param nb_exclussion: if True, non-bonded interactions between the restrained atoms are suppressed, as for a chemical bond :type nb_exclussion: bool """ if isinstance(obj1, int) and isinstance(obj2, int): # Older MMTK versions admitted only single atoms and # stored single indices. Support this mode for opening # trajectories made with those versions self.atom_indices_1 = [obj1] self.atom_indices_2 = [obj2] if isChemicalObject(obj1) or isCollection(obj1): obj1 = obj1.atomList() if isChemicalObject(obj2) or isCollection(obj2): obj2 = obj2.atomList() self.atom_indices_1 = self.getAtomParameterIndices(obj1) self.atom_indices_2 = self.getAtomParameterIndices(obj2) if nb_exclusion and (len(self.atom_indices_1) > 1 or len(self.atom_indices_2) > 1): raise ValueError("Non-bonded exclusion possible only " "between single-atom objects") self.arguments = (self.atom_indices_1, self.atom_indices_2, distance, force_constant, nb_exclusion) self.distance = distance self.force_constant = force_constant self.nb_exclusion = nb_exclusion ForceField.__init__(self, 'harmonic distance restraint') def declareDependencies(self, global_data): global_data.add('nb_exclusions', self.__class__) def evaluatorParameters(self, universe, subset1, subset2, global_data): if universe.is_periodic and \ (len(self.atom_indices_1) > 1 or len(self.atom_indices_1) > 1): raise ValueError("Center-of-mass restraints not implemented" " for periodic universes") ok = False for s1, s2 in [(subset1, subset2), (subset2, subset1)]: if s1 is None and s1 is None: ok = True break s1 = set(a.index for a in s1.atomIterator()) diff1 = set(self.atom_indices_1).difference(s1) s2 = set(a.index for a in s2.atomIterator()) diff2 = set(self.atom_indices_2).difference(s2) if not diff1 and not diff2: # Each object is in one of the subsets ok = True break if (diff1 and len(diff1) != len(self.atom_indices_1)) \ or (diff2 and len(diff2) != len(self.atom_indices_2)): # The subset contains some but not all of the # restrained atoms. raise ValueError("Restrained atoms partially " "in a subset") global_data.add('initialized', self.__class__) if not ok: # The objects are not in the subsets, so there is no # contribution to the total energy. return {'harmonic_distance_cm': []} if self.nb_exclusion: assert len(self.atom_indices_1) == 1 assert len(self.atom_indices_2) == 1 i1 = self.atom_indices_1[0] i2 = self.atom_indices_2[0] global_data.add('excluded_pairs', Utility.normalizePair((i1, i2))) if len(self.atom_indices_1) == 1 and len(self.atom_indices_2) == 1: # Keep the old format for the single-atom case for best # compatibility with older MMTK versions. return {'harmonic_distance_term': [(self.atom_indices_1[0], self.atom_indices_2[0], self.distance, self.force_constant)]} else: return {'harmonic_distance_cm': [(self.atom_indices_1, self.atom_indices_2, self.distance, self.force_constant)]} def evaluatorTerms(self, universe, subset1, subset2, global_data): params = self.evaluatorParameters(universe, subset1, subset2, global_data) if params.has_key('harmonic_distance_term'): params = params['harmonic_distance_term'] return [HarmonicDistanceTerm(universe._spec, N.array([p[:2]]), N.array([p[2:]]), self.name) for p in params] else: params = params['harmonic_distance_cm'] return [HarmonicCMDistanceTerm(universe, N.array(p[0]), N.array(p[1]), universe.masses().array, p[2], p[3]) for p in params] def description(self): return 'ForceFields.Restraints.' + self.__class__.__name__ + \ `self.arguments`
[docs]class HarmonicAngleRestraint(ForceField): """ Harmonic angle restraint between three atoms """ def __init__(self, atom1, atom2, atom3, angle, force_constant): """ :param atom1: first atom :type atom1: :class:`~MMTK.ChemicalObjects.Atom` :param atom2: second (central) atom :type atom2: :class:`~MMTK.ChemicalObjects.Atom` :param atom3: third atom :type atom3: :class:`~MMTK.ChemicalObjects.Atom` :param angle: the angle at which the restraint is zero :type angle: float :param force_constant: the force constant of the restraint term. The functional form of the restraint is force_constant*(phi-angle)**2, where phi is the angle atom1-atom2-atom3. """ self.index1, self.index2, self.index3 = \ self.getAtomParameterIndices((atom1, atom2, atom3)) self.arguments = (self.index1, self.index2, self.index3, angle, force_constant) self.angle = angle self.force_constant = force_constant ForceField.__init__(self, 'harmonic angle restraint') def evaluatorParameters(self, universe, subset1, subset2, global_data): return {'harmonic_angle_term': [(self.index1, self.index2, self.index3, self.angle, self.force_constant)]} def evaluatorTerms(self, universe, subset1, subset2, global_data): params = self.evaluatorParameters(universe, subset1, subset2, global_data)['harmonic_angle_term'] assert len(params) == 1 indices = N.array([params[0][:3]]) parameters = N.array([params[0][3:]]) return [HarmonicAngleTerm(universe._spec, indices, parameters, self.name)]
[docs]class HarmonicDihedralRestraint(ForceField): """ Harmonic dihedral angle restraint between four atoms """ def __init__(self, atom1, atom2, atom3, atom4, dihedral, force_constant): """ :param atom1: first atom :type atom1: :class:`~MMTK.ChemicalObjects.Atom` :param atom2: second (axis) atom :type atom2: :class:`~MMTK.ChemicalObjects.Atom` :param atom3: third (axis)atom :type atom3: :class:`~MMTK.ChemicalObjects.Atom` :param atom4: fourth atom :type atom4: :class:`~MMTK.ChemicalObjects.Atom` :param dihedral: the dihedral angle at which the restraint is zero :type dihedral: float :param force_constant: the force constant of the restraint term. The functional form of the restraint is force_constant*(phi-abs(dihedral))**2, where phi is the dihedral angle atom1-atom2-atom3-atom4. """ self.index1, self.index2, self.index3, self.index4 = \ self.getAtomParameterIndices((atom1, atom2, atom3, atom4)) self.dihedral = dihedral self.force_constant = force_constant self.arguments = (self.index1, self.index2, self.index3, self.index4, dihedral, force_constant) ForceField.__init__(self, 'harmonic dihedral restraint') def evaluatorParameters(self, universe, subset1, subset2, global_data): return {'cosine_dihedral_term': [(self.index1, self.index2, self.index3, self.index4, 0., self.dihedral, 0., self.force_constant)]} def evaluatorTerms(self, universe, subset1, subset2, global_data): params = self.evaluatorParameters(universe, subset1, subset2, global_data)['cosine_dihedral_term'] assert len(params) == 1 indices = N.array([params[0][:4]]) parameters = N.array([params[0][4:]]) return [CosineDihedralTerm(universe._spec, indices, parameters, self.name)]
[docs]class HarmonicTrapForceField(ForceField): """ Harmonic potential with respect to a fixed point in space """ def __init__(self, obj, center, force_constant): """ :param obj: the object on whose center of mass the force field acts :type obj: :class:`~MMTK.Collections.GroupOfAtoms` :param center: the point to which the atom is attached by the harmonic potential :type center: Scientific.Geometry.Vector :param force_constant: the force constant of the harmonic potential :type force_constant: float """ if isChemicalObject(obj): obj = obj.atomList() self.atom_indices = self.getAtomParameterIndices(obj) self.arguments = (self.atom_indices, center, force_constant) ForceField.__init__(self, 'harmonic_trap') self.center = center self.force_constant = force_constant def ready(self, global_data): return True def evaluatorParameters(self, universe, subset1, subset2, global_data): if universe.is_periodic and len(self.atom_indices) > 1: raise ValueError("Center-of-mass restraints not implemented" " for periodic universes") for subset in [subset1, subset2]: if subset is not None: subset = set(a.index for a in subset.atomIterator()) diff = set(self.atom_indices).difference(subset) if diff: if len(diff) == len(self.atom_indices): # The subset doesn't contain the restrained atoms. return {'harmonic_trap_cm': []} else: # The subset contains some but not all of the # restrained atoms. raise ValueError("Restrained atoms partially " "in a subset") return {'harmonic_trap_cm': [(self.atom_indices, self.center, self.force_constant)]} def evaluatorTerms(self, universe, subset1, subset2, global_data): params = self.evaluatorParameters(universe, subset1, subset2, global_data)['harmonic_trap_cm'] return [HarmonicCMTrapTerm(universe, N.array(p[0]), universe.masses().array, p[1], p[2]) for p in params]